ryankask · August 29, 2015 14:02
diff --git a/db.py b/db.py
 """
 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 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

 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).

        """
        nested = self._in_atomic
        self.begin(nested=nested)
        self._in_atomic = True

        try:
            yield
        except:
            self.rollback()
            raise
        else:
            self.commit()
        finally:
            if not nested:
                self._in_atomic = False


 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)

        # Keep track of which DBAPI connection(s) had autocommit turned off for
        # a particular transaction object.
        dconns_by_trans = {}

        @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
            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.Session, 'after_transaction_end')
        def receive_after_transaction_end(session, transaction):
            """Restore autocommit anywhere this transaction turned it off."""
            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]
	"""
	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 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

	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).

	"""
	nested = self._in_atomic
	self.begin(nested=nested)
	self._in_atomic = True

	try:
	yield
	except:
	self.rollback()
	raise
	else:
	self.commit()
	finally:
	if not nested:
	self._in_atomic = False


	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)

	# Keep track of which DBAPI connection(s) had autocommit turned off for
	# a particular transaction object.
	dconns_by_trans = {}

	@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
	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.Session, 'after_transaction_end')
	def receive_after_transaction_end(session, transaction):
	"""Restore autocommit anywhere this transaction turned it off."""
	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]
No results found