SQLAlchemy通过@contextmanager简化回滚操作

问题引出

　　通过SQLAlchemy来修改数据库时，经常需要使用db.session.commit()和db.session.rollback()结合异常处理语句来实现事务回滚。如果在所有需要修改数据库的地方都加上异常处理和回滚语句就过于麻烦了，而通过with语句结合@contextmanager我们就可以方便地实现这一操作。具体的@contextmanager用法已在Python中@contextmanager的用法一文中写出。

 

简化前

from flask_sqlalchemy import SQLAlchemy
db = SQLAlchemy()
def setData():
    try:
        user = User()
        user.name = 'xiao'
        db.session.add(user)
        db.session.commit()
    except Exception as e:
        db.session.rollback()
        raise e

 

简化过程

　　我们需要先为SQLAlchemy添加@contextmanager，由于SQLAlchemy不能直接被修改，所以我们创建子类继承SQLAlchemy来代替其使用：

from flask_sqlalchemy import SQLAlchemy as _SQLAlchemy
from contextlib import contextmanager

class SQLAlchemy(_SQLAlchemy):
    @contextmanager
    def auto_commit(self):
        try:
            yield
            self.session.commit()
        except Exception as e:
            self.session.rollback()
            raise e

　　

　　当我们得到了这个添加了@contextmanager的SQLAlchemy后，就可以让其他模型继承这个SQLAlchemy：

db = SQLAlchemy()  # 新的SQLAlchemy

class Base(db.Model):
    pass

 

　　然后在需要修改数据库的地方就可以用with语句来实现提交和回滚了：

from flask_sqlalchemy import SQLAlchemy
db = SQLAlchemy()
def setData():
    with db.auto_commit():
        user = User()
        user.name = 'xiao'
        db.session.add(user)