2 contextmanager helpers to run PostgreSQL queries with savepoints and deferred constraints

postgres_contexthelpers.py

@contextlib.contextmanager
def with_savepoint(postgres, name="temp_savepoint"):
    """
    A context manager that wraps database queries in a PostgreSQL savepoint

    Usage::

        with with_savepoint(postgres, name='my_savepoint) as rollback:
            do_something()
            if something:
                rollback()  # rolls back to the start of the with block

    """

    with postgres.cursor() as cur:
        cur.execute("SAVEPOINT " + name)

        def rollback():
            cur.execute("ROLLBACK TO SAVEPOINT " + name)

        try:
            yield rollback
        except Exception:
            raising = True
            raise
        else:
            raising = False
        finally:
            try:
                cur.execute("RELEASE SAVEPOINT " + name)
            except Exception:
                # don't squash the original exception
                if not raising:
                    raise
                else:
                    logger.exception("suppressing exception")

@contextlib.contextmanager
def with_deferred_constraint(postgres, name):
    """
    A context manager that deferrs a constraint within a with: block

    Usage::

        with with_deferred_constraint(postgres, 'my_constraint'):
            some_queries()

    Note: assumes the constraints were set to IMMEDIATE beforehand
    (the default).
    """

    with postgres.cursor() as cur:
        cur.execute("SET CONSTRAINTS " + name + " DEFERRED")

        try:
            yield
        except Exception:
            raising = True
            raise
        else:
            raising = False
        finally:
            try:
                cur.execute("SET CONSTRAINTS " + name + " IMMEDIATE")
            except Exception:
                # don't squash the original exception
                if not raising:
                    raise
                else:
                    logger.exception("suppressing exception")