Named DictCursor/RealDictCursor honour itersize

Closes ticket #80.

3 files changed, 49 insertions, 8 deletions

diff --git a/NEWS b/NEWS
index c255b8f..4bb218e 100644
--- a/NEWS
+++ b/NEWS
@@ -1,3 +1,9 @@
+What's new in psycopg 2.4.4
+---------------------------
+
+  - Named DictCursor/RealDictCursor honour itersize (ticket #80).
+
+
 What's new in psycopg 2.4.3
 ---------------------------
 
diff --git a/lib/extras.py b/lib/extras.py
index 78f962e..491a390 100644
--- a/lib/extras.py
+++ b/lib/extras.py
@@ -86,18 +86,28 @@ class DictCursorBase(_cursor):
             res = _cursor.fetchall(self)
         return res
 
-    def next(self):
+    def __iter__(self):
         if self._prefetch:
-            res = _cursor.fetchone(self)
-            if res is None:
-                raise StopIteration()
+            res = _cursor.fetchmany(self, self.itersize)
+            if not res:
+                return
         if self._query_executed:
             self._build_index()
         if not self._prefetch:
-            res = _cursor.fetchone(self)
-            if res is None:
-                raise StopIteration()
-        return res
+            res = _cursor.fetchmany(self, self.itersize)
+
+        for r in res:
+            yield r
+
+        # the above was the first itersize record. the following are
+        # in a repeated loop.
+        while 1:
+            res = _cursor.fetchmany(self, self.itersize)
+            if not res:
+                return
+            for r in res:
+                yield r
+
 
 class DictConnection(_connection):
     """A connection that uses `DictCursor` automatically."""
diff --git a/tests/test_extras_dictcursor.py b/tests/test_extras_dictcursor.py
index 0af5afd..3346c11 100755
--- a/tests/test_extras_dictcursor.py
+++ b/tests/test_extras_dictcursor.py
@@ -65,6 +65,7 @@ class ExtrasDictCursorTests(unittest.TestCase):
                 return row
         self._testWithPlainCursorReal(getter)
 
+
     def testDictCursorWithNamedCursorFetchOne(self):
         self._testWithNamedCursor(lambda curs: curs.fetchone())
 
@@ -80,6 +81,12 @@ class ExtrasDictCursorTests(unittest.TestCase):
                 return row
         self._testWithNamedCursor(getter)
 
+    @skip_before_postgres(8, 2)
+    def testDictCursorWithNamedCursorNotGreedy(self):
+        curs = self.conn.cursor('tmp', cursor_factory=psycopg2.extras.DictCursor)
+        self._testNamedCursorNotGreedy(curs)
+
+
     def testDictCursorRealWithNamedCursorFetchOne(self):
         self._testWithNamedCursorReal(lambda curs: curs.fetchone())
 
@@ -95,6 +102,12 @@ class ExtrasDictCursorTests(unittest.TestCase):
                 return row
         self._testWithNamedCursorReal(getter)
 
+    @skip_before_postgres(8, 2)
+    def testDictCursorRealWithNamedCursorNotGreedy(self):
+        curs = self.conn.cursor('tmp', cursor_factory=psycopg2.extras.RealDictCursor)
+        self._testNamedCursorNotGreedy(curs)
+
+
     def _testWithPlainCursor(self, getter):
         curs = self.conn.cursor(cursor_factory=psycopg2.extras.DictCursor)
         curs.execute("SELECT * FROM ExtrasDictCursorTests")
@@ -128,6 +141,18 @@ class ExtrasDictCursorTests(unittest.TestCase):
         self.failUnless(row['foo'] == 'qux')
         self.failUnless(row[0] == 'qux')
 
+    def _testNamedCursorNotGreedy(self, curs):
+        curs.itersize = 2
+        curs.execute("""select clock_timestamp() as ts from generate_series(1,3)""")
+        recs = []
+        for t in curs:
+            time.sleep(0.01)
+            recs.append(t)
+
+        # check that the dataset was not fetched in a single gulp
+        self.assert_(recs[1]['ts'] - recs[0]['ts'] < timedelta(seconds=0.005))
+        self.assert_(recs[2]['ts'] - recs[1]['ts'] > timedelta(seconds=0.0099))
+
 
 class NamedTupleCursorTest(unittest.TestCase):
     def setUp(self):