"""Plugin migration management using Alembic.""" from pathlib import Path from typing import Optional import logging import subprocess import sys logger = logging.getLogger(__name__) class PluginMigrationManager: """ Manages database migrations for plugins. Each plugin has its own migrations directory. """ def __init__(self, plugins_dir: Path, database_url: str): self.plugins_dir = plugins_dir self.database_url = database_url # Set by PluginManager.init_app; a PluginVerifier or None (no policy). self.verifier = None def _verify_ok(self, plugin_name: str) -> bool: """Strip bytecode + verify the tree before ANY alembic import. Alembic imports env.py and every versions/*.py module body (running their top-level code) with full DB rights, so every method that reaches alembic for a plugin must pass through here first (finding #3). """ if self.verifier is None or not self.verifier.require_signed: return True from .packaging import strip_bytecode strip_bytecode(self.plugins_dir / plugin_name) ok, reason = self.verifier.check(plugin_name) if not ok: logger.error( "Refusing plugin migration code for %s: %s", plugin_name, reason) return ok def get_migrations_dir(self, plugin_name: str) -> Optional[Path]: """Get migrations directory for a plugin.""" migrations_dir = self.plugins_dir / plugin_name / 'migrations' if migrations_dir.exists(): return migrations_dir return None def run_plugin_migrations( self, plugin_name: str, revision: str = 'head' ) -> bool: """ Run migrations for a plugin. Uses flask db upgrade with the plugin's migrations directory. """ # verify-at-migrate: never run a plugin's migration code (full DB rights) # from an unverified tree when the site enforces signing. if not self._verify_ok(plugin_name): return False migrations_dir = self.get_migrations_dir(plugin_name) if not migrations_dir: logger.info(f"No migrations directory for plugin {plugin_name}") return True # No migrations to run try: # Use alembic directly with plugin's migrations from alembic.config import Config from alembic import command from .packaging import no_bytecode config = Config() config.set_main_option('script_location', str(migrations_dir)) config.set_main_option('sqlalchemy.url', self.database_url) # Use plugin-specific version table config.set_main_option( 'version_table', f'alembic_version_{plugin_name}' ) with no_bytecode(): command.upgrade(config, revision) logger.info(f"Migrations completed for {plugin_name}") return True except ImportError: # Fallback to subprocess if alembic not available in context logger.warning("Using subprocess for migrations") return self._run_migrations_subprocess(plugin_name, revision) except Exception as e: logger.error(f"Migration failed for {plugin_name}: {e}") return False def _run_migrations_subprocess( self, plugin_name: str, revision: str = 'head' ) -> bool: """Run migrations via subprocess as fallback.""" if not self._verify_ok(plugin_name): return False migrations_dir = self.get_migrations_dir(plugin_name) if not migrations_dir: return True try: result = subprocess.run( [ sys.executable, '-B', '-m', 'alembic', '-c', str(migrations_dir / 'alembic.ini'), 'upgrade', revision ], capture_output=True, text=True, env={ **dict(__import__('os').environ), 'DATABASE_URL': self.database_url, 'PYTHONDONTWRITEBYTECODE': '1', } ) if result.returncode != 0: logger.error(f"Migration error: {result.stderr}") return False return True except Exception as e: logger.error(f"Migration subprocess failed: {e}") return False def downgrade_plugin( self, plugin_name: str, revision: str = 'base' ) -> bool: """ Downgrade/rollback plugin migrations. """ # Downgrade runs the same env.py + version module code as upgrade, so it # gets the same fail-closed verification (finding #3: this path had none). if not self._verify_ok(plugin_name): return False migrations_dir = self.get_migrations_dir(plugin_name) if not migrations_dir: return True try: from alembic.config import Config from alembic import command from .packaging import no_bytecode config = Config() config.set_main_option('script_location', str(migrations_dir)) config.set_main_option('sqlalchemy.url', self.database_url) config.set_main_option( 'version_table', f'alembic_version_{plugin_name}' ) with no_bytecode(): command.downgrade(config, revision) logger.info(f"Downgrade completed for {plugin_name}") return True except Exception as e: logger.error(f"Downgrade failed for {plugin_name}: {e}") return False def get_current_revision(self, plugin_name: str) -> Optional[str]: """Get current migration revision for a plugin. ScriptDirectory imports every versions/*.py to build the revision map, so this reads (executes) plugin code and must verify first (finding #3). """ if not self._verify_ok(plugin_name): return None migrations_dir = self.get_migrations_dir(plugin_name) if not migrations_dir: return None try: from alembic.config import Config from alembic.script import ScriptDirectory from .packaging import no_bytecode config = Config() config.set_main_option('script_location', str(migrations_dir)) with no_bytecode(): script = ScriptDirectory.from_config(config) return script.get_current_head() except Exception: return None def has_pending_migrations(self, plugin_name: str) -> bool: """Check if plugin has any migration scripts on disk. File-level check only (no DB): True when the plugin ships a migrations/versions dir with at least one script. Used by upgrade-all to decide whether the plugin participates in the per-plugin Alembic flow at all. """ migrations_dir = self.get_migrations_dir(plugin_name) if not migrations_dir: return False versions_dir = migrations_dir / 'versions' if not versions_dir.exists(): return False # Any migration files on disk? migration_files = list(versions_dir.glob('*.py')) return len(migration_files) > 0 def get_applied_revision(self, plugin_name: str) -> Optional[str]: """Return the revision stamped in alembic_version_ in the DB, or None if the plugin chain has never been stamped (or on any error).""" migrations_dir = self.get_migrations_dir(plugin_name) if not migrations_dir or not self.database_url: return None try: from sqlalchemy import create_engine from alembic.runtime.migration import MigrationContext engine = create_engine(self.database_url) with engine.connect() as connection: context = MigrationContext.configure( connection, opts={'version_table': f'alembic_version_{plugin_name}'}, ) return context.get_current_revision() except Exception: return None def has_unapplied_migrations(self, plugin_name: str) -> bool: """True when the plugin's on-disk chain head is ahead of what the DB has stamped. Compares the script head against alembic_version_. Best-effort: returns False when it cannot tell (no chain / no DB).""" head = self.get_current_revision(plugin_name) # script head on disk if not head: return False return self.get_applied_revision(plugin_name) != head