New migrations system, ported from sqlite-migrate

Refs #752

sqlite_utils/__init__.py

@@ -2,5 +2,6 @@ from .utils import suggest_column_types
 from .hookspecs import hookimpl
 from .hookspecs import hookspec
 from .db import Database
+from .migrations import Migrations
 
-__all__ = ["Database", "suggest_column_types", "hookimpl", "hookspec"]
+__all__ = ["Database", "Migrations", "suggest_column_types", "hookimpl", "hookspec"]

sqlite_utils/cli.py

@@ -1,4 +1,5 @@
 import base64
+import difflib
 from typing import Any
 import click
 from click_default_group import DefaultGroup
@@ -3252,6 +3253,125 @@ def create_spatial_index(db_path, table, column_name, load_extension):
     db.table(table).create_spatial_index(column_name)
 
 
+def _find_migration_files(migrations):
+    if not migrations:
+        migrations = [pathlib.Path(".").resolve()]
+    files = set()
+    for path_str in migrations:
+        path = pathlib.Path(path_str)
+        if path.is_dir():
+            files.update(path.rglob("migrations.py"))
+        else:
+            files.add(path)
+    return sorted(files)
+
+
+def _compatible_migration_set(obj):
+    return isinstance(obj, sqlite_utils.Migrations) or all(
+        hasattr(obj, attr) for attr in ("name", "applied", "pending", "apply")
+    )
+
+
+def _load_migration_sets(files):
+    migration_sets = []
+    for filepath in files:
+        code = filepath.read_text()
+        namespace = {
+            "__file__": str(filepath),
+            "__name__": "__sqlite_utils_migration__",
+        }
+        exec(code, namespace)
+        migration_sets.extend(
+            obj for obj in namespace.values() if _compatible_migration_set(obj)
+        )
+    return migration_sets
+
+
+def _display_migration_list(db, migration_sets):
+    for migration_set in migration_sets:
+        click.echo("Migrations for: {}".format(migration_set.name))
+        click.echo()
+        click.echo("  Applied:")
+        for migration in migration_set.applied(db):
+            click.echo("    {} - {}".format(migration.name, migration.applied_at))
+        click.echo()
+        click.echo("  Pending:")
+        output = False
+        for migration in migration_set.pending(db):
+            output = True
+            click.echo("    {}".format(migration.name))
+        if not output:
+            click.echo("    (none)")
+        click.echo()
+
+
+@click.command()
+@click.argument(
+    "db_path", type=click.Path(dir_okay=False, readable=True, writable=True)
+)
+@click.argument("migrations", type=click.Path(dir_okay=True, exists=True), nargs=-1)
+@click.option("--stop-before", help="Stop before applying this migration")
+@click.option(
+    "list_", "--list", is_flag=True, help="List migrations without running them"
+)
+@click.option("-v", "--verbose", is_flag=True, help="Show verbose output")
+def migrate(db_path, migrations, stop_before, list_, verbose):
+    """
+    Apply pending database migrations.
+
+    Usage:
+
+        sqlite-utils migrate database.db
+
+    This will find the migrations.py file in the current directory
+    or subdirectories and apply any pending migrations.
+
+    Or pass paths to one or more migrations.py files directly:
+
+        sqlite-utils migrate database.db path/to/migrations.py
+
+    Pass --list to see a list of applied and pending migrations
+    without applying them.
+    """
+    files = _find_migration_files(migrations)
+    migration_sets = _load_migration_sets(files)
+    if not migration_sets:
+        raise click.ClickException("No migrations.py files found")
+
+    if stop_before and len(migration_sets) > 1:
+        raise click.ClickException(
+            "--stop-before can only be used with a single migrations.py file"
+        )
+
+    db = sqlite_utils.Database(db_path)
+    _register_db_for_cleanup(db)
+
+    if list_:
+        _display_migration_list(db, migration_sets)
+        return
+
+    prev_schema = db.schema
+    if verbose:
+        click.echo("Migrating {}".format(db_path))
+        click.echo("\nSchema before:\n")
+        click.echo(textwrap.indent(prev_schema, "  ") or "  (empty)")
+        click.echo()
+    for migration_set in migration_sets:
+        migration_set.apply(db, stop_before=stop_before)
+    if verbose:
+        click.echo("Schema after:\n")
+        post_schema = db.schema
+        if post_schema == prev_schema:
+            click.echo("  (unchanged)")
+        else:
+            click.echo(textwrap.indent(post_schema, "  "))
+            click.echo("\nSchema diff:\n")
+            diff = list(
+                difflib.unified_diff(prev_schema.splitlines(), post_schema.splitlines())
+            )
+            click.echo("\n".join(diff[3:]))
+
+
 @cli.command(name="plugins")
 def plugins_list():
     "List installed plugins"
@@ -3259,6 +3379,7 @@ def plugins_list():
 
 
 pm.hook.register_commands(cli=cli)
+cli.add_command(migrate)
 
 
 def _render_common(title, values):

sqlite_utils/migrations.py

@@ -0,0 +1,119 @@
+from dataclasses import dataclass
+import datetime
+from typing import Callable, cast, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from sqlite_utils.db import Database, Table
+
+
+class Migrations:
+    migrations_table = "_sqlite_migrations"
+
+    @dataclass
+    class _Migration:
+        name: str
+        fn: Callable
+
+    @dataclass
+    class _AppliedMigration:
+        name: str
+        applied_at: datetime.datetime
+
+    def __init__(self, name: str):
+        """
+        :param name: The name of the migration set. This should be unique.
+        """
+        self.name = name
+        self._migrations: list[Migrations._Migration] = []
+
+    def __call__(self, *, name: str | None = None) -> Callable:
+        """
+        :param name: The name to use for this migration - if not provided,
+          the name of the function will be used.
+        """
+
+        def inner(func: Callable) -> Callable:
+            self._migrations.append(
+                self._Migration(name or getattr(func, "__name__"), func)
+            )
+            return func
+
+        return inner
+
+    def pending(self, db: "Database") -> list["Migrations._Migration"]:
+        """
+        Return a list of pending migrations.
+        """
+        self.ensure_migrations_table(db)
+        already_applied = {
+            r["name"]
+            for r in db[self.migrations_table].rows_where(
+                "migration_set = ?", [self.name]
+            )
+        }
+        return [
+            migration
+            for migration in self._migrations
+            if migration.name not in already_applied
+        ]
+
+    def applied(self, db: "Database") -> list["Migrations._AppliedMigration"]:
+        """
+        Return a list of applied migrations.
+        """
+        self.ensure_migrations_table(db)
+        return [
+            self._AppliedMigration(name=row["name"], applied_at=row["applied_at"])
+            for row in db[self.migrations_table].rows_where(
+                "migration_set = ?", [self.name]
+            )
+        ]
+
+    def apply(self, db: "Database", *, stop_before: str | None = None):
+        """
+        Apply any pending migrations to the database.
+        """
+        self.ensure_migrations_table(db)
+        for migration in self.pending(db):
+            name = migration.name
+            if name == stop_before:
+                return
+            migration.fn(db)
+            _table(db, self.migrations_table).insert(
+                {
+                    "migration_set": self.name,
+                    "name": name,
+                    "applied_at": str(datetime.datetime.now(datetime.timezone.utc)),
+                }
+            )
+
+    def ensure_migrations_table(self, db: "Database"):
+        """
+        Ensure the _sqlite_migrations table exists and has the correct schema.
+        """
+        table = _table(db, self.migrations_table)
+        if not table.exists():
+            table.create(
+                {
+                    "id": int,
+                    "migration_set": str,
+                    "name": str,
+                    "applied_at": str,
+                },
+                pk="id",
+            )
+            table.create_index(["migration_set", "name"], unique=True)
+        elif table.pks != ["id"]:
+            table.transform(pk="id")
+            unique_indexes = {tuple(index.columns) for index in table.indexes}
+            if ("migration_set", "name") not in unique_indexes:
+                table.create_index(["migration_set", "name"], unique=True)
+
+    def __repr__(self):
+        return "<Migrations '{}': [{}]>".format(
+            self.name, ", ".join(m.name for m in self._migrations)
+        )
+
+
+def _table(db: "Database", name: str) -> "Table":
+    return cast("Table", db[name])