Release v0.1.7

src/govoplan_core/db/migrations.py

@@ -1,19 +1,24 @@
 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.db.session import configure_database
+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
@@ -33,6 +38,9 @@ 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"),
@@ -160,12 +168,31 @@ class MigrationResult:
     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:
@@ -187,7 +214,130 @@ def registered_module_migration_plan(
         active_enabled_modules,
         manifest_factories=active_manifest_factories,
     )
-    return migration_metadata_plan(registry)
+    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:
@@ -518,6 +668,7 @@ 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
@@ -526,14 +677,19 @@ def migrate_database(
         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,
+    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,
-        ),
-        "heads",
-    )
+        )
+        _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)
@@ -542,3 +698,43 @@ def migrate_database(
         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}