from __future__ import annotations from collections.abc import Mapping from dataclasses import dataclass import json import os from pathlib import Path from typing import Any from alembic import command from alembic.config import Config from alembic.runtime.migration import MigrationContext from alembic.script import ScriptDirectory from sqlalchemy import create_engine, inspect, text from govoplan_core.core.migrations import MigrationMetadataPlan, migration_metadata_plan from govoplan_core.core import change_sequence as core_change_sequence_models # noqa: F401 - populate core metadata from govoplan_core.core.change_sequence import ChangeSequenceEntry, ChangeSequenceRetentionFloor from govoplan_core.core.module_management import load_startup_enabled_modules, startup_candidate_module_ids from govoplan_core.core.modules import ModuleMigrationTaskContext, ModuleMigrationTaskResult from govoplan_core.db.session import configure_database, get_database from govoplan_core.server.default_config import get_server_config from govoplan_core.server.config import ManifestFactory from govoplan_core.server.registry import available_module_manifests, build_platform_registry from govoplan_core.settings import settings # Historic development databases could be created partly through Alembic and # partly through Base.metadata.create_all(). In that state Alembic still says # "2c..." while the 3d/4e file-storage tables already exist, so a normal # upgrade attempts to create file_blobs again. This reconciliation is kept # deliberately narrow and only advances the marker when the complete expected # schema for the skipped revisions is already present. REVISION_AUTH_RBAC = "2c3d4e5f6a7b" REVISION_FILE_STORAGE = "3d4e5f6a7b8c" REVISION_FILE_FOLDERS = "4e5f6a7b8c9d" REVISION_NAMESPACE_PLATFORM_TABLES = "2e3f4a5b6c7d" REVISION_CORE_CHANGE_SEQUENCE = "3f4a5b6c7d8e" REVISION_HIERARCHICAL_SETTINGS = "f5a6b7c8d9e0" LEGACY_SCOPE_TABLE = "tenancy_tenants" CORE_SCOPE_TABLE = "core_scopes" PRE_MIGRATION_TASK_PHASES = ("pre_migration_check", "pre_migration_prepare") POST_MIGRATION_TASK_PHASES = ("post_migration_backfill", "post_migration_verify") MIGRATION_TASK_PHASES = (*PRE_MIGRATION_TASK_PHASES, *POST_MIGRATION_TASK_PHASES) _NAMESPACE_TABLE_RENAMES = ( ("tenants", "tenancy_tenants"), ("accounts", "access_accounts"), ("users", "access_users"), ("groups", "access_groups"), ("roles", "access_roles"), ("system_role_assignments", "access_system_role_assignments"), ("user_group_memberships", "access_user_group_memberships"), ("user_role_assignments", "access_user_role_assignments"), ("group_role_assignments", "access_group_role_assignments"), ("api_keys", "access_api_keys"), ("auth_sessions", "access_auth_sessions"), ("system_settings", "core_system_settings"), ("governance_templates", "admin_governance_templates"), ("governance_template_assignments", "admin_governance_template_assignments"), ) _FILE_STORAGE_TABLES = { "file_blobs", "file_assets", "file_versions", "file_shares", "campaign_attachment_uses", } _FILE_FOLDER_TABLES = {"file_folders"} _FILE_STORAGE_COLUMNS = { "file_blobs": { "id", "tenant_id", "storage_backend", "storage_key", "checksum_sha256", "size_bytes", }, "file_assets": { "id", "tenant_id", "owner_type", "display_path", "filename", "current_version_id", }, "file_versions": { "id", "file_asset_id", "blob_id", "version_number", "checksum_sha256", }, "file_shares": {"id", "file_asset_id", "target_type", "target_id", "permission"}, "campaign_attachment_uses": { "id", "campaign_id", "campaign_version_id", "file_asset_id", "file_version_id", "file_blob_id", }, "file_folders": {"id", "tenant_id", "owner_type", "path"}, } _CREATE_ALL_THROUGH_HIERARCHICAL_TABLES = { "access_accounts", "access_auth_sessions", "audit_log", "campaign_versions", "campaign_jobs", "send_attempts", "access_system_role_assignments", "core_system_settings", "admin_governance_templates", "admin_governance_template_assignments", "mail_server_profiles", } _CREATE_ALL_THROUGH_HIERARCHICAL_COLUMNS = { "access_auth_sessions": {"account_id", "csrf_token_hash"}, "audit_log": {"scope", "tenant_id", "user_id", "api_key_id"}, "campaign_versions": { "workflow_state", "current_flow", "user_lock_state", "user_locked_at", "user_locked_by_user_id", "execution_snapshot", "execution_snapshot_hash", "execution_snapshot_at", }, "campaign_jobs": {"claimed_at", "claim_token", "smtp_started_at", "outcome_unknown_at", "eml_sha256"}, "send_attempts": {"status", "claim_token"}, "access_users": {"account_id", "settings", "mail_profile_policy"}, "access_groups": {"system_template_id", "system_required", "settings", "mail_profile_policy"}, "access_roles": {"system_template_id", "system_required"}, "campaigns": {"settings", "mail_profile_policy"}, "mail_server_profiles": {"scope_type", "scope_id"}, "core_system_settings": {"settings", "allow_tenant_custom_groups", "allow_tenant_custom_roles", "allow_tenant_api_keys"}, } _SCOPE_TABLE_COLUMNS = {"settings", "allow_custom_groups", "allow_custom_roles", "allow_api_keys"} _CREATE_ALL_THROUGH_HIERARCHICAL_INDEXES = { "file_folders": {"uq_file_folders_active_user_path", "uq_file_folders_active_group_path"}, "campaign_versions": { "ix_campaign_versions_user_lock_state", "ix_campaign_versions_user_locked_by_user_id", "ix_campaign_versions_execution_snapshot_hash", }, "campaign_jobs": {"ix_campaign_jobs_claim_token", "ix_campaign_jobs_eml_sha256"}, "send_attempts": {"ix_send_attempts_status", "ix_send_attempts_claim_token"}, "audit_log": {"ix_audit_log_scope_created_at", "ix_audit_log_tenant_scope_created_at"}, "access_auth_sessions": {"ix_access_auth_sessions_account_id"}, "access_users": {"ix_access_users_account_id"}, "access_groups": {"ix_access_groups_system_template_id"}, "access_roles": {"ix_access_roles_system_template_id"}, "mail_server_profiles": {"ix_mail_server_profiles_scope_type", "ix_mail_server_profiles_scope_id", "ix_mail_server_profiles_scope"}, } @dataclass(frozen=True, slots=True) class MigrationResult: previous_revision: str | None reconciled_revision: str | None current_revision: str | None class ModuleMigrationTaskExecutionError(RuntimeError): def __init__(self, message: str, *, records: tuple[dict[str, object], ...]) -> None: super().__init__(message) self.records = records def registered_module_migration_plan( database_url: str | None = None, *, enabled_modules: tuple[str, ...] | list[str] | None = None, manifest_factories: tuple[ManifestFactory, ...] = (), ) -> MigrationMetadataPlan: return migration_metadata_plan(_registered_module_registry( database_url=database_url, enabled_modules=enabled_modules, manifest_factories=manifest_factories, )) def _registered_module_registry( *, database_url: str | None = None, enabled_modules: tuple[str, ...] | list[str] | None = None, manifest_factories: tuple[ManifestFactory, ...] = (), ): server_config = get_server_config() active_database_url = database_url or settings.database_url if active_database_url: configure_database(active_database_url) active_manifest_factories = manifest_factories or tuple(server_config.manifest_factories) raw_enabled_modules = tuple(enabled_modules) if enabled_modules is not None else load_startup_enabled_modules(server_config.enabled_modules) candidate_modules = startup_candidate_module_ids(server_config.enabled_modules, raw_enabled_modules) available_modules = available_module_manifests( active_manifest_factories, enabled_modules=candidate_modules, ignore_load_errors=True, ) active_enabled_modules = ( tuple(enabled_modules) if enabled_modules is not None else load_startup_enabled_modules(server_config.enabled_modules, available=available_modules) ) registry = build_platform_registry( active_enabled_modules, manifest_factories=active_manifest_factories, ) return registry def run_registered_module_migration_tasks( *, database_url: str | None = None, enabled_modules: tuple[str, ...] | list[str] | None = None, migration_module_order: tuple[str, ...] | list[str] | None = None, phases: tuple[str, ...] | list[str] = MIGRATION_TASK_PHASES, dry_run: bool = False, manifest_factories: tuple[ManifestFactory, ...] = (), ) -> tuple[dict[str, object], ...]: active_phases = tuple(dict.fromkeys(str(item).strip() for item in phases if str(item).strip())) invalid_phases = tuple(phase for phase in active_phases if phase not in MIGRATION_TASK_PHASES) if invalid_phases: raise ValueError("Unsupported module migration task phase(s): " + ", ".join(invalid_phases)) url = database_url or settings.database_url registry = _registered_module_registry( database_url=url, enabled_modules=enabled_modules, manifest_factories=manifest_factories, ) manifests = {manifest.id: manifest for manifest in registry.manifests()} ordered_ids = tuple(dict.fromkeys([ *(str(item).strip() for item in (migration_module_order or ()) if str(item).strip()), *manifests.keys(), ])) records: list[dict[str, object]] = [] database = get_database() with database.SessionLocal() as session: for phase in active_phases: for module_id in ordered_ids: manifest = manifests.get(module_id) if manifest is None or manifest.migration_spec is None: continue for task in manifest.migration_spec.migration_tasks: if task.phase != phase: continue record: dict[str, object] = { "module_id": manifest.id, "task_id": task.task_id, "phase": task.phase, "summary": task.summary, "task_version": task.task_version, "safety": task.safety, "idempotent": task.idempotent, "dry_run": dry_run, } if task.timeout_seconds is not None: record["timeout_seconds"] = task.timeout_seconds if not task.idempotent: record.update({"status": "blocked", "message": "Task is not idempotent."}) records.append(record) raise ModuleMigrationTaskExecutionError( f"Module migration task {manifest.id}/{task.task_id} is not idempotent.", records=tuple(records), ) if dry_run: record.update({"status": "skipped", "message": "Dry run; executor was not called."}) records.append(record) continue if task.executor is None: record.update({"status": "blocked", "message": "Task has no executor."}) records.append(record) raise ModuleMigrationTaskExecutionError( f"Module migration task {manifest.id}/{task.task_id} has no executor.", records=tuple(records), ) context = ModuleMigrationTaskContext( module_id=manifest.id, task_id=task.task_id, phase=task.phase, database_url=url, target_version=manifest.version, session=session, dry_run=dry_run, metadata=task.metadata, ) try: result = task.executor(context) normalized = _normalize_migration_task_result(result) except Exception as exc: session.rollback() record.update({ "status": "blocked", "message": f"{type(exc).__name__}: {exc}", }) records.append(record) raise ModuleMigrationTaskExecutionError( f"Module migration task {manifest.id}/{task.task_id} failed.", records=tuple(records), ) from exc record.update({ "status": normalized.status, "message": normalized.message, "details": _jsonable_migration_task_details(normalized.details), }) records.append(record) if normalized.status == "blocked": session.rollback() raise ModuleMigrationTaskExecutionError( normalized.message or f"Module migration task {manifest.id}/{task.task_id} blocked migration.", records=tuple(records), ) session.commit() return tuple(records) def _normalize_migration_task_result(result: ModuleMigrationTaskResult | None) -> ModuleMigrationTaskResult: if result is None: return ModuleMigrationTaskResult() if not isinstance(result, ModuleMigrationTaskResult): raise TypeError("Module migration task executors must return ModuleMigrationTaskResult or None.") if result.status not in {"ok", "warning", "blocked", "skipped"}: raise ValueError(f"Unsupported module migration task status: {result.status!r}") return result def _jsonable_migration_task_details(value: Mapping[str, Any]) -> Mapping[str, Any] | str: try: json.dumps(value) except (TypeError, ValueError): return str(dict(value)) return value def _repo_root() -> Path: packaged_root = Path(__file__).resolve().parents[3] configured = os.environ.get("GOVOPLAN_CORE_SOURCE_ROOT") candidates = [ Path(configured).expanduser() if configured else None, Path.cwd(), packaged_root, ] for candidate in candidates: if candidate is None: continue resolved = candidate.resolve() if (resolved / "alembic.ini").exists() and (resolved / "alembic").is_dir(): return resolved return packaged_root def alembic_config( *, database_url: str | None = None, enabled_modules: tuple[str, ...] | list[str] | None = None, manifest_factories: tuple[ManifestFactory, ...] = (), ) -> Config: root = _repo_root() config = Config(str(root / "alembic.ini")) config.set_main_option("script_location", str(root / "alembic")) active_database_url = database_url or settings.database_url plan = registered_module_migration_plan( active_database_url, enabled_modules=enabled_modules, manifest_factories=manifest_factories, ) version_locations = [str(root / "alembic" / "versions")] version_locations.extend(location for location in plan.script_locations if location) config.set_main_option("version_locations", os.pathsep.join(dict.fromkeys(version_locations))) config.set_main_option("version_path_separator", "os") config.attributes["database_url"] = active_database_url if enabled_modules is not None: config.attributes["enabled_modules"] = tuple(enabled_modules) if manifest_factories: config.attributes["manifest_factories"] = tuple(manifest_factories) return config def database_revision(database_url: str | None = None) -> str | None: url = database_url or settings.database_url engine = create_engine(url) try: with engine.connect() as connection: heads = MigrationContext.configure(connection).get_current_heads() if not heads: return None if len(heads) == 1: return heads[0] return ",".join(sorted(heads)) finally: engine.dispose() def _has_columns(inspector, table_name: str, required: set[str]) -> bool: try: actual = {column["name"] for column in inspector.get_columns(table_name)} except Exception: return False return required.issubset(actual) def _has_indexes(inspector, table_name: str, required: set[str]) -> bool: try: actual = {index["name"] for index in inspector.get_indexes(table_name)} except Exception: return False return required.issubset(actual) def _has_create_all_schema_through_hierarchical_settings(inspector, tables: set[str]) -> bool: if not _CREATE_ALL_THROUGH_HIERARCHICAL_TABLES.issubset(tables): return False if not any( table_name in tables and _has_columns(inspector, table_name, _SCOPE_TABLE_COLUMNS) for table_name in (CORE_SCOPE_TABLE, LEGACY_SCOPE_TABLE) ): return False for table_name, required_columns in _CREATE_ALL_THROUGH_HIERARCHICAL_COLUMNS.items(): if table_name not in tables or not _has_columns(inspector, table_name, required_columns): return False for table_name, required_indexes in _CREATE_ALL_THROUGH_HIERARCHICAL_INDEXES.items(): if table_name not in tables or not _has_indexes(inspector, table_name, required_indexes): return False return True def _backfill_user_lock_state_for_create_all_schema(database_url: str) -> None: engine = create_engine(database_url) try: with engine.begin() as connection: connection.execute(text(""" UPDATE campaign_versions SET user_lock_state = 'permanent', user_locked_at = published_at, user_locked_by_user_id = NULL WHERE published_at IS NOT NULL AND user_lock_state IS NULL """)) finally: engine.dispose() def _row_count(connection, table_name: str) -> int: quoted = connection.dialect.identifier_preparer.quote(table_name) return int(connection.execute(text(f"SELECT COUNT(*) FROM {quoted}")).scalar_one()) def _drop_table(connection, table_name: str) -> None: quoted = connection.dialect.identifier_preparer.quote(table_name) connection.execute(text(f"DROP TABLE {quoted}")) def _rename_table(connection, old_name: str, new_name: str) -> None: quoted_old = connection.dialect.identifier_preparer.quote(old_name) quoted_new = connection.dialect.identifier_preparer.quote(new_name) connection.execute(text(f"ALTER TABLE {quoted_old} RENAME TO {quoted_new}")) def _reconcile_scope_table_names(connection, tables: set[str]) -> bool: if LEGACY_SCOPE_TABLE not in tables: return False changed = False if CORE_SCOPE_TABLE in tables: if _row_count(connection, CORE_SCOPE_TABLE) == 0: _drop_table(connection, CORE_SCOPE_TABLE) tables.remove(CORE_SCOPE_TABLE) changed = True elif _row_count(connection, LEGACY_SCOPE_TABLE) == 0: _drop_table(connection, LEGACY_SCOPE_TABLE) tables.remove(LEGACY_SCOPE_TABLE) changed = True else: raise RuntimeError(f"Cannot reconcile non-empty {LEGACY_SCOPE_TABLE} over non-empty {CORE_SCOPE_TABLE}") if LEGACY_SCOPE_TABLE in tables and CORE_SCOPE_TABLE not in tables: _rename_table(connection, LEGACY_SCOPE_TABLE, CORE_SCOPE_TABLE) tables.remove(LEGACY_SCOPE_TABLE) tables.add(CORE_SCOPE_TABLE) changed = True return changed def reconcile_scope_table_retirement_drift(database_url: str | None = None) -> bool: """Repair databases that are stamped after the scope-table rename. The real Alembic migration performs the same rename. This helper covers development databases that were manually stamped or partially created via create_all(), where the version marker can be newer than the table name. It runs after Alembic upgrades so older migrations can still use the historical table name while they execute. """ url = database_url or settings.database_url changed = False engine = create_engine(url) try: with engine.begin() as connection: tables = set(inspect(connection).get_table_names()) changed = _reconcile_scope_table_names(connection, tables) finally: engine.dispose() return changed def reconcile_namespace_table_drift(database_url: str | None = None) -> bool: """Repair dev databases stamped past the namespace-table migration. During the repository split some development databases were stamped at the newer Alembic heads while still carrying the old platform table names. The real migration only renames tables, so replay that idempotent operation before startup bootstrap code touches the ORM. """ url = database_url or settings.database_url changed = False engine = create_engine(url) try: with engine.begin() as connection: schema = inspect(connection) tables = set(schema.get_table_names()) if not any(old_name in tables and new_name not in tables for old_name, new_name in _NAMESPACE_TABLE_RENAMES): return False old_tables_to_drop: set[str] = set() new_tables_to_drop: set[str] = set() for old_name, new_name in _NAMESPACE_TABLE_RENAMES: if old_name not in tables or new_name not in tables: continue if _row_count(connection, old_name) == 0: old_tables_to_drop.add(old_name) elif _row_count(connection, new_name) == 0: new_tables_to_drop.add(new_name) else: raise RuntimeError(f"Cannot reconcile non-empty {old_name} over non-empty {new_name}") for old_name, _new_name in reversed(_NAMESPACE_TABLE_RENAMES): if old_name in old_tables_to_drop: _drop_table(connection, old_name) tables.remove(old_name) changed = True for _old_name, new_name in reversed(_NAMESPACE_TABLE_RENAMES): if new_name in new_tables_to_drop: _drop_table(connection, new_name) tables.remove(new_name) changed = True for old_name, new_name in _NAMESPACE_TABLE_RENAMES: if old_name not in tables or new_name in tables: continue _rename_table(connection, old_name, new_name) tables.remove(old_name) tables.add(new_name) changed = True finally: engine.dispose() return changed def reconcile_change_sequence_retention_floor_drift(database_url: str | None = None) -> bool: """Repair databases stamped after the change-sequence migration changed. Early development databases may have applied the change-sequence revision before the retention-floor table was added to that migration. The main change-sequence table is enough for full snapshots, but incremental delta requests need the retention floor to decide whether a watermark is stale. """ url = database_url or settings.database_url changed = False engine = create_engine(url) try: with engine.begin() as connection: schema = inspect(connection) tables = set(schema.get_table_names()) if ChangeSequenceEntry.__tablename__ not in tables: return False if ChangeSequenceRetentionFloor.__tablename__ not in tables: ChangeSequenceRetentionFloor.__table__.create(bind=connection, checkfirst=True) changed = True else: indexes = { index["name"] for index in schema.get_indexes(ChangeSequenceRetentionFloor.__tablename__) } for index in ChangeSequenceRetentionFloor.__table__.indexes: if index.name not in indexes: index.create(bind=connection) changed = True finally: engine.dispose() return changed def reconcile_legacy_create_all_schema(database_url: str | None = None) -> str | None: """Repair the known Alembic/create_all drift without modifying application data. Returns the revision stamped during reconciliation, or ``None`` when no repair was necessary. A partial/unknown schema is intentionally left alone so Alembic can fail visibly instead of guessing. """ url = database_url or settings.database_url engine = create_engine(url) try: with engine.connect() as connection: heads = MigrationContext.configure(connection).get_current_heads() current = heads[0] if len(heads) == 1 else None has_no_revision = len(heads) == 0 schema = inspect(connection) tables = set(schema.get_table_names()) has_file_storage = _FILE_STORAGE_TABLES.issubset(tables) and all( _has_columns(schema, table, _FILE_STORAGE_COLUMNS[table]) for table in _FILE_STORAGE_TABLES ) has_file_folders = _FILE_FOLDER_TABLES.issubset(tables) and _has_columns( schema, "file_folders", _FILE_STORAGE_COLUMNS["file_folders"], ) has_create_all_hierarchical_schema = _has_create_all_schema_through_hierarchical_settings(schema, tables) finally: engine.dispose() target: str | None = None if current == REVISION_AUTH_RBAC and has_file_storage and has_file_folders: target = REVISION_FILE_FOLDERS elif current == REVISION_AUTH_RBAC and has_file_storage: target = REVISION_FILE_STORAGE elif current == REVISION_FILE_STORAGE and has_file_folders: target = REVISION_FILE_FOLDERS elif current == REVISION_FILE_FOLDERS and has_create_all_hierarchical_schema: # Development DBs may be stamped at 4e after earlier create_all # reconciliation even though later tables/columns were already created # by a newer model set. Skip only when that complete known schema is # present, then let newer migrations such as mail credential usernames # run normally. _backfill_user_lock_state_for_create_all_schema(url) target = REVISION_HIERARCHICAL_SETTINGS elif has_no_revision and has_create_all_hierarchical_schema: _backfill_user_lock_state_for_create_all_schema(url) target = REVISION_HIERARCHICAL_SETTINGS elif has_no_revision and has_file_storage and has_file_folders: # This is the other create_all-only development shape. The strict # column checks above ensure that we only stamp a complete known schema. target = REVISION_FILE_FOLDERS if target is None: return None command.stamp(alembic_config(database_url=url), target) return target def migrate_database( *, database_url: str | None = None, reconcile_legacy_schema: bool = True, enabled_modules: tuple[str, ...] | list[str] | None = None, migration_module_order: tuple[str, ...] | list[str] | None = None, manifest_factories: tuple[ManifestFactory, ...] = (), ) -> MigrationResult: url = database_url or settings.database_url if reconcile_legacy_schema: reconcile_namespace_table_drift(url) reconcile_change_sequence_retention_floor_drift(url) previous = database_revision(url) reconciled = reconcile_legacy_create_all_schema(url) if reconcile_legacy_schema else None config = alembic_config( database_url=url, enabled_modules=enabled_modules, manifest_factories=manifest_factories, ) if migration_module_order: plan = registered_module_migration_plan( url, enabled_modules=enabled_modules, manifest_factories=manifest_factories, ) _upgrade_ordered_module_heads(config, plan, tuple(migration_module_order)) command.upgrade(config, "heads") if reconcile_legacy_schema: reconcile_scope_table_retirement_drift(url) current = database_revision(url) return MigrationResult( previous_revision=previous, reconciled_revision=reconciled, current_revision=current, ) def _upgrade_ordered_module_heads( config: Config, plan: MigrationMetadataPlan, module_order: tuple[str, ...], ) -> None: module_heads = _module_heads_by_id(config, plan) upgraded: set[str] = set() for module_id in dict.fromkeys(module_order): for head in module_heads.get(module_id, ()): if head in upgraded: continue command.upgrade(config, head) upgraded.add(head) def _module_heads_by_id(config: Config, plan: MigrationMetadataPlan) -> dict[str, tuple[str, ...]]: script = ScriptDirectory.from_config(config) heads = script.get_heads() module_locations = { item.module_id: Path(item.script_location).resolve() for item in plan.modules if item.script_location } grouped: dict[str, list[str]] = {module_id: [] for module_id in module_locations} for head in heads: revision = script.get_revision(head) path_text = getattr(revision, "path", None) if revision is None or not path_text: continue revision_path = Path(path_text).resolve() for module_id, location in module_locations.items(): try: revision_path.relative_to(location) except ValueError: continue grouped[module_id].append(head) break return {module_id: tuple(sorted(values)) for module_id, values in grouped.items() if values}