Adding a Database to a Flask App

In this tutorial I will document the process of adding a database to your flask app.

We'll use the following tools:

Flask-SQLAlchemy: ORM mapping
Flask-Migrate: DB Migrations for SQLAlchemy
Flask-Testing: Testing for SQLAlchemy

Let's add all of these to requirements.txt to get started:

...
Flask-SQLAlchemy
Flask-Script
Flask-Migrate
Flask-Testing
...

Step 1: Add Flask-SQLAlchemy to our app

This part is easy:

app.py


from flask import Flask
from flask.ext.sqlalchemy import SQLAlchemy

app = Flask(__name__)
app.config['SQLALCHEMY_DATABASE_URI'] = 'sqlite:///db.db'
db = SQLAlchemy(app)

Step 2: Let's set up testing from the get-go:

We're using Flask-Testing to test our database. Here is an example TestCase which we'll use to test our code:

tests.py


from flask.ext.testing import TestCase
from app import db, app

TEST_SQLALCHEMY_DATABASE_URI = "sqlite:///test.sqlite"
    
class MyTest(TestCase):

    def create_app(self):
        app.config['SQLALCHEMY_DATABASE_URI'] = TEST_SQLALCHEMY_DATABASE_URI
        return app

    def setUp(self):
        db.create_all()

    def tearDown(self):
        db.session.remove()
        db.drop_all()

What's going on here:

We import our app from our flask project and we configure a test database uri. We don't want to be running our tests against a production database by mistake!
In setUp() we create a fresh database (this is run before each test in the TestCase)
In tearDown() we destroy our database again. This ensures that each test is running against a clean database in a predictable state.

You should be able to run this successfully with your test runner of choice. I like nose, so I would run this with:

$ nosetests
..
----------------------------------------------------------------------
Ran 0 tests in 0.000s

OK

Step 3: Create a model:

models.py


from app import db 

class User(db.Model):
    id = db.Column(db.Integer, primary_key=True)
    username = db.Column(db.String(80), unique=True)
    email = db.Column(db.String(120), unique=True)

    def __init__(self, username, email):
        self.username = username
        self.email = email

    def __repr__(self):
        return '<User %r>' % self.username

Step 4: Test our model

In tests.py we can add a couple of tests.

Firstly, in setUp(), let's create some initial users:


def setUp(self):
    db.create_all()
    ## create users:
    user = User('joe', 'joe@soap.com')
    user2 = User('jane', 'jane@soap.com')
    db.session.add(user)
    db.session.add(user2)
    db.session.commit()

Let's test querying for all users:


def test_get_all_users(self):
    users = User.query.all()
    assert len(users) == 2, 'Expect all users to be returned'

and fetching a specific user:


def test_get_user(self):
    user = User.query.filter_by(username='joe').first()
    assert user.email == 'joe@soap.com', 'Expect the correct user to be returned'

Again, we can run this with our test runner of choice:

$ nosetests tests/db_tests.py
..
----------------------------------------------------------------------
Ran 2 tests in 0.037s

OK

Excellent! We've setup Flask to use a simple database. We've created our first ORM model and we've got some unit tests making sure it's all working as expected. Last step: we need to get migrations in place.

Step 5: Setup migrations

The last part of the puzzle is to setup our DB migrations

What are DB migrations?

Part of the purpose of using an ORM is to abstract away the database. This means that your ORM mappings need to be tightly coupled to the db schema. Migrations are a way to manage this coupling with code in a mostly painless way. Migrations make sure that your database is in a state that matches your code.

We're going to use: Flask-Migrate for this.

Adding Flask-Migrate to our app is easy. Create a file called manage.py

manage.py


from flask.ext.script import Manager
from flask.ext.migrate import Migrate, MigrateCommand
from app import app, db

manager = Manager(app)
migrate = Migrate(app, db)
manager.add_command('db', MigrateCommand)

if __name__ == "__main__":
    manager.run()

Notes:

If you already have a mange.py file, you can simply add the relevant parts.
See Flask-Script for more info on running commands with Flask

We're now ready to add migrations to our project. To start managing the project with migrations run:


$ python manage.py db init

This will create a migrations folder which will contain all the migrations we need for our project.

We then need to generate our initial migration. Use:


$ python manage.py db migrate

Then, when we make changes to our models, we need to generate a migration to apply these changes to our database. We do this with:


$ python manage.py db upgrade

You will need to run upgrade every time you make changes to your models.

Now, when checking out the project from fresh, a new user can simply run: python app.py db upgrade and all migrations will be run so that their application's database is in the correct state.

Menu