shopdb-flask/shopdb/plugins/alembic_template.py

"""Shared Alembic env.py logic for bundled plugins.

Every bundled plugin that owns tables (computers, employees, knowledgebase,
machines, network, notifications, printers, slides, usb, warranty) has a
`migrations/env.py` that does the minimum:

    import os
    os.environ['PLUGIN_NAME'] = 'computers'
    from shopdb.plugins.alembic_template import run_migrations
    run_migrations()

This module wires the plugin's models into a MetaData object filtered to
only the tables that belong to that plugin, then runs Alembic in either
offline or online mode against the Flask app's configured engine.

Ownership cutover (see ADR-008): the core Alembic chain created every table
that exists through its head (`7d16_directoryemployees`), including the plugin
tables. Each plugin's `0001` migration is therefore a stamp-only no-op that
just records the anchor revision in `alembic_version_<plugin>`. NEW plugin
schema changes land as `plugins/<name>/migrations/000N` from here on, never in
the core chain.

Plugin tables must be importable via `plugins.<name>.models`. Plugins
register their `__tablename__` set in PLUGIN_TABLE_OWNERS below so the
filter is explicit (avoids depending on import-side-effect global state).
"""
from __future__ import annotations

import importlib
import logging
import os
from typing import Iterable

from alembic import context
from sqlalchemy import MetaData, pool

logger = logging.getLogger('alembic.env.plugin')

# Explicit table-ownership map. Adding tables to a plugin requires updating
# this dict so the per-plugin migration knows which tables to include.
PLUGIN_TABLE_OWNERS: dict[str, Iterable[str]] = {
    'computers':     ('computertypes', 'computers', 'computerinstalledapps',
                      'accessprotocols', 'computeraccess'),
    'employees':     ('directoryemployees',),
    'geenforce':     ('manifestscopes', 'manifestentries', 'manifestentrypctypes',
                      'manifestentryhostnames', 'manifestentrymachinenumbers',
                      'manifestinusechecks', 'manifestinusecheckprocesses',
                      'manifestpublishedversions', 'manifestpayloads',
                      'manifestenforcementreports', 'manifestenforcementresults',
                      'pctypealiases'),
    'knowledgebase': ('knowledgebase',),
    'machines':      ('machinetypes', 'machines'),
    'measuringtools': ('measuringtooltypes', 'measuringtools'),
    'network':       ('networkdevicetypes', 'networkdevices', 'vlans', 'subnets'),
    'notifications': ('notificationtypes', 'notifications'),
    'printers':      ('printertypes', 'printers', 'modelsupplies', 'printerdrivers'),
    'slides':        ('tvslides',),
    'usb':           ('usbdevicetypes', 'usbdevices', 'usbcheckouts'),
    'warranty':      ('warranties', 'warrantyassets'),
}


def _get_plugin_metadata(plugin_name: str) -> MetaData:
    """Import the plugin's models and return a MetaData containing only its
    declared tables (filtered via PLUGIN_TABLE_OWNERS)."""
    owned = set(PLUGIN_TABLE_OWNERS.get(plugin_name, ()))
    if not owned:
        raise RuntimeError(
            f"PLUGIN_TABLE_OWNERS has no entry for plugin '{plugin_name}'. "
            f"Update shopdb/plugins/alembic_template.py."
        )

    # Importing models attaches them to the global db.metadata.
    importlib.import_module(f'plugins.{plugin_name}.models')
    from shopdb.extensions import db
    full = db.metadata

    plugin_md = MetaData()
    for table in list(full.tables.values()):
        if table.name in owned:
            table.to_metadata(plugin_md)
    return plugin_md


def create_plugin_tables(plugin_name: str):
    """Emit CreateTable DDL for every table this plugin owns. Idempotent
    via Alembic's batch_op.create_table behavior (raises if exists; the
    baseline migration is meant to run against an empty schema).

    Called from each plugin's 0001_baseline.py upgrade() so the table
    definitions stay sourced from the SQLAlchemy models rather than being
    duplicated in handwritten Alembic ops.
    """
    from alembic import op
    from sqlalchemy.schema import CreateTable

    md = _get_plugin_metadata(plugin_name)
    bind = op.get_bind()
    # Sort by FK dependency so parent tables are created first.
    for table in md.sorted_tables:
        op.execute(str(CreateTable(table).compile(dialect=bind.dialect)))


def drop_plugin_tables(plugin_name: str):
    """Mirror of create_plugin_tables for downgrade(). Drops in reverse FK
    order."""
    from alembic import op

    md = _get_plugin_metadata(plugin_name)
    for table in reversed(md.sorted_tables):
        op.execute(f'DROP TABLE IF EXISTS "{table.name}"')


def run_migrations():
    """Entry point called by each plugin's migrations/env.py."""
    plugin_name = os.environ.get('PLUGIN_NAME')
    if not plugin_name:
        raise RuntimeError("PLUGIN_NAME env var must be set before run_migrations()")

    config = context.config
    target_metadata = _get_plugin_metadata(plugin_name)

    # Per-plugin version table so each plugin's chain is independent of core
    # Alembic's alembic_version table.
    version_table = f'alembic_version_{plugin_name}'

    db_url = config.get_main_option('sqlalchemy.url')
    if not db_url:
        # Pull from the Flask app config if running inside an app context
        # (e.g. via flask plugin migrate <name>).
        try:
            from flask import current_app
            db_url = current_app.config['SQLALCHEMY_DATABASE_URI']
            config.set_main_option('sqlalchemy.url', db_url.replace('%', '%%'))
        except Exception as ex:
            raise RuntimeError(
                "sqlalchemy.url not set and no Flask app context available. "
                f"Original error: {ex}"
            )

    if context.is_offline_mode():
        context.configure(
            url=db_url,
            target_metadata=target_metadata,
            literal_binds=True,
            version_table=version_table,
            include_schemas=False,
        )
        with context.begin_transaction():
            context.run_migrations()
    else:
        # Build the engine straight from the resolved URL. The plugin manager
        # drives this via a programmatic alembic Config (no ini file), so
        # config.get_section returns an empty dict and engine_from_config would
        # find no sqlalchemy.url. db_url is already resolved above.
        from sqlalchemy import create_engine
        connectable = create_engine(db_url, poolclass=pool.NullPool)
        with connectable.connect() as connection:
            context.configure(
                connection=connection,
                target_metadata=target_metadata,
                version_table=version_table,
                include_schemas=False,
            )
            with context.begin_transaction():
                context.run_migrations()