Allow specification of alternate return formats

thobbs · thobbs · commit 2c498ff5d291 · 2013-05-20T18:53:46.000-05:00
diff --git a/cassandra/cluster.py b/cassandra/cluster.py
@@ -13,7 +13,8 @@
                                WriteTimeoutErrorMessage,
                                UnavailableErrorMessage,
                                OverloadedErrorMessage,
-                               IsBootstrappingErrorMessage)
+                               IsBootstrappingErrorMessage, named_tuple_factory,
+                               dict_factory)
 from cassandra.metadata import Metadata
 from cassandra.policies import (RoundRobinPolicy, SimpleConvictionPolicy,
                                 ExponentialReconnectionPolicy, HostDistance,
@@ -430,6 +431,19 @@ class Session(object):
     keyspace = None
     is_shutdown = False
 
+    row_factory = staticmethod(named_tuple_factory)
+    """
+    The format to return row results in.  By default, each
+    returned row will be a named tuple.  You can alternatively
+    use any of the following:
+
+      - :func:`cassandra.decoder.tuple_factory`
+      - :func:`cassandra.decoder.named_tuple_factory`
+      - :func:`cassandra.decoder.dict_factory`
+      - :func:`cassandra.decoder.ordered_dict_factory`
+
+    """
+
     _lock = None
     _pools = None
     _load_balancer = None
@@ -798,9 +812,14 @@ def _refresh_schema(self, connection, keyspace=None, table=None):
 
         if ks_query:
             ks_result, cf_result, col_result = connection.wait_for_responses(ks_query, cf_query, col_query)
+            ks_result = dict_factory(*ks_result.results)
+            cf_result = dict_factory(*cf_result.results)
+            col_result = dict_factory(*col_result.results)
         else:
             ks_result = None
             cf_result, col_result = connection.wait_for_responses(cf_query, col_query)
+            cf_result = dict_factory(*cf_result.results)
+            col_result = dict_factory(*col_result.results)
 
         self._cluster.metadata.rebuild_schema(keyspace, table, ks_result, cf_result, col_result)
 
@@ -817,12 +836,14 @@ def _refresh_node_list_and_token_map(self, connection):
         peers_query = QueryMessage(query=self._SELECT_PEERS, consistency_level=cl)
         local_query = QueryMessage(query=self._SELECT_LOCAL, consistency_level=cl)
         peers_result, local_result = connection.wait_for_responses(peers_query, local_query)
+        peers_result = dict_factory(*peers_result.results)
 
         partitioner = None
         token_map = {}
 
         if local_result.results:
-            local_row = local_result.results[0]
+            local_rows = dict_factory(*(local_result.results))
+            local_row = local_rows[0]
             cluster_name = local_row["cluster_name"]
             self._cluster.metadata.cluster_name = cluster_name
 
@@ -836,7 +857,7 @@ def _refresh_node_list_and_token_map(self, connection):
                 token_map[host] = tokens
 
         found_hosts = set()
-        for row in peers_result.results:
+        for row in peers_result:
             addr = row.get("rpc_address")
 
             # TODO handle ipv6 equivalent
@@ -909,14 +930,15 @@ def wait_for_schema_agreement(self, connection=None):
             peers_query = QueryMessage(query=self._SELECT_SCHEMA_PEERS, consistency_level=cl)
             local_query = QueryMessage(query=self._SELECT_SCHEMA_LOCAL, consistency_level=cl)
             peers_result, local_result = connection.wait_for_responses(peers_query, local_query)
+            peers_result = dict_factory(*peers_result.results)
 
             versions = set()
             if local_result.results:
-                local_row = local_result.results[0]
+                local_row = dict_factory(*local_result.results)[0]
                 if local_row.get("schema_version"):
                     versions.add(local_row.get("schema_version"))
 
-            for row in peers_result.results:
+            for row in peers_result:
                 if not row.get("rpc_address") or not row.get("schema_version"):
                     continue
 
@@ -1074,6 +1096,7 @@ class ResponseFuture(object):
 
     def __init__(self, session, message, query):
         self.session = session
+        self.row_factory = session.row_factory
         self.message = message
         self.query = query
 
@@ -1149,7 +1172,10 @@ def _set_result(self, response):
                     self.session.cluster.control_connection,
                     self)
             else:
-                self._set_final_result(getattr(response, 'results', None))
+                results = getattr(response, 'results', None)
+                if results:
+                    results = self.row_factory(*results)
+                self._set_final_result(results)
         elif isinstance(response, ErrorMessage):
             retry_policy = self.query.retry_policy
             if not retry_policy:
diff --git a/cassandra/decoder.py b/cassandra/decoder.py
@@ -14,6 +14,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from collections import namedtuple, OrderedDict
 import socket
 try:
     from cStringIO import StringIO
@@ -46,6 +47,23 @@ def warn(msg):
     print msg
 
 
+def tuple_factory(colnames, rows):
+    return rows
+
+
+def named_tuple_factory(colnames, rows):
+    Row = namedtuple('Row', colnames)
+    return [Row(*row) for row in rows]
+
+
+def dict_factory(colnames, rows):
+    return [dict(zip(colnames, row)) for row in rows]
+
+
+def ordered_dict_factory(colnames, rows):
+    return [OrderedDict(zip(colnames, row)) for row in rows]
+
+
 class PreparedResult:
     def __init__(self, queryid, param_metadata):
         self.queryid = queryid
@@ -393,8 +411,8 @@ def recv_results_rows(cls, f):
         rows = [cls.recv_row(f, len(colspecs)) for x in xrange(rowcount)]
         colnames = [c[2] for c in colspecs]
         coltypes = [c[3] for c in colspecs]
-        return [dict(zip(colnames, [ctype.from_binary(val) for ctype, val in zip(coltypes, row)]))
-                for row in rows]
+        return (colnames, [tuple(ctype.from_binary(val) for ctype, val in zip(coltypes, row))
+                           for row in rows])
 
     @classmethod
     def recv_results_prepared(cls, f):
diff --git a/cassandra/metadata.py b/cassandra/metadata.py
@@ -63,10 +63,10 @@ def rebuild_schema(self, keyspace, table, ks_results, cf_results, col_results):
         cf_def_rows = defaultdict(list)
         col_def_rows = defaultdict(lambda: defaultdict(list))
 
-        for row in cf_results.results:
+        for row in cf_results:
             cf_def_rows[row["keyspace_name"]].append(row)
 
-        for row in col_results.results:
+        for row in col_results:
             ksname = row["keyspace_name"]
             cfname = row["columnfamily_name"]
             col_def_rows[ksname][cfname].append(row)
@@ -75,7 +75,7 @@ def rebuild_schema(self, keyspace, table, ks_results, cf_results, col_results):
         if not table:
             # ks_results is not None
             added_keyspaces = set()
-            for row in ks_results.results:
+            for row in ks_results:
                 keyspace_meta = self._build_keyspace_metadata(row)
                 for table_row in cf_def_rows.get(keyspace_meta.name, []):
                     table_meta = self._build_table_metadata(
diff --git a/example.py b/example.py
@@ -15,7 +15,7 @@ def main():
     s = c.connect()
 
     rows = s.execute("SELECT keyspace_name FROM system.schema_keyspaces")
-    if KEYSPACE in [row.values()[0] for row in rows]:
+    if KEYSPACE in [row[0] for row in rows]:
         print "dropping existing keyspace..."
         s.execute("DROP KEYSPACE " + KEYSPACE)
 
@@ -55,7 +55,7 @@ def main():
         log.exeception()
 
     for row in rows:
-        print '\t'.join(row.values())
+        print '\t'.join(row)
 
     s.execute("DROP KEYSPACE " + KEYSPACE)
 
