254 lines
8.8 KiB
Python
254 lines
8.8 KiB
Python
"""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_<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
|