Make ResultSet handle more than just list results

aholmberg · aholmberg · commit d4c380cb010f · 2015-10-16T14:45:42.000-05:00
for use with custom ProtocolHandlers providing different ResultMessage deserialization

PYTHON-430
diff --git a/cassandra/cluster.py b/cassandra/cluster.py
@@ -61,7 +61,7 @@
                                 BatchMessage, RESULT_KIND_PREPARED,
                                 RESULT_KIND_SET_KEYSPACE, RESULT_KIND_ROWS,
                                 RESULT_KIND_SCHEMA_CHANGE, MIN_SUPPORTED_VERSION,
-                                ProtocolHandler)
+                                ProtocolHandler, _RESULT_SEQUENCE_TYPES)
 from cassandra.metadata import Metadata, protect_name, murmur3
 from cassandra.policies import (TokenAwarePolicy, DCAwareRoundRobinPolicy, SimpleConvictionPolicy,
                                 ExponentialReconnectionPolicy, HostDistance,
@@ -3345,7 +3345,7 @@ class ResultSet(object):
 
     def __init__(self, response_future, initial_response):
         self.response_future = response_future
-        self._current_rows = initial_response or []
+        self._set_current_rows(initial_response)
         self._page_iter = None
         self._list_mode = False
 
@@ -3386,10 +3386,16 @@ def fetch_next_page(self):
         if self.response_future.has_more_pages:
             self.response_future.start_fetching_next_page()
             result = self.response_future.result()
-            self._current_rows = result._current_rows
+            self._current_rows = result._current_rows  # ResultSet has already _set_current_rows to the appropriate form
         else:
             self._current_rows = []
 
+    def _set_current_rows(self, result):
+        if isinstance(result, _RESULT_SEQUENCE_TYPES):
+            self._current_rows = result
+        else:
+            self._current_rows = [result] if result else []
+
     def _fetch_all(self):
         self._current_rows = list(self)
         self._page_iter = None
diff --git a/cassandra/obj_parser.pyx b/cassandra/obj_parser.pyx
@@ -38,13 +38,17 @@ cdef class LazyParser(ColumnParser):
         # supported in cpdef methods
         return parse_rows_lazy(reader, desc)
 
+    cpdef get_cython_generator_type(self):
+        return get_cython_generator_type()
 
 def parse_rows_lazy(BytesIOReader reader, ParseDesc desc):
     cdef Py_ssize_t i, rowcount
     rowcount = read_int(reader)
     cdef RowParser rowparser = TupleRowParser()
     return (rowparser.unpack_row(reader, desc) for i in range(rowcount))
 
+def get_cython_generator_type():
+    return type((i for i in range(0)))
 
 cdef class TupleRowParser(RowParser):
     """
diff --git a/cassandra/protocol.py b/cassandra/protocol.py
@@ -1000,6 +1000,7 @@ def decode_message(cls, protocol_version, user_type_map, stream_id, flags, opcod
 
         return msg
 
+_RESULT_SEQUENCE_TYPES = (list, tuple)  # types retuned by ResultMessages
 
 def cython_protocol_handler(colparser):
     """
@@ -1045,6 +1046,9 @@ class CythonProtocolHandler(ProtocolHandler):
 if HAVE_CYTHON:
     from cassandra.obj_parser import ListParser, LazyParser
     ProtocolHandler = cython_protocol_handler(ListParser())
+
+    lazy_parser = LazyParser()
+    _RESULT_SEQUENCE_TYPES += (lazy_parser.get_cython_generator_type(),)
     LazyProtocolHandler = cython_protocol_handler(LazyParser())
 else:
     # Use Python-based ProtocolHandler
