Add dict and namedtuple cursor factories #290

clickhouse_driver/__init__.py

@@ -3,7 +3,7 @@
 from .dbapi import connect
 
 
-VERSION = (0, 2, 3)
+VERSION = (0, 2, 4)
 __version__ = '.'.join(str(x) for x in VERSION)
 
 __all__ = ['Client', 'connect']

clickhouse_driver/dbapi/connection.py

@@ -79,9 +79,11 @@ def rollback(self):
         """
         pass
 
-    def cursor(self):
+    def cursor(self, cursor_factory=None):
         """
-        :return: a new Cursor Object using the connection.
+        :param cursor_factory: Argument can be used to create non-standard
+                               cursors.
+        :return: a new cursor object using the connection.
         """
         if self.is_closed:
             raise InterfaceError('connection already closed')
@@ -91,6 +93,7 @@ def cursor(self):
             self._hosts = client.connection.hosts
         else:
             client.connection.hosts = self._hosts
-        cursor = Cursor(client, self)
+        cursor_factory = cursor_factory or Cursor
+        cursor = cursor_factory(client, self)
         self.cursors.append(cursor)
         return cursor

clickhouse_driver/dbapi/extras.py

@@ -0,0 +1,73 @@
+import re
+from collections import namedtuple
+from functools import lru_cache
+
+from .cursor import Cursor
+
+
+class DictCursor(Cursor):
+    """
+    A cursor that generates results as :class:`dict`.
+
+    ``fetch*()`` methods will return dicts instead of tuples.
+    """
+
+    def fetchone(self):
+        rv = super(DictCursor, self).fetchone()
+        if rv is not None:
+            rv = dict(zip(self._columns, rv))
+        return rv
+
+    def fetchmany(self, size=None):
+        rv = super(DictCursor, self).fetchmany(size=size)
+        return [dict(zip(self._columns, x)) for x in rv]
+
+    def fetchall(self):
+        rv = super(DictCursor, self).fetchall()
+        return [dict(zip(self._columns, x)) for x in rv]
+
+
+class NamedTupleCursor(Cursor):
+    """
+    A cursor that generates results as named tuples created by
+    :func:`~collections.namedtuple`.
+
+    ``fetch*()`` methods will return named tuples instead of regular tuples, so
+    their elements can be accessed both as regular numeric items as well as
+    attributes.
+    """
+
+    # ascii except alnum and underscore
+    _re_clean = re.compile(
+        '[' + re.escape(' !"#$%&\'()*+,-./:;<=>?@[\\]^`{|}~') + ']')
+
+    @classmethod
+    @lru_cache(512)
+    def _make_nt(self, key):
+        fields = []
+        for s in key:
+            s = self._re_clean.sub('_', s)
+            # Python identifier cannot start with numbers, namedtuple fields
+            # cannot start with underscore.
+            if s[0] == '_' or '0' <= s[0] <= '9':
+                s = 'f' + s
+            fields.append(s)
+
+        return namedtuple('Record', fields)
+
+    def fetchone(self):
+        rv = super(NamedTupleCursor, self).fetchone()
+        if rv is not None:
+            nt = self._make_nt(self._columns)
+            rv = nt(*rv)
+        return rv
+
+    def fetchmany(self, size=None):
+        rv = super(NamedTupleCursor, self).fetchmany(size=size)
+        nt = self._make_nt(self._columns)
+        return [nt(*x) for x in rv]
+
+    def fetchall(self):
+        rv = super(NamedTupleCursor, self).fetchall()
+        nt = self._make_nt(self._columns)
+        return [nt(*x) for x in rv]

docs/dbapi.rst

@@ -27,3 +27,12 @@ Cursor
 .. autoclass:: clickhouse_driver.dbapi.cursor.Cursor
    :members:
    :inherited-members:
+
+Extras
+------
+
+.. _dbapi-extras:
+
+.. autoclass:: clickhouse_driver.dbapi.extras.DictCursor
+
+.. autoclass:: clickhouse_driver.dbapi.extras.NamedTupleCursor

docs/features.rst

@@ -352,6 +352,26 @@ managers:
         >>>        print(cursor.fetchall())
 
 
+You can use ``cursor_factory`` argument to get results as dicts or named tuples
+(since version 0.2.4):
+
+    .. code-block:: python
+
+        >>> from clickhouse_driver.dbapi.extras import DictCursor
+        >>> with connect('clickhouse://localhost') as conn:
+        >>>     with conn.cursor(cursor_factory=DictCursor) as cursor:
+        >>>        cursor.execute('SELECT * FROM system.tables')
+        >>>        print(cursor.fetchall())
+
+    .. code-block:: python
+
+        >>> from clickhouse_driver.dbapi.extras import NamedTupleCursor
+        >>> with connect('clickhouse://localhost') as conn:
+        >>>     with conn.cursor(cursor_factory=NamedTupleCursor) as cursor:
+        >>>        cursor.execute('SELECT * FROM system.tables')
+        >>>        print(cursor.fetchall())
+
+
 NumPy/Pandas support
 --------------------
 

tests/test_dbapi.py

@@ -1,4 +1,5 @@
 import types
+from collections import namedtuple
 from contextlib import contextmanager
 import socket
 from unittest.mock import patch
@@ -7,6 +8,7 @@
 from clickhouse_driver.dbapi import (
     ProgrammingError, InterfaceError, OperationalError
 )
+from clickhouse_driver.dbapi.extras import DictCursor, NamedTupleCursor
 from tests.testcase import BaseTestCase
 
 
@@ -30,8 +32,9 @@ def created_connection(self, **kwargs):
 
     @contextmanager
     def created_cursor(self, **kwargs):
+        cursor_kwargs = kwargs.pop('cursor_kwargs', {})
         with self.created_connection(**kwargs) as connection:
-            cursor = connection.cursor()
+            cursor = connection.cursor(**cursor_kwargs)
 
             try:
                 yield cursor
@@ -368,3 +371,126 @@ def test_execute_fetch_before_query(self):
             with self.assertRaises(ProgrammingError) as e:
                 cursor.fetchall()
             self.assertEqual(str(e.exception), 'no results to fetch')
+
+
+class DictCursorFactoryTestCase(DBAPITestCaseBase):
+    def test_execute_fetchone(self):
+        cursor_kwargs = {'cursor_factory': DictCursor}
+
+        with self.created_cursor(cursor_kwargs=cursor_kwargs) as cursor:
+            cursor.execute('SELECT number FROM system.numbers LIMIT 4')
+
+            self.assertIsInstance(cursor._rows, list)
+            self.assertEqual(cursor.fetchone(), {'number': 0})
+            self.assertEqual(cursor.fetchone(), {'number': 1})
+            self.assertEqual(cursor.fetchone(), {'number': 2})
+            self.assertEqual(cursor.fetchone(), {'number': 3})
+            self.assertEqual(cursor.fetchone(), None)
+
+    def test_execute_fetchmany(self):
+        cursor_kwargs = {'cursor_factory': DictCursor}
+
+        with self.created_cursor(cursor_kwargs=cursor_kwargs) as cursor:
+            cursor.execute('SELECT number FROM system.numbers LIMIT 4')
+
+            self.assertIsInstance(cursor._rows, list)
+            self.assertEqual(cursor.fetchmany(), [{'number': 0}])
+            self.assertEqual(cursor.fetchmany(None), [{'number': 1}])
+            self.assertEqual(cursor.fetchmany(0), [])
+            self.assertEqual(
+                cursor.fetchmany(-1), [{'number': 2}, {'number': 3}]
+            )
+
+            cursor.execute('SELECT number FROM system.numbers LIMIT 4')
+            self.assertEqual(cursor.fetchmany(1), [{'number': 0}])
+            self.assertEqual(
+                cursor.fetchmany(2), [{'number': 1}, {'number': 2}]
+            )
+            self.assertEqual(cursor.fetchmany(3), [{'number': 3}])
+            self.assertEqual(cursor.fetchmany(3), [])
+
+            cursor.arraysize = 2
+            cursor.execute('SELECT number FROM system.numbers LIMIT 4')
+            self.assertEqual(
+                cursor.fetchmany(), [{'number': 0}, {'number': 1}])
+            self.assertEqual(
+                cursor.fetchmany(), [{'number': 2}, {'number': 3}]
+            )
+
+    def test_execute_fetchall(self):
+        cursor_kwargs = {'cursor_factory': DictCursor}
+
+        with self.created_cursor(cursor_kwargs=cursor_kwargs) as cursor:
+            cursor.execute('SELECT number FROM system.numbers LIMIT 4')
+            self.assertEqual(cursor.rowcount, 4)
+            self.assertEqual(
+                cursor.fetchall(), [
+                    {'number': 0}, {'number': 1}, {'number': 2}, {'number': 3}
+                ])
+
+
+class NamedTupleCursorFactoryTestCase(DBAPITestCaseBase):
+    def test_execute_fetchone(self):
+        cursor_kwargs = {'cursor_factory': NamedTupleCursor}
+
+        with self.created_cursor(cursor_kwargs=cursor_kwargs) as cursor:
+            cursor.execute('SELECT number FROM system.numbers LIMIT 4')
+
+            self.assertIsInstance(cursor._rows, list)
+            nt = namedtuple('Record', cursor._columns)
+
+            self.assertEqual(cursor.fetchone(), nt(0))
+            self.assertEqual(cursor.fetchone(), nt(1))
+            self.assertEqual(cursor.fetchone(), nt(2))
+            self.assertEqual(cursor.fetchone(), nt(3))
+            self.assertEqual(cursor.fetchone(), None)
+
+    def test_execute_fetchmany(self):
+        cursor_kwargs = {'cursor_factory': NamedTupleCursor}
+
+        with self.created_cursor(cursor_kwargs=cursor_kwargs) as cursor:
+            cursor.execute('SELECT number FROM system.numbers LIMIT 4')
+
+            self.assertIsInstance(cursor._rows, list)
+            nt = namedtuple('Record', cursor._columns)
+
+            self.assertEqual(cursor.fetchmany(), [nt(0)])
+            self.assertEqual(cursor.fetchmany(None), [nt(1)])
+            self.assertEqual(cursor.fetchmany(0), [])
+            self.assertEqual(cursor.fetchmany(-1), [nt(2), nt(3)])
+
+            cursor.execute('SELECT number FROM system.numbers LIMIT 4')
+            self.assertEqual(cursor.fetchmany(1), [nt(0)])
+            self.assertEqual(cursor.fetchmany(2), [nt(1), nt(2)])
+            self.assertEqual(cursor.fetchmany(3), [nt(3)])
+            self.assertEqual(cursor.fetchmany(3), [])
+
+            cursor.arraysize = 2
+            cursor.execute('SELECT number FROM system.numbers LIMIT 4')
+            self.assertEqual(cursor.fetchmany(), [nt(0), nt(1)])
+            self.assertEqual(cursor.fetchmany(), [nt(2), nt(3)])
+
+    def test_execute_fetchall(self):
+        cursor_kwargs = {'cursor_factory': NamedTupleCursor}
+
+        with self.created_cursor(cursor_kwargs=cursor_kwargs) as cursor:
+            cursor.execute('SELECT number FROM system.numbers LIMIT 4')
+            self.assertEqual(cursor.rowcount, 4)
+            nt = namedtuple('Record', cursor._columns)
+
+            self.assertEqual(cursor.fetchall(), [nt(0), nt(1), nt(2), nt(3)])
+
+    def test_make_nt_caching(self):
+        cursor_kwargs = {'cursor_factory': NamedTupleCursor}
+
+        with self.created_cursor(cursor_kwargs=cursor_kwargs) as cursor:
+            cursor.execute('SELECT number FROM system.numbers LIMIT 4')
+
+            self.assertIsInstance(cursor._rows, list)
+            nt = namedtuple('Record', cursor._columns)
+
+            self.assertEqual(cursor.fetchone(), nt(0))
+
+            with patch('clickhouse_driver.dbapi.extras.namedtuple') as nt_mock:
+                nt_mock.side_effect = ValueError
+                self.assertEqual(cursor.fetchone(), nt(1))