Flask Migrate

Part 3, Lesson 2



In this lesson, we'll utilize Flask Migrate to handle database migrations...


Navigate to flask-microservices-users, activate the virtual environment, add the environment variables, and then run the tests to ensure they pass:

$ source env/bin/activate
(env)$ export APP_SETTINGS=project.config.DevelopmentConfig
(env)$ export DATABASE_URL=postgres://postgres:[email protected]:5432/users_dev
(env)$ export DATABASE_TEST_URL=postgres://postgres:[email protected]:5432/users_test
(env)$ python manage.py test

You may need to change the username and password depending on your local Postgres config.

Let's make a few changes to the schema in flask-microservices-users/project/api/models.py:

  1. username must be unique
  2. email must be unique
  3. active should default to True

We'll also add a password field, which will be hashed before it's added to the database:

password = db.Column(db.String(255), nullable=False)

With that, let's start with some tests. Add a new file to "flask-microservices-users/project/tests" called test_user_model.py. This file will hold tests related to our database model:

# project/tests/test_user_model.py


from project import db
from project.api.models import User
from project.tests.base import BaseTestCase


class TestUserModel(BaseTestCase):

    def test_add_user(self):
        user = User(
            username='justatest',
            email='[email protected]',
        )
        db.session.add(user)
        db.session.commit()
        self.assertTrue(user.id)
        self.assertEqual(user.username, 'justatest')
        self.assertEqual(user.email, '[email protected]')
        self.assertTrue(user.active)
        self.assertTrue(user.created_at)

Run the tests. You should see a single failure - AssertionError: False is not true - since user.active is False.

Flask Migrate

Since we need to make a schema change, add Flask-Migrate:

(env)$ pip install flask-migrate==2.0.4
(env)$ pip freeze > requirements.txt

In flask-microservices-users/project/__init__.py add the import, create a new instance, and update create_app():

# project/__init__.py


import os

from flask import Flask, jsonify
from flask_cors import CORS
from flask_sqlalchemy import SQLAlchemy
from flask_migrate import Migrate


# instantiate the db
db = SQLAlchemy()
# instantiate flask migrate
migrate = Migrate()

def create_app():

    # instantiate the app
    app = Flask(__name__)

    # enable CORS
    CORS(app)

    # set config
    app_settings = os.getenv('APP_SETTINGS')
    app.config.from_object(app_settings)

    # set up extensions
    db.init_app(app)
    migrate.init_app(app, db)

    # register blueprints
    from project.api.views import users_blueprint
    app.register_blueprint(users_blueprint)

    return app

Then, add a new manager command to flask-microservices-users/manage.py, just below manager = Manager(app):

manager.add_command('db', MigrateCommand)

Add the import as well:

from flask_migrate import MigrateCommand

Generate the migrations folder, add the initial migration, and then apply it to the database:

(env)$ python manage.py db init
(env)$ python manage.py db migrate
(env)$ python manage.py db upgrade

Review the documentation for more info.

Now, we can make the changes to the schema:

class User(db.Model):
    __tablename__ = "users"
    id = db.Column(db.Integer, primary_key=True, autoincrement=True)
    username = db.Column(db.String(128), unique=True, nullable=False)
    email = db.Column(db.String(128), unique=True, nullable=False)
    active = db.Column(db.Boolean, default=True, nullable=False)
    created_at = db.Column(db.DateTime, nullable=False)

    def __init__(self, username, email, created_at=datetime.datetime.utcnow()):
        self.username = username
        self.email = email
        self.created_at = created_at

Again, run:

(env)$ python manage.py db migrate
(env)$ python manage.py db upgrade

Keep in mind that if you have any duplicate usernames and/or emails already in your database, you will get an error when trying to apply the migration to the database. You can either update the data or drop the db and start over.

Run the tests again. They should pass! Before moving on, let's add a few more tests to flask-microservices-users/project/tests/test_user_model.py:

def test_add_user_duplicate_username(self):
    user = User(
        username='justatest',
        email='[email protected]',
    )
    db.session.add(user)
    db.session.commit()
    duplicate_user = User(
        username='justatest',
        email='[email protected]',
    )
    db.session.add(duplicate_user)
    self.assertRaises(IntegrityError, db.session.commit)

def test_add_user_duplicate_email(self):
    user = User(
        username='justatest',
        email='[email protected]',
    )
    db.session.add(user)
    db.session.commit()
    duplicate_user = User(
        username='justanothertest',
        email='[email protected]',
    )
    db.session.add(duplicate_user)
    self.assertRaises(IntegrityError, db.session.commit)

Notice how we didn't invoke db.session.commit the second time, when adding a user. Instead, we passed it to assertRaises() and let it invoke it and assert the exception was raised.

Add the import:

from sqlalchemy.exc import IntegrityError

Test again.

Refactor

Now is a good time to do some refactoring...

First, in flask-microservices-users/project/tests/test_users.py, rename test_add_user_duplicate_user to test_add_user_duplicate_email.

Also, did you notice that we added a new user a number of times in the test_user_model.py tests? Let's abstract out the add_user helper function from test_users.py to a utility file so we can use it in both test files.

Add a new file called "utils.py" to "tests":

# project/tests/utils.py


import datetime


from project import db
from project.api.models import User


def add_user(username, email, created_at=datetime.datetime.utcnow()):
    user = User(username=username, email=email, created_at=created_at)
    db.session.add(user)
    db.session.commit()
    return user

Then remove the helper from test_users.py and add the import to the same file:

from project.tests.utils import add_user

Refactor test_user_model.py like so:

# project/tests/test_user_model.py


from sqlalchemy.exc import IntegrityError

from project import db
from project.api.models import User
from project.tests.base import BaseTestCase
from project.tests.utils import add_user

class TestUserModel(BaseTestCase):

    def test_add_user(self):
        user = add_user('justatest', '[email protected]')
        self.assertTrue(user.id)
        self.assertEqual(user.username, 'justatest')
        self.assertEqual(user.email, '[email protected]')
        self.assertTrue(user.active)
        self.assertTrue(user.created_at)

    def test_add_user_duplicate_username(self):
        add_user('justatest', '[email protected]')
        duplicate_user = User(
            username='justatest',
            email='[email protected]',
        )
        db.session.add(duplicate_user)
        self.assertRaises(IntegrityError, db.session.commit)

    def test_add_user_duplicate_email(self):
        add_user('justatest', '[email protected]')
        duplicate_user = User(
            username='justatest2',
            email='[email protected]',
        )
        db.session.add(duplicate_user)
        self.assertRaises(IntegrityError, db.session.commit)

Run the tests again to ensure nothing broke.


Flask Migrate

In this lesson, we'll utilize Flask Migrate to handle database migrations...


Navigate to flask-microservices-users, activate the virtual environment, add the environment variables, and then run the tests to ensure they pass:

$ source env/bin/activate
(env)$ export APP_SETTINGS=project.config.DevelopmentConfig
(env)$ export DATABASE_URL=postgres://postgres:[email protected]:5432/users_dev
(env)$ export DATABASE_TEST_URL=postgres://postgres:[email protected]:5432/users_test
(env)$ python manage.py test

You may need to change the username and password depending on your local Postgres config.

Let's make a few changes to the schema in flask-microservices-users/project/api/models.py:

  1. username must be unique
  2. email must be unique
  3. active should default to True

We'll also add a password field, which will be hashed before it's added to the database:

password = db.Column(db.String(255), nullable=False)

With that, let's start with some tests. Add a new file to "flask-microservices-users/project/tests" called test_user_model.py. This file will hold tests related to our database model:

# project/tests/test_user_model.py


from project import db
from project.api.models import User
from project.tests.base import BaseTestCase


class TestUserModel(BaseTestCase):

    def test_add_user(self):
        user = User(
            username='justatest',
            email='[email protected]',
        )
        db.session.add(user)
        db.session.commit()
        self.assertTrue(user.id)
        self.assertEqual(user.username, 'justatest')
        self.assertEqual(user.email, 'tes[email protected]')
        self.assertTrue(user.active)
        self.assertTrue(user.created_at)

Run the tests. You should see a single failure - AssertionError: False is not true - since user.active is False.

Flask Migrate

Since we need to make a schema change, add Flask-Migrate:

(env)$ pip install flask-migrate==2.0.4
(env)$ pip freeze > requirements.txt

In flask-microservices-users/project/__init__.py add the import, create a new instance, and update create_app():

# project/__init__.py


import os

from flask import Flask, jsonify
from flask_cors import CORS
from flask_sqlalchemy import SQLAlchemy
from flask_migrate import Migrate


# instantiate the db
db = SQLAlchemy()
# instantiate flask migrate
migrate = Migrate()

def create_app():

    # instantiate the app
    app = Flask(__name__)

    # enable CORS
    CORS(app)

    # set config
    app_settings = os.getenv('APP_SETTINGS')
    app.config.from_object(app_settings)

    # set up extensions
    db.init_app(app)
    migrate.init_app(app, db)

    # register blueprints
    from project.api.views import users_blueprint
    app.register_blueprint(users_blueprint)

    return app

Then, add a new manager command to flask-microservices-users/manage.py, just below manager = Manager(app):

manager.add_command('db', MigrateCommand)

Add the import as well:

from flask_migrate import MigrateCommand

Generate the migrations folder, add the initial migration, and then apply it to the database:

(env)$ python manage.py db init
(env)$ python manage.py db migrate
(env)$ python manage.py db upgrade

Review the documentation for more info.

Now, we can make the changes to the schema:

class User(db.Model):
    __tablename__ = "users"
    id = db.Column(db.Integer, primary_key=True, autoincrement=True)
    username = db.Column(db.String(128), unique=True, nullable=False)
    email = db.Column(db.String(128), unique=True, nullable=False)
    active = db.Column(db.Boolean, default=True, nullable=False)
    created_at = db.Column(db.DateTime, nullable=False)

    def __init__(self, username, email, created_at=datetime.datetime.utcnow()):
        self.username = username
        self.email = email
        self.created_at = created_at

Again, run:

(env)$ python manage.py db migrate
(env)$ python manage.py db upgrade

Keep in mind that if you have any duplicate usernames and/or emails already in your database, you will get an error when trying to apply the migration to the database. You can either update the data or drop the db and start over.

Run the tests again. They should pass! Before moving on, let's add a few more tests to flask-microservices-users/project/tests/test_user_model.py:

def test_add_user_duplicate_username(self):
    user = User(
        username='justatest',
        email='[email protected]',
    )
    db.session.add(user)
    db.session.commit()
    duplicate_user = User(
        username='justatest',
        email='[email protected]',
    )
    db.session.add(duplicate_user)
    self.assertRaises(IntegrityError, db.session.commit)

def test_add_user_duplicate_email(self):
    user = User(
        username='justatest',
        email='[email protected]',
    )
    db.session.add(user)
    db.session.commit()
    duplicate_user = User(
        username='justanothertest',
        email='[email protected]',
    )
    db.session.add(duplicate_user)
    self.assertRaises(IntegrityError, db.session.commit)

Notice how we didn't invoke db.session.commit the second time, when adding a user. Instead, we passed it to assertRaises() and let it invoke it and assert the exception was raised.

Add the import:

from sqlalchemy.exc import IntegrityError

Test again.

Refactor

Now is a good time to do some refactoring...

First, in flask-microservices-users/project/tests/test_users.py, rename test_add_user_duplicate_user to test_add_user_duplicate_email.

Also, did you notice that we added a new user a number of times in the test_user_model.py tests? Let's abstract out the add_user helper function from test_users.py to a utility file so we can use it in both test files.

Add a new file called "utils.py" to "tests":

# project/tests/utils.py


import datetime


from project import db
from project.api.models import User


def add_user(username, email, created_at=datetime.datetime.utcnow()):
    user = User(username=username, email=email, created_at=created_at)
    db.session.add(user)
    db.session.commit()
    return user

Then remove the helper from test_users.py and add the import to the same file:

from project.tests.utils import add_user

Refactor test_user_model.py like so:

# project/tests/test_user_model.py


from sqlalchemy.exc import IntegrityError

from project import db
from project.api.models import User
from project.tests.base import BaseTestCase
from project.tests.utils import add_user

class TestUserModel(BaseTestCase):

    def test_add_user(self):
        user = add_user('justatest', '[email protected]')
        self.assertTrue(user.id)
        self.assertEqual(user.username, 'justatest')
        self.assertEqual(user.email, '[email protected]')
        self.assertTrue(user.active)
        self.assertTrue(user.created_at)

    def test_add_user_duplicate_username(self):
        add_user('justatest', '[email protected]')
        duplicate_user = User(
            username='justatest',
            email='[email protected]',
        )
        db.session.add(duplicate_user)
        self.assertRaises(IntegrityError, db.session.commit)

    def test_add_user_duplicate_email(self):
        add_user('justatest', '[email protected]')
        duplicate_user = User(
            username='justatest2',
            email='[email protected]',
        )
        db.session.add(duplicate_user)
        self.assertRaises(IntegrityError, db.session.commit)

Run the tests again to ensure nothing broke.