Python decorator that can accept an SQLAlchemy engine and make the function context transactional and automatically rollback. Great for integration tests.

rollback.py

from functools import wraps

def rollback(engine):
    def rollback_decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            with engine.connect() as connection:
                with connection.begin() as tx:
                    try:
                        kwargs['connection'] = connection
                        func(*args, **kwargs)
                    finally:
                        tx.rollback()

        return wrapper

    return rollback_decorator