Merge branch 'cursors-rownumber' into devel

NEWS

@@ -9,6 +9,8 @@ What's new in psycopg 2.4.5
     interaction (ticket #90).
   - Better efficiency and formatting of timezone offset objects thanks
     to Menno Smits (tickets #94, #95).
+  - Fixed 'rownumber' during iteration on cursor subclasses.
+    Regression introduced in 2.4.4 (ticket #100).
 
 
 What's new in psycopg 2.4.4

lib/extras.py

@@ -88,25 +88,17 @@ class DictCursorBase(_cursor):
 
     def __iter__(self):
         if self._prefetch:
-            res = _cursor.fetchmany(self, self.itersize)
+            res = _cursor.__iter__(self)
-            if not res:
+            first = res.next()
-                return
         if self._query_executed:
             self._build_index()
         if not self._prefetch:
-            res = _cursor.fetchmany(self, self.itersize)
+            res = _cursor.__iter__(self)
+            first = res.next()
 
-        for r in res:
+        yield first
-            yield r
-
-        # the above was the first itersize record. the following are
-        # in a repeated loop.
         while 1:
-            res = _cursor.fetchmany(self, self.itersize)
+            yield res.next()
-            if not res:
-                return
-            for r in res:
-                yield r
 
 
 class DictConnection(_connection):
@@ -318,14 +310,17 @@ class NamedTupleCursor(_cursor):
         return [nt(*t) for t in ts]
 
     def __iter__(self):
-        # Invoking _cursor.__iter__(self) goes to infinite recursion,
+        it = _cursor.__iter__(self)
-        # so we do pagination by hand
+        t = it.next()
+
+        nt = self.Record
+        if nt is None:
+            nt = self.Record = self._make_nt()
+
+        yield nt(*t)
+
         while 1:
-            recs = self.fetchmany(self.itersize)
+            yield nt(*it.next())
-            if not recs:
-                return
-            for rec in recs:
-                yield rec
 
     try:
         from collections import namedtuple

tests/test_cursor.py

@@ -234,6 +234,17 @@ class CursorTests(unittest.TestCase):
         # everything swallowed in two gulps
         self.assertEqual(rv, [(i,((i - 1) % 30) + 1) for i in range(1,51)])
 
+    @skip_before_postgres(8, 0)
+    def test_iter_named_cursor_rownumber(self):
+        curs = self.conn.cursor('tmp')
+        # note: this fails if itersize < dataset: internally we check
+        # rownumber == rowcount to detect when to read anoter page, so we
+        # would need an extra attribute to have a monotonic rownumber.
+        curs.itersize = 20
+        curs.execute('select generate_series(1,10)')
+        for i, rec in enumerate(curs):
+            self.assertEqual(i + 1, curs.rownumber)
+
     @skip_if_no_namedtuple
     def test_namedtuple_description(self):
         curs = self.conn.cursor()

tests/test_extras_dictcursor.py

@@ -35,6 +35,7 @@ class ExtrasDictCursorTests(unittest.TestCase):
     def tearDown(self):
         self.conn.close()
 
+
     def testDictCursorWithPlainCursorFetchOne(self):
         self._testWithPlainCursor(lambda curs: curs.fetchone())
 
@@ -53,6 +54,26 @@ class ExtrasDictCursorTests(unittest.TestCase):
                 return row
         self._testWithPlainCursor(getter)
 
+    def testUpdateRow(self):
+        row = self._testWithPlainCursor(lambda curs: curs.fetchone())
+        row['foo'] = 'qux'
+        self.failUnless(row['foo'] == 'qux')
+        self.failUnless(row[0] == 'qux')
+
+    @skip_before_postgres(8, 0)
+    def testDictCursorWithPlainCursorIterRowNumber(self):
+        curs = self.conn.cursor(cursor_factory=psycopg2.extras.DictCursor)
+        self._testIterRowNumber(curs)
+
+    def _testWithPlainCursor(self, getter):
+        curs = self.conn.cursor(cursor_factory=psycopg2.extras.DictCursor)
+        curs.execute("SELECT * FROM ExtrasDictCursorTests")
+        row = getter(curs)
+        self.failUnless(row['foo'] == 'bar')
+        self.failUnless(row[0] == 'bar')
+        return row
+
+
     def testDictCursorWithPlainCursorRealFetchOne(self):
         self._testWithPlainCursorReal(lambda curs: curs.fetchone())
 
@@ -71,6 +92,17 @@ class ExtrasDictCursorTests(unittest.TestCase):
                 return row
         self._testWithPlainCursorReal(getter)
 
+    @skip_before_postgres(8, 0)
+    def testDictCursorWithPlainCursorRealIterRowNumber(self):
+        curs = self.conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor)
+        self._testIterRowNumber(curs)
+
+    def _testWithPlainCursorReal(self, getter):
+        curs = self.conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor)
+        curs.execute("SELECT * FROM ExtrasDictCursorTests")
+        row = getter(curs)
+        self.failUnless(row['foo'] == 'bar')
+
 
     def testDictCursorWithNamedCursorFetchOne(self):
         self._testWithNamedCursor(lambda curs: curs.fetchone())
@@ -95,6 +127,18 @@ class ExtrasDictCursorTests(unittest.TestCase):
         curs = self.conn.cursor('tmp', cursor_factory=psycopg2.extras.DictCursor)
         self._testNamedCursorNotGreedy(curs)
 
+    @skip_before_postgres(8, 0)
+    def testDictCursorWithNamedCursorIterRowNumber(self):
+        curs = self.conn.cursor('tmp', cursor_factory=psycopg2.extras.DictCursor)
+        self._testIterRowNumber(curs)
+
+    def _testWithNamedCursor(self, getter):
+        curs = self.conn.cursor('aname', cursor_factory=psycopg2.extras.DictCursor)
+        curs.execute("SELECT * FROM ExtrasDictCursorTests")
+        row = getter(curs)
+        self.failUnless(row['foo'] == 'bar')
+        self.failUnless(row[0] == 'bar')
+
 
     def testDictCursorRealWithNamedCursorFetchOne(self):
         self._testWithNamedCursorReal(lambda curs: curs.fetchone())
@@ -119,27 +163,10 @@ class ExtrasDictCursorTests(unittest.TestCase):
         curs = self.conn.cursor('tmp', cursor_factory=psycopg2.extras.RealDictCursor)
         self._testNamedCursorNotGreedy(curs)
 
-
+    @skip_before_postgres(8, 0)
-    def _testWithPlainCursor(self, getter):
+    def testDictCursorRealWithNamedCursorIterRowNumber(self):
-        curs = self.conn.cursor(cursor_factory=psycopg2.extras.DictCursor)
+        curs = self.conn.cursor('tmp', cursor_factory=psycopg2.extras.RealDictCursor)
-        curs.execute("SELECT * FROM ExtrasDictCursorTests")
+        self._testIterRowNumber(curs)
-        row = getter(curs)
-        self.failUnless(row['foo'] == 'bar')
-        self.failUnless(row[0] == 'bar')
-        return row
-
-    def _testWithNamedCursor(self, getter):
-        curs = self.conn.cursor('aname', cursor_factory=psycopg2.extras.DictCursor)
-        curs.execute("SELECT * FROM ExtrasDictCursorTests")
-        row = getter(curs)
-        self.failUnless(row['foo'] == 'bar')
-        self.failUnless(row[0] == 'bar')
-
-    def _testWithPlainCursorReal(self, getter):
-        curs = self.conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor)
-        curs.execute("SELECT * FROM ExtrasDictCursorTests")
-        row = getter(curs)
-        self.failUnless(row['foo'] == 'bar')
 
     def _testWithNamedCursorReal(self, getter):
         curs = self.conn.cursor('aname', cursor_factory=psycopg2.extras.RealDictCursor)
@@ -147,11 +174,6 @@ class ExtrasDictCursorTests(unittest.TestCase):
         row = getter(curs)
         self.failUnless(row['foo'] == 'bar')
 
-    def testUpdateRow(self):
-        row = self._testWithPlainCursor(lambda curs: curs.fetchone())
-        row['foo'] = 'qux'
-        self.failUnless(row['foo'] == 'qux')
-        self.failUnless(row[0] == 'qux')
 
     def _testNamedCursorNotGreedy(self, curs):
         curs.itersize = 2
@@ -165,6 +187,14 @@ class ExtrasDictCursorTests(unittest.TestCase):
         self.assert_(recs[1]['ts'] - recs[0]['ts'] < timedelta(seconds=0.005))
         self.assert_(recs[2]['ts'] - recs[1]['ts'] > timedelta(seconds=0.0099))
 
+    def _testIterRowNumber(self, curs):
+        # Only checking for dataset < itersize:
+        # see CursorTests.test_iter_named_cursor_rownumber
+        curs.itersize = 20
+        curs.execute("""select * from generate_series(1,10)""")
+        for i, r in enumerate(curs):
+            self.assertEqual(i + 1, curs.rownumber)
+
 
 class NamedTupleCursorTest(unittest.TestCase):
     def setUp(self):
@@ -192,12 +222,14 @@ class NamedTupleCursorTest(unittest.TestCase):
     @skip_if_no_namedtuple
     def test_fetchone(self):
         curs = self.conn.cursor()
-        curs.execute("select * from nttest where i = 1")
+        curs.execute("select * from nttest order by 1")
         t = curs.fetchone()
         self.assertEqual(t[0], 1)
         self.assertEqual(t.i, 1)
         self.assertEqual(t[1], 'foo')
         self.assertEqual(t.s, 'foo')
+        self.assertEqual(curs.rownumber, 1)
+        self.assertEqual(curs.rowcount, 3)
 
     @skip_if_no_namedtuple
     def test_fetchmany_noarg(self):
@@ -210,6 +242,8 @@ class NamedTupleCursorTest(unittest.TestCase):
         self.assertEqual(res[0].s, 'foo')
         self.assertEqual(res[1].i, 2)
         self.assertEqual(res[1].s, 'bar')
+        self.assertEqual(curs.rownumber, 2)
+        self.assertEqual(curs.rowcount, 3)
 
     @skip_if_no_namedtuple
     def test_fetchmany(self):
@@ -221,6 +255,8 @@ class NamedTupleCursorTest(unittest.TestCase):
         self.assertEqual(res[0].s, 'foo')
         self.assertEqual(res[1].i, 2)
         self.assertEqual(res[1].s, 'bar')
+        self.assertEqual(curs.rownumber, 2)
+        self.assertEqual(curs.rowcount, 3)
 
     @skip_if_no_namedtuple
     def test_fetchall(self):
@@ -234,6 +270,8 @@ class NamedTupleCursorTest(unittest.TestCase):
         self.assertEqual(res[1].s, 'bar')
         self.assertEqual(res[2].i, 3)
         self.assertEqual(res[2].s, 'baz')
+        self.assertEqual(curs.rownumber, 3)
+        self.assertEqual(curs.rowcount, 3)
 
     @skip_if_no_namedtuple
     def test_executemany(self):
@@ -251,16 +289,26 @@ class NamedTupleCursorTest(unittest.TestCase):
         curs = self.conn.cursor()
         curs.execute("select * from nttest order by 1")
         i = iter(curs)
+        self.assertEqual(curs.rownumber, 0)
+
         t = i.next()
         self.assertEqual(t.i, 1)
         self.assertEqual(t.s, 'foo')
+        self.assertEqual(curs.rownumber, 1)
+        self.assertEqual(curs.rowcount, 3)
+
         t = i.next()
         self.assertEqual(t.i, 2)
         self.assertEqual(t.s, 'bar')
+        self.assertEqual(curs.rownumber, 2)
+        self.assertEqual(curs.rowcount, 3)
+
         t = i.next()
         self.assertEqual(t.i, 3)
         self.assertEqual(t.s, 'baz')
         self.assertRaises(StopIteration, i.next)
+        self.assertEqual(curs.rownumber, 3)
+        self.assertEqual(curs.rowcount, 3)
 
     def test_error_message(self):
         try:
@@ -385,6 +433,17 @@ class NamedTupleCursorTest(unittest.TestCase):
         self.assert_(recs[1].ts - recs[0].ts < timedelta(seconds=0.005))
         self.assert_(recs[2].ts - recs[1].ts > timedelta(seconds=0.0099))
 
+    @skip_if_no_namedtuple
+    @skip_before_postgres(8, 0)
+    def test_named_rownumber(self):
+        curs = self.conn.cursor('tmp')
+        # Only checking for dataset < itersize:
+        # see CursorTests.test_iter_named_cursor_rownumber
+        curs.itersize = 4
+        curs.execute("""select * from generate_series(1,3)""")
+        for i, t in enumerate(curs):
+            self.assertEqual(i + 1, curs.rownumber)
+
 
 def test_suite():
     return unittest.TestLoader().loadTestsFromName(__name__)