Adding "Streaming" cursor functionality

pyhs2/connections.py

@@ -16,7 +16,7 @@ class Connection(object):
     client = None
     session = None
 
-    def __init__(self, host=None, port=10000, authMechanism=None, user=None, password=None, database=None, configuration=None):
+    def __init__(self, host=None, port=10000, authMechanism=None, user=None, password=None, database=None, configuration=None, cursorclass=Cursor):
         authMechanisms = set(['NOSASL', 'PLAIN', 'KERBEROS', 'LDAP'])
         if authMechanism not in authMechanisms:
             raise NotImplementedError('authMechanism is either not supported or not implemented')
@@ -69,8 +69,10 @@ def _get_krb_settings(self, default_host, config):
 
         return host, service
 
-    def cursor(self):
-        return Cursor(self.client, self.session)
+    def cursor(self, cursor = None):
+        if cursor:
+            return cursor(self.client, self.session)
+        return self.cursorclass(self.client, self.session)
 
     def close(self):
         req = TCloseSessionReq(sessionHandle=self.session)

pyhs2/cursor.py

@@ -105,3 +105,23 @@ def close(self):
         if self.operationHandle is not None:
             req = TCloseOperationReq(operationHandle=self.operationHandle)
             self.client.CloseOperation(req) 
+
+
+class SSCursor(Cursor):
+    """
+    Unbuffered Cursor, mainly useful for queries that return a lot of data,
+    or for connections to remote servers over a slow network.
+    """
+    def fetch(self):    
+        fetchReq = TFetchResultsReq(operationHandle=self.operationHandle,
+                                    orientation=TFetchOrientation.FETCH_NEXT,
+                                    maxRows=10000)
+        while True:
+            resultsRes = self.client.FetchResults(fetchReq)
+            for row in resultsRes.results.rows:
+                rowData= []
+                for col in row.colVals:
+                    rowData.append(get_value(col))
+                yield rowData
+            if len(resultsRes.results.rows) == 0:
+                break