Commit 6cfd220a by Lukáš Lalinský

Refactor the code to use the cursor_factory pattern used in psycopg2 and other drivers

parent 3aa2f11b
......@@ -29,6 +29,7 @@ The library implements the standard DB API 2.0 interface, so it can be
used the same way you would use any other SQL database from Python, for example::
import phoenixdb
import phoenixdb.cursor
database_url = 'http://localhost:8765/'
conn = phoenixdb.connect(database_url, autocommit=True)
......@@ -39,6 +40,11 @@ used the same way you would use any other SQL database from Python, for example:
cursor.execute("SELECT * FROM users")
print cursor.fetchall()
cursor = conn.cursor(cursor_factory=phoenixdb.cursor.DictCursor)
cursor.execute("SELECT * FROM users WHERE id=1")
print cursor.fetchone()['USERNAME']
Setting up a development environment
------------------------------------
......
......@@ -56,6 +56,9 @@ def connect(url, max_retries=None, **kwargs):
:param max_retries:
The maximum number of retries in case there is a connection error.
:param cursor_factory:
If specified, the connection's :attr:`~phoenixdb.connection.Connection.cursor_factory` is set to it.
:returns:
:class:`~phoenixdb.connection.Connection` object.
"""
......
......@@ -31,9 +31,18 @@ class Connection(object):
You should not construct this object manually, use :func:`~phoenixdb.connect` instead.
"""
def __init__(self, client, **kwargs):
cursor_factory = None
"""
The default cursor factory used by :meth:`cursor` if the parameter is not specified.
"""
def __init__(self, client, cursor_factory=None, **kwargs):
self._client = client
self._closed = False
if cursor_factory is not None:
self.cursor_factory = cursor_factory
else:
self.cursor_factory = Cursor
self._cursors = []
# Extract properties to pass to OpenConnectionRequest
self._connection_args = {}
......@@ -97,15 +106,22 @@ class Connection(object):
if self._closed:
raise ProgrammingError('the connection is already closed')
def cursor(self):
def cursor(self, cursor_factory=None):
"""Creates a new cursor.
:param cursor_factory:
This argument can be used to create non-standard cursors.
The class returned must be a subclass of
:class:`~phoenixdb.cursor.Cursor` (for example :class:`~phoenixdb.cursor.DictCursor`).
A default factory for the connection can also be specified using the
:attr:`cursor_factory` attribute.
:returns:
A :class:`~phoenixdb.cursor.Cursor` object.
"""
if self._closed:
raise ProgrammingError('the connection is already closed')
cursor = Cursor(self)
cursor = (cursor_factory or self.cursor_factory)(self)
self._cursors.append(weakref.ref(cursor, self._cursors.remove))
return cursor
......@@ -165,5 +181,6 @@ class Connection(object):
props = self._client.connection_sync(self._id, {'transactionIsolation': bool(value)})
self._transactionisolation = props.transaction_isolation
for name in errors.__all__:
setattr(Connection, name, getattr(errors, name))
......@@ -18,7 +18,7 @@ from phoenixdb.types import TypeHelper
from phoenixdb.errors import OperationalError, NotSupportedError, ProgrammingError, InternalError
from phoenixdb.calcite import common_pb2
__all__ = ['Cursor', 'ColumnDescription']
__all__ = ['Cursor', 'ColumnDescription', 'DictCursor']
logger = logging.getLogger(__name__)
......@@ -257,7 +257,7 @@ class Cursor(object):
tmp_row.append(value)
return tmp_row
def fetchone(self, doc=False):
def fetchone(self):
if self._frame is None:
raise ProgrammingError('no select statement was executed')
if self._pos is None:
......@@ -269,37 +269,28 @@ class Cursor(object):
self._pos = None
if not self._frame.done:
self._fetch_next_frame()
if doc:
return self.docResult(row=row)
else:
return row
return row
def fetchmany(self, size=None, doc=False):
def fetchmany(self, size=None):
if size is None:
size = self.arraysize
rows = []
while size > 0:
row = self.fetchone(doc=doc)
row = self.fetchone()
if row is None:
break
rows.append(row)
size -= 1
return rows
def fetchall(self, doc=False):
def fetchall(self):
rows = []
while True:
row = self.fetchone(doc=doc)
row = self.fetchone()
if row is None:
break
rows.append(row)
return rows
def docResult(self, row):
d = {}
for ind, val in enumerate(row):
d[self._signature.columns.__getitem__(ind).column_name] = val
return d
def setinputsizes(self, sizes):
pass
......@@ -336,3 +327,14 @@ class Cursor(object):
if self._frame is not None and self._pos is not None:
return self._frame.offset + self._pos
return self._pos
class DictCursor(Cursor):
"""A cursor which returns results as a dictionary"""
def _transform_row(self, row):
row = super(DictCursor, self)._transform_row(row)
d = {}
for ind, val in enumerate(row):
d[self._signature.columns[ind].column_name] = val
return d
import unittest
import phoenixdb
import phoenixdb.cursor
from phoenixdb.tests import TEST_DB_URL
......@@ -33,3 +34,32 @@ class PhoenixDatabaseTest(unittest.TestCase):
cursor.itersize = 4
cursor.execute("SELECT * FROM test WHERE id>? ORDER BY id", [1])
self.assertEqual(cursor.fetchall(), [[i, 'text {}'.format(i)] for i in range(2, 10)])
def _check_dict_cursor(self, cursor):
cursor.execute("DROP TABLE IF EXISTS test")
cursor.execute("CREATE TABLE test (id INTEGER PRIMARY KEY, text VARCHAR)")
cursor.execute("UPSERT INTO test VALUES (?, ?)", [1, 'text 1'])
cursor.execute("SELECT * FROM test ORDER BY id")
self.assertEqual(cursor.fetchall(), [{'ID': 1, 'TEXT': 'text 1'}])
def test_dict_cursor_default_parameter(self):
db = phoenixdb.connect(TEST_DB_URL, autocommit=True, cursor_factory=phoenixdb.cursor.DictCursor)
self.addCleanup(db.close)
with db.cursor() as cursor:
self._check_dict_cursor(cursor)
def test_dict_cursor_default_attribute(self):
db = phoenixdb.connect(TEST_DB_URL, autocommit=True)
db.cursor_factory = phoenixdb.cursor.DictCursor
self.addCleanup(db.close)
with db.cursor() as cursor:
self._check_dict_cursor(cursor)
def test_dict_cursor(self):
db = phoenixdb.connect(TEST_DB_URL, autocommit=True)
self.addCleanup(db.close)
with db.cursor(cursor_factory=phoenixdb.cursor.DictCursor) as cursor:
self._check_dict_cursor(cursor)
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment