Inheritance, yield, contextmanager and rollback

Python

Publish Date: 2020-04-11

Author: Isaac Zhou

Word Count: 426

Read Times: 2 Min

Building a webapp with flask could involve a couple of more advanced topics in Python, Flask and SQLAlchemy

Event and Rollback

Event and rollback are database concepts. With ORM, they are also relevant for the “Model” layer. E.g. I have some code for user registration

@web.route("/register", methods=['GET', 'POST'])
def register():
    form = RegisterForm(request.form)
    if request.method == 'POST' and form.validate():
        user = User()
        user.set_attrs(form.data)
        db.session.add(user)
        db.session.commit()
    return render_template('auth/register.html', form=form)

Here, SQLAlchemy has secured data consistency with db.session.commit(), however rollback is missing, we may use try for some refactor

@web.route("/register", methods=['GET', 'POST'])
def register():
    form = RegisterForm(request.form)
    if request.method == 'POST' and form.validate():
        try:
            user = User()
            user.set_attrs(form.data)
            db.session.add(user)
            db.session.commit()
        except Exception as e:
            db.rollback()
            raise e
    return render_template('auth/register.html', form=form)

Adding db.rollback() will avoid any mistakes caused by db data inconsistency here. It’s a good practice to keep here.

ContextManager and Yield

The code above we can separate into two parts: core and common part

Core Logic

To instantiate a user and save it into the db

user = User()
user.set_attrs(form.data)
db.session.add(user)

Common Part

Basically try, except part, or commit and rollback part

try:
    # Core Code
    db.session.commit()
except Exception as e:
    db.rollback()
    raise e

We can apply contextmanager and yield to rewrite the above logic as below

@contextmanager
def auto_commit(db):
    try:
        yield
        db.session.commit()
    except Exception as e:
        db.session.rollback()
        raise e

with auto_commit(db):
    user = User()
    user.set_attrs(form.data)
    db.session.add(user)

Inheritance and add to SQLAlchemy

The above code can be gracefully converted to OOP by inheriting SQLAlchemy class

from contextlib import contextmanager

from flask_sqlalchemy import SQLAlchemy

# inherit SQLAlchemy add context manager
class MySQLAlch(SQLAlchemy):
    @contextmanager
    def auto_commit(self):
        try:
            yield
            self.session.commit()
        except Exception as e:
            self.session.rollback()
            raise e


# instantiate updated MySQLAlch
db = MySQLAlch()

To use it in the view func:

@web.route("/register", methods=['GET', 'POST'])
def register():
    form = RegisterForm(request.form)
    if request.method == 'POST' and form.validate():
        with db.auto_commit():
            user = User()
            user.set_attrs(form.data)
            db.session.add(user)
        return redirect(url_for('web.login'))
    return render_template('auth/register.html', form=form)