ryankask · August 29, 2015 14:02 · Jun 17, 2014 · Jun 16, 2014 · Jun 16, 2014 · Jun 16, 2014
diff --git a/db.py b/db.py
@@ -68,5 +68,6 @@ def receive_after_transaction_end(session, transaction):
             if transaction in dconns_by_trans:
                 for dbapi_connection in dconns_by_trans[transaction]:
                     assert not dbapi_connection.autocommit
-                    dbapi_connection.autocommit = True
+                    if not dbapi_connection.closed:
+                        dbapi_connection.autocommit = True
                 del dconns_by_trans[transaction]
diff --git a/db.py b/db.py
@@ -1,65 +1,7 @@
 """
 SQLAlchemy, PostgreSQL (psycopg2), and autocommit
 
-The goal: 
-- DB-level autocommit by default.
-- No unnecessary transaction around reads.
-- No implicit transaction start on any query.
-- Support for transactions (and nested transactions using savepoints) when explicitly requested.
-
-The below code seems to achieve all that, though I'm not very good at SQLAlchemy so it might be terrible in
-some way. Thanks to Mike Bayer for pointing me to the new-ish AUTOCOMMIT transaction isolation level
-setting for SA's psycopg2 dialect. Thanks to Christophe Pettus and Aymeric Augustin for implementing
-this style of transaction handling in Django and inspiring this API. Any errors are mine, of course.
-
-What this does and why:
-
-Postgres itself is always in AUTOCOMMIT mode unless you explicitly start a transaction with BEGIN,
-but in order to comply with PEP 249 (DBAPI), psycopg2's default mode starts a transaction implicitly
-(issues a BEGIN) anytime you send a query. SQLAlchemy's default behavior follows PEP 249. But psycopg2
-does expose an "autocommit" property on connections which can be set to True to prevent it from
-implicitly starting transactions.
-
-In order to get the desired behavior, we first need to ensure that all psycopg2 connections are
-initially set up with their autocommit property set to True. This is easy since SA 0.8.2 - the psycopg2
-SA dialect exposes an AUTOCOMMIT transaction isolation level, and selecting this sets autocommit=True
-on all the psycopg2 connections.
-
-(Sidebar: Postgres doesn't actually have AUTOCOMMIT as a transaction isolation level, but psycopg2
-[and thus SQLAlchemy] choose to represent it as one. This is weird, since you can have autocommit on
-and still start transactions, which still need to have a real isolation level. And when you set the
-AUTOCOMMIT isolation level in psycopg2, it actually resets the real Postgres isolation level to the
-default, "read committed". So that's an odd API choice, but whatever.)
-
-But we also want to be able to start a transaction, ideally via SQLAlchemy's session.begin().
-Unfortunately, that method doesn't actually issue a BEGIN (because SA relies on the underlying DBAPI
-module to do that implicitly). So when we create a transaction in SA via session.begin(), we need
-to turn off the autocommit property on the connection for that transaction. This still doesn't
-explicitly issue a BEGIN to the database, but it ensures that psycopg2 will automatically issue
-one on the next query.
-
-We accomplish this using a listener to SQLAlchemy's "after_begin" event. (We have to use "after_begin"
-instead of "after_transaction_start" because a single SQLAlchemy "SessionTransaction" object can
-involve multiple database connections, and "after_begin" is the only event fired per-connection; 
-"after_transaction_*" are only fired per-SessionTransaction and don't give our receiver an actual
-connection.)
-
-Lastly, we need to restore autocommit = True on that connection (or possibly connections) when
-the transaction ends. Sadly, the after_commit and after_rollback events fire after the connection
-has been given back to the pool, and don't give us access to it. Likewise the after_transaction_end
-event only gives us the transaction object, not a connection. So we have to get a little clever
-here and keep track of which connections have had autocommit turned off, on behalf of which
-transaction object. And then when that transaction object fires its after_transaction_end event,
-we turn autocommit back on for all those connections.
-
-We use the autocommit=True argument to Session, which doesn't actually implement database-level
-autocommit at all, but it does prevent the Session from automatically starting a transaction
-immediately on initialization and after any other transaction ends, so we need it in conjunction
-with our real autocommit stuff.
-
-The session.atomic() context manager below isn't really part of the autocommit implementation, just
-an example of a context manager API for transactions that is much easier to implement correctly in
-an autocommit context (with explicit start points for a transaction).
+See blog post: http://oddbird.net/2014/06/14/sqlalchemy-postgres-autocommit/
 
 """
 from contextlib import contextmanager

diff --git a/db.py b/db.py
@@ -14,44 +14,52 @@
 
 What this does and why:
 
-Postgres itself is always in AUTOCOMMIT mode by default, but in order to comply with PEP 249 (DBAPI),
-psycopg2's default mode starts a transaction implicitly (issues a BEGIN) anytime you send a query.
-SQLAlchemy's default behavior follows PEP 249. But psycopg2 does expose an "autocommit" property on
-connections which can be set to True to prevent it from implicitly starting transactions.
+Postgres itself is always in AUTOCOMMIT mode unless you explicitly start a transaction with BEGIN,
+but in order to comply with PEP 249 (DBAPI), psycopg2's default mode starts a transaction implicitly
+(issues a BEGIN) anytime you send a query. SQLAlchemy's default behavior follows PEP 249. But psycopg2
+does expose an "autocommit" property on connections which can be set to True to prevent it from
+implicitly starting transactions.
 
-In order to get the desired behavior, we first need to ensure that all psycopg2 connections are initially set
-up with their autocommit property set to True. This is easy since SA 0.8.2 - the psycopg2 SA dialect exposes an
-AUTOCOMMIT transaction isolation level, and selecting this sets autocommit=True on all the psycopg2 connections.
+In order to get the desired behavior, we first need to ensure that all psycopg2 connections are
+initially set up with their autocommit property set to True. This is easy since SA 0.8.2 - the psycopg2
+SA dialect exposes an AUTOCOMMIT transaction isolation level, and selecting this sets autocommit=True
+on all the psycopg2 connections.
 
 (Sidebar: Postgres doesn't actually have AUTOCOMMIT as a transaction isolation level, but psycopg2
 [and thus SQLAlchemy] choose to represent it as one. This is weird, since you can have autocommit on
-and still start transactions, which still need to have a real isolation level. And when you set the AUTOCOMMIT
-isolation level in psycopg2, it actually resets the real Postgres isolation level to the default,
-"read committed". So that's an odd API choice, but whatever.)
-
-But we also want to be able to start a transaction, ideally via SQLAlchemy's session.begin(). Unfortunately,
-that method doesn't actually issue a BEGIN (because SA relies on the underlying DBAPI module to do that
-implicitly). So when we create a transaction in SA via session.begin(), we need to turn off the
-autocommit property on the connection for that transaction. This still doesn't explicitly issue a BEGIN to
-the database, but it ensures that psycopg2 will automatically issue one on the next query.
-
-We accomplish this using a listener to SQLAlchemy's "after_begin" event. (We have to use "after_begin" instead of
-"after_transaction_start" because a single SQLAlchemy "SessionTransaction" object can involve multiple
-database connections, and "after_begin" is the only event fired per-connection; "after_transaction_*" are only
-fired per-SessionTransaction and don't give our receiver an actual connection.)
-
-Lastly, we need to restore autocommit = True on that connection (or possibly connections) when the transaction
-ends. Sadly, the after_commit and after_rollback events fire after the connection has been given back to the pool,
-and don't give us access to it. Likewise the after_transaction_end event only gives us the transaction object,
-not a connection. So we have to get a little clever here and keep track of which connections have had autocommit
-turned off, on behalf of which transaction object. And then when that transaction object fires its after_transaction_end
-event, we turn autocommit back on for all those connections.
-
-We use the autocommit=True argument to Session, which doesn't actually implement database-level autocommit
-at all, but it does prevent the Session from automatically starting a transaction immediately on
-initialization and after any other transaction ends, so we need it in conjunction with our real autocommit stuff.
-
-And that's about it. Code below.
+and still start transactions, which still need to have a real isolation level. And when you set the
+AUTOCOMMIT isolation level in psycopg2, it actually resets the real Postgres isolation level to the
+default, "read committed". So that's an odd API choice, but whatever.)
+
+But we also want to be able to start a transaction, ideally via SQLAlchemy's session.begin().
+Unfortunately, that method doesn't actually issue a BEGIN (because SA relies on the underlying DBAPI
+module to do that implicitly). So when we create a transaction in SA via session.begin(), we need
+to turn off the autocommit property on the connection for that transaction. This still doesn't
+explicitly issue a BEGIN to the database, but it ensures that psycopg2 will automatically issue
+one on the next query.
+
+We accomplish this using a listener to SQLAlchemy's "after_begin" event. (We have to use "after_begin"
+instead of "after_transaction_start" because a single SQLAlchemy "SessionTransaction" object can
+involve multiple database connections, and "after_begin" is the only event fired per-connection; 
+"after_transaction_*" are only fired per-SessionTransaction and don't give our receiver an actual
+connection.)
+
+Lastly, we need to restore autocommit = True on that connection (or possibly connections) when
+the transaction ends. Sadly, the after_commit and after_rollback events fire after the connection
+has been given back to the pool, and don't give us access to it. Likewise the after_transaction_end
+event only gives us the transaction object, not a connection. So we have to get a little clever
+here and keep track of which connections have had autocommit turned off, on behalf of which
+transaction object. And then when that transaction object fires its after_transaction_end event,
+we turn autocommit back on for all those connections.
+
+We use the autocommit=True argument to Session, which doesn't actually implement database-level
+autocommit at all, but it does prevent the Session from automatically starting a transaction
+immediately on initialization and after any other transaction ends, so we need it in conjunction
+with our real autocommit stuff.
+
+The session.atomic() context manager below isn't really part of the autocommit implementation, just
+an example of a context manager API for transactions that is much easier to implement correctly in
+an autocommit context (with explicit start points for a transaction).
 
 """
 from contextlib import contextmanager

diff --git a/db.py b/db.py
@@ -115,8 +115,8 @@ def receive_after_begin(session, transaction, connection):
         @event.listens_for(self.Session, 'after_transaction_end')
         def receive_after_transaction_end(session, transaction):
             """Restore autocommit anywhere this transaction turned it off."""
-            dconns = dconns_by_trans.get(transaction, set())
-            while dconns:
-                dbapi_connection = dconns.pop()
-                assert not dbapi_connection.autocommit
-                dbapi_connection.autocommit = True
+            if transaction in dconns_by_trans:
+                for dbapi_connection in dconns_by_trans[transaction]:
+                    assert not dbapi_connection.autocommit
+                    dbapi_connection.autocommit = True
+                del dconns_by_trans[transaction]
diff --git a/db.py b/db.py
@@ -100,7 +100,7 @@ def __init__(self, db_uri):
         # a particular transaction object.
         dconns_by_trans = {}
 
-        @event.listens_for(self.Parley, 'after_begin')
+        @event.listens_for(self.Session, 'after_begin')
         def receive_after_begin(session, transaction, connection):
             """When a (non-nested) transaction begins, turn autocommit off."""
             dbapi_connection = connection.connection.connection
@@ -112,7 +112,7 @@ def receive_after_begin(session, transaction, connection):
             dconns_by_trans.setdefault(transaction, set()).add(
                 dbapi_connection)
 
-        @event.listens_for(self.Parley, 'after_transaction_end')
+        @event.listens_for(self.Session, 'after_transaction_end')
         def receive_after_transaction_end(session, transaction):
             """Restore autocommit anywhere this transaction turned it off."""
             dconns = dconns_by_trans.get(transaction, set())

diff --git a/db.py b/db.py
@@ -1,12 +1,57 @@
 """
-SQLAlchemy, PostgreSQL, and autocommit
+SQLAlchemy, PostgreSQL (psycopg2), and autocommit
 
 The goal: 
 - DB-level autocommit by default.
 - No unnecessary transaction around reads.
-- Explicit transactions (and nested transactions using savepoints) when desired.
-
-The below seems to achieve that, though I'm not very good at SQLAlchemy so it might be terrible.
+- No implicit transaction start on any query.
+- Support for transactions (and nested transactions using savepoints) when explicitly requested.
+
+The below code seems to achieve all that, though I'm not very good at SQLAlchemy so it might be terrible in
+some way. Thanks to Mike Bayer for pointing me to the new-ish AUTOCOMMIT transaction isolation level
+setting for SA's psycopg2 dialect. Thanks to Christophe Pettus and Aymeric Augustin for implementing
+this style of transaction handling in Django and inspiring this API. Any errors are mine, of course.
+
+What this does and why:
+
+Postgres itself is always in AUTOCOMMIT mode by default, but in order to comply with PEP 249 (DBAPI),
+psycopg2's default mode starts a transaction implicitly (issues a BEGIN) anytime you send a query.
+SQLAlchemy's default behavior follows PEP 249. But psycopg2 does expose an "autocommit" property on
+connections which can be set to True to prevent it from implicitly starting transactions.
+
+In order to get the desired behavior, we first need to ensure that all psycopg2 connections are initially set
+up with their autocommit property set to True. This is easy since SA 0.8.2 - the psycopg2 SA dialect exposes an
+AUTOCOMMIT transaction isolation level, and selecting this sets autocommit=True on all the psycopg2 connections.
+
+(Sidebar: Postgres doesn't actually have AUTOCOMMIT as a transaction isolation level, but psycopg2
+[and thus SQLAlchemy] choose to represent it as one. This is weird, since you can have autocommit on
+and still start transactions, which still need to have a real isolation level. And when you set the AUTOCOMMIT
+isolation level in psycopg2, it actually resets the real Postgres isolation level to the default,
+"read committed". So that's an odd API choice, but whatever.)
+
+But we also want to be able to start a transaction, ideally via SQLAlchemy's session.begin(). Unfortunately,
+that method doesn't actually issue a BEGIN (because SA relies on the underlying DBAPI module to do that
+implicitly). So when we create a transaction in SA via session.begin(), we need to turn off the
+autocommit property on the connection for that transaction. This still doesn't explicitly issue a BEGIN to
+the database, but it ensures that psycopg2 will automatically issue one on the next query.
+
+We accomplish this using a listener to SQLAlchemy's "after_begin" event. (We have to use "after_begin" instead of
+"after_transaction_start" because a single SQLAlchemy "SessionTransaction" object can involve multiple
+database connections, and "after_begin" is the only event fired per-connection; "after_transaction_*" are only
+fired per-SessionTransaction and don't give our receiver an actual connection.)
+
+Lastly, we need to restore autocommit = True on that connection (or possibly connections) when the transaction
+ends. Sadly, the after_commit and after_rollback events fire after the connection has been given back to the pool,
+and don't give us access to it. Likewise the after_transaction_end event only gives us the transaction object,
+not a connection. So we have to get a little clever here and keep track of which connections have had autocommit
+turned off, on behalf of which transaction object. And then when that transaction object fires its after_transaction_end
+event, we turn autocommit back on for all those connections.
+
+We use the autocommit=True argument to Session, which doesn't actually implement database-level autocommit
+at all, but it does prevent the Session from automatically starting a transaction immediately on
+initialization and after any other transaction ends, so we need it in conjunction with our real autocommit stuff.
+
+And that's about it. Code below.
 
 """
 from contextlib import contextmanager

diff --git a/db.py b/db.py
@@ -51,6 +51,27 @@ def __init__(self, db_uri):
         self.engine = create_engine(db_uri, isolation_level="AUTOCOMMIT")
         self.Session = sessionmaker(bind=self.engine, class_=Session, autocommit=True)
 
-        @event.listens_for(self.Session, 'after_begin')
+        # Keep track of which DBAPI connection(s) had autocommit turned off for
+        # a particular transaction object.
+        dconns_by_trans = {}
+
+        @event.listens_for(self.Parley, 'after_begin')
         def receive_after_begin(session, transaction, connection):
-            self.engine.dialect._set_connection_isolation(connection, 'READ COMMITTED')
+            """When a (non-nested) transaction begins, turn autocommit off."""
+            dbapi_connection = connection.connection.connection
+            if transaction.nested:
+                assert not dbapi_connection.autocommit
+                return
+            assert dbapi_connection.autocommit
+            dbapi_connection.autocommit = False
+            dconns_by_trans.setdefault(transaction, set()).add(
+                dbapi_connection)
+
+        @event.listens_for(self.Parley, 'after_transaction_end')
+        def receive_after_transaction_end(session, transaction):
+            """Restore autocommit anywhere this transaction turned it off."""
+            dconns = dconns_by_trans.get(transaction, set())
+            while dconns:
+                dbapi_connection = dconns.pop()
+                assert not dbapi_connection.autocommit
+                dbapi_connection.autocommit = True
diff --git a/db.py b/db.py
@@ -30,12 +30,9 @@ def atomic(self):
         Supports nested usage (via savepoints).
 
         """
-        outer = not self._in_atomic
-        if outer:
-            self.begin()
-            self._in_atomic = True
-        else:
-            self.begin_nested()
+        nested = self._in_atomic
+        self.begin(nested=nested)
+        self._in_atomic = True
 
         try:
             yield
@@ -45,7 +42,7 @@ def atomic(self):
         else:
             self.commit()
         finally:
-            if outer:
+            if not nested:
                 self._in_atomic = False
 
 

diff --git a/db.py b/db.py
@@ -52,10 +52,8 @@ def atomic(self):
 class Database(object):
     def __init__(self, db_uri):
         self.engine = create_engine(db_uri, isolation_level="AUTOCOMMIT")
-        self.Session = sessionmaker(
-            bind=self.engine, class_=Session, autocommit=True)
+        self.Session = sessionmaker(bind=self.engine, class_=Session, autocommit=True)
 
         @event.listens_for(self.Session, 'after_begin')
         def receive_after_begin(session, transaction, connection):
-            self.engine.dialect._set_connection_isolation(
-                connection, 'READ COMMITTED')
+            self.engine.dialect._set_connection_isolation(connection, 'READ COMMITTED')
diff --git a/db.py b/db.py
@@ -48,16 +48,6 @@ def atomic(self):
             if outer:
                 self._in_atomic = False
 
-    def xadd(self, *objs):
-        with self.atomic():
-            for obj in objs:
-                self.add(obj)
-
-    def is_deleted(self, obj):
-        obj_id = obj.id
-        self.expire(obj)
-        return self.query(type(obj)).get(obj_id) is None
-
 
 class Database(object):
     def __init__(self, db_uri):

diff --git a/db.py b/db.py
@@ -1,7 +1,10 @@
 """
 SQLAlchemy, PostgreSQL, and autocommit
 
-The goal: DB-level autocommit by default. No unnecessary transaction around reads. Explicit transactions (and nested transactions using savepoints) when desired.
+The goal: 
+- DB-level autocommit by default.
+- No unnecessary transaction around reads.
+- Explicit transactions (and nested transactions using savepoints) when desired.
 
 The below seems to achieve that, though I'm not very good at SQLAlchemy so it might be terrible.
 

diff --git a/db.py b/db.py
@@ -0,0 +1,68 @@
+"""
+SQLAlchemy, PostgreSQL, and autocommit
+
+The goal: DB-level autocommit by default. No unnecessary transaction around reads. Explicit transactions (and nested transactions using savepoints) when desired.
+
+The below seems to achieve that, though I'm not very good at SQLAlchemy so it might be terrible.
+
+"""
+from contextlib import contextmanager
+
+from sqlalchemy import create_engine, event
+from sqlalchemy.orm import sessionmaker, Session as BaseSession
+
+
+class Session(BaseSession):
+    def __init__(self, *a, **kw):
+        super(Session, self).__init__(*a, **kw)
+        self._in_atomic = False
+
+    @contextmanager
+    def atomic(self):
+        """Transaction context manager.
+
+        Will commit the transaction on successful completion of the block, or
+        roll it back on error.
+
+        Supports nested usage (via savepoints).
+
+        """
+        outer = not self._in_atomic
+        if outer:
+            self.begin()
+            self._in_atomic = True
+        else:
+            self.begin_nested()
+
+        try:
+            yield
+        except:
+            self.rollback()
+            raise
+        else:
+            self.commit()
+        finally:
+            if outer:
+                self._in_atomic = False
+
+    def xadd(self, *objs):
+        with self.atomic():
+            for obj in objs:
+                self.add(obj)
+
+    def is_deleted(self, obj):
+        obj_id = obj.id
+        self.expire(obj)
+        return self.query(type(obj)).get(obj_id) is None
+
+
+class Database(object):
+    def __init__(self, db_uri):
+        self.engine = create_engine(db_uri, isolation_level="AUTOCOMMIT")
+        self.Session = sessionmaker(
+            bind=self.engine, class_=Session, autocommit=True)
+
+        @event.listens_for(self.Session, 'after_begin')
+        def receive_after_begin(session, transaction, connection):
+            self.engine.dialect._set_connection_isolation(
+                connection, 'READ COMMITTED')