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 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 DDL (full DB rights) from an
        # unverified tree when the site enforces signing.
        if self.verifier is not None:
            ok, reason = self.verifier.check(plugin_name)
            if not ok:
                logger.error(
                    "Refusing migrations for %s: signature verification failed "
                    "(%s)", plugin_name, reason)
                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

            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}'
            )

            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."""
        migrations_dir = self.get_migrations_dir(plugin_name)
        if not migrations_dir:
            return True

        try:
            result = subprocess.run(
                [
                    sys.executable, '-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
                }
            )

            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.
        """
        migrations_dir = self.get_migrations_dir(plugin_name)

        if not migrations_dir:
            return True

        try:
            from alembic.config import Config
            from alembic import command

            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}'
            )

            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."""
        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

            config = Config()
            config.set_main_option('script_location', str(migrations_dir))

            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