1"""Decide what to do with a baseline migration: run it, record it, or refuse.
  2
  3A baseline stands in for a package's deleted history (`Migration.supersedes`).
  4The loader computes this once, next to the records it observed, and stores it
  5as `loader.baseline_status`; the planner, preflight, and plain-dev all read
  6that. Nothing here writes.
  7"""
  8
  9from dataclasses import dataclass, field
 10from typing import TYPE_CHECKING, Any
 11
 12from plain.postgres.options import default_db_table
 13
 14from .exceptions import (
 15    MigrationHistoryError,
 16    ResetBoundaryError,
 17    StaleMigrationRecordsError,
 18    UnrecordedTablesError,
 19)
 20
 21if TYPE_CHECKING:
 22    from plain.postgres.connection import DatabaseConnection
 23
 24    from .loader import MigrationLoader
 25    from .migration import Migration
 26
 27
 28@dataclass
 29class BaselineStatus:
 30    # Baselines to record without running, before anything else runs.
 31    adopt: list[Migration] = field(default_factory=list)
 32    # Why the planner must stop, per package. `str(error)` is the operator message.
 33    refusals: list[MigrationHistoryError] = field(default_factory=list)
 34
 35    @property
 36    def adopt_keys(self) -> set[tuple[str, str]]:
 37        return {(b.package_label, b.name) for b in self.adopt}
 38
 39
 40def classify_baselines(
 41    loader: MigrationLoader,
 42    applied: dict[tuple[str, str], Any],
 43    connection: DatabaseConnection,
 44) -> BaselineStatus:
 45    """
 46    For each package that ships a baseline:
 47
 48    - baseline recorded -> nothing to do
 49    - no records for the package, no tables -> the planner runs it
 50    - no records for the package, tables present -> refuse: fake it or drop them
 51    - the superseded migration recorded -> adopt (if the tables are there)
 52    - records but not the superseded migration -> refuse, too old
 53    - records but tables missing -> refuse, stale records
 54    """
 55    status = BaselineStatus()
 56    if not loader.baselines:
 57        return status
 58
 59    recorded_packages = {label for label, _name in applied}
 60    existing_tables: set[str] | None = None
 61
 62    def tables_present(tables: list[str]) -> list[str]:
 63        nonlocal existing_tables
 64        if existing_tables is None:
 65            existing_tables = set(connection.table_names())
 66        return [table for table in tables if table in existing_tables]
 67
 68    for package_label, baseline in loader.baselines.items():
 69        key = (package_label, baseline.name)
 70        if key in applied:
 71            continue
 72        if package_label not in recorded_packages:
 73            present = tables_present(applied_tables(loader, applied, baseline))
 74            if present:
 75                status.refusals.append(
 76                    UnrecordedTablesError(package_label, baseline.name, present)
 77                )
 78            continue
 79        assert baseline.supersedes is not None
 80        if (package_label, baseline.supersedes) not in applied:
 81            status.refusals.append(
 82                ResetBoundaryError(
 83                    package_label, baseline.supersedes, baseline.shipped_in
 84                )
 85            )
 86            continue
 87
 88        tables = applied_tables(loader, applied, baseline)
 89        missing = [t for t in tables if t not in tables_present(tables)]
 90        if missing:
 91            status.refusals.append(
 92                StaleMigrationRecordsError(package_label, missing, tables)
 93            )
 94            continue
 95
 96        status.adopt.append(baseline)
 97
 98    return status
 99
100
101def applied_tables(
102    loader: MigrationLoader, applied: dict[tuple[str, str], Any], baseline: Migration
103) -> list[str]:
104    """The package's tables as this database should have them right now.
105
106    Replays the baseline plus whatever recorded migrations follow it, so a
107    kept migration that renamed a table counts by its new name, and a pending
108    one that hasn't run yet doesn't count at all.
109    """
110    package_label = baseline.package_label
111    nodes = [(package_label, baseline.name)] + [
112        key for key in applied if key[0] == package_label and key in loader.graph.nodes
113    ]
114    state = loader.graph.make_state(
115        nodes=nodes, real_packages=loader.unmigrated_packages
116    )
117    return [
118        model_state.options.get("db_table")
119        or default_db_table(package_label, model_state.name)
120        for (label, _name), model_state in state.models.items()
121        if label == package_label
122    ]