shopdb-flask/shopdb/plugins/migrations.py

"""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)
            # The lock-wait cap lives with the engine in alembic_template, which
            # builds its own; see MIGRATION_LOCK_WAIT_SECONDS there.
            _unused = """Bound how long this migration will wait for a metadata lock.
            #
            # THIS connection is the one that blocks. A plugin table with a
            # foreign key to a core table (users, assets) needs an EXCLUSIVE
            # metadata lock on it, and MySQL's lock_wait_timeout defaults to a
            # YEAR - so a DDL behind any long-running transaction waits
            # effectively forever. Worse, MySQL grants metadata locks in request
            # order, so every later reader queues behind the pending exclusive
            # and the whole site stops answering, which is what took this
            # instance down twice. Failing the migration after a bounded wait
            # turns a site-wide outage into one failed enable the operator can
            # retry.
            """  # noqa: F841

            # 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_<plugin> 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_<plugin>.
        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