Expose cursor_factory to configuration; #10

Author: chadwhitacre

postgres.py

@@ -209,7 +209,8 @@ class Postgres(object):
 
     :param unicode url: A ``postgres://`` URL or a `PostgreSQL connection string <http://www.postgresql.org/docs/current/static/libpq-connect.html>`_
     :param int minconn: The minimum size of the connection pool
-    :param int maxconn: The minimum size of the connection pool
+    :param int maxconn: The maximum size of the connection pool
+    :param cursor_factory: Defaults to :py:class:`~psycopg2.extras.RealDictCursor`
     :param strict_one: The default :py:attr:`strict` parameter for :py:meth:`~postgres.Postgres.one`
     :type strict_one: :py:class:`bool`
 
@@ -223,6 +224,11 @@ class Postgres(object):
     :py:class:`~postgres.Postgres` instance is that it runs everything through
     its connection pool.
 
+    Check the :py:mod:`psycopg2` `docs
+    <http://initd.org/psycopg/docs/extras.html#connection-and-cursor-subclasses>`_
+    for additional :py:attr:`cursor_factories`, such as
+    :py:class:`NamedTupleCursor`.
+
     The names in our simple API, :py:meth:`~postgres.Postgres.run`,
     :py:meth:`~postgres.Postgres.one`, and :py:meth:`~postgres.Postgres.rows`,
     were chosen to be short and memorable, and to not conflict with the DB-API
@@ -242,10 +248,13 @@ class Postgres(object):
 
     """
 
-    def __init__(self, url, minconn=1, maxconn=10, strict_one=None):
+    def __init__(self, url, minconn=1, maxconn=10, \
+                               cursor_factory=RealDictCursor, strict_one=None):
         if url.startswith("postgres://"):
             dsn = url_to_dsn(url)
 
+        Connection.cursor_factory = cursor_factory
+
         self.pool = ConnectionPool( minconn=minconn
                                   , maxconn=maxconn
                                   , dsn=dsn
@@ -400,18 +409,20 @@ class Connection(psycopg2.extensions.connection):
 
         - We set :py:attr:`autocommit` to :py:const:`True`.
         - We set the client encoding to ``UTF-8``.
-        - We use :py:class:`psycopg2.extras.RealDictCursor`.
+        - We use :py:attr:`self.cursor_factory`.
 
     """
 
+    cursor_factory = None   # set this before using this object
+
     def __init__(self, *a, **kw):
         psycopg2.extensions.connection.__init__(self, *a, **kw)
         self.set_client_encoding('UTF-8')
         self.autocommit = True
 
     def cursor(self, *a, **kw):
         if 'cursor_factory' not in kw:
-            kw['cursor_factory'] = RealDictCursor
+            kw['cursor_factory'] = self.cursor_factory
         return psycopg2.extensions.connection.cursor(self, *a, **kw)
 
 

tests.py

@@ -1,9 +1,11 @@
 from __future__ import unicode_literals
 
 import os
+from collections import namedtuple
 from unittest import TestCase
 
 from postgres import Postgres, TooFew, TooMany
+from psycopg2.extras import NamedTupleCursor
 
 
 DATABASE_URL = os.environ['DATABASE_URL']
@@ -259,3 +261,29 @@ def test_get_connection_gets_a_connection(self):
             cursor.execute("SELECT * FROM foo ORDER BY bar")
             actual = cursor.fetchall()
         assert actual == [{"bar": "baz"}, {"bar": "buz"}]
+
+
+# cursor_factory
+# ==============
+
+class TestCursorFactory(WithData):
+
+    def setUp(self):                    # override
+        self.db = Postgres(DATABASE_URL, cursor_factory=NamedTupleCursor)
+        self.db.run("DROP SCHEMA IF EXISTS public CASCADE")
+        self.db.run("CREATE SCHEMA public")
+        self.db.run("CREATE TABLE foo (bar text)")
+        self.db.run("INSERT INTO foo VALUES ('baz')")
+        self.db.run("INSERT INTO foo VALUES ('buz')")
+
+    def test_NamedDictCursor_results_in_namedtuples(self):
+        Record = namedtuple("Record", ["bar"])
+        expected = [Record(bar="baz"), Record(bar="buz")]
+        actual = self.db.rows("SELECT * FROM foo ORDER BY bar")
+        assert actual == expected
+
+    def test_NamedDictCursor_results_in_namedtuples(self):
+        Record = namedtuple("Record", ["bar"])
+        expected = [Record(bar="baz"), Record(bar="buz")]
+        actual = self.db.rows("SELECT * FROM foo ORDER BY bar")
+        assert actual == expected