danielrichman · July 8, 2024 12:46 · Oct 2, 2014 · Oct 19, 2013 · Oct 19, 2013 · Oct 19, 2013
diff --git a/flask_postgres.py b/flask_postgres.py
@@ -1,6 +1,7 @@
 from __future__ import unicode_literals
 
 import logging
+import threading
 
 import flask
 from werkzeug.local import LocalProxy
@@ -37,7 +38,7 @@ def __init__(self, app=None, pool_size=2):
         self._pool = []
         self.pool_size = pool_size
         self._lock = threading.RLock()
-        self.logger = getLogger(__name__ + ".PostgreSQL")
+        self.logger = logging.getLogger(__name__ + ".PostgreSQL")
 
         if app is not None:
             self.init_app(app)

diff --git a/connection_factory_python2.py → connection_factory.py b/connection_factory_python2.py → connection_factory.py
diff --git a/nosetests.py b/nosetests.py
@@ -0,0 +1,339 @@
+import psycopg2
+import psycopg2.extras
+import psycopg2.extensions
+import flask
+
+class MockConnectionBase(object):
+    def __init__(self):
+        # prevent connection setup
+        self.registered_types = []
+        pass
+    def cursor(self, cursor_factory=None):
+        assert cursor_factory is None or \
+                cursor_factory == psycopg2.extras.RealDictCursor
+        if cursor_factory is None:
+            return "stubbed cursor"
+        else:
+            return "stubbed dict cursor"
+
+class ConnectionRebaser(type):
+    def __new__(mcs, name, bases, dict):
+        bases += (MockConnectionBase, )
+        return type.__new__(mcs, name, bases, dict)
+    def mro(cls):
+        return (cls, utils.PostgreSQLConnection, MockConnectionBase) + \
+                utils.PostgreSQLConnection.__mro__[1:]
+
+class RebasedPostgreSQLConnection(utils.PostgreSQLConnection):
+    __metaclass__ = ConnectionRebaser
+
+class FakeExtensions(object):
+    UNICODE = psycopg2.extensions.UNICODE
+    UNICODEARRAY = psycopg2.extensions.UNICODEARRAY
+
+    @classmethod
+    def register_type(self, what, connection):
+        connection.registered_types.append(what)
+
+class TestPostgreSQLConnection(object):
+    def setup(self):
+        assert psycopg2.extensions.__name__ == "psycopg2.extensions"
+        self.fakes = FakeExtensions()
+        self.original_extensions = psycopg2.extensions
+        psycopg2.extensions = self.fakes
+
+    def teardown(self):
+        assert isinstance(psycopg2.extensions, FakeExtensions)
+        psycopg2.extensions = self.original_extensions
+
+    def test_only_affects_cursor(self):
+        assert [x for x in utils.PostgreSQLConnection.__dict__
+                if not x.startswith("__")] == ["cursor"]
+
+    def test_cursor(self):
+        c = RebasedPostgreSQLConnection()
+        assert c.cursor() == "stubbed cursor"
+        assert c.cursor(False) == "stubbed cursor"
+        assert c.cursor(True) == "stubbed dict cursor"
+
+    def test_register_types(self):
+        c = RebasedPostgreSQLConnection()
+        assert c.registered_types == [psycopg2.extensions.UNICODE,
+                                      psycopg2.extensions.UNICODEARRAY]
+
+
+class FakePsycopg2(object):
+    class connection(object):
+        class _cursor(object):
+            def __init__(self, connection, real_dict_cursor):
+                self.queries = []
+                self.connection = connection
+                self.real_dict_cursor = real_dict_cursor
+            def __enter__(self):
+                return self
+            def __exit__(self, *args):
+                pass
+            def execute(self, query, args=None):
+                if self.connection.calls["close"] != 0:
+                    raise psycopg2.OperationalError
+                self.queries.append((query, args))
+
+        close_error = False
+        autocommit = False
+
+        def __init__(self, **settings):
+            self.settings = settings
+            self.types = []
+            self.cursors = []
+            self.calls = {"commit": 0, "reset": 0, "close": 0}
+        def cursor(self, real_dict_cursor=False):
+            c = self._cursor(self, real_dict_cursor)
+            self.cursors.append(c)
+            return c
+        def commit(self):
+            assert self.calls["close"] == 0
+            self.calls["commit"] += 1
+        def reset(self):
+            self.calls["reset"] += 1
+            if self.calls["close"] != 0:
+                raise psycopg2.OperationalError
+        def close(self):
+            self.calls["close"] += 1
+            if self.close_error:
+                raise psycopg2.OperationalError
+
+    connections = 0
+
+    def connect(self, **settings):
+        self.connections += 1
+        return self.connection(**settings)
+
+    OperationalError = psycopg2.OperationalError
+    extras = psycopg2.extras
+
+class TestPostgreSQL(object):
+    def setup(self):
+        assert utils.psycopg2 is psycopg2
+        self.fakes = utils.psycopg2 = FakePsycopg2()
+        self.app = flask.Flask(__name__)
+        self.app.config["POSTGRES"] = {"database": "mydb", "user": "steve"}
+        self.postgres = utils.PostgreSQL(self.app)
+
+    def teardown(self):
+        assert isinstance(utils.psycopg2, FakePsycopg2)
+        utils.psycopg2 = psycopg2
+
+    def test_adds_hooks(self):
+        assert self.app.after_request_funcs == {None: [self.postgres.commit]}
+        assert self.app.teardown_appcontext_funcs == [self.postgres.teardown]
+
+    def test_connect_new(self):
+        with self.app.test_request_context("/"):
+            c = self.postgres.connection
+            assert isinstance(c, self.fakes.connection)
+            assert c.settings == \
+                    {"database": "mydb", "user": "steve",
+                     "connection_factory": utils.PostgreSQLConnection}
+            assert c.calls == {"commit": 0, "reset": 0, "close": 0}
+            assert c.autocommit is False
+
+    def test_connect_once(self):
+        with self.app.test_request_context("/"):
+            c = self.postgres.connection
+            d = self.postgres.connection
+            assert c is d
+            assert self.fakes.connections == 1
+            assert c.calls == {"commit": 0, "reset": 0, "close": 0}
+
+    def test_teardown_resets_before_store(self):
+        with self.app.test_request_context("/"):
+            c = self.postgres.connection
+        assert c.calls == {"commit": 0, "reset": 1, "close": 0}
+
+    def test_connect_from_pool(self):
+        with self.app.test_request_context("/"):
+            c = self.postgres.connection
+        assert c.calls == {"commit": 0, "reset": 1, "close": 0}
+
+        with self.app.test_request_context("/"):
+            d = self.postgres.connection
+            assert d is c
+            assert d.calls == {"commit": 0, "reset": 2, "close": 0}
+        assert d.calls == {"commit": 0, "reset": 3, "close": 0}
+
+    def test_removes_from_pool(self):
+        # put a connection in the pool
+        with self.app.test_request_context("/"):
+            c = self.postgres.connection
+
+        # now get two connections from the pool
+
+        # must explicitly create two app contexts.
+        # in normal usage, flask promises to never share an app context
+        # between requests. When testing, it will only create an app context
+        # when test_request_context is __enter__'d and there is no existing
+        # app context
+        with self.app.app_context(), self.app.test_request_context("/1"):
+            d = self.postgres.connection
+
+            with self.app.app_context(), self.app.test_request_context("/2"):
+                e = self.postgres.connection
+        assert d is c
+        assert e is not d
+
+        assert d.calls == {"commit": 0, "reset": 3, "close": 0}
+        assert e.calls == {"commit": 0, "reset": 1, "close": 0}
+
+    def test_connect_from_pool_bad(self):
+        # put two distinct connections in the pool
+        with self.app.app_context(), self.app.test_request_context("/1"):
+            c = self.postgres.connection
+
+            with self.app.app_context(), self.app.test_request_context("/2"):
+                d = self.postgres.connection
+        assert c is not d
+
+        assert c.calls == d.calls == {"commit": 0, "reset": 1, "close": 0}
+        c.close()
+
+        with self.app.test_request_context("/"):
+            e = self.postgres.connection
+
+            # it should try c.reset, which will fail, and then destroy the
+            # pool by closing d as well
+            # one close call by uut, one close call from above
+            assert c.calls == {"commit": 0, "reset": 2, "close": 1 + 1}
+            assert d.calls == {"commit": 0, "reset": 1, "close": 1}
+
+            # e should be a new connection
+            assert e is not c and e is not d
+            assert e.calls == {"commit": 0, "reset": 0, "close": 0}
+
+        assert e.calls == {"commit": 0, "reset": 1, "close": 0}
+
+    def test_absorbs_close_errors(self):
+        with self.app.app_context(), self.app.test_request_context("/1"):
+            c = self.postgres.connection
+            with self.app.app_context(), self.app.test_request_context("/2"):
+                d = self.postgres.connection
+
+        c.close()
+        d.close_error = True
+
+        with self.app.test_request_context("/"):
+            e = self.postgres.connection
+
+    def test_teardown_closes_if_pool_full(self):
+        # default pool size is 2
+        with self.app.app_context(), self.app.test_request_context("/1"):
+            c = self.postgres.connection
+
+            with self.app.app_context(), self.app.test_request_context("/2"):
+                d = self.postgres.connection
+
+                with self.app.app_context(), \
+                        self.app.test_request_context("/3"):
+                    e = self.postgres.connection
+
+        assert len(set([c, d, e])) == 3
+
+        assert c.calls == {"commit": 0, "reset": 0, "close": 1}
+        assert d.calls == {"commit": 0, "reset": 1, "close": 0}
+        assert e.calls == {"commit": 0, "reset": 1, "close": 0}
+
+        with self.app.app_context(), self.app.test_request_context("/1"):
+            f = self.postgres.connection
+
+            with self.app.app_context(), self.app.test_request_context("/2"):
+                g = self.postgres.connection
+
+                with self.app.app_context(), \
+                        self.app.test_request_context("/3"):
+                    h = self.postgres.connection
+
+        assert f is d
+        assert g is e
+        assert len(set([d, e, h])) == 3
+
+    def test_cursor(self):
+        with self.app.test_request_context("/"):
+            c = self.postgres.connection
+            x = self.postgres.cursor()
+            assert isinstance(x, c._cursor)
+            assert x.connection is c
+            assert len(c.cursors) == 1
+
+            y = self.postgres.cursor()
+            assert y.connection is c
+            assert len(c.cursors) == 2
+
+        assert c.calls == {"commit": 0, "reset": 1, "close": 0}
+
+        with self.app.test_request_context("/"):
+            # cursor first without asking for connection explicitly
+            x = self.postgres.cursor()
+            assert isinstance(x, utils.psycopg2.connection._cursor)
+            c = x.connection
+            assert isinstance(c, utils.psycopg2.connection)
+            assert c.calls == {"commit": 0, "reset": 2, "close": 0}
+            assert c.settings == \
+                    {"database": "mydb", "user": "steve",
+                     "connection_factory": utils.PostgreSQLConnection}
+            assert c is self.postgres.connection
+
+    def test_dict_cursor(self):
+        with self.app.test_request_context("/"):
+            c = self.postgres.cursor(True)
+            assert len(self.postgres.connection.cursors) == 1
+            assert c.real_dict_cursor
+            c = self.postgres.cursor()
+            assert not c.real_dict_cursor
+
+    def test_commit(self):
+        with self.app.test_request_context("/"):
+            c = self.postgres.connection
+            self.postgres.commit()
+            assert c.calls == {"commit": 1, "reset": 0, "close": 0}
+
+    def test_commit_as_hook(self):
+        # as an after request hook, commit must return the response object
+        # it is passed
+        response = object()
+        with self.app.test_request_context("/"):
+            c = self.postgres.connection
+            assert self.postgres.commit(response) is response
+
+        # now check it works as a hook
+        with self.app.test_request_context("/"):
+            d = self.postgres.connection
+            assert d is c
+            assert self.app.process_response(response) is response
+
+        assert c.calls == {"commit": 2, "reset": 3, "close": 0}
+
+    def test_commit_nop_if_no_connection(self):
+        with self.app.test_request_context("/"):
+            self.postgres.commit()
+        assert utils.psycopg2.connections == 0
+
+        with self.app.test_request_context("/"):
+            self.app.process_response(None)
+        assert utils.psycopg2.connections == 0
+
+        # should nop if teardown puts the connection in the pool
+        with self.app.test_request_context("/"):
+            c = self.postgres.connection
+            self.postgres.teardown(None)
+            self.postgres.commit()
+            assert c.calls == {"commit": 0, "reset": 1, "close": 0}
+
+    def test_teardown_nop_if_no_connection(self):
+        with self.app.test_request_context("/"):
+            self.postgres.teardown(None)
+        assert utils.psycopg2.connections == 0
+
+        with self.app.test_request_context("/"):
+            c = self.postgres.connection
+            self.postgres.teardown(None)
+            assert c.calls == {"commit": 0, "reset": 1, "close": 0}
+        assert c.calls == {"commit": 0, "reset": 1, "close": 0}
diff --git a/connection_factory_python2.py b/connection_factory_python2.py
@@ -0,0 +1,29 @@
+class PostgreSQLConnection(psycopg2.extensions.connection):
+    """
+    A custom `connection_factory` for :func:`psycopg2.connect`.
+
+    This
+    * puts the connection into unicode mode (for text)
+    * modifies the :meth:`cursor` method of a :class:`psycopg2.connection`,
+      facilitating easy acquiring of cursors made from
+      :cls:`psycopg2.extras.RealDictCursor`.
+    """
+
+    # this may be omitted in py3k
+    def __init__(self, *args, **kwargs):
+        super(PostgreSQLConnection, self).__init__(*args, **kwargs)
+        for type in (psycopg2.extensions.UNICODE,
+                     psycopg2.extensions.UNICODEARRAY):
+            psycopg2.extensions.register_type(type, self)
+
+    def cursor(self, real_dict_cursor=False):
+        """
+        Get a new cursor.
+
+        If real_dict_cursor is set, a RealDictCursor is returned
+        """
+
+        kwargs = {}
+        if real_dict_cursor:
+            kwargs["cursor_factory"] = psycopg2.extras.RealDictCursor
+        return super(PostgreSQLConnection, self).cursor(**kwargs)
diff --git a/flask_postgres.py b/flask_postgres.py
@@ -0,0 +1,164 @@
+from __future__ import unicode_literals
+
+import logging
+
+import flask
+from werkzeug.local import LocalProxy
+import psycopg2
+import psycopg2.extras
+import psycopg2.extensions
+
+
+postgres = LocalProxy(lambda: flask.current_app.postgres)
+
+
+class PostgreSQL(object):
+    """
+    A PostgreSQL helper extension for Flask apps
+
+    On initialisation it adds an after_request function that commits the
+    transaction (so that if the transaction rolls back the request will
+    fail) and a app context teardown function that disconnects any active
+    connection.
+
+    You can of course (and indeed should) use :meth:`commit` if you need to
+    ensure some changes have made it to the database before performing
+    some other action. :meth:`teardown` is also available to be called
+    directly.
+
+    Connections are created by ``psycopg2.connect(**app.config["POSTGRES"])``
+    (e.g., ``app.config["POSTGRES"] = {"database": "mydb"}``),
+    are pooled (you can adjust the pool size with `pool`) and are tested for
+    server shutdown before being given to the request.
+    """
+
+    def __init__(self, app=None, pool_size=2):
+        self.app = app
+        self._pool = []
+        self.pool_size = pool_size
+        self._lock = threading.RLock()
+        self.logger = getLogger(__name__ + ".PostgreSQL")
+
+        if app is not None:
+            self.init_app(app)
+
+    def init_app(self, app):
+        """
+        Initialises the app by adding hooks
+
+        * Hook: ``app.after_request(self.commit)``
+        * Hook: ``app.teardown_appcontext(self.teardown)``
+        """
+
+        app.after_request(self.commit)
+        app.teardown_appcontext(self.teardown)
+        app.postgresql = self
+
+    def _connect(self):
+        """Returns a connection to the database"""
+
+        with self._lock:
+            c = None
+
+            if len(self._pool):
+                c = self._pool.pop()
+                try:
+                    # This tests if the connection is still alive.
+                    c.reset()
+                except psycopg2.OperationalError:
+                    self.logger.debug("assuming pool dead", exc_info=True)
+
+                    # assume that the entire pool is dead
+                    try:
+                        c.close()
+                    except psycopg2.OperationalError:
+                        pass
+
+                    for c in self._pool:
+                        try:
+                            c.close()
+                        except psycopg2.OperationalError:
+                            pass
+
+                    self._pool = []
+                    c = None
+                else:
+                    self.logger.debug("got connection from pool")
+
+            if c is None:
+                c = self._new_connection()
+
+        return c
+
+    def _new_connection(self):
+        """Create a new connection to the database"""
+        s = flask.current_app.config["POSTGRES"]
+        summary = ' '.join(k + "=" + v for k, v in s.iteritems())
+        self.logger.debug("connecting (%s)", summary)
+        c = psycopg2.connect(connection_factory=PostgreSQLConnection, **s)
+        return c
+
+    @property
+    def connection(self):
+        """
+        Gets the PostgreSQL connection for this Flask request
+
+        If no connection has been used in this request, it connects to the
+        database. Further use of this property will reference the same
+        connection
+
+        The connection is committed and closed at the end of the request.
+        """
+
+        g = flask.g
+        if not hasattr(g, '_postgresql'):
+            g._postgresql = self._connect()
+        return g._postgresql
+
+    def cursor(self, real_dict_cursor=False):
+        """
+        Get a new postgres cursor for immediate use during a request
+
+        If a cursor has not yet been used in this request, it connects to the
+        database. Further cursors re-use the per-request connection.
+
+        The connection is committed and closed at the end of the request.
+
+        If real_dict_cursor is set, a RealDictCursor is returned
+        """
+
+        return self.connection.cursor(real_dict_cursor)
+
+    def commit(self, response=None):
+        """
+        (Almost an) alias for self.connection.commit()
+
+        ... except if self.connection has never been used this is a noop
+        (i.e., it does nothing)
+
+        Returns `response` unmodified, so that this may be used as an
+        :meth:`flask.after_request` function.
+        """
+        g = flask.g
+        if hasattr(g, '_postgresql'):
+            self.logger.debug("committing")
+            g._postgresql.commit()
+        return response
+
+    def teardown(self, exception):
+        """Either return the connection to the pool or close it"""
+        g = flask.g
+        if hasattr(g, '_postgresql'):
+            c = g._postgresql
+            del g._postgresql
+
+            with self._lock:
+                s = len(self._pool)
+                if s >= self.pool_size:
+                    self.logger.debug("teardown: pool size %i - closing", s)
+                    c.close()
+                else:
+                    self.logger.debug("teardown: adding to pool, new size %i",
+                                      s + 1)
+                    c.reset()
+                    self._pool.append(c)
diff --git a/flask_postgres_python2.py b/flask_postgres_python2.py
@@ -1,111 +0,0 @@
-from __future__ import unicode_literals
-
-import logging
-
-import flask
-from werkzeug.local import LocalProxy
-import psycopg2
-import psycopg2.extras
-import psycopg2.extensions
-
-
-postgres = LocalProxy(lambda: flask.current_app.postgres)
-
-
-class PostgreSQL(object):
-    """
-    A PostgreSQL helper extension for Flask apps
-
-    On initialisation it adds an after_request function that commits the
-    transaction (so that if the transaction rolls back the request will
-    fail) and a app context teardown function that disconnects any active
-    connection.
-
-    You can of course (and indeed should) use :meth:`commit` if you need to
-    ensure some changes have made it to the database before performing
-    some other action. :meth:`teardown` is also available to be called
-    directly.
-
-    Connections are created by ``psycopg2.connenct(app.config["POSTGRES"])``.
-    """
-
-    def __init__(self, app=None):
-        self.app = app
-        self.logger = logging.getLogger(__module__ + ".PostgreSQL")
-
-        if app is not None:
-            self.init_app(app)
-
-    def init_app(self, app):
-        """
-        Initialises the app by adding hooks
-
-        * Hook: ``app.after_request(self.commit)``
-        * Hook: ``app.teardown_appcontext(self.teardown)``
-        """
-
-        app.after_request(self.commit)
-        app.teardown_appcontext(self.teardown)
-        app.postgres = self
-
-    def _connect(self):
-        """Returns (always) a new connection to the database"""
-        s = flask.current_app.config["POSTGRES"]
-        self.logger.debug("connecting (%r)", s)
-        c = psycopg2.connect(s)
-        psycopg2.extensions.register_type(psycopg2.extensions.UNICODE, c)
-        psycopg2.extensions.register_type(psycopg2.extensions.UNICODEARRAY, c)
-        return c
-
-    @property
-    def connection(self):
-        """
-        Gets the PostgreSQL connection for this Flask request
-     
-        If no connection has been used in this request, it connects to the
-        database. Further use of this property will reference the same
-        connection
-     
-        The connection is committed and closed at the end of the request.
-        """
-
-        g = flask.current_app.g
-        if not hasattr(g, '_postgresql'):
-            g._postgresql = self._connect()
-        return g._postgresql
-
-    def cursor(self, real_dict_cursor=False):
-        """
-        Get a new postgres cursor for immediate use during a request
-     
-        If a cursor has not yet been used in this request, it connects to the
-        database. Further cursors re-use the per-request connection.
-     
-        The connection is committed and closed at the end of the request.
-     
-        If real_dict_cursor is set, a RealDictCursor is returned
-        """
-
-        kwargs = {}
-        if real_dict_cursor:
-            kwargs["cursor_factory"] = psycopg2.extras.RealDictCursor
-        return self.connection.cursor(**kwargs)
-
-    def commit(self):
-        """
-        (Almost an) alias for self.connection.commit()
-
-        ... except if self.connection has never been used (i.e., it does
-        nothing this is a noop)
-        """
-        g = flask.current_app.g
-        if hasattr(g, '_postgresql'):
-            self.logger.debug("committing")
-            g._postgresql.commit()
-
-    def teardown(self, exception):
-        """Tears down any connection on current_app.g"""
-        g = flask.current_app.g
-        if hasattr(g, '_postgresql'):
-            self.logger("teardown: closing")
-            g._postgresql.close()

diff --git a/flask_postgres_python3.py b/flask_postgres_python3.py
@@ -1,106 +0,0 @@
-import logging
-
-import flask
-from werkzeug.local import LocalProxy
-import psycopg2
-import psycopg2.extras
-
-
-postgres = LocalProxy(lambda: flask.current_app.postgres)
-
-
-class PostgreSQL(object):
-    """
-    A PostgreSQL helper extension for Flask apps
-
-    On initialisation it adds an after_request function that commits the
-    transaction (so that if the transaction rolls back the request will
-    fail) and a app context teardown function that disconnects any active
-    connection.
-
-    You can of course (and indeed should) use :meth:`commit` if you need to
-    ensure some changes have made it to the database before performing
-    some other action. :meth:`teardown` is also available to be called
-    directly.
-
-    Connections are created by ``psycopg2.connenct(app.config["POSTGRES"])``.
-    """
-
-    def __init__(self, app=None):
-        self.app = app
-        self.logger = logging.getLogger(__module__ + ".PostgreSQL")
-
-        if app is not None:
-            self.init_app(app)
-
-    def init_app(self, app):
-        """
-        Initialises the app by adding hooks
-
-        * Hook: ``app.after_request(self.commit)``
-        * Hook: ``app.teardown_appcontext(self.teardown)``
-        """
-
-        app.after_request(self.commit)
-        app.teardown_appcontext(self.teardown)
-        app.postgres = self
-
-    def _connect(self):
-        """Returns (always) a new connection to the database"""
-        s = flask.current_app.config["POSTGRES"]
-        self.logger.debug("connecting (%r)", s)
-        c = psycopg2.connect(s)
-        return c
-
-    @property
-    def connection(self):
-        """
-        Gets the PostgreSQL connection for this Flask request
-     
-        If no connection has been used in this request, it connects to the
-        database. Further use of this property will reference the same
-        connection
-     
-        The connection is committed and closed at the end of the request.
-        """
-
-        g = flask.current_app.g
-        if not hasattr(g, '_postgresql'):
-            g._postgresql = self._connect()
-        return g._postgresql
-
-    def cursor(self, real_dict_cursor=False):
-        """
-        Get a new postgres cursor for immediate use during a request
-     
-        If a cursor has not yet been used in this request, it connects to the
-        database. Further cursors re-use the per-request connection.
-     
-        The connection is committed and closed at the end of the request.
-     
-        If real_dict_cursor is set, a RealDictCursor is returned
-        """
-
-        kwargs = {}
-        if real_dict_cursor:
-            kwargs["cursor_factory"] = psycopg2.extras.RealDictCursor
-        return self.connection.cursor(**kwargs)
-
-    def commit(self):
-        """
-        (Almost an) alias for self.connection.commit()
-
-        ... except if self.connection has never been used (i.e., it does
-        nothing this is a noop)
-        """
-        g = flask.current_app.g
-        if hasattr(g, '_postgresql'):
-            self.logger.debug("committing")
-            g._postgresql.commit()
-
-    def teardown(self, exception):
-        """Tears down any connection on current_app.g"""
-        g = flask.current_app.g
-        if hasattr(g, '_postgresql'):
-            self.logger("teardown: closing")
-            g._postgresql.close()

diff --git a/flask_postgres_python2.py b/flask_postgres_python2.py
@@ -3,11 +3,15 @@
 import logging
 
 import flask
+from werkzeug.local import LocalProxy
 import psycopg2
 import psycopg2.extras
 import psycopg2.extensions
 
 
+postgres = LocalProxy(lambda: flask.current_app.postgres)
+
+
 class PostgreSQL(object):
     """
     A PostgreSQL helper extension for Flask apps
@@ -42,6 +46,7 @@ def init_app(self, app):
 
         app.after_request(self.commit)
         app.teardown_appcontext(self.teardown)
+        app.postgres = self
 
     def _connect(self):
         """Returns (always) a new connection to the database"""

diff --git a/flask_postgres_python3.py b/flask_postgres_python3.py
@@ -1,10 +1,14 @@
 import logging
 
 import flask
+from werkzeug.local import LocalProxy
 import psycopg2
 import psycopg2.extras
 
 
+postgres = LocalProxy(lambda: flask.current_app.postgres)
+
+
 class PostgreSQL(object):
     """
     A PostgreSQL helper extension for Flask apps
@@ -39,6 +43,7 @@ def init_app(self, app):
 
         app.after_request(self.commit)
         app.teardown_appcontext(self.teardown)
+        app.postgres = self
 
     def _connect(self):
         """Returns (always) a new connection to the database"""

diff --git a/flask_postgres.py → flask_postgres_python2.py b/flask_postgres.py → flask_postgres_python2.py
@@ -5,6 +5,7 @@
 import flask
 import psycopg2
 import psycopg2.extras
+import psycopg2.extensions
 
 
 class PostgreSQL(object):

diff --git a/flask_postgres_python3.py b/flask_postgres_python3.py
@@ -0,0 +1,101 @@
+import logging
+
+import flask
+import psycopg2
+import psycopg2.extras
+
+
+class PostgreSQL(object):
+    """
+    A PostgreSQL helper extension for Flask apps
+
+    On initialisation it adds an after_request function that commits the
+    transaction (so that if the transaction rolls back the request will
+    fail) and a app context teardown function that disconnects any active
+    connection.
+
+    You can of course (and indeed should) use :meth:`commit` if you need to
+    ensure some changes have made it to the database before performing
+    some other action. :meth:`teardown` is also available to be called
+    directly.
+
+    Connections are created by ``psycopg2.connenct(app.config["POSTGRES"])``.
+    """
+
+    def __init__(self, app=None):
+        self.app = app
+        self.logger = logging.getLogger(__module__ + ".PostgreSQL")
+
+        if app is not None:
+            self.init_app(app)
+
+    def init_app(self, app):
+        """
+        Initialises the app by adding hooks
+
+        * Hook: ``app.after_request(self.commit)``
+        * Hook: ``app.teardown_appcontext(self.teardown)``
+        """
+
+        app.after_request(self.commit)
+        app.teardown_appcontext(self.teardown)
+
+    def _connect(self):
+        """Returns (always) a new connection to the database"""
+        s = flask.current_app.config["POSTGRES"]
+        self.logger.debug("connecting (%r)", s)
+        c = psycopg2.connect(s)
+        return c
+
+    @property
+    def connection(self):
+        """
+        Gets the PostgreSQL connection for this Flask request
+     
+        If no connection has been used in this request, it connects to the
+        database. Further use of this property will reference the same
+        connection
+     
+        The connection is committed and closed at the end of the request.
+        """
+
+        g = flask.current_app.g
+        if not hasattr(g, '_postgresql'):
+            g._postgresql = self._connect()
+        return g._postgresql
+
+    def cursor(self, real_dict_cursor=False):
+        """
+        Get a new postgres cursor for immediate use during a request
+     
+        If a cursor has not yet been used in this request, it connects to the
+        database. Further cursors re-use the per-request connection.
+     
+        The connection is committed and closed at the end of the request.
+     
+        If real_dict_cursor is set, a RealDictCursor is returned
+        """
+
+        kwargs = {}
+        if real_dict_cursor:
+            kwargs["cursor_factory"] = psycopg2.extras.RealDictCursor
+        return self.connection.cursor(**kwargs)
+
+    def commit(self):
+        """
+        (Almost an) alias for self.connection.commit()
+
+        ... except if self.connection has never been used (i.e., it does
+        nothing this is a noop)
+        """
+        g = flask.current_app.g
+        if hasattr(g, '_postgresql'):
+            self.logger.debug("committing")
+            g._postgresql.commit()
+
+    def teardown(self, exception):
+        """Tears down any connection on current_app.g"""
+        g = flask.current_app.g
+        if hasattr(g, '_postgresql'):
+            self.logger("teardown: closing")
+            g._postgresql.close()
diff --git a/flask_postgres.py b/flask_postgres.py
@@ -46,7 +46,10 @@ def _connect(self):
         """Returns (always) a new connection to the database"""
         s = flask.current_app.config["POSTGRES"]
         self.logger.debug("connecting (%r)", s)
-        return psycopg2.connect(s)
+        c = psycopg2.connect(s)
+        psycopg2.extensions.register_type(psycopg2.extensions.UNICODE, c)
+        psycopg2.extensions.register_type(psycopg2.extensions.UNICODEARRAY, c)
+        return c
 
     @property
     def connection(self):

diff --git a/flask_postgres.py b/flask_postgres.py
@@ -1,5 +1,3 @@
-"""utils - miscellaneous utilities"""
-
 from __future__ import unicode_literals
 
 import logging
@@ -47,7 +45,7 @@ def init_app(self, app):
     def _connect(self):
         """Returns (always) a new connection to the database"""
         s = flask.current_app.config["POSTGRES"]
-        self.logger.debug("Connecting to %s", s)
+        self.logger.debug("connecting (%r)", s)
         return psycopg2.connect(s)
 
     @property
@@ -63,9 +61,9 @@ def connection(self):
         """
 
         g = flask.current_app.g
-        if not hasattr(g, '_database'):
-            g._database = self._connect()
-        return g._database
+        if not hasattr(g, '_postgresql'):
+            g._postgresql = self._connect()
+        return g._postgresql
 
     def cursor(self, real_dict_cursor=False):
         """
@@ -85,14 +83,20 @@ def cursor(self, real_dict_cursor=False):
         return self.connection.cursor(**kwargs)
 
     def commit(self):
-        """Alias for self.connection.commit()"""
-        self.connection.commit()
+        """
+        (Almost an) alias for self.connection.commit()
+
+        ... except if self.connection has never been used (i.e., it does
+        nothing this is a noop)
+        """
+        g = flask.current_app.g
+        if hasattr(g, '_postgresql'):
+            self.logger.debug("committing")
+            g._postgresql.commit()
 
     def teardown(self, exception):
         """Tears down any connection on current_app.g"""
         g = flask.current_app.g
         if hasattr(g, '_postgresql'):
-            try:
-                g._database.commit()
-            finally:
-                g._database.close()
+            self.logger("teardown: closing")
+            g._postgresql.close()
diff --git a/flask_postgres.py b/flask_postgres.py
@@ -1,48 +1,98 @@
+"""utils - miscellaneous utilities"""
+
+from __future__ import unicode_literals
+
+import logging
+
 import flask
 import psycopg2
-from psycopg2.extras import DateTimeRange, RealDictCursor
+import psycopg2.extras
 
 
-def connection():
+class PostgreSQL(object):
     """
-    Get a connection to use in this request
+    A PostgreSQL helper extension for Flask apps
 
-    If no connection has been used in this request, it connects to the
-    database. Further calls to connection() in this request context will
-    get the same connection.
+    On initialisation it adds an after_request function that commits the
+    transaction (so that if the transaction rolls back the request will
+    fail) and a app context teardown function that disconnects any active
+    connection.
 
-    The connection is committed and closed at the end of the request.
+    You can of course (and indeed should) use :meth:`commit` if you need to
+    ensure some changes have made it to the database before performing
+    some other action. :meth:`teardown` is also available to be called
+    directly.
+
+    Connections are created by ``psycopg2.connenct(app.config["POSTGRES"])``.
     """
 
-    assert flask.has_request_context()
-    if not hasattr(g, '_database'):
-        g._database = psycopg2.connect(app.config["POSTGRES"])
-    return g._database
+    def __init__(self, app=None):
+        self.app = app
+        self.logger = logging.getLogger(__module__ + ".PostgreSQL")
 
-def cursor(real_dict_cursor=False):
-    """
-    Get a postgres cursor for immediate use during a request
+        if app is not None:
+            self.init_app(app)
 
-    If a cursor has not yet been used in this request, it connects to the
-    database. Further cursors re-use the per-request connection.
+    def init_app(self, app):
+        """
+        Initialises the app by adding hooks
 
-    The connection is committed and closed at the end of the request.
+        * Hook: ``app.after_request(self.commit)``
+        * Hook: ``app.teardown_appcontext(self.teardown)``
+        """
 
-    If real_dict_cursor is set, a RealDictCursor is returned
-    """
+        app.after_request(self.commit)
+        app.teardown_appcontext(self.teardown)
+
+    def _connect(self):
+        """Returns (always) a new connection to the database"""
+        s = flask.current_app.config["POSTGRES"]
+        self.logger.debug("Connecting to %s", s)
+        return psycopg2.connect(s)
+
+    @property
+    def connection(self):
+        """
+        Gets the PostgreSQL connection for this Flask request
+     
+        If no connection has been used in this request, it connects to the
+        database. Further use of this property will reference the same
+        connection
+     
+        The connection is committed and closed at the end of the request.
+        """
+
+        g = flask.current_app.g
+        if not hasattr(g, '_database'):
+            g._database = self._connect()
+        return g._database
+
+    def cursor(self, real_dict_cursor=False):
+        """
+        Get a new postgres cursor for immediate use during a request
+     
+        If a cursor has not yet been used in this request, it connects to the
+        database. Further cursors re-use the per-request connection.
+     
+        The connection is committed and closed at the end of the request.
+     
+        If real_dict_cursor is set, a RealDictCursor is returned
+        """
+
+        kwargs = {}
+        if real_dict_cursor:
+            kwargs["cursor_factory"] = psycopg2.extras.RealDictCursor
+        return self.connection.cursor(**kwargs)
+
+    def commit(self):
+        """Alias for self.connection.commit()"""
+        self.connection.commit()
 
-    if real_dict_cursor:
-        f = RealDictCursor
-        return connection().cursor(cursor_factory=f)
-    else:
-        return connection().cursor()
-
-@app.teardown_appcontext
-def close_db_connection(exception):
-    """Commit and close the per-request postgres connection"""
-
-    if hasattr(g, '_database'):
-        try:
-            g._database.commit()
-        finally:
-            g._database.close()
+    def teardown(self, exception):
+        """Tears down any connection on current_app.g"""
+        g = flask.current_app.g
+        if hasattr(g, '_postgresql'):
+            try:
+                g._database.commit()
+            finally:
+                g._database.close()
diff --git a/flask_postgres_pool.py b/flask_postgres_pool.py
@@ -1,76 +0,0 @@
-import flask
-import psycopg2
-from psycopg2.extras import DateTimeRange, RealDictCursor
-
-
-postgres_pool = None
-
-@app.before_first_request
-def setup_postgres_pool():
-    """Initialise the postgres connection pool"""
-
-    # Happens "before_first_request" rather than at module init since
-    # app.config could change
-
-    global postgres_pool
-    postgres_pool = ThreadedConnectionPool(1, 10, app.config["POSTGRES"])
-
-def connection():
-    """
-    Get a connection to use in this request
-
-    If no connection has been used in this request, it connects to the
-    database. Further calls to connection() in this request context will
-    get the same connection.
-
-    The connection is committed and closed at the end of the request.
-    """
-
-    assert flask.has_request_context()
-    if not hasattr(g, '_database'):
-        for attempt in range(postgres_pool.maxconn, -1, -1):
-            conn = postgres_pool.getconn()
-            try:
-                conn.cursor().execute('SELECT 1')
-            except psycopg2.OperationalError:
-                # thrown if the connection is dead due to server restart
-                # putting a broken connection back into the pool will clean up
-                # some internal keys and then discard it.
-                postgres_pool.putconn(conn)
-                if attempt == 0:
-                    raise
-                else:
-                    continue
-            else:
-                g._database = conn
-                break
-
-    return g._database
-
-def cursor(real_dict_cursor=False):
-    """
-    Get a postgres cursor for immediate use during a request
-
-    If a cursor has not yet been used in this request, it connects to the
-    database. Further cursors re-use the per-request connection.
-
-    The connection is committed and closed at the end of the request.
-
-    If real_dict_cursor is set, a RealDictCursor is returned
-    """
-
-    if real_dict_cursor:
-        f = RealDictCursor
-        return connection().cursor(cursor_factory=f)
-    else:
-        return connection().cursor()
-
-@app.teardown_appcontext
-def close_db_connection(exception):
-    """Commit and close the per-request postgres connection"""
-
-    if hasattr(g, '_database'):
-        try:
-            g._database.commit()
-        finally:
-            postgres_pool.putconn(g._database)

diff --git a/flask_postgres.py b/flask_postgres.py
@@ -1,3 +1,8 @@
+import flask
+import psycopg2
+from psycopg2.extras import DateTimeRange, RealDictCursor
+
+
 def connection():
     """
     Get a connection to use in this request

diff --git a/flask_postgres_pool.py b/flask_postgres_pool.py
@@ -1,3 +1,8 @@
+import flask
+import psycopg2
+from psycopg2.extras import DateTimeRange, RealDictCursor
+
+
 postgres_pool = None
 
 @app.before_first_request

diff --git a/flask_postgres.py b/flask_postgres.py
@@ -1,15 +1,3 @@
-postgres_pool = None
-
-@app.before_first_request
-def setup_postgres_pool():
-    """Initialise the postgres connection pool"""
-
-    # Happens "before_first_request" rather than at module init since
-    # app.config could change
-
-    global postgres_pool
-    postgres_pool = ThreadedConnectionPool(1, 10, app.config["POSTGRES"])
-
 def connection():
     """
     Get a connection to use in this request
@@ -23,23 +11,7 @@ def connection():
 
     assert flask.has_request_context()
     if not hasattr(g, '_database'):
-        for attempt in range(postgres_pool.maxconn, -1, -1):
-            conn = postgres_pool.getconn()
-            try:
-                conn.cursor().execute('SELECT 1')
-            except psycopg2.OperationalError:
-                # thrown if the connection is dead due to server restart
-                # putting a broken connection back into the pool will clean up
-                # some internal keys and then discard it.
-                postgres_pool.putconn(conn)
-                if attempt == 0:
-                    raise
-                else:
-                    continue
-            else:
-                g._database = conn
-                break
-
+        g._database = psycopg2.connect(app.config["POSTGRES"])
     return g._database
 
 def cursor(real_dict_cursor=False):
@@ -68,4 +40,4 @@ def close_db_connection(exception):
         try:
             g._database.commit()
         finally:
-            postgres_pool.putconn(g._database)
+            g._database.close()
diff --git a/flask_postgres_pool.py b/flask_postgres_pool.py
@@ -0,0 +1,71 @@
+postgres_pool = None
+
+@app.before_first_request
+def setup_postgres_pool():
+    """Initialise the postgres connection pool"""
+
+    # Happens "before_first_request" rather than at module init since
+    # app.config could change
+
+    global postgres_pool
+    postgres_pool = ThreadedConnectionPool(1, 10, app.config["POSTGRES"])
+
+def connection():
+    """
+    Get a connection to use in this request
+
+    If no connection has been used in this request, it connects to the
+    database. Further calls to connection() in this request context will
+    get the same connection.
+
+    The connection is committed and closed at the end of the request.
+    """
+
+    assert flask.has_request_context()
+    if not hasattr(g, '_database'):
+        for attempt in range(postgres_pool.maxconn, -1, -1):
+            conn = postgres_pool.getconn()
+            try:
+                conn.cursor().execute('SELECT 1')
+            except psycopg2.OperationalError:
+                # thrown if the connection is dead due to server restart
+                # putting a broken connection back into the pool will clean up
+                # some internal keys and then discard it.
+                postgres_pool.putconn(conn)
+                if attempt == 0:
+                    raise
+                else:
+                    continue
+            else:
+                g._database = conn
+                break
+
+    return g._database
+
+def cursor(real_dict_cursor=False):
+    """
+    Get a postgres cursor for immediate use during a request
+
+    If a cursor has not yet been used in this request, it connects to the
+    database. Further cursors re-use the per-request connection.
+
+    The connection is committed and closed at the end of the request.
+
+    If real_dict_cursor is set, a RealDictCursor is returned
+    """
+
+    if real_dict_cursor:
+        f = RealDictCursor
+        return connection().cursor(cursor_factory=f)
+    else:
+        return connection().cursor()
+
+@app.teardown_appcontext
+def close_db_connection(exception):
+    """Commit and close the per-request postgres connection"""
+
+    if hasattr(g, '_database'):
+        try:
+            g._database.commit()
+        finally:
+            postgres_pool.putconn(g._database)
diff --git a/flask_postgres.py b/flask_postgres.py
@@ -0,0 +1,71 @@
+postgres_pool = None
+
+@app.before_first_request
+def setup_postgres_pool():
+    """Initialise the postgres connection pool"""
+
+    # Happens "before_first_request" rather than at module init since
+    # app.config could change
+
+    global postgres_pool
+    postgres_pool = ThreadedConnectionPool(1, 10, app.config["POSTGRES"])
+
+def connection():
+    """
+    Get a connection to use in this request
+
+    If no connection has been used in this request, it connects to the
+    database. Further calls to connection() in this request context will
+    get the same connection.
+
+    The connection is committed and closed at the end of the request.
+    """
+
+    assert flask.has_request_context()
+    if not hasattr(g, '_database'):
+        for attempt in range(postgres_pool.maxconn, -1, -1):
+            conn = postgres_pool.getconn()
+            try:
+                conn.cursor().execute('SELECT 1')
+            except psycopg2.OperationalError:
+                # thrown if the connection is dead due to server restart
+                # putting a broken connection back into the pool will clean up
+                # some internal keys and then discard it.
+                postgres_pool.putconn(conn)
+                if attempt == 0:
+                    raise
+                else:
+                    continue
+            else:
+                g._database = conn
+                break
+
+    return g._database
+
+def cursor(real_dict_cursor=False):
+    """
+    Get a postgres cursor for immediate use during a request
+
+    If a cursor has not yet been used in this request, it connects to the
+    database. Further cursors re-use the per-request connection.
+
+    The connection is committed and closed at the end of the request.
+
+    If real_dict_cursor is set, a RealDictCursor is returned
+    """
+
+    if real_dict_cursor:
+        f = RealDictCursor
+        return connection().cursor(cursor_factory=f)
+    else:
+        return connection().cursor()
+
+@app.teardown_appcontext
+def close_db_connection(exception):
+    """Commit and close the per-request postgres connection"""
+
+    if hasattr(g, '_database'):
+        try:
+            g._database.commit()
+        finally:
+            postgres_pool.putconn(g._database)
No results found