PYTHON migrate.py

#!/usr/bin/env python3
"""Simple versioned SQL migration runner for Command Sovereignty.

Migrations live in migrations/ as NNN_description.sql files (e.g.
001_add_index_users_email.sql). Applied migrations are tracked in a
schema_migrations table. Each migration runs inside a transaction.

Usage:
    .venv/bin/python scripts/migrate.py            # apply pending migrations
    .venv/bin/python scripts/migrate.py --status   # show applied/pending
    .venv/bin/python scripts/migrate.py --dry-run  # list what would run
    .venv/bin/python scripts/migrate.py --create "add index on users.email"
"""
import argparse
import datetime
import os
import re
import sys

BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
MIGRATIONS_DIR = os.path.join(BASE_DIR, 'migrations')
sys.path.insert(0, BASE_DIR)

# Load env file the same way the service does
ENV_FILE = os.path.expanduser('~/.config/command-sovereignty/env')
if os.path.exists(ENV_FILE):
    with open(ENV_FILE) as f:
        for line in f:
            line = line.strip()
            if line and not line.startswith('#') and '=' in line:
                key, _, value = line.partition('=')
                os.environ.setdefault(key.strip(), value.strip())

from sqlalchemy import create_engine, text  # noqa: E402

MIGRATION_RE = re.compile(r'^(\d{3,})_[\w-]+\.sql$')


def get_engine():
    url = os.environ.get('DATABASE_URL')
    if not url:
        print('ERROR: DATABASE_URL not set (checked env and ~/.config/command-sovereignty/env)')
        sys.exit(1)
    if url.startswith('postgres://'):
        url = url.replace('postgres://', 'postgresql://', 1)
    return create_engine(url)


def ensure_table(engine):
    with engine.begin() as conn:
        conn.execute(text(
            'CREATE TABLE IF NOT EXISTS schema_migrations ('
            'version VARCHAR(255) PRIMARY KEY, '
            'applied_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP)'
        ))


def applied_versions(engine):
    with engine.connect() as conn:
        rows = conn.execute(text('SELECT version FROM schema_migrations')).fetchall()
    return {r[0] for r in rows}


def discover_migrations():
    if not os.path.isdir(MIGRATIONS_DIR):
        return []
    files = []
    for name in sorted(os.listdir(MIGRATIONS_DIR)):
        m = MIGRATION_RE.match(name)
        if m:
            files.append((name, os.path.join(MIGRATIONS_DIR, name)))
        elif name.endswith('.sql'):
            print(f'WARNING: skipping {name} (must match NNN_description.sql)')
    return files


def cmd_status(engine):
    applied = applied_versions(engine)
    migrations = discover_migrations()
    if not migrations:
        print('No migration files found in migrations/')
        return
    for name, _ in migrations:
        marker = 'applied' if name in applied else 'PENDING'
        print(f'  [{marker:>7}] {name}')


def cmd_migrate(engine, dry_run=False):
    applied = applied_versions(engine)
    pending = [(n, p) for n, p in discover_migrations() if n not in applied]
    if not pending:
        print('Database is up to date. No pending migrations.')
        return
    for name, path in pending:
        if dry_run:
            print(f'Would apply: {name}')
            continue
        print(f'Applying {name} ...', end=' ', flush=True)
        with open(path) as f:
            sql = f.read()
        try:
            with engine.begin() as conn:
                # Split on semicolons to handle multi-statement migrations (SQLite limitation)
                statements = [s.strip() for s in sql.split(';') if s.strip() and not s.strip().startswith('--')]
                for stmt in statements:
                    conn.execute(text(stmt))
                conn.execute(
                    text('INSERT INTO schema_migrations (version) VALUES (:v)'),
                    {'v': name},
                )
            print('OK')
        except Exception as e:
            print('FAILED')
            print(f'ERROR in {name}: {e}')
            print('Transaction rolled back. Fix the migration and re-run.')
            sys.exit(1)
    if not dry_run:
        print(f'Applied {len(pending)} migration(s).')


def cmd_create(description):
    os.makedirs(MIGRATIONS_DIR, exist_ok=True)
    existing = discover_migrations()
    next_num = 1
    if existing:
        next_num = max(int(MIGRATION_RE.match(n).group(1)) for n, _ in existing) + 1
    slug = re.sub(r'[^\w]+', '_', description.lower()).strip('_')[:60]
    name = f'{next_num:03d}_{slug}.sql'
    path = os.path.join(MIGRATIONS_DIR, name)
    with open(path, 'w') as f:
        f.write(f'-- Migration: {description}\n'
                f'-- Created: {datetime.date.today().isoformat()}\n\n')
    print(f'Created {path}')


def main():
    parser = argparse.ArgumentParser(description='Run versioned SQL migrations')
    parser.add_argument('--status', action='store_true', help='show migration status')
    parser.add_argument('--dry-run', action='store_true', help='list pending without applying')
    parser.add_argument('--create', metavar='DESC', help='create a new empty migration file')
    args = parser.parse_args()

    if args.create:
        cmd_create(args.create)
        return

    engine = get_engine()
    ensure_table(engine)
    if args.status:
        cmd_status(engine)
    else:
        cmd_migrate(engine, dry_run=args.dry_run)


if __name__ == '__main__':
    main()