Python上下文管理器实现数据库事务
使用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')