Paul BecotteAdmin

Opinionated Flask Setup

There are a LOT of ways to put together a Python api service. I've done it quite a few times at this point, and thought it may be helpful to write down the mix that I personally feel is the best-in-class setup. Our goals-

  • Modular code
  • JSON API
  • Postgres
  • Repeatable database migrations
  • Easy to write tests for functionality
  • Fully automated deployments

Database

The data access layer has, in my opinion, the best Python library period- SqlAlchemy. The best part about it is that it can just provide a simple query builder, all the way to a full blown ORM. However, I am NOT a fan of `Flask-Sqlalchemy`, which tries to abstract away session management- at the cost of making individual sections of the code much less independent, and often harder to test. I always have a healthy suspicion of global objects anyway, in opposition to the recommended pattern there. Instead, I like to provide a context manager to scope out sessions that I can use when necessary.

from typing import Optional

from flask import Flask
from sqlalchemy import create_engine
from sqlalchemy.engine import Engine
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker

class Alchemy:
    engine: Engine
    Session: sessionmaker

    def __init__(
        self,
        app: Optional[Flask] = None,
        engine: Optional[Engine] = None,
        session: Optional[sessionmaker] = None,
    ):
        if engine:
            self.engine = engine
        if session:
            self.Session = session
        if app:
            self.register_app(app)

    def register_app(self, app: Flask):
        db_string = app.config["SQLALCHEMY_DATABASE_URI"]
        if not hasattr(self, "engine"):
            self.engine = create_engine(db_string)
        if not hasattr(self, "Session"):
            self.Session = sessionmaker(bind=self.engine)
        app.alchemy = self

    @contextmanager
    def session(self):
        session = self.Session()
        yield session
        try:
            session.commit()
        except Exception:
            session.rollback()
            raise
        finally:
            session.close()
            

Now, we can make database queries like this-

from flask import current_app

from models import MyModel


def some_view():
    with current_app.alchemy.session() as session:
        blue_models = session.query(MyModel).filter(MyModel.color == "blue").all()
        red_model = MyModel(color="red")
        session.add(red_model)
        session.commit()  
        # Leaving the block will attempt to auto-commit, but if you
        # need to do it earlier, you can just call the method directly
        
        

Alembic is a great library that is built on top of SqlAlchemy that gives you atomic migrations. You define migration scripts, and you can use `alembic upgrade head` to take a db from one state to the next. This is ALWAYS the way to set up database schemas so that your app just works everywhere.

Testing

For testing, Pytest has become the industry standard for good reason. The fixture paradigm works well with SqlAlchemy- we can build a fixture that will keep the database in a nicely cleaned state in between each test case without individual tests having to think about it.

import os
from pathlib import Path
from unittest.mock import patch

import alembic.config
import pytest
from sqlalchemy import event

from package.application import create_app


@pytest.fixture(scope="session")
def _migrations():
    # At the beginning of our tests, we should have an empty db- use
    # our migrations to set it up. Session scoped means we just do
    # this once- since we use a transaction around our tests, we dont
    # need any further cleanup. HOWEVER, we could run this before every
    # test case as an alternative way to get isolation, but this tends 
    # to be the quicker option
    folder = Path(__file__).parent.parent
    os.chdir(folder)
    ini = folder.joinpath("alembic.ini")
    _app = create_app()
    args = ["--raiseerr", "-c", str(ini)]
    alembic.config.main(argv=args + ["upgrade", "head"])
    yield
    # pytest fixtures that use `yield` let you do cleanup steps
    # after the test finishes
    alembicArgs = args + ["downgrade", "base"]
    alembic.config.main(argv=alembicArgs)


@pytest.fixture(name="_app")
def app_fixture():
    # We can tweak the config for testing here-
    app = create_app()
    app.config["TESTING"] = True
    app.config["BCRYPT_ROUNDS"] = 4
    with app.test_request_context():
        yield app


@pytest.fixture
def client(_app):
    with _app.test_client() as c:
        yield c


@pytest.fixture
def _session(_app, _migrations):
    # This magic prevents any test from accidentily changing
    # the database. Even if the test calls `commit`, this
    # will roll that change back (so long as your db engine
    # supports save points)
    # First, we start a connection and transaction entirely
    # outside the ORM
    conn = _app.alchemy.engine.connect()
    transaction = conn.begin()    
    original = _app.alchemy.Session

    # Next, we patch our session function to return a new session
    # bound to the same connection as the transaction we just
    # created, and begin a nested transaction for the test to use
    def make_session():
        session = original(bind=conn)
        session.begin_nested()

    # Use SqlAlchemy events- if the nested transaction ends, we
    # go ahead and start it back up
    @event.listens_for(session, "after_transaction_end")
    def _restart_savepoint(sess, trans):
        if trans.nested and not trans._parent.nested:
            sess.expire_all()
            sess.begin_nested()
        return session
    
    with patch.object(_app.alchemy, "Session", side_effect=make_session):
        yield make_session()
    
    # Finally, rollback the outer transaction and close our connection
    transaction.rollback()
    conn.close()
    

There is another common problem - building objects to use in your tests. I like the library Factory Boy for this purpose- it lets you setup factory functions that can build object where you only specify the important stuff and let the library handle the rest. We can combine this with sqlalchemy and pytest fixtures like this-

import factory
import pytest

from myapp.models import MyModel


class ModelFactory(factory.Factory):
    class Meta:
        model = MyModel
        
    color = "red"
    name = "my_red_model"
    

@pytest.fixture(name="_model")
def model(_session):
    _model = ModelFactory()
    _session.add(_model)
    _session.commit()
    return _model

# Now, for any test that depends on a red model, we can
# just use the fixture-

def test_model(_client, _model):
    response = _client.get("/model?color=red")
    assert response.json["name"] == _model.name
    
    

I'll come back in the next installment to go over the structure of the Flask app and the API.