"""Plugin DB migration runner with tenant_id column validation.""" from __future__ import annotations import os from pathlib import Path from typing import Any from sqlalchemy import inspect, text from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession from app.models.plugin import PluginMigration class MigrationValidationError(Exception): """Raised when a plugin migration creates a table without tenant_id column.""" class MigrationRunner: """Runs plugin SQL migrations and validates that all created tables have tenant_id.""" def __init__(self, engine: AsyncEngine, migrations_dir: str | None = None) -> None: self._engine = engine self._migrations_dir = Path( migrations_dir or os.path.join(os.path.dirname(__file__), "migrations") ) def _resolve_migration_path(self, filename: str, plugin_name: str | None = None) -> Path: """Resolve a migration filename to an absolute path. Search order: 1. Plugin-specific migrations dir: app/plugins/builtins//migrations/ 2. Global migrations_dir (default: app/plugins/migrations/) 3. Shared builtins migrations dir: app/plugins/builtins/migrations/ """ # 1. Plugin-specific migrations directory if plugin_name: plugin_migrations_dir = Path(__file__).parent / "builtins" / plugin_name / "migrations" candidate = plugin_migrations_dir / filename if candidate.exists(): return candidate # 2. Global migrations_dir candidate = self._migrations_dir / filename if candidate.exists(): return candidate # 3. Shared builtins migrations directory builtins_dir = Path(__file__).parent / "builtins" / "migrations" candidate2 = builtins_dir / filename if candidate2.exists(): return candidate2 raise FileNotFoundError(f"Migration file not found: {filename}") async def run_migration( self, db: AsyncSession, plugin_name: str, migration_filename: str, tenant_id: Any | None = None, ) -> PluginMigration: """Run a single SQL migration file for a plugin. Reads the SQL file, executes it, validates created tables have tenant_id, and records the migration in plugin_migrations table. Args: db: Async database session. plugin_name: Name of the plugin being migrated. migration_filename: Filename of the migration SQL file. tenant_id: Optional tenant UUID for tenant-scoped migrations. Returns: PluginMigration record. Raises: FileNotFoundError: If migration file doesn't exist. MigrationValidationError: If a created table lacks tenant_id column. """ migration_path = self._resolve_migration_path(migration_filename, plugin_name) sql_content = migration_path.read_text(encoding="utf-8") # Capture existing table names before migration (static analysis of SQL) existing_tables = await self._get_table_names_via_session(db) # Execute the migration SQL # Split on semicolons but handle potential function definitions statements = self._split_sql(sql_content) for stmt in statements: stmt = stmt.strip() if stmt: await db.execute(text(stmt)) await db.flush() # Capture table names after migration (using same session connection) new_tables = await self._get_table_names_via_session(db) created_tables = new_tables - existing_tables # Validate: all newly created tables must have tenant_id column # Exception: tables marked as global in the SQL with `-- GLOBAL TABLE` comment global_tables = self._extract_global_table_names(sql_content) tables_without_tenant = [] for table_name in created_tables: if table_name in global_tables: continue # Global tables are exempt from tenant_id requirement if not await self._table_has_column_via_session(db, table_name, "tenant_id"): tables_without_tenant.append(table_name) if tables_without_tenant: # Rollback the migration await db.rollback() raise MigrationValidationError( f"Plugin '{plugin_name}' migration '{migration_filename}' created tables without tenant_id column: " f"{', '.join(tables_without_tenant)}. All plugin tables must have tenant_id for multi-tenant isolation. " f"If a table is intentionally global, add `-- GLOBAL TABLE: table_name` comment to the migration SQL." ) # Enable RLS on all newly created tenant tables. # Plugin migrations run AFTER core alembic migration 0085, which sets up RLS # for all known core tables. Without this step, plugin-created tables would # be left without RLS — a critical multi-tenant isolation gap. await self._enable_rls_for_tables(db, created_tables, sql_content) # Record the migration in plugin_migrations table migration_record = PluginMigration( plugin_name=plugin_name, migration_file=migration_filename, status="applied", ) db.add(migration_record) await db.flush() return migration_record async def run_all_migrations( self, db: AsyncSession, plugin_name: str, migration_files: list[str], tenant_id: Any | None = None, ) -> list[PluginMigration]: """Run all migration files for a plugin in order. Skips already-applied migrations (idempotent). """ # Check which migrations are already applied from sqlalchemy import select result = await db.execute( select(PluginMigration).where( PluginMigration.plugin_name == plugin_name, PluginMigration.status == "applied", ) ) applied_files = {row.migration_file for row in result.scalars().all()} records: list[PluginMigration] = [] for filename in migration_files: if filename in applied_files: continue # Already applied — idempotent skip record = await self.run_migration(db, plugin_name, filename, tenant_id) records.append(record) return records async def drop_plugin_tables(self, db: AsyncSession, plugin_name: str) -> list[str]: """Drop all tables that were created by a plugin's migrations. Args: db: Async database session. plugin_name: Name of the plugin to remove tables for. Returns: List of dropped table names. """ from sqlalchemy import select # Get all migration records for this plugin result = await db.execute( select(PluginMigration).where(PluginMigration.plugin_name == plugin_name) ) migration_records = result.scalars().all() dropped_tables: list[str] = [] for record in migration_records: migration_path = self._resolve_migration_path(record.migration_file, plugin_name) sql_content = migration_path.read_text(encoding="utf-8") # Parse table names from CREATE TABLE statements table_names = self._extract_table_names(sql_content) for table_name in table_names: try: await db.execute(text(f'DROP TABLE IF EXISTS "{table_name}" CASCADE')) dropped_tables.append(table_name) except Exception: pass # Table may already be gone await db.flush() # Remove migration records for record in migration_records: await db.delete(record) await db.flush() return dropped_tables async def get_applied_migrations( self, db: AsyncSession, plugin_name: str, ) -> list[str]: """Get all applied migration filenames for a plugin, sorted by application order. Args: db: Async database session. plugin_name: Name of the plugin. Returns: List of migration filenames sorted by application order (oldest first). """ from sqlalchemy import select result = await db.execute( select(PluginMigration) .where( PluginMigration.plugin_name == plugin_name, PluginMigration.status == "applied", ) .order_by(PluginMigration.id) ) return [row.migration_file for row in result.scalars().all()] async def _find_down_sql( self, migration_filename: str, plugin_name: str | None = None, ) -> str | None: """Find rollback SQL for a migration. Search order: 1. A dedicated down file: _down.sql 2. A `-- DOWN:` block inside the original migration file Returns the rollback SQL string, or None if no rollback is found. """ # 1. Try dedicated down file base, ext = os.path.splitext(migration_filename) down_filename = f"{base}_down{ext}" try: down_path = self._resolve_migration_path(down_filename, plugin_name) return down_path.read_text(encoding="utf-8") except FileNotFoundError: pass # 2. Try parsing -- DOWN: block from the original migration file try: up_path = self._resolve_migration_path(migration_filename, plugin_name) content = up_path.read_text(encoding="utf-8") except FileNotFoundError: return None down_marker = "-- DOWN:" if down_marker in content: parts = content.split(down_marker, 1) if len(parts) == 2: down_sql = parts[1].strip() return down_sql if down_sql else None return None async def run_migration_down( self, db: AsyncSession, plugin_name: str, migration_filename: str, ) -> None: """Roll back a single migration. Searches for rollback SQL (dedicated _down.sql file or -- DOWN: block in the original migration), executes it, and removes the migration record from plugin_migrations. Args: db: Async database session. plugin_name: Name of the plugin. migration_filename: Filename of the migration to roll back. Raises: FileNotFoundError: If no rollback SQL can be found. """ down_sql = await self._find_down_sql(migration_filename, plugin_name) if down_sql is None: raise FileNotFoundError( f"No rollback SQL found for migration '{migration_filename}'. " f"Create a '{migration_filename.replace('.sql', '_down.sql')}' file " f"or add a '-- DOWN:' section to the migration file." ) # Execute the rollback SQL statements = self._split_sql(down_sql) for stmt in statements: stmt = stmt.strip() if stmt: await db.execute(text(stmt)) await db.flush() # Remove the migration record from sqlalchemy import select result = await db.execute( select(PluginMigration).where( PluginMigration.plugin_name == plugin_name, PluginMigration.migration_file == migration_filename, PluginMigration.status == "applied", ) ) record = result.scalar_one_or_none() if record: await db.delete(record) await db.flush() async def rollback_to_version( self, db: AsyncSession, plugin_name: str, target_version: str, ) -> list[str]: """Roll back all applied migrations after a target version. Migrations are rolled back in reverse order (newest first) until the target version is reached. The plugin version in the database is updated to the target version. Args: db: Async database session. plugin_name: Name of the plugin. target_version: Target version string (e.g. '0002'). Migrations with filenames greater than this will be rolled back. Returns: List of migration filenames that were rolled back. """ applied = await self.get_applied_migrations(db, plugin_name) # Filter migrations after target_version (by filename sort order) migrations_to_rollback = [ m for m in applied if m > target_version ] if not migrations_to_rollback: return [] # Roll back in reverse order (newest first) rolled_back: list[str] = [] for migration_filename in reversed(migrations_to_rollback): await self.run_migration_down(db, plugin_name, migration_filename) rolled_back.append(migration_filename) # Update plugin version in DB (if a plugin_versions table exists) try: from sqlalchemy import select, update as sa_update result = await db.execute( text("SELECT 1 FROM information_schema.tables " "WHERE table_schema = 'public' AND table_name = 'plugin_versions'") ) if result.fetchone(): await db.execute( text("UPDATE plugin_versions SET version = :version " "WHERE plugin_name = :plugin_name"), {"version": target_version, "plugin_name": plugin_name}, ) await db.flush() except Exception: pass # plugin_versions table may not exist — that's ok return rolled_back async def _get_table_names_via_session(self, db: AsyncSession) -> set[str]: """Get current table names using the session's own connection. This ensures we see uncommitted DDL changes within the same transaction. """ result = await db.execute( text("SELECT tablename FROM pg_tables WHERE schemaname = 'public'") ) return {row[0] for row in result.fetchall()} async def _table_has_column_via_session( self, db: AsyncSession, table_name: str, column_name: str ) -> bool: """Check if a table has a specific column using the session's connection.""" result = await db.execute( text( "SELECT column_name FROM information_schema.columns " "WHERE table_schema = 'public' AND table_name = :table_name AND column_name = :column_name" ), {"table_name": table_name, "column_name": column_name}, ) return result.fetchone() is not None async def _get_table_names(self) -> set[str]: """Get current table names from the database (separate connection).""" def _get_names(sync_conn): insp = inspect(sync_conn) return set(insp.get_table_names()) async with self._engine.connect() as conn: return await conn.run_sync(_get_names) async def _table_has_column(self, table_name: str, column_name: str) -> bool: """Check if a table has a specific column (separate connection).""" def _has_col(sync_conn): insp = inspect(sync_conn) if table_name not in insp.get_table_names(): return False columns = [col["name"] for col in insp.get_columns(table_name)] return column_name in columns async with self._engine.connect() as conn: return await conn.run_sync(_has_col) @staticmethod def _split_sql(sql: str) -> list[str]: """Split SQL into individual statements, respecting basic semicolon boundaries.""" statements: list[str] = [] current: list[str] = [] in_dollar_quote = False dollar_tag = "" for line in sql.splitlines(): stripped = line.strip() # Handle dollar-quoted blocks (PostgreSQL function bodies) if "$$" in stripped and not in_dollar_quote: parts = stripped.split("$$", 1) if len(parts) > 1: dollar_tag = parts[0].strip() if dollar_tag: in_dollar_quote = True elif stripped.count("$$") >= 2: # Single-line $$ block current.append(line) if stripped.rstrip().endswith(";"): statements.append("\n".join(current)) current = [] continue else: in_dollar_quote = True current.append(line) continue if in_dollar_quote: current.append(line) if "$$" in stripped: in_dollar_quote = False # After closing $$, check if the line ends with ; to flush statement after_dollar = stripped.split("$$", 1)[-1] if "$$" in stripped else "" if after_dollar.rstrip().endswith(";"): statements.append("\n".join(current)) current = [] continue current.append(line) if stripped.endswith(";"): statements.append("\n".join(current)) current = [] if current: remaining = "\n".join(current).strip() if remaining: statements.append(remaining) return statements @staticmethod def _extract_table_names(sql: str) -> list[str]: """Extract table names from CREATE TABLE statements in SQL.""" import re # Match CREATE TABLE [IF NOT EXISTS] "table_name" or CREATE TABLE table_name pattern = r'CREATE\s+TABLE\s+(?:IF\s+NOT\s+EXISTS\s+)?["\']?(\w+)["\']?' return re.findall(pattern, sql, re.IGNORECASE) def _extract_global_table_names(self, sql: str) -> set[str]: """Extract table names marked as global in SQL comments. Looks for `-- GLOBAL TABLE: table_name` comments in the SQL. These tables are exempt from the tenant_id column requirement. """ import re pattern = r'--\s*GLOBAL\s+TABLE:\s*(\w+)' return set(re.findall(pattern, sql, re.IGNORECASE)) async def _enable_rls_for_tables( self, db: AsyncSession, created_tables: set[str], sql_content: str, ) -> None: """Enable Row Level Security on all newly created tenant tables. Plugin migrations run AFTER core alembic migration 0085, which sets up RLS for all known core tables. Without this step, plugin-created tables would be left without RLS — a critical multi-tenant isolation gap. This method replicates the exact RLS pattern from 0085 for every new tenant table: 1. ALTER TABLE ... ENABLE ROW LEVEL SECURITY 2. ALTER TABLE ... FORCE ROW LEVEL SECURITY 3. DROP old policies (idempotent) + CREATE fail-closed tenant isolation policy 4. GRANT CRUD to crm_api and crm_worker 5. ALTER TABLE ... OWNER TO crm_migration Global tables (marked with `-- GLOBAL TABLE` in SQL) are skipped. """ global_tables = self._extract_global_table_names(sql_content) tenant_tables = created_tables - global_tables for table_name in tenant_tables: # Idempotent: DROP old policies first, then ENABLE + FORCE RLS await db.execute( text( f"DO $$ BEGIN " f"IF EXISTS (SELECT 1 FROM information_schema.tables " f"WHERE table_schema = 'public' AND table_name = '{table_name}') THEN " f"DROP POLICY IF EXISTS tenant_isolation ON public.{table_name}; " f"DROP POLICY IF EXISTS {table_name}_tenant_isolation ON public.{table_name}; " f"ALTER TABLE public.{table_name} ENABLE ROW LEVEL SECURITY; " f"ALTER TABLE public.{table_name} FORCE ROW LEVEL SECURITY; " f"END IF; END $$" ) ) # Create fail-closed tenant isolation policy (same pattern as 0085) await db.execute( text( f"DO $$ BEGIN " f"IF EXISTS (SELECT 1 FROM information_schema.tables " f"WHERE table_schema = 'public' AND table_name = '{table_name}') THEN " f"CREATE POLICY {table_name}_tenant_isolation " f"ON public.{table_name} " f"FOR ALL " f"TO crm_api, crm_worker " f"USING (tenant_id = NULLIF(current_setting('app.current_tenant_id', true), '')::uuid) " f"WITH CHECK (tenant_id = NULLIF(current_setting('app.current_tenant_id', true), '')::uuid); " f"END IF; END $$" ) ) # Grant CRUD to runtime roles await db.execute( text( f"DO $$ BEGIN " f"IF EXISTS (SELECT 1 FROM information_schema.tables " f"WHERE table_schema = 'public' AND table_name = '{table_name}') THEN " f"GRANT SELECT, INSERT, UPDATE, DELETE ON public.{table_name} TO crm_api; " f"GRANT SELECT, INSERT, UPDATE, DELETE ON public.{table_name} TO crm_worker; " f"END IF; END $$" ) ) # Transfer ownership to crm_migration (same as 0085) await db.execute( text( f"DO $$ BEGIN " f"IF EXISTS (SELECT 1 FROM information_schema.tables " f"WHERE table_schema = 'public' AND table_name = '{table_name}') THEN " f"ALTER TABLE public.{table_name} OWNER TO crm_migration; " f"END IF; END $$" ) ) await db.flush()