diff --git a/ChangeLog b/ChangeLog index d1bd3ded..ffaab92e 100644 --- a/ChangeLog +++ b/ChangeLog @@ -1,5 +1,9 @@ 2009-03-02 Federico Di Gregorio + * Applied patch from Menno Smits to avoid problems with DictCursor + when the query is executed by a named cursor. Also added Menno's + tests and uniformed other DictCursor tests. + * psycopg/psycopgmodule.c: fixed unwanted exception when passing an explicit None to register_type(). diff --git a/lib/extras.py b/lib/extras.py index 297d039e..78156467 100644 --- a/lib/extras.py +++ b/lib/extras.py @@ -19,12 +19,14 @@ and classes untill a better place in the distribution is found. import os import time +import re as regex try: import logging except: logging = None +from psycopg2 import DATETIME, DataError from psycopg2 import extensions as _ext from psycopg2.extensions import cursor as _cursor from psycopg2.extensions import connection as _connection @@ -46,26 +48,29 @@ class DictCursorBase(_cursor): self.row_factory = row_factory def fetchone(self): + res = _cursor.fetchone(self) if self._query_executed: self._build_index() - return _cursor.fetchone(self) + return res def fetchmany(self, size=None): + res = _cursor.fetchmany(self, size) if self._query_executed: self._build_index() - return _cursor.fetchmany(self, size) + return res def fetchall(self): + res = _cursor.fetchall(self) if self._query_executed: self._build_index() - return _cursor.fetchall(self) + return res def next(self): - if self._query_executed: - self._build_index() res = _cursor.fetchone(self) if res is None: raise StopIteration() + if self._query_executed: + self._build_index() return res class DictConnection(_connection): @@ -74,7 +79,7 @@ class DictConnection(_connection): if name is None: return _connection.cursor(self, cursor_factory=DictCursor) else: - return _connection.cursor(self, name, cursor_factory=DictCursor) + return _connection.cursor(self, name, cursor_factory=DictCursor) class DictCursor(DictCursorBase): """A cursor that keeps a list of column name -> index mappings.""" @@ -302,12 +307,12 @@ try: __str__ = getquoted - def register_uuid(oid=None): + def register_uuid(oid=None, conn_or_curs=None): """Create the UUID type and an uuid.UUID adapter.""" if not oid: oid = 2950 _ext.UUID = _ext.new_type((oid, ), "UUID", lambda data, cursor: data and uuid.UUID(data) or None) - _ext.register_type(_ext.UUID) + _ext.register_type(_ext.UUID, conn_or_curs) _ext.register_adapter(uuid.UUID, UUID_adapter) return _ext.UUID @@ -346,12 +351,12 @@ class Inet(object): def __str__(self): return str(self.addr) -def register_inet(oid=None): +def register_inet(oid=None, conn_or_curs=None): """Create the INET type and an Inet adapter.""" if not oid: oid = 869 _ext.INET = _ext.new_type((oid, ), "INET", lambda data, cursor: data and Inet(data) or None) - _ext.register_type(_ext.INET) + _ext.register_type(_ext.INET, conn_or_curs) return _ext.INET diff --git a/tests/__init__.py b/tests/__init__.py index b8c39304..bce49187 100755 --- a/tests/__init__.py +++ b/tests/__init__.py @@ -24,6 +24,7 @@ import test_quote import test_connection import test_transaction import types_basic +import types_extras import test_lobject def test_suite(): @@ -36,6 +37,7 @@ def test_suite(): suite.addTest(test_connection.test_suite()) suite.addTest(test_transaction.test_suite()) suite.addTest(types_basic.test_suite()) + suite.addTest(types_extras.test_suite()) suite.addTest(test_lobject.test_suite()) return suite diff --git a/tests/extras_dictcursor.py b/tests/extras_dictcursor.py old mode 100755 new mode 100644 index 323cc93f..11579ba0 --- a/tests/extras_dictcursor.py +++ b/tests/extras_dictcursor.py @@ -27,18 +27,55 @@ class ExtrasDictCursorTests(unittest.TestCase): self.conn = psycopg2.connect(tests.dsn) curs = self.conn.cursor() curs.execute("CREATE TEMPORARY TABLE ExtrasDictCursorTests (foo text)") + curs.execute("INSERT INTO ExtrasDictCursorTests VALUES ('bar')") + self.conn.commit() def tearDown(self): self.conn.close() - def testDictCursor(self): + def testDictCursorWithPlainCursorFetchOne(self): + self._testWithPlainCursor(lambda curs: curs.fetchone()) + + def testDictCursorWithPlainCursorFetchMany(self): + self._testWithPlainCursor(lambda curs: curs.fetchmany(100)[0]) + + def testDictCursorWithPlainCursorFetchAll(self): + self._testWithPlainCursor(lambda curs: curs.fetchall()[0]) + + def testDictCursorWithPlainCursorIter(self): + def getter(curs): + for row in curs: + return row + self._testWithPlainCursor(getter) + + def testDictCursorWithNamedCursorFetchOne(self): + self._testWithNamedCursor(lambda curs: curs.fetchone()) + + def testDictCursorWithNamedCursorFetchMany(self): + self._testWithNamedCursor(lambda curs: curs.fetchmany(100)[0]) + + def testDictCursorWithNamedCursorFetchAll(self): + self._testWithNamedCursor(lambda curs: curs.fetchall()[0]) + + def testDictCursorWithNamedCursorIter(self): + def getter(curs): + for row in curs: + return row + self._testWithNamedCursor(getter) + + def _testWithPlainCursor(self, getter): curs = self.conn.cursor(cursor_factory=psycopg2.extras.DictCursor) - curs.execute("INSERT INTO ExtrasDictCursorTests VALUES ('bar')") curs.execute("SELECT * FROM ExtrasDictCursorTests") - row = curs.fetchone() + row = getter(curs) self.failUnless(row['foo'] == 'bar') self.failUnless(row[0] == 'bar') + 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 test_suite(): return unittest.TestLoader().loadTestsFromName(__name__) diff --git a/tests/types_extras.py b/tests/types_extras.py index 1c0d0d7e..1342dd15 100644 --- a/tests/types_extras.py +++ b/tests/types_extras.py @@ -25,7 +25,7 @@ import psycopg2.extras import tests -class TypesBasicTests(unittest.TestCase): +class TypesExtrasTests(unittest.TestCase): """Test that all type conversions are working.""" def setUp(self):