6

I am using python 3 and flask, with flask-migrate (which uses alembic) to handle my SQL migrations. When I run local integration tests, I want to rebuild the database each time so I can run my API calls against a clean db for each api call i'm testing (yes, i could use sqlite, but i want to check constraints are correct).

I can do the following on the command line easily:

mysql -uroot -e 'drop database DBNAME; create database DBNAME;'
FLASK_APP=flask_app.py flask db upgrade

But I would rather run it in the python code for 2 reasons:

  1. I don't want to have to worry about the mysql client being installed on the CI machines that will (eventually) run this code (they should just need to the python mysql packages).
  2. I want to manipulate the flask settings to force the database name to avoid accidents (so it needs to run in the same thread/memory space as the script which invokes it).

The app object (created with app = Flask(__name__)) has a cli property, but it requires a context object, and it doesn't feel like i'm using the right tool. I expected app.cli.invoke('db', 'upgrade') or similar...

Any suggestions on how to invoke flask commands from the code without a child cli process?

Sam Adams
  • 5,327
  • 2
  • 18
  • 17

3 Answers3

4

I use the following pattern (see below). An alternate approach can be seen at https://flask.palletsprojects.com/en/1.1.x/cli/?highlight=click#application-context

# file: commands.py
import click
from click import pass_context
from flask.cli import AppGroup, with_appcontext
from flask import current_app
from flask_migrate import Migrate
from alembic import command

from extensions import flask_db as db

db_cli = AppGroup('db', help='Various database management commands.')

@db_cli.command('init')
def db_init():
    """Initialize the database."""
    db.create_all()
    click.echo("Create all tables.")


@db_cli.command('drop')
def db_drop():
    """Drop the database."""
    db.engine.execute("SET FOREIGN_KEY_CHECKS=0;")
    db.drop_all()
    db.engine.execute("SET FOREIGN_KEY_CHECKS=1;")
    click.echo("Drop all tables.")

@db_cli.command('migrate')
def db_migrate():
    "Migrate with alembic."

    config = Migrate(current_app, db).get_config()
    command.upgrade(config, 'head')


@db_cli.command('db_upgrade')
@pass_context
def db_upgrade(ctx):
    """Alias for 'db reset'."""
    db_drop.invoke(ctx)
    db_init.invoke(ctx)
    db_migrate.invoke(ctx)

# file: extensions.py
# Keep your extenstions separate to allow importing without import loops.

from flask_sqlalchemy import SQLAlchemy

flask_db = SQLAlchemy()

# file: app.py (app/__init__.py) wherever your app is built
from extensions import flask_db

app = Flask(__name__)

flask_db.init_app(app)  # I'm not sure if the order matters here.
app.cli.add_command(db_cli)

# file: wsgi.py (top level file)
# This file lets you run 'flask' commands (e.g. flask routes)

# noinspection PyUnresolvedReferences
from app import app as application  # noqa

# file layout
- /
  - app/  (or app.py)
    - __init__.py  (optional)
  - commands.py
  - extensions.py
  - wsgi.py

Usage: flask db upgrade

Marlen T. B.
  • 664
  • 7
  • 12
2

It's not great, but in the end I avoided using flask commands directly and this seems to do what i need:

from my.app import app, db, initialize_app
from flask_migrate import Migrate
from alembic import command
from my.settings import settings
from sqlalchemy_utils.functions import drop_database, create_database, database_exists

test_db_name = 'test_db'
db_url = f'mysql+pymysql://mysqluser@127.0.0.1/{test_db_name}'
settings.SQLALCHEMY_DATABASE_URI = db_url


def reset():
    if database_exists(db_url):
        drop_database(db_url)
    create_database(db_url)
    initialize_app(app) # sets flask config SQLALCHEMY_DATABASE_URI to include test_db
    with app.app_context():
        config = Migrate(app, db).get_config()
        command.upgrade(config, 'head')
Sam Adams
  • 5,327
  • 2
  • 18
  • 17
  • 1
    This has a disadvantage though. Normally you'd want the database user assigned to your application to not have full admin privileges, but only enough permissions to work with the one database. Using an admin user for the application is a security risk, if someone got hold of your db admin credentials then it can wipe out your database server. The standard practice is to use a limited user for the application that has only read and write access to the database. In these cases you can use `db.drop_all()` to drop all the tables in your database, without dropping the database itself. – Miguel Grinberg Jun 21 '18 at 15:30
0

Is there any reason that you can't simply wrap the functions that have all the logic for the CLI?

something like

# cli.py

def db_init(cli=False):
    db.create_all()
    if cli:
        print("Created all tables")

@app.cli.command('db-init')
def _db_init(): # I'm just a wrapper function
    db_init(cli=True)

If you do it that way, there should be no reason why you can't simply do this afterwards:

# some_module.py

from .cli import db_init

db_init()
JackofSpades
  • 113
  • 1
  • 8