New migrations system, ported from sqlite-migrate (#754)

Closes #752
This commit is contained in:
Simon Willison 2026-06-21 09:40:21 -07:00 committed by GitHub
commit 3cc27d69bc
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
10 changed files with 963 additions and 19 deletions

View file

@ -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"]

View file

@ -1,4 +1,5 @@
import base64
import difflib
from typing import Any
import click
from click_default_group import DefaultGroup
@ -10,6 +11,7 @@ import sqlite_utils
from sqlite_utils.db import (
AlterError,
BadMultiValues,
DEFAULT,
DescIndex,
NoTable,
quote_identifier,
@ -2578,7 +2580,6 @@ def transform(
_register_db_for_cleanup(db)
_load_extensions(db, load_extension)
types = {}
kwargs = {}
for column, ctype in type:
if ctype.upper() not in VALID_COLUMN_TYPES:
raise click.ClickException(
@ -2598,29 +2599,46 @@ def transform(
for column in default_none:
default_dict[column] = None
kwargs["types"] = types
kwargs["drop"] = set(drop)
kwargs["rename"] = dict(rename)
kwargs["column_order"] = column_order or None
kwargs["not_null"] = not_null_dict
drop_set = set(drop)
rename_dict = dict(rename)
column_order_list = list(column_order) or None
drop_foreign_keys_value = drop_foreign_keys or None
add_foreign_keys_value = add_foreign_keys or None
pk_value = DEFAULT
if pk:
if len(pk) == 1:
kwargs["pk"] = pk[0]
pk_value = pk[0]
else:
kwargs["pk"] = pk
pk_value = pk
elif pk_none:
kwargs["pk"] = None
kwargs["defaults"] = default_dict
if drop_foreign_keys:
kwargs["drop_foreign_keys"] = drop_foreign_keys
if add_foreign_keys:
kwargs["add_foreign_keys"] = add_foreign_keys
pk_value = None
table_obj = db.table(table)
if sql:
for line in db.table(table).transform_sql(**kwargs):
for line in table_obj.transform_sql(
types=types,
drop=drop_set,
rename=rename_dict,
column_order=column_order_list,
not_null=not_null_dict,
pk=pk_value,
defaults=default_dict,
drop_foreign_keys=drop_foreign_keys_value,
add_foreign_keys=add_foreign_keys_value,
):
click.echo(line)
else:
db.table(table).transform(**kwargs)
table_obj.transform(
types=types,
drop=drop_set,
rename=rename_dict,
column_order=column_order_list,
not_null=not_null_dict,
pk=pk_value,
defaults=default_dict,
drop_foreign_keys=drop_foreign_keys_value,
add_foreign_keys=add_foreign_keys_value,
)
@cli.command()
@ -3252,6 +3270,142 @@ 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()
def _stop_before_for_migration_set(stop_before, migration_set_name):
matches = []
for value in stop_before:
set_name, separator, migration_name = value.partition(":")
if separator:
if set_name == migration_set_name:
matches.append(migration_name)
else:
matches.append(value)
return matches
@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",
multiple=True,
help="Stop before applying this migration. Use set:name to target a migration set.",
)
@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.
Use --stop-before migration_set:name to stop before a
migration. This option can be used multiple times.
"""
files = _find_migration_files(migrations)
migration_sets = _load_migration_sets(files)
if not migration_sets:
raise click.ClickException("No migrations.py files found")
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_for_migration_set(stop_before, migration_set.name),
)
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 +3413,7 @@ def plugins_list():
pm.hook.register_commands(cli=cli)
cli.add_command(migrate)
def _render_common(title, values):

126
sqlite_utils/migrations.py Normal file
View file

@ -0,0 +1,126 @@
from collections.abc import Iterable
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 | Iterable[str] | None = None):
"""
Apply any pending migrations to the database.
"""
self.ensure_migrations_table(db)
if stop_before is None:
stop_before_names = set()
elif isinstance(stop_before, str):
stop_before_names = {stop_before}
else:
stop_before_names = set(stop_before)
for migration in self.pending(db):
name = migration.name
if name in stop_before_names:
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])