A single file ORM-less SQLite3 migration framework in Python.

sqlite-migrator.py

#!/usr/bin/env python3
import sqlite3
import argparse

db = sqlite3.connect("test.db")


class Migration:
    """
    The base migration class. Migrations that inherit from here can
    mix Python code and SQL queries. An example, would be importing
    data from a CSV or remote API endpoint.
    """

    def __init__(self, name):
        self.name = name

    def migrate_up(self, db):
        raise NotImplementedError()

    def migrate_down(self, db):
        raise NotImplementedError()


class SQLMigration(Migration):
    """
    A migration consisting only of SQL commands. Use this for
    adding tables and columns.
    """

    def __init__(self, name, up_sql, down_sql):
        super().__init__(name)
        self.up_sql = up_sql
        self.down_sql = down_sql

    def migrate_up(self, db):
        db.execute(self.up_sql)

    def migrate_down(self, db):
        db.execute(self.down_sql)


# Define all migrations here in this list.
MIGRATIONS = [
    SQLMigration(
        "create_user_table",
        """CREATE TABLE IF NOT EXISTS user (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            username TEXT NON NULL
        );""",
        "DROP TABLE IF EXISTS user",
    ),
]


def bootstrap():
    """Creates the table used for tracking migrations."""

    db.execute(
        """CREATE TABLE IF NOT EXISTS migrations (
        name TEXT NOT NULL PRIMARY KEY,
        applied TEXT NOT NULL
    )"""
    )
    db.commit()


def migrate_up():
    """Migrates the database up to the latest migration."""

    bootstrap()
    for migration in MIGRATIONS:
        cursor = db.execute(
            "SELECT * FROM migrations WHERE name = ?", (migration.name,)
        )
        if cursor.fetchone():
            print("Skipping migration:", migration.name)
        else:
            print("Running migration:", migration.name)
            migration.migrate_up(db)
            db.execute(
                "INSERT INTO migrations VALUES (?, datetime('now'))", (migration.name,)
            )
            db.commit()


def migrate_dev():
    """Drop and re-run the latest migration."""
    last_migration = MIGRATIONS[-1]

    bootstrap()
    cursor = db.execute(
        "SELECT * FROM migrations WHERE name = ?", (last_migration.name,)
    )
    if cursor.fetchone() == None:
        # Migration has never been applied. Just migrate up.
        migrate_up()
    else:
        (last_applied,) = db.execute(
            "SELECT name FROM migrations ORDER BY applied DESC LIMIT 1"
        ).fetchone()
        assert last_applied == last_migration.name

        print("Dropping and re-running migration:", last_migration.name)
        last_migration.migrate_down(db)
        last_migration.migrate_up(db)

        db.execute(
            "UPDATE migrations SET applied = datetime('now') WHERE name = ?",
            (last_migration.name,),
        )
        db.commit()


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Migrate the sqlite3 database.")
    subparsers = parser.add_subparsers(dest="subparser")
    parser_dev = subparsers.add_parser("dev")
    args = parser.parse_args()

    if args.subparser == "dev":
        migrate_dev()
    else:
        migrate_up()