Build reusable bot framework
This commit is contained in:
Chelsea Lee
202
core/migrations/__init__.py
Normal file
202
core/migrations/__init__.py
Normal file
@@ -0,0 +1,202 @@
|
||||
"""Versioned PostgreSQL migrations for the framework and feature modules."""
|
||||
|
||||
from dataclasses import dataclass
|
||||
from hashlib import sha256
|
||||
from pathlib import Path
|
||||
import re
|
||||
|
||||
import psycopg2.extras
|
||||
|
||||
from core import postgres
|
||||
|
||||
|
||||
MIGRATION_NAME = re.compile(r"^(?P<version>\d+)(?:[_-].*)?\.sql$")
|
||||
LOCK_NAME = "llm-bot-framework:schema-migrations"
|
||||
|
||||
|
||||
class MigrationError(RuntimeError):
|
||||
"""Raised when migration history is inconsistent or cannot be applied."""
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Migration:
|
||||
namespace: str
|
||||
version: int
|
||||
name: str
|
||||
path: Path
|
||||
checksum: str
|
||||
sql: str
|
||||
|
||||
|
||||
def _project_root(project_root=None):
|
||||
if project_root is not None:
|
||||
return Path(project_root).resolve()
|
||||
return Path(__file__).resolve().parents[2]
|
||||
|
||||
|
||||
def _read_namespace(namespace, directory):
|
||||
migrations = []
|
||||
if not directory.is_dir():
|
||||
return migrations
|
||||
|
||||
for path in sorted(directory.glob("*.sql")):
|
||||
match = MIGRATION_NAME.match(path.name)
|
||||
if not match:
|
||||
raise MigrationError(f"Invalid migration filename: {path}")
|
||||
sql = path.read_text(encoding="utf-8")
|
||||
migrations.append(
|
||||
Migration(
|
||||
namespace=namespace,
|
||||
version=int(match.group("version")),
|
||||
name=path.name,
|
||||
path=path,
|
||||
checksum=sha256(sql.encode("utf-8")).hexdigest(),
|
||||
sql=sql,
|
||||
)
|
||||
)
|
||||
return migrations
|
||||
|
||||
|
||||
def discover_migrations(project_root=None):
|
||||
"""Return core and feature migrations in deterministic application order."""
|
||||
root = _project_root(project_root)
|
||||
migrations = _read_namespace("core", root / "config" / "migrations")
|
||||
|
||||
modules_root = root / "modules"
|
||||
if modules_root.is_dir():
|
||||
for module_path in sorted(modules_root.iterdir(), key=lambda path: path.name):
|
||||
if module_path.is_dir() and not module_path.name.startswith("_"):
|
||||
migrations.extend(
|
||||
_read_namespace(module_path.name, module_path / "migrations")
|
||||
)
|
||||
|
||||
seen = set()
|
||||
for migration in migrations:
|
||||
key = (migration.namespace, migration.version)
|
||||
if key in seen:
|
||||
raise MigrationError(
|
||||
f"Duplicate migration {migration.namespace}:{migration.version}"
|
||||
)
|
||||
seen.add(key)
|
||||
|
||||
return sorted(
|
||||
migrations,
|
||||
key=lambda item: (
|
||||
item.namespace != "core",
|
||||
item.namespace,
|
||||
item.version,
|
||||
item.name,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _ensure_history_table(cursor):
|
||||
cursor.execute(
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS schema_migrations (
|
||||
namespace VARCHAR(255) NOT NULL,
|
||||
version INTEGER NOT NULL,
|
||||
checksum CHAR(64) NOT NULL,
|
||||
applied_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
|
||||
PRIMARY KEY (namespace, version)
|
||||
)
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
def _applied_migrations(cursor):
|
||||
cursor.execute(
|
||||
"""
|
||||
SELECT namespace, version, checksum, applied_at
|
||||
FROM schema_migrations
|
||||
ORDER BY namespace, version
|
||||
"""
|
||||
)
|
||||
return {
|
||||
(row["namespace"], row["version"]): dict(row) for row in cursor.fetchall()
|
||||
}
|
||||
|
||||
|
||||
def upgrade(project_root=None):
|
||||
"""Apply pending migrations transactionally under a PostgreSQL advisory lock."""
|
||||
migrations = discover_migrations(project_root)
|
||||
applied_now = []
|
||||
|
||||
with postgres.get_connection() as connection:
|
||||
with connection.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cursor:
|
||||
cursor.execute("SELECT pg_advisory_xact_lock(hashtext(%s))", (LOCK_NAME,))
|
||||
_ensure_history_table(cursor)
|
||||
applied = _applied_migrations(cursor)
|
||||
|
||||
for migration in migrations:
|
||||
key = (migration.namespace, migration.version)
|
||||
previous = applied.get(key)
|
||||
if previous:
|
||||
if previous["checksum"] != migration.checksum:
|
||||
raise MigrationError(
|
||||
"Applied migration checksum changed: "
|
||||
f"{migration.namespace}:{migration.version}"
|
||||
)
|
||||
continue
|
||||
|
||||
cursor.execute(migration.sql)
|
||||
cursor.execute(
|
||||
"""
|
||||
INSERT INTO schema_migrations (namespace, version, checksum)
|
||||
VALUES (%s, %s, %s)
|
||||
""",
|
||||
(migration.namespace, migration.version, migration.checksum),
|
||||
)
|
||||
applied_now.append(migration)
|
||||
|
||||
return applied_now
|
||||
|
||||
|
||||
def migration_status(project_root=None):
|
||||
"""Return applied, pending, changed, and source-missing migration records."""
|
||||
migrations = discover_migrations(project_root)
|
||||
with postgres.get_connection() as connection:
|
||||
with connection.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cursor:
|
||||
cursor.execute(
|
||||
"SELECT to_regclass('public.schema_migrations') AS table_name"
|
||||
)
|
||||
if cursor.fetchone()["table_name"] is None:
|
||||
applied = {}
|
||||
else:
|
||||
applied = _applied_migrations(cursor)
|
||||
|
||||
status = []
|
||||
source_keys = set()
|
||||
for migration in migrations:
|
||||
key = (migration.namespace, migration.version)
|
||||
source_keys.add(key)
|
||||
record = applied.get(key)
|
||||
state = "pending"
|
||||
applied_at = None
|
||||
if record:
|
||||
state = "applied" if record["checksum"] == migration.checksum else "changed"
|
||||
applied_at = record["applied_at"]
|
||||
status.append(
|
||||
{
|
||||
"namespace": migration.namespace,
|
||||
"version": migration.version,
|
||||
"name": migration.name,
|
||||
"checksum": migration.checksum,
|
||||
"applied_at": applied_at,
|
||||
"state": state,
|
||||
}
|
||||
)
|
||||
|
||||
for key, record in sorted(applied.items()):
|
||||
if key not in source_keys:
|
||||
status.append(
|
||||
{
|
||||
"namespace": key[0],
|
||||
"version": key[1],
|
||||
"name": None,
|
||||
"checksum": record["checksum"],
|
||||
"applied_at": record["applied_at"],
|
||||
"state": "missing",
|
||||
}
|
||||
)
|
||||
return status
|
||||
37
core/migrations/__main__.py
Normal file
37
core/migrations/__main__.py
Normal file
@@ -0,0 +1,37 @@
|
||||
"""Command-line entry point for database migrations."""
|
||||
|
||||
import argparse
|
||||
|
||||
from core.migrations import MigrationError, migration_status, upgrade
|
||||
|
||||
|
||||
def main(argv=None):
|
||||
parser = argparse.ArgumentParser(description="Manage database migrations")
|
||||
parser.add_argument("command", choices=("upgrade", "status"))
|
||||
args = parser.parse_args(argv)
|
||||
|
||||
try:
|
||||
if args.command == "upgrade":
|
||||
applied = upgrade()
|
||||
for migration in applied:
|
||||
print(f"applied {migration.namespace}:{migration.version} {migration.name}")
|
||||
if not applied:
|
||||
print("database is up to date")
|
||||
return 0
|
||||
|
||||
records = migration_status()
|
||||
if not records:
|
||||
print("no migrations found")
|
||||
for record in records:
|
||||
name = record["name"] or "<source missing>"
|
||||
print(
|
||||
f"{record['state']:<8} "
|
||||
f"{record['namespace']}:{record['version']} {name}"
|
||||
)
|
||||
return 1 if any(row["state"] == "changed" for row in records) else 0
|
||||
except MigrationError as error:
|
||||
parser.exit(1, f"migration error: {error}\n")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
Reference in New Issue
Block a user