CodeWalk

Python上下文管理器实现数据库事务

作者:孤独的心 · 2026-05-30 12:55

使用Python上下文管理器实现一个数据库事务封装器。要求:1)begin/commit/rollback通过__enter__/__exit__管理;2)__exit__根据异常类型决定commit还是rollback;3)支持嵌套事务(使用计数或savepoint);4)给出一个转账操作的示例。

回答

孤独的心

class Transaction:
    def __init__(self, conn):
        self.conn = conn
        self.depth = 0
    
    def __enter__(self):
        if self.depth == 0:
            self.conn.execute('BEGIN')
        else:
            self.conn.execute(f'SAVEPOINT sp_{self.depth}')
        self.depth += 1
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        self.depth -= 1
        if exc_type is None:
            if self.depth == 0:
                self.conn.execute('COMMIT')
            else:
                self.conn.execute(f'RELEASE SAVEPOINT sp_{self.depth}')
        else:
            if self.depth == 0:
                self.conn.execute('ROLLBACK')
            else:
                self.conn.execute(f'ROLLBACK TO SAVEPOINT sp_{self.depth}')
        return False  # 不压制异常

# 使用示例
with Transaction(conn) as tx:
    conn.execute('UPDATE accounts SET balance = balance - 100 WHERE id=1')
    conn.execute('UPDATE accounts SET balance = balance + 100 WHERE id=2')
    # 如果任何一条SQL失败,自动ROLLBACK

@contextmanager版本

@contextmanager
def transaction(conn):
    conn.execute('BEGIN')
    try:
        yield conn
    except:
        conn.execute('ROLLBACK')
        raise
    else:
        conn.execute('COMMIT')