psycopg/psycopg2
1
1
Fork 0
mirror of https://github.com/psycopg/psycopg2.git synced 2025-01-31 17:34:08 +03:00
Code Issues Packages Projects Releases Wiki Activity

Merge branch 'cursors-rownumber' into devel

Browse Source
This commit is contained in:
Daniele Varrazzo 2012-02-23 23:05:40 +00:00
commit e57f3284eb
4 changed files with 115 additions and 48 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
NEWS
View File

@ -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

37
lib/extras.py
View File

@ -88,25 +88,17 @@ class DictCursorBase(_cursor):
def __iter__(self):
if self._prefetch:
res = _cursor.fetchmany(self, self.itersize)
if not res:
return
res = _cursor.__iter__(self)
first = res.next()
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 r
# the above was the first itersize record. the following are
# in a repeated loop.
yield first
while 1:
res = _cursor.fetchmany(self, self.itersize)
if not res:
return
for r in res:
yield r
yield res.next()
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,
# so we do pagination by hand
it = _cursor.__iter__(self)
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)
if not recs:
return
for rec in recs:
yield rec
yield nt(*it.next())
try:
from collections import namedtuple

11
tests/test_cursor.py
View File

@ -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()

113
tests/test_extras_dictcursor.py
View File

@ -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)
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 _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')
@skip_before_postgres(8, 0)
def testDictCursorRealWithNamedCursorIterRowNumber(self):
curs = self.conn.cursor('tmp', cursor_factory=psycopg2.extras.RealDictCursor)
self._testIterRowNumber(curs)
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__)