Skip to content

Instantly share code, notes, and snippets.

@carljm
Last active May 27, 2024 06:54
Show Gist options
  • Save carljm/57bfb8616f11bceaf865 to your computer and use it in GitHub Desktop.
Save carljm/57bfb8616f11bceaf865 to your computer and use it in GitHub Desktop.
SQLAlchemy and Postgres autocommit
"""
SQLAlchemy, PostgreSQL (psycopg2), and autocommit
See blog post: http://oddbird.net/2014/06/14/sqlalchemy-postgres-autocommit/
"""
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]
@willvousden
Copy link

@carljm Thanks so much for the blog post that accompanied this – very helpful! Did anything ever come of this? Is this still your approach to transaction management, or is there a better solution now?

@z00sts
Copy link

z00sts commented Feb 8, 2018

Nice research. But isn't it some sort of hacking SQLAlchemy 'by-design-behaviour'?

@lukecyca
Copy link

@carljm I just wanted to add to the chorus of Thank Yous for the well-researched and clearly-written blog post. I can finally make sense of all the layers at play.

@carljm
Copy link
Author

carljm commented Aug 19, 2019

Just saw the latest comment notification and realized I missed some earlier comments. Thanks for the kind words about the blog post! A couple responses:

Did anything ever come of this? Is this still your approach to transaction management, or is there a better solution now?

In my current work I'm no longer using either Postgres or SQLAlchemy. The project I did this on is still running in prod and there haven't been issues caused by this approach that I know of. But I only ever did it on one project.

isn't it some sort of hacking SQLAlchemy 'by-design-behaviour'

Yup. I was curious to see if it could be done, and as I mentioned in the blog post, I'm not super happy with how hacky the solution feels, and wish there were a well-supported way to get this behavior in SQLAlchemy.

@MartinSchmidt123
Copy link

Thanks for the nice blog post and all the explanations.
But please make sure to not use this code as a dependency in a fastapi project.
Doing a commit in a dependency with yield is only executed after returning data to the client.
From https://fastapi.tiangolo.com/tutorial/dependencies/dependencies-with-yield/:
Only the code prior to and including the yield statement is executed before creating a response

For more discussions see e.g. fastapi/fastapi#3620

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment