shopdb-flask/shopdb/plugins/alembic_template.py

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

Every bundled plugin that owns tables (backups, computers, endpoints, guides,
knowledgebase, machines, network, notifications, printers, slides, tools,
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).

ORPHAN_TABLE_OWNERS below is the other half of that map, and it exists because
the core chain creates tables that no bundled plugin owns any more. Those rows
are not a plugin's, so they must not be in PLUGIN_TABLE_OWNERS - nothing can
import their models and no chain owns them - but they still land on every site,
and `flask plugin prune-schema` only knows how to drop a table it can attribute
to an owner. Drop the entry and the table becomes undroppable.
"""
from __future__ import annotations

import importlib
import logging
import os
from typing import Iterable

from alembic import context
from sqlalchemy import MetaData, pool

# Registers the compiler hook that forces utf8mb4 + DYNAMIC on MySQL. Imported
# for the side effect. It used to live inline in migrations/env.py, so it applied
# to the CORE chain only: a plugin's baseline tables were created at the server's
# default charset while core's were utf8mb4, on the same database. On a server
# defaulting to latin1 that difference is invisible until a join between the two
# stops using an index, or a character comes back mangled.
import shopdb.utils.mysql_charset  # noqa: F401,E402

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

# Seconds a migration will wait for a metadata lock before failing. Long enough
# to ride out a normal request, far short of MySQL's one-year default.
MIGRATION_LOCK_WAIT_SECONDS = 20

# 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]] = {
    'backups':       ('backuprevisions',),
    'computers':     ('computertypes', 'computers', 'computerinstalledapps',
                      'accessprotocols', 'computeraccess',
                      'configurationpolicies', 'configurationpolicyactions',
                      'businesspolicies', 'businesspolicyitems',
                      'policyactionholds'),
    'endpoints':     ('endpointstates', 'endpointevents', 'endpointcommands',
                      'endpointapplications', 'endpointmonitors',
                      'operatorreports', 'endpointhealthstates',
                      'endpointhealthtransitions', 'endpointmonitorimpacts',
                      'endpointlocks', 'endpointlogcollections',
                      'endpointcredentials', 'endpointenrollrequests',
                      'endpointmigrationwaves', 'endpointbootstraptokens'),
    'guides':        ('guides', 'guiderevisions', 'guideimages'),
    'knowledgebase': ('knowledgebase',),
    'machines':      ('machinetypes', 'machines'),
    'network':       ('networkdevicetypes', 'networkdevices', 'vlans', 'subnets'),
    'notifications': ('notificationtypes', 'notifications', 'notificationacknowledgements'),
    'printers':      ('printertypes', 'printers', 'modelsupplies', 'printerdrivers',
                      'printersupplyalerts', 'printerobservedqueues'),
    'slides':        ('tvslides',),
    # Tech Tools was schemaless until secure send: every other tool in it
    # runs in the browser and owns nothing.
    'tools':         ('securesendsecrets',),
    'warranty':      ('warranties', 'warrantyassets'),
}


# Tables the CORE Alembic chain creates that NO bundled plugin owns.
#
# They are not dead weight by accident: the core baseline (68b3947ae14f) creates
# usbdevicetypes/usbdevices/usbcheckouts and 7d16_directoryemployees creates
# directoryemployees, and neither migration can be removed - the whole plugin
# chain set anchors on that head, so editing it breaks `flask db upgrade` on a
# fresh database. The tables therefore arrive at every site whether or not
# anything reads them.
#
# Keyed by the name of the plugin that used to own them, because that is the
# name an operator recognises in prune-schema output. A key here is NOT a plugin
# this build ships, and nothing imports models for it.
ORPHAN_TABLE_OWNERS: dict[str, Iterable[str]] = {
    'employees': ('directoryemployees',),
    'usb':       ('usbdevicetypes', 'usbdevices', 'usbcheckouts'),
}


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

    NOT actually called from anywhere: every post-cutover baseline emits
    explicit Alembic ops instead, because this helper's per-plugin MetaData
    cannot resolve a foreign key to a CORE table at CreateTable-compile time.
    Kept because it is still the reference for what idempotent means here, and
    because create_table_if_missing below is the same rule in a form the
    baselines can use. See that function for what went wrong without it.
    """
    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 create_table_if_missing(table_name: str, *columns, **kwargs) -> bool:
    """op.create_table, skipped when the table is already there.

    WHY EVERY POST-CUTOVER BASELINE NEEDS THIS. A baseline is the sole creator
    of its plugin's tables, so it was written as a plain op.create_table - which
    is correct only if the table's existence and the plugin's Alembic stamp can
    never disagree. They can, and routinely do:

      * `flask db-utils create-all` runs db.create_all(), which builds EVERY
        table registered on the metadata - plugin tables included - and stamps
        NOTHING. `flask plugin upgrade-all` then runs the baseline against
        tables that already exist.
      * A run that creates the table and then fails before its stamp commits
        leaves the same split, and re-running is the obvious thing to try.

    Reported as `flask plugin upgrade-all` failing the tools plugin with
    "table securesendsecrets already exists". Nothing caught it because the
    migration test drops the post-cutover tables db.create_all() made before
    running upgrade-all, which arranges the one state where an unguarded create
    works and never the state a real database is in.

    Returns True when it created the table, so a caller can skip the indexes
    that belong to it.
    """
    from alembic import op
    from sqlalchemy import inspect

    if inspect(op.get_bind()).has_table(table_name):
        logger.info('%s already exists, skipping create', table_name)
        return False
    op.create_table(table_name, *columns, **kwargs)
    return True


def create_index_if_missing(index_name: str, table_name: str, columns, **kwargs):
    """op.create_index, skipped when an equivalent index already exists.

    Matched on the COLUMN LIST rather than the name, because the two creators
    do not agree on names. db.create_all() takes them from SQLAlchemy's default
    convention (`ix_securesendsecrets_token`) while a baseline spells them the
    way this repo's naming rules require (`ixsecuresendsecretstoken`). Matching
    on the name would call the index missing and add a second one over the same
    column - MySQL permits that, with only a warning, so the duplicate would
    accumulate silently on every affected site.
    """
    from alembic import op
    from sqlalchemy import inspect

    wanted = [str(column) for column in columns]
    inspector = inspect(op.get_bind())
    if not inspector.has_table(table_name):
        op.create_index(index_name, table_name, columns, **kwargs)
        return
    for existing in inspector.get_indexes(table_name):
        if list(existing.get('column_names') or []) == wanted:
            logger.info('index over %s(%s) already exists as %s, skipping',
                        table_name, ', '.join(wanted), existing.get('name'))
            return
    op.create_index(index_name, table_name, columns, **kwargs)


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

        # Bound how long a migration waits 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 the DDL
        # waits behind any open transaction effectively forever. Worse, MySQL
        # grants metadata locks in request order: every later reader queues
        # behind the pending exclusive, so one stuck migration stops the whole
        # site answering. Failing after a bounded wait turns a site-wide outage
        # into one failed enable the operator can retry.
        #
        # MySQL only - SQLite has no such statement and would reject it.
        connect_args = {}
        if db_url.startswith('mysql'):
            connect_args['init_command'] = (
                f'SET SESSION lock_wait_timeout = {MIGRATION_LOCK_WAIT_SECONDS}')
        connectable = create_engine(db_url, poolclass=pool.NullPool,
                                    connect_args=connect_args)
        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()
            # SAME FIX AS migrations/env.py. MySQL DDL commits implicitly, which
            # flushes everything queued before it - including the PREVIOUS
            # migration's version stamp. The LAST migration of a run has no DDL
            # after it, so its stamp is rolled back when the connection closes:
            # the schema change survives, alembic_version stays one revision
            # behind, and `flask plugin upgrade-all` exits 0 having silently
            # re-run that migration - again on every deploy after.
            #
            # Core got this fix; the per-plugin chains run through THIS file and
            # did not, so every plugin head has been re-running on MySQL. That is
            # only invisible while every head happens to be idempotent.
            connection.commit()