"""Database migration service using Flask-Migrate (Alembic).""" import os import logging import shutil from datetime import datetime from alembic import command from alembic.config import Config as AlembicConfig from alembic.script import ScriptDirectory from alembic.migration import MigrationContext from sqlalchemy import inspect as sa_inspect, text logger = logging.getLogger(__name__) class MigrationService: """Handles database migration detection, backup, and execution.""" _needs_migration = False _current_revision = None _head_revision = None _pending_migrations = [] # True when the database is stamped at a revision this codebase no longer # contains (a reverted branch, a downgraded install). Upgrading is impossible # until it's stamped back to a known revision, so the UI must say that rather # than offer a Continue button that cannot work. _orphaned_revision = False @classmethod def _get_alembic_config(cls, app): """Build an Alembic config pointing at the migrations directory.""" migrations_dir = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(__file__))), 'migrations') cfg = AlembicConfig() cfg.set_main_option('script_location', migrations_dir) cfg.set_main_option('sqlalchemy.url', app.config['SQLALCHEMY_DATABASE_URI']) return cfg # Map SQLAlchemy type names to SQLite-compatible type strings _TYPE_MAP = { 'INTEGER': 'INTEGER', 'BIGINTEGER': 'INTEGER', 'SMALLINTEGER': 'INTEGER', 'FLOAT': 'REAL', 'NUMERIC': 'REAL', 'BOOLEAN': 'BOOLEAN', 'DATETIME': 'DATETIME', 'DATE': 'DATE', 'TIME': 'TIME', 'TEXT': 'TEXT', 'STRING': 'TEXT', 'VARCHAR': 'TEXT', 'JSON': 'TEXT', } @classmethod def _sqlite_type(cls, sa_type): """Convert a SQLAlchemy column type to a SQLite type string.""" type_name = type(sa_type).__name__.upper() return cls._TYPE_MAP.get(type_name, 'TEXT') @classmethod def _sqlite_default(cls, col): """Render a column's literal server_default as a SQLite DEFAULT value.""" default = getattr(col, 'server_default', None) arg = getattr(default, 'arg', None) if isinstance(arg, bool): return '1' if arg else '0' if isinstance(arg, (int, float)): return repr(arg) if isinstance(arg, str): return "'" + arg.replace("'", "''") + "'" return None @classmethod def _fix_missing_columns(cls, db): """Sync database schema with ORM models. Compares every model column against the actual database and adds any that are missing. Also creates tables that don't exist yet. Runs raw SQL before any ORM queries to prevent crashes when models reference columns that don't exist in the database yet. """ inspector = sa_inspect(db.engine) existing_tables = set(inspector.get_table_names()) added = 0 for table_name, table_obj in db.metadata.tables.items(): if table_name not in existing_tables: # Entire table is missing — create_all will handle it later continue existing_cols = {c['name'] for c in inspector.get_columns(table_name)} for col in table_obj.columns: if col.name in existing_cols: continue sqlite_type = cls._sqlite_type(col.type) sql = f'ALTER TABLE {table_name} ADD COLUMN {col.name} {sqlite_type}' # Honour a declared server_default so existing rows are # backfilled instead of left NULL — this runs BEFORE alembic, # so a migration that only adds the column when missing would # otherwise skip its own default (users.auth_version → no login). default = cls._sqlite_default(col) if default is not None: sql += f' DEFAULT {default}' try: with db.engine.begin() as conn: conn.execute(text(sql)) logger.info(f'Auto-added missing column: {table_name}.{col.name} ({sqlite_type})') added += 1 except Exception as e: logger.warning(f'Failed to add {table_name}.{col.name}: {e}') # Create any entirely new tables db.create_all() if added: logger.info(f'Schema sync complete: {added} column(s) added') @classmethod def check_and_prepare(cls, app): """Called on startup to detect migration state. Three scenarios: 1. Fresh install (no DB / no tables) -> upgrade head to create everything 2. Existing install (tables exist, no alembic_version) -> stamp baseline, then check 3. Normal state (alembic_version exists) -> compare current vs head """ from app import db try: # Fix any missing columns before ORM queries can fail cls._fix_missing_columns(db) cfg = cls._get_alembic_config(app) script = ScriptDirectory.from_config(cfg) head = script.get_current_head() cls._head_revision = head inspector = sa_inspect(db.engine) existing_tables = inspector.get_table_names() has_alembic = 'alembic_version' in existing_tables # Count real application tables (exclude alembic_version) app_tables = [t for t in existing_tables if t != 'alembic_version'] if not app_tables and not has_alembic: # Scenario 1: Fresh install — create everything via Alembic logger.info('Fresh install detected — running alembic upgrade head') with app.app_context(): command.upgrade(cfg, 'head') cls._needs_migration = False cls._current_revision = head cls._pending_migrations = [] return if app_tables and not has_alembic: # Scenario 2: Existing install upgrading to Alembic # Run upgrade so migrations can add missing columns to existing tables logger.info('Existing install detected — running alembic upgrade head') with app.app_context(): command.upgrade(cfg, 'head') # Scenario 3 (or after stamping): Check current vs head with db.engine.connect() as conn: context = MigrationContext.configure(conn) current_heads = context.get_current_heads() cls._current_revision = current_heads[0] if current_heads else None if cls._current_revision != head: # Is the stamped revision one this codebase still knows about? # A branch that added a migration and was then reverted leaves the # database stamped at a revision whose file is gone. The walk below # breaks on `current_revision`, so an unknown one never matches and # every revision in history gets reported as "pending" — which is # both wrong and alarming (and `upgrade` would fail anyway with # "Can't locate revision"). Detect it and say so instead. known_revisions = {rev.revision for rev in script.walk_revisions()} cls._orphaned_revision = bool( cls._current_revision and cls._current_revision not in known_revisions ) cls._pending_migrations = [] if cls._orphaned_revision: logger.warning( f'Database is stamped at unknown revision ' f'{cls._current_revision!r} (head={head}). It is not in this ' f'codebase — the schema cannot be upgraded until it is ' f'stamped back to a known revision.' ) cls._needs_migration = True return for rev in script.walk_revisions(): if cls._current_revision and rev.revision == cls._current_revision: break cls._pending_migrations.append({ 'revision': rev.revision, 'description': rev.doc or '', 'down_revision': rev.down_revision, }) cls._pending_migrations.reverse() if cls._pending_migrations: logger.info( f'Database migration needed: {len(cls._pending_migrations)} pending ' f'(current={cls._current_revision}, head={head}) — running upgrade' ) with app.app_context(): command.upgrade(cfg, 'head') cls._current_revision = head cls._needs_migration = False else: cls._needs_migration = False cls._pending_migrations = [] except Exception as e: logger.error(f'Migration check failed: {e}') # Don't block startup on migration check failure — fall back to no-migration state cls._needs_migration = False @classmethod def get_status(cls): """Return current migration status. ``needs_migration`` is derived from the live revision vs head rather than the cached boot flag, so the login gate (App.jsx) prompts whenever the database isn't at head — including when the boot auto-upgrade couldn't complete (e.g. an orphaned revision). When current == head it stays false, so a normally up-to-date instance never sees the gate. """ needs = bool( cls._current_revision and cls._head_revision and cls._current_revision != cls._head_revision ) return { 'needs_migration': needs or cls._orphaned_revision, 'current_revision': cls._current_revision, 'head_revision': cls._head_revision, 'pending_count': len(cls._pending_migrations), 'pending_migrations': cls._pending_migrations, # When true, `pending_migrations` is empty because it is unknowable, # not because the database is up to date. 'orphaned_revision': cls._orphaned_revision, } @classmethod def create_backup(cls, app): """Create a database backup before migration. SQLite: file copy. PostgreSQL: pg_dump. """ db_url = app.config['SQLALCHEMY_DATABASE_URI'] timestamp = datetime.utcnow().strftime('%Y%m%d_%H%M%S') try: if db_url.startswith('sqlite'): # Extract file path from sqlite:///path or sqlite:////path db_path = db_url.replace('sqlite:///', '') if not os.path.isabs(db_path): db_path = os.path.join(os.path.dirname(os.path.dirname(__file__)), db_path) if not os.path.exists(db_path): return {'success': False, 'error': 'Database file not found'} backup_dir = os.path.join(os.path.dirname(db_path), 'backups') os.makedirs(backup_dir, exist_ok=True) backup_name = f'serverkit_pre_migration_{timestamp}.db' backup_path = os.path.join(backup_dir, backup_name) shutil.copy2(db_path, backup_path) logger.info(f'Database backup created: {backup_path}') return {'success': True, 'path': backup_path} elif 'postgresql' in db_url: from app.utils.system import run_checked backup_dir = '/var/serverkit/backups/db' os.makedirs(backup_dir, exist_ok=True) backup_name = f'serverkit_pre_migration_{timestamp}.sql' backup_path = os.path.join(backup_dir, backup_name) result = run_checked(['pg_dump', db_url, '-f', backup_path], timeout=300) if not result['success']: return {'success': False, 'error': result['error']} logger.info(f'Database backup created: {backup_path}') return {'success': True, 'path': backup_path} else: return {'success': False, 'error': f'Unsupported database type'} except Exception as e: logger.error(f'Backup failed: {e}') return {'success': False, 'error': str(e)} @classmethod def apply_migrations(cls, app): """Run all pending Alembic migrations. Self-heals an orphaned ``alembic_version``: if the DB records a revision that no longer exists in the migration scripts (renamed/removed during development), Alembic can't compute an upgrade path. Because the boot schema sync (``_fix_missing_columns`` + ``create_all``) already brings the live schema up to the ORM models (== head), we re-stamp to head to repair the bookkeeping instead of failing. """ try: cfg = cls._get_alembic_config(app) script = ScriptDirectory.from_config(cfg) from app import db # What revision does the DB think it's at, and does Alembic still know it? with db.engine.connect() as conn: context = MigrationContext.configure(conn) db_heads = list(context.get_current_heads()) def _known(rev): try: script.get_revision(rev) return True except Exception: return False orphaned = bool(db_heads) and not all(_known(r) for r in db_heads) with app.app_context(): if orphaned: # command.stamp() can't help here: it still resolves the # (unknown) current revision and fails the same way. The live # schema already matches head via the boot sync, so reset the # alembic_version table directly to head. head = script.get_current_head() logger.warning( f'Recorded revision {db_heads} not found in migration scripts — ' f'resetting alembic_version to head ({head})' ) with db.engine.begin() as conn: conn.execute(text('DELETE FROM alembic_version')) if head: conn.execute( text('INSERT INTO alembic_version (version_num) VALUES (:v)'), {'v': head}, ) else: command.upgrade(cfg, 'head') # Update internal state with db.engine.connect() as conn: context = MigrationContext.configure(conn) current_heads = context.get_current_heads() cls._current_revision = current_heads[0] if current_heads else None cls._needs_migration = False cls._pending_migrations = [] # Record in SystemSettings try: from app.models import SystemSettings from app import db as _db SystemSettings.set('schema_version', cls._current_revision) SystemSettings.set('last_migration_at', datetime.utcnow().isoformat()) _db.session.commit() except Exception as e: logger.warning(f'Failed to record migration in settings: {e}') logger.info(f'Migrations applied successfully (now at {cls._current_revision})') return {'success': True, 'revision': cls._current_revision} except Exception as e: logger.error(f'Migration failed: {e}') return {'success': False, 'error': str(e)} @classmethod def get_migration_history(cls, app): """Return list of all Alembic revisions with descriptions.""" try: cfg = cls._get_alembic_config(app) script = ScriptDirectory.from_config(cfg) revisions = [] for rev in script.walk_revisions(): revisions.append({ 'revision': rev.revision, 'down_revision': rev.down_revision, 'description': rev.doc or '', 'is_current': rev.revision == cls._current_revision, 'is_head': rev.revision == cls._head_revision, }) revisions.reverse() return revisions except Exception as e: logger.error(f'Failed to get migration history: {e}') return []