From 9667e9892263a91e237ebaba955326d19fa0fc54 Mon Sep 17 00:00:00 2001
From: Gene Peters <epeters@kixeye.com>
Date: Tue, 10 Jun 2014 11:33:48 -0700
Subject: [PATCH] Adding stream cursor functionality, similar to MySQLdb

---
 pyhs2/connections.py |  8 +++++---
 pyhs2/cursor.py      | 20 ++++++++++++++++++++
 2 files changed, 25 insertions(+), 3 deletions(-)

diff --git a/pyhs2/connections.py b/pyhs2/connections.py
index 0760ddb..d9d7ec5 100644
--- a/pyhs2/connections.py
+++ b/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)
diff --git a/pyhs2/cursor.py b/pyhs2/cursor.py
index 8dfe8f3..9af1c58 100644
--- a/pyhs2/cursor.py
+++ b/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