"""Audit and atomically baseline the known unversioned platform schema. The command is read-only unless ``--apply`` is supplied together with the schema fingerprint printed by a preceding dry run. It never accepts a database URL on the command line, so credentials are not exposed through process arguments. """ import argparse from collections.abc import Mapping from dataclasses import dataclass import hashlib import json from pathlib import Path import re from typing import Any from alembic import command from alembic.autogenerate import compare_metadata from alembic.config import Config from alembic.migration import MigrationContext from alembic.script import ScriptDirectory from sqlalchemy import MetaData, Table, create_engine, inspect, select, text from sqlalchemy.engine import Connection, Engine from sqlalchemy.pool import NullPool from app.core.config import get_settings from app.core.database import Base from app.core.database.safety import validate_platform_migration_target from app.modules.ai_memory import models as ai_memory_models from app.modules.audit import models as audit_models from app.modules.business import models as business_models from app.modules.events import models as event_models from app.modules.feishu import models as feishu_models from app.modules.feishu_users import models as feishu_user_models from app.modules.observability import models as observability_models from app.modules.personalization import models as personalization_models from app.modules.subscriptions import models as subscription_models from app.modules.workflows import models as workflow_models PREVIOUS_REVISION = "202607260005" TARGET_REVISION = "202607270002" APPROVAL_TABLE = "approval_requests" SCHEMA_FINGERPRINT_PATTERN = re.compile(r"^[0-9a-f]{64}$") ADVISORY_LOCK_KEY = int.from_bytes(b"CAIPSCHE", byteorder="big", signed=True) MINIMUM_POSTGRESQL_VERSION_NUM = 150000 # Keep imports referenced so every model is registered with Base.metadata. _REGISTERED_MODEL_MODULES = ( ai_memory_models, audit_models, business_models, event_models, feishu_models, feishu_user_models, observability_models, personalization_models, subscription_models, workflow_models, ) _ADD_COLUMNS = { "ai_memory_entries": {"kind", "owner_id"}, "attendance_records": { "attendance_scope", "is_active", "last_seen_at", "source_location_status", "source_status", "source_updated_at", }, "audit_logs": {"request_id"}, "feishu_event_receipts": { "attempt_count", "auto_reply", "event_type", "last_error", "locked_by", "locked_until", "max_attempts", "next_attempt_at", "payload", "processed_at", "reply_attempt_count", "reply_last_error", "reply_locked_by", "reply_locked_until", "reply_next_attempt_at", "reply_payload", "reply_sent_at", "reply_status", "status", }, "market_watchlists": {"owner_id"}, "projects": { "department_code", "department_name", "display_code", "is_active", "last_seen_at", "owner_employee_code", "source_archived", "source_contract_amount", "source_created_at", "source_project_investment_amount", "source_stage", "source_stage_label", "source_updated_at", }, "risk_events": { "assigned_to", "closed_at", "closed_reason", "resolved_at", "review_summary", }, "work_reports": { "employee_code", "external_id", "is_active", "is_draft", "is_late", "last_seen_at", "source_updated_at", }, "work_tasks": { "employee_code", "external_id", "is_active", "last_seen_at", "source_created_at", "source_system", "source_updated_at", }, } _ADD_INDEXES = { "ai_memory_entries": { "ix_ai_memory_entries_fingerprint", "ix_ai_memory_entries_kind", "ix_ai_memory_entries_owner_id", }, "attendance_records": { "ix_attendance_records_attendance_scope", "ix_attendance_records_is_active", "ix_attendance_records_last_seen_at", "ix_attendance_records_source_status", "ix_attendance_records_source_updated_at", }, "audit_logs": {"ix_audit_logs_request_id"}, "feishu_admin_bootstrap_tombstones": { "ix_feishu_admin_bootstrap_tombstones_created_at", "ix_feishu_admin_bootstrap_tombstones_identity_hash", }, "feishu_event_receipts": { "ix_feishu_event_receipts_event_type", "ix_feishu_event_receipts_locked_by", "ix_feishu_event_receipts_locked_until", "ix_feishu_event_receipts_next_attempt_at", "ix_feishu_event_receipts_processed_at", "ix_feishu_event_receipts_reply_locked_by", "ix_feishu_event_receipts_reply_locked_until", "ix_feishu_event_receipts_reply_next_attempt_at", "ix_feishu_event_receipts_reply_sent_at", "ix_feishu_event_receipts_reply_status", "ix_feishu_event_receipts_status", }, "feishu_app_tickets": { "ix_feishu_app_tickets_app_id", "ix_feishu_app_tickets_received_at", }, "market_watchlists": {"ix_market_watchlists_owner_id"}, "projects": { "ix_projects_department_code", "ix_projects_department_name", "ix_projects_display_code", "ix_projects_is_active", "ix_projects_last_seen_at", "ix_projects_owner_employee_code", "ix_projects_source_archived", "ix_projects_source_stage", "ix_projects_source_updated_at", }, "risk_events": {"ix_risk_events_assigned_to"}, "work_reports": { "ix_work_reports_employee_code", "ix_work_reports_external_id", "ix_work_reports_is_active", "ix_work_reports_is_draft", "ix_work_reports_is_late", "ix_work_reports_last_seen_at", "ix_work_reports_source_updated_at", }, "work_tasks": { "ix_work_tasks_employee_code", "ix_work_tasks_external_id", "ix_work_tasks_is_active", "ix_work_tasks_last_seen_at", "ix_work_tasks_source_updated_at", }, } _REMOVE_INDEXES = { "ai_memory_entries": {"ix_ai_memory_entries_fingerprint"}, APPROVAL_TABLE: { "ix_approval_requests_action", "ix_approval_requests_applicant", "ix_approval_requests_approver", "ix_approval_requests_created_at", "ix_approval_requests_domain", "ix_approval_requests_record_id", "ix_approval_requests_status", "ix_approval_requests_ticket_id", }, } _MODIFY_DEFAULTS = { "attendance_records": {"attendance_scope", "is_active"}, "projects": {"is_active", "source_archived"}, "work_reports": {"is_active", "is_draft", "is_late"}, "work_tasks": {"is_active"}, } _IMPACTED_TABLES = frozenset( { *_ADD_COLUMNS, *_ADD_INDEXES, *_MODIFY_DEFAULTS, *_REMOVE_INDEXES, APPROVAL_TABLE, "feishu_users", } ) _CONTENT_FINGERPRINT_TABLES = _IMPACTED_TABLES - {APPROVAL_TABLE} class ReconciliationError(RuntimeError): """Raised when a baseline precondition is not satisfied.""" @dataclass(frozen=True, order=True) class DriftKey: operation: str table: str object_name: str def render(self) -> str: suffix = f":{self.object_name}" if self.object_name else "" return f"{self.operation}:{self.table}{suffix}" @dataclass(frozen=True) class BaselineAudit: dialect: str fingerprint: str impacted_table_row_counts: tuple[tuple[str, int], ...] impacted_table_content_digests: tuple[tuple[str, str], ...] postgresql_server_version_num: int | None postgresql_version_supported: bool | None revision_rows: tuple[str, ...] | None observed_drift: frozenset[DriftKey] unexpected_drift: frozenset[DriftKey] approval_rows: int | None approval_inbound_foreign_keys: int schema_privileges_ok: bool | None table_ownership_ok: bool | None @property def eligible(self) -> bool: return ( self.revision_rows == () and not self.unexpected_drift and (self.approval_rows is None or self.approval_rows == 0) and self.approval_inbound_foreign_keys == 0 and self.schema_privileges_ok is not False and self.table_ownership_ok is not False and self.postgresql_version_supported is not False ) def public_dict(self) -> dict[str, Any]: return { "eligible": self.eligible, "dialect": self.dialect, "schema_fingerprint_sha256": self.fingerprint, "impacted_table_row_counts": dict(self.impacted_table_row_counts), "postgresql_server_version_num": self.postgresql_server_version_num, "postgresql_version_supported": self.postgresql_version_supported, "alembic_version_state": ( "missing" if self.revision_rows is None else "empty" if not self.revision_rows else "versioned" ), "alembic_revision_count": ( None if self.revision_rows is None else len(self.revision_rows) ), "observed_drift_count": len(self.observed_drift), "unexpected_drift": [ item.render() for item in sorted(self.unexpected_drift) ], "approval_rows": self.approval_rows, "approval_inbound_foreign_keys": self.approval_inbound_foreign_keys, "schema_privileges_ok": self.schema_privileges_ok, "table_ownership_ok": self.table_ownership_ok, } def _allowed_drift() -> frozenset[DriftKey]: items = { DriftKey("add_table", "feishu_app_tickets", ""), DriftKey("add_table", "feishu_admin_bootstrap_tombstones", ""), DriftKey("remove_table", APPROVAL_TABLE, ""), DriftKey( "add_constraint", "ai_memory_entries", "uq_ai_memory_owner_fingerprint", ), DriftKey( "add_constraint", "market_watchlists", "uq_market_watchlist_owner_symbol", ), DriftKey( "remove_constraint", "ai_memory_entries", "uq_ai_memory_owner_fingerprint", ), DriftKey( "remove_constraint", "market_watchlists", "uq_market_watchlist_owner_symbol", ), DriftKey( "remove_constraint", "market_watchlists", "uq_market_watchlist_actor_symbol", ), DriftKey( "add_fk", "ai_memory_entries", "owner_id->feishu_users.id", ), DriftKey( "add_fk", "market_watchlists", "owner_id->feishu_users.id", ), DriftKey("modify_nullable", "work_tasks", "source_system"), } for table_name, columns in _ADD_COLUMNS.items(): items.update( DriftKey("add_column", table_name, column_name) for column_name in columns ) for table_name, indexes in _ADD_INDEXES.items(): items.update( DriftKey("add_index", table_name, index_name) for index_name in indexes ) for table_name, indexes in _REMOVE_INDEXES.items(): items.update( DriftKey("remove_index", table_name, index_name) for index_name in indexes ) for table_name, columns in _MODIFY_DEFAULTS.items(): items.update( DriftKey("modify_default", table_name, column_name) for column_name in columns ) return frozenset(items) ALLOWED_DRIFT = _allowed_drift() def _foreign_key_name(constraint: Any) -> str: local_columns = ",".join(column.name for column in constraint.columns) remote_columns = ",".join( element.target_fullname for element in constraint.elements ) return f"{local_columns}->{remote_columns}" def _normalize_diff(diff: tuple[Any, ...]) -> DriftKey: operation = str(diff[0]) if operation in {"add_table", "remove_table"}: return DriftKey(operation, str(diff[1].name), "") if operation in {"add_column", "remove_column"}: return DriftKey(operation, str(diff[2]), str(diff[3].name)) if operation in {"add_index", "remove_index"}: index = diff[1] return DriftKey(operation, str(index.table.name), str(index.name)) if operation in {"add_constraint", "remove_constraint"}: constraint = diff[1] return DriftKey( operation, str(constraint.table.name), str(constraint.name), ) if operation in {"add_fk", "remove_fk"}: constraint = diff[1] return DriftKey( operation, str(constraint.table.name), _foreign_key_name(constraint), ) if operation.startswith("modify_"): return DriftKey(operation, str(diff[2]), str(diff[3])) return DriftKey(operation, "", "") def _metadata_drift(connection: Connection) -> frozenset[DriftKey]: context = MigrationContext.configure( connection, opts={ "compare_type": True, "compare_server_default": True, }, ) raw_diffs = compare_metadata(context, Base.metadata) flattened: list[tuple[Any, ...]] = [] for item in raw_diffs: if isinstance(item, list): flattened.extend(item) else: flattened.append(item) return frozenset(_normalize_diff(diff) for diff in flattened) def _canonical_reflection_value(value: Any) -> Any: """Convert SQLAlchemy reflection values into stable JSON-compatible data.""" if value is None or isinstance(value, (bool, int, float, str)): return value if isinstance(value, bytes): return {"bytes_hex": value.hex()} if isinstance(value, Mapping): return { str(key): _canonical_reflection_value(value[key]) for key in sorted(value, key=lambda item: str(item)) } if isinstance(value, (list, tuple)): return [_canonical_reflection_value(item) for item in value] if isinstance(value, (set, frozenset)): items = [_canonical_reflection_value(item) for item in value] return sorted(items, key=_stable_json) return str(value) def _stable_json(value: Any) -> str: return json.dumps( value, sort_keys=True, separators=(",", ":"), ensure_ascii=True, ) def _sorted_reflection_items(items: list[dict[str, Any]]) -> list[dict[str, Any]]: normalized = [_canonical_reflection_value(item) for item in items] return sorted(normalized, key=_stable_json) def _schema_snapshot(connection: Connection) -> list[dict[str, Any]]: inspector = inspect(connection) snapshot: list[dict[str, Any]] = [] for table_name in sorted(inspector.get_table_names()): snapshot.append( { "table": table_name, "columns": [ _canonical_reflection_value(column) for column in inspector.get_columns(table_name) ], "pk": _canonical_reflection_value( inspector.get_pk_constraint(table_name) ), "uniques": _sorted_reflection_items( inspector.get_unique_constraints(table_name) ), "checks": _sorted_reflection_items( inspector.get_check_constraints(table_name) ), "fks": _sorted_reflection_items( inspector.get_foreign_keys(table_name) ), "indexes": _sorted_reflection_items( inspector.get_indexes(table_name) ), } ) return snapshot def _impacted_table_row_counts( connection: Connection, ) -> tuple[tuple[str, int], ...]: existing_tables = set(inspect(connection).get_table_names()) preparer = connection.dialect.identifier_preparer rows: list[tuple[str, int]] = [] for table_name in sorted(_IMPACTED_TABLES & existing_tables): quoted_table = preparer.quote_identifier(table_name) count = connection.execute( text(f"SELECT COUNT(*) FROM {quoted_table}") ).scalar_one() rows.append((table_name, int(count))) return tuple(rows) def _row_content_digest(row: Mapping[str, Any]) -> str: canonical_row = { str(column_name): _canonical_reflection_value(value) for column_name, value in row.items() } return hashlib.sha256(_stable_json(canonical_row).encode()).hexdigest() def _table_content_digest( connection: Connection, table_name: str, primary_key_columns: tuple[str, ...], ) -> str: table = Table( table_name, MetaData(), autoload_with=connection, resolve_fks=False, ) statement = select(table) has_stable_primary_key = bool(primary_key_columns) and all( column_name in table.c for column_name in primary_key_columns ) if has_stable_primary_key: statement = statement.order_by( *(table.c[column_name].asc() for column_name in primary_key_columns) ) row_digests = ( _row_content_digest(row) for row in connection.execute(statement).mappings() ) if not has_stable_primary_key: row_digests = iter(sorted(row_digests)) table_digest = hashlib.sha256() for row_digest in row_digests: table_digest.update(row_digest.encode("ascii")) table_digest.update(b"\n") return table_digest.hexdigest() def _impacted_table_content_digests( connection: Connection, ) -> tuple[tuple[str, str], ...]: inspector = inspect(connection) existing_tables = set(inspector.get_table_names()) digests: list[tuple[str, str]] = [] for table_name in sorted(_CONTENT_FINGERPRINT_TABLES & existing_tables): primary_key = inspector.get_pk_constraint(table_name) primary_key_columns = tuple( str(column_name) for column_name in primary_key.get("constrained_columns") or () ) digests.append( ( table_name, _table_content_digest( connection, table_name, primary_key_columns, ), ) ) return tuple(digests) def _schema_fingerprint( connection: Connection, impacted_table_row_counts: tuple[tuple[str, int], ...] | None = None, impacted_table_content_digests: tuple[tuple[str, str], ...] | None = None, ) -> str: row_counts = ( impacted_table_row_counts if impacted_table_row_counts is not None else _impacted_table_row_counts(connection) ) content_digests = ( impacted_table_content_digests if impacted_table_content_digests is not None else _impacted_table_content_digests(connection) ) payload = _stable_json( { "schema": _schema_snapshot(connection), "impacted_table_row_counts": [ {"table": table_name, "row_count": row_count} for table_name, row_count in row_counts ], "impacted_table_content_digests": [ {"table": table_name, "sha256": digest} for table_name, digest in content_digests ], } ).encode() return hashlib.sha256(payload).hexdigest() def _revision_rows(connection: Connection) -> tuple[str, ...] | None: if "alembic_version" not in inspect(connection).get_table_names(): return None rows = connection.execute( text("SELECT version_num FROM alembic_version ORDER BY version_num") ).scalars() return tuple(str(row) for row in rows) def _approval_state(connection: Connection) -> tuple[int | None, int]: inspector = inspect(connection) tables = inspector.get_table_names() if APPROVAL_TABLE not in tables: return None, 0 rows = int( connection.execute( text("SELECT COUNT(*) FROM approval_requests") ).scalar_one() ) inbound = 0 for table_name in tables: inbound += sum( 1 for foreign_key in inspector.get_foreign_keys(table_name) if foreign_key.get("referred_table") == APPROVAL_TABLE ) return rows, inbound def _postgres_privileges( connection: Connection, ) -> tuple[bool | None, bool | None]: if connection.dialect.name != "postgresql": return None, None schema_ok = bool( connection.execute( text( """ SELECT has_schema_privilege(current_schema(), 'USAGE') AND has_schema_privilege(current_schema(), 'CREATE') """ ) ).scalar_one() ) ownership_ok = bool( connection.execute( text( """ SELECT COALESCE(bool_and(pg_has_role(c.relowner, 'USAGE')), true) FROM pg_class AS c JOIN pg_namespace AS n ON n.oid = c.relnamespace WHERE n.nspname = current_schema() AND c.relname = ANY(:table_names) """ ), {"table_names": sorted(_IMPACTED_TABLES)}, ).scalar_one() ) return schema_ok, ownership_ok def _postgres_version_state( connection: Connection, ) -> tuple[int | None, bool | None]: if connection.dialect.name != "postgresql": return None, None version_num = int( connection.execute( text("SELECT current_setting('server_version_num')") ).scalar_one() ) return version_num, version_num >= MINIMUM_POSTGRESQL_VERSION_NUM def audit_connection(connection: Connection) -> BaselineAudit: """Return a non-secret, read-only assessment of baseline eligibility.""" observed = _metadata_drift(connection) approval_rows, approval_inbound = _approval_state(connection) schema_ok, ownership_ok = _postgres_privileges(connection) postgres_version_num, postgres_version_supported = _postgres_version_state( connection ) impacted_table_row_counts = _impacted_table_row_counts(connection) impacted_table_content_digests = _impacted_table_content_digests(connection) return BaselineAudit( dialect=connection.dialect.name, fingerprint=_schema_fingerprint( connection, impacted_table_row_counts, impacted_table_content_digests, ), impacted_table_row_counts=impacted_table_row_counts, impacted_table_content_digests=impacted_table_content_digests, postgresql_server_version_num=postgres_version_num, postgresql_version_supported=postgres_version_supported, revision_rows=_revision_rows(connection), observed_drift=observed, unexpected_drift=observed - ALLOWED_DRIFT, approval_rows=approval_rows, approval_inbound_foreign_keys=approval_inbound, schema_privileges_ok=schema_ok, table_ownership_ok=ownership_ok, ) def audit_engine(engine: Engine) -> BaselineAudit: """Run the default read-only audit.""" with engine.connect() as connection: transaction = connection.begin() try: if connection.dialect.name == "postgresql": connection.exec_driver_sql("SET TRANSACTION READ ONLY") return audit_connection(connection) finally: transaction.rollback() def _alembic_config(connection: Connection) -> Config: project_root = Path(__file__).resolve().parents[2] config = Config(str(project_root / "alembic.ini")) config.set_main_option("script_location", str(project_root / "alembic")) config.attributes["connection"] = connection return config def _validate_expected_head(config: Config) -> None: heads = tuple(ScriptDirectory.from_config(config).get_heads()) if heads != (TARGET_REVISION,): raise ReconciliationError("Alembic head changed after this baseline was prepared") def _acquire_postgres_lock(connection: Connection) -> None: connection.exec_driver_sql("SET LOCAL lock_timeout = '5s'") connection.exec_driver_sql("SET LOCAL statement_timeout = '5min'") connection.execute( text("SELECT pg_advisory_xact_lock(:lock_key)"), {"lock_key": ADVISORY_LOCK_KEY}, ) def _lock_impacted_postgres_tables(connection: Connection) -> None: if connection.dialect.name != "postgresql": return existing_tables = set(inspect(connection).get_table_names()) table_names = sorted(_IMPACTED_TABLES & existing_tables) if not table_names: return preparer = connection.dialect.identifier_preparer quoted_tables = ", ".join( preparer.quote_identifier(table_name) for table_name in table_names ) connection.exec_driver_sql( f"LOCK TABLE {quoted_tables} IN ACCESS EXCLUSIVE MODE" ) def _validate_apply_preconditions( audit: BaselineAudit, expected_fingerprint: str, ) -> None: if audit.postgresql_version_supported is False: raise ReconciliationError( "PostgreSQL 15 or newer is required for NULLS NOT DISTINCT constraints" ) if audit.fingerprint != expected_fingerprint: raise ReconciliationError( "Schema fingerprint changed after dry run; run the audit again" ) if audit.revision_rows is None: raise ReconciliationError("alembic_version table is missing") if audit.revision_rows: raise ReconciliationError("Database already has an Alembic revision") if audit.unexpected_drift: raise ReconciliationError("Schema contains drift outside the approved allowlist") if audit.approval_rows not in {None, 0}: raise ReconciliationError("approval_requests is not empty") if audit.approval_inbound_foreign_keys: raise ReconciliationError("approval_requests has dependent foreign keys") if audit.schema_privileges_ok is False or audit.table_ownership_ok is False: raise ReconciliationError("Database role lacks required schema ownership privileges") def apply_baseline( engine: Engine, expected_fingerprint: str, *, require_postgresql: bool = True, ) -> BaselineAudit: """Atomically stamp, reconcile, upgrade, and verify the approved schema.""" if not SCHEMA_FINGERPRINT_PATTERN.fullmatch(expected_fingerprint): raise ReconciliationError("Expected fingerprint must be a lowercase SHA-256 value") if require_postgresql and engine.dialect.name != "postgresql": raise ReconciliationError("Baseline apply is supported only on PostgreSQL") with engine.begin() as connection: if connection.dialect.name == "postgresql": _acquire_postgres_lock(connection) _lock_impacted_postgres_tables(connection) before = audit_connection(connection) _validate_apply_preconditions(before, expected_fingerprint) config = _alembic_config(connection) _validate_expected_head(config) command.stamp(config, PREVIOUS_REVISION) command.upgrade(config, "head") after = audit_connection(connection) if after.revision_rows != (TARGET_REVISION,): raise ReconciliationError("Alembic revision was not advanced atomically") if after.observed_drift: raise ReconciliationError("Schema still differs from SQLAlchemy metadata") if after.approval_rows is not None: raise ReconciliationError("Obsolete approval_requests table still exists") return after def _build_parser() -> argparse.ArgumentParser: parser = argparse.ArgumentParser( description=( "Audit the platform schema, or atomically apply the approved one-time baseline" ) ) parser.add_argument( "--apply", action="store_true", help="Apply the baseline; without this flag the command is read-only", ) parser.add_argument( "--expected-fingerprint", help="Exact SHA-256 fingerprint printed by the immediately preceding dry run", ) return parser def _validated_cli_database_url() -> str: settings = get_settings() try: target = validate_platform_migration_target( settings.database_url, settings.legacy_database_url, ) except RuntimeError as exc: raise ReconciliationError(str(exc)) from None if target[0] != "postgresql": raise ReconciliationError( "Schema reconciliation CLI requires a PostgreSQL platform database" ) return settings.database_url def main() -> None: parser = _build_parser() arguments = parser.parse_args() if arguments.apply and not arguments.expected_fingerprint: parser.error("--apply requires --expected-fingerprint") if not arguments.apply and arguments.expected_fingerprint: parser.error("--expected-fingerprint is only valid with --apply") engine: Engine | None = None try: engine = create_engine( _validated_cli_database_url(), poolclass=NullPool, ) if arguments.apply: report = apply_baseline( engine, str(arguments.expected_fingerprint), ) payload = { "applied": True, "target_revision": TARGET_REVISION, **report.public_dict(), } else: report = audit_engine(engine) payload = {"applied": False, **report.public_dict()} print(json.dumps(payload, ensure_ascii=False, indent=2, sort_keys=True)) if not arguments.apply and not report.eligible: raise SystemExit(2) except ReconciliationError as exc: print( json.dumps( { "applied": False, "error": str(exc), }, ensure_ascii=False, sort_keys=True, ) ) raise SystemExit(2) from None except Exception as exc: print( json.dumps( { "applied": False, "error": f"Unexpected {type(exc).__name__}", }, ensure_ascii=False, sort_keys=True, ) ) raise SystemExit(1) from None finally: if engine is not None: engine.dispose() if __name__ == "__main__": main()