An exploration of transaction integration for PonyORM

integration.py

import transaction

from pony.orm import Database, PrimaryKey, Required
from pony.orm.core import DBSessionContextManager
from pony.orm.core import commit, flush, rollback


class CustomDBSessionContextManager(DBSessionContextManager):

    def _enter(self, *args, **kwargs):
        super()._enter(*args, **kwargs)
        transaction.get().join(PonyDataManager(self))


db_session = CustomDBSessionContextManager


class PonyDataManager(object):
    """ Example integration of pony into transaction. """

    def __init__(self, session_context):
        self.session_context = session_context

    def abort(self, trans):
        rollback()

    def tpc_begin(self, trans):
        flush()

    def commit(self, trans):
        commit()

    def tpc_vote(self, trans):
        pass

    def tpc_finish(self, trans):
        pass

    def tpc_abort(self, trans):
        pass

    def sortKey(self):
        # Try to sort last, so that we vote last - we may commit in tpc_vote(),
        # which allows Zope to roll back its transaction if the RDBMS
        # threw a conflict error.
        return "~ponydatamanager:%d" % id(self.session_context)

    @property
    def savepoint(self):
        raise NotImplementedError

    def should_retry(self, error):
        if not self.session_context.retry_exceptions:
            return False

        if callable(self.session_context.retry_exceptions):
            return self.session_context.retry_exceptions(error)

        return error in self.session_context.retry_exceptions


if __name__ == '__main__':
    db = Database()

    class Record(db.Entity):
        id = PrimaryKey(int, auto=True)
        name = Required(str)

    db.bind('sqlite', ':memory:')
    db.generate_mapping(create_tables=True)

    record_id = None

    with db_session():
        record = Record(name='foo')
        transaction.abort()

    with db_session():
        assert sum(1 for r in Record.select()) == 0

    with db_session():
        record = Record(name='foo')
        assert not record.id

        transaction.commit()
        assert record.id

    with db_session():
        assert sum(1 for r in Record.select()) == 1