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',
                      'manifestblobs',
                      'manifestenforcementreports', 'manifestenforcementresults',
                      'pctypealiases'),
    'knowledgebase': ('knowledgebase',),
    'machines':      ('machinetypes', 'machines'),
    'measuringtools': ('measuringtooltypes', 'measuringtools'),
    'network':       ('networkdevicetypes', 'networkdevices', 'vlans', 'subnets'),
    'notifications': ('notificationtypes', 'notifications'),
    'printedparts':  ('printeditems', 'printeditemtransactions',
                      'printeditemfiles'),
    'printers':      ('printertypes', 'printers', 'modelsupplies', 'printerdrivers',
                      'printersupplyalerts'),
    '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):
    """Create every table this plugin owns, sourced from the SQLAlchemy models
    (not duplicated DDL). IDEMPOTENT: a table that already exists is skipped, so
    this is safe on an existing database that has the table from the pre-cutover
    core baseline as well as on a fresh install (ADR-014 Phase 2).

    Called from each plugin's 0001 baseline.py upgrade().
    """
    from alembic import op
    from sqlalchemy import inspect
    from sqlalchemy.schema import CreateTable

    md = _get_plugin_metadata(plugin_name)
    bind = op.get_bind()
    existing = set(inspect(bind).get_table_names())
    # Sort by FK dependency so parent tables are created first.
    for table in md.sorted_tables:
        if table.name in existing:
            continue
        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()