from __future__ import annotations from dataclasses import dataclass import os from pathlib import Path from alembic import command from alembic.config import Config from alembic.runtime.migration import MigrationContext 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.db.session import configure_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" _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"}, "tenancy_tenants": {"settings", "allow_custom_groups", "allow_custom_roles", "allow_api_keys"}, "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"}, } _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 def registered_module_migration_plan( database_url: str | None = None, *, enabled_modules: tuple[str, ...] | list[str] | None = None, manifest_factories: tuple[ManifestFactory, ...] = (), ) -> MigrationMetadataPlan: 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 migration_metadata_plan(registry) def _repo_root() -> Path: return Path(__file__).resolve().parents[3] 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 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_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, 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 command.upgrade( alembic_config( database_url=url, enabled_modules=enabled_modules, manifest_factories=manifest_factories, ), "heads", ) current = database_revision(url) return MigrationResult( previous_revision=previous, reconciled_revision=reconciled, current_revision=current, )