fix(janitor): keep the origin of a pending model migration

Author: kvlonge

src/sqlbuild/.agents/skills/sqlbuild/references/docs/cli/janitor.md

@@ -12,6 +12,7 @@ Online: https://sqlbuild.com/docs/cli/janitor/
 - Flags
 - Lifecycle
 - Compatibility views
+- Pending migration origins
 - Relation age
 - Audit record
 - Restoring an archive
@@ -74,6 +75,10 @@ A compatibility view is never archived as a stale relation. Once it expires, or
 
 Run the same SQLBuild version for `sqb janitor` as for builds. An older janitor does not know compatibility views and archives them as stale relations.
 
+## Pending migration origins
+
+A relation that a project model still has to migrate from is never archived. While a model declares `migrate_from` and its own relation does not exist yet in the target, the janitor skips the origin and lists it as `pending migration origin for <model>`. A bare `migrate_from` name is matched in each schema the project builds into; a qualified name only at that location. Once the model has built, the origin is handled as usual: after a migration, the old name holds a compatibility view, which expires as described above.
+
 ## Relation age
 
 `retention_days` is measured from the last time a relation was written, as reported by the warehouse. For a relation the project no longer builds, that is roughly when it stopped being built.

src/sqlbuild/executor/janitor/_helpers/plan.py

@@ -8,12 +8,14 @@
 
 from sqlbuild.adapter.contract.classes.base_adapter import BaseAdapter
 from sqlbuild.adapter.contract.models import RelationInfo, RelationLookup
+from sqlbuild.compiler.compile.constants import MIGRATE_FROM_CONFIG_KEY
 from sqlbuild.compiler.compile.models import (
     CompiledModel,
     CompiledProject,
     CompiledSeed,
     CompiledSource,
 )
+from sqlbuild.executor.janitor.constants import QUALIFIED_ORIGIN_PARTS
 from sqlbuild.executor.janitor.models import JanitorRelationKey
 
 
@@ -42,6 +44,59 @@ def collect_desired_keys(project: CompiledProject) -> set[JanitorRelationKey]:
     return keys
 
 
+def pending_migration_origins(
+    *,
+    project: CompiledProject,
+    relations_by_schema: dict[tuple[str | None, str | None], tuple[RelationInfo, ...]],
+) -> dict[JanitorRelationKey, str]:
+    """Return listed origins of migrations whose destination relation does not exist yet."""
+
+    listed: dict[tuple[str | None, str | None, str], JanitorRelationKey] = {}
+    relations: tuple[RelationInfo, ...]
+    for relations in relations_by_schema.values():
+        listed.update({_folded(relation_key(item)): relation_key(item) for item in relations})
+    schemas: set[tuple[str | None, str | None]] = collect_target_schemas(project)
+    pending: dict[JanitorRelationKey, str] = {}
+    model: CompiledModel
+    for model in project.models:
+        raw: object | None = model.config.values.get(MIGRATE_FROM_CONFIG_KEY)
+        destination: JanitorRelationKey = JanitorRelationKey(
+            database=model.destination.database,
+            schema=model.destination.schema,
+            name=model.destination.name,
+        )
+        if not isinstance(raw, str) or _folded(destination) in listed:
+            continue
+        origin: JanitorRelationKey
+        for origin in _origin_keys(raw=raw, destination=destination, schemas=schemas):
+            found: JanitorRelationKey | None = listed.get(_folded(origin))
+            if found is not None:
+                pending.setdefault(found, f"pending migration origin for {model.name}")
+    return pending
+
+
+def _origin_keys(
+    *, raw: str, destination: JanitorRelationKey, schemas: set[tuple[str | None, str | None]]
+) -> tuple[JanitorRelationKey, ...]:
+    parts: list[str] = [part.strip().strip('"`') for part in raw.split(".")]
+    if len(parts) == 1:
+        return tuple(
+            JanitorRelationKey(database=database, schema=schema, name=parts[0])
+            for database, schema in sorted(schemas, key=_schema_sort_key)
+        )
+    return (
+        JanitorRelationKey(
+            database=parts[0] if len(parts) == QUALIFIED_ORIGIN_PARTS else destination.database,
+            schema=parts[-2],
+            name=parts[-1],
+        ),
+    )
+
+
+def _folded(key: JanitorRelationKey) -> tuple[str | None, str | None, str]:
+    return RelationLookup.key(database=key.database, schema=key.schema, name=key.name)
+
+
 def collect_target_schemas(project: CompiledProject) -> set[tuple[str | None, str | None]]:
     """Collect schemas where the compiled project writes relations, one spelling per schema."""
 

src/sqlbuild/executor/janitor/constants.py

@@ -43,3 +43,4 @@
 MALFORMED_ARCHIVE_REASON: str = (
     "name resembles a janitor archive but does not match the strict archive grammar"
 )
+QUALIFIED_ORIGIN_PARTS: int = 3

src/sqlbuild/executor/janitor/main/plan.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+from dataclasses import replace
 from datetime import UTC, datetime
 from typing import Any
 
@@ -14,7 +15,11 @@
     collect_query_diff_artifact_candidates,
     gather_janitor_warehouse_facts,
 )
-from sqlbuild.executor.janitor._helpers.plan import collect_scan_schemas, collect_target_schemas
+from sqlbuild.executor.janitor._helpers.plan import (
+    collect_scan_schemas,
+    collect_target_schemas,
+    pending_migration_origins,
+)
 from sqlbuild.executor.janitor._helpers.schema_planning import classify_target_schemas
 from sqlbuild.executor.janitor.classes.relation_age_reader import JanitorRelationAgeReader
 from sqlbuild.executor.janitor.models import (
@@ -23,6 +28,7 @@
     JanitorDirectStatePruneCandidate,
     JanitorOldNameViewPlanning,
     JanitorPlan,
+    JanitorRelationKey,
     JanitorRelationScope,
     JanitorSchemaClassification,
     JanitorWarehouseFacts,
@@ -107,12 +113,11 @@ def build_janitor_plan(
         )
         inspection.completed(metadata={"item_count": len(target_schemas)})
     old_names: JanitorOldNameViewPlanning
-    old_names, scope = plan_old_name_view_cleanup(
+    old_names, scope = _plan_protected_names(
         adapter=adapter,
         connection=connection,
         managed_target_schemas=managed_target_schemas,
-        relations_by_schema=facts.relations_by_schema,
-        target_name=project.effective_target_name,
+        facts=facts,
         early_drops=early_old_name_view_drops,
         project=project,
         scope=scope,
@@ -173,3 +178,56 @@ def build_janitor_plan(
         age_metadata_supported=age_supported,
         planned_at=now,
     )
+
+
+def _plan_protected_names(
+    *,
+    adapter: BaseAdapter,
+    connection: Any,
+    managed_target_schemas: set[tuple[str | None, str | None]],
+    facts: JanitorWarehouseFacts,
+    early_drops: tuple[str, ...],
+    project: CompiledProject,
+    scope: JanitorRelationScope,
+    now: datetime,
+) -> tuple[JanitorOldNameViewPlanning, JanitorRelationScope]:
+    """Plan old-name views, then keep general cleanup away from pending migration origins."""
+
+    old_names: JanitorOldNameViewPlanning
+    old_names, scope = plan_old_name_view_cleanup(
+        adapter=adapter,
+        connection=connection,
+        managed_target_schemas=managed_target_schemas,
+        relations_by_schema=facts.relations_by_schema,
+        target_name=project.effective_target_name,
+        early_drops=early_drops,
+        project=project,
+        scope=scope,
+        now=now,
+    )
+    return old_names, _protect_pending_origins(
+        project=project, facts=facts, scope=scope, old_names=old_names
+    )
+
+
+def _protect_pending_origins(
+    *,
+    project: CompiledProject,
+    facts: JanitorWarehouseFacts,
+    scope: JanitorRelationScope,
+    old_names: JanitorOldNameViewPlanning,
+) -> JanitorRelationScope:
+    pending: dict[JanitorRelationKey, str] = {
+        key: reason
+        for key, reason in pending_migration_origins(
+            project=project, relations_by_schema=facts.relations_by_schema
+        ).items()
+        if key not in old_names.keys
+    }
+    if not pending:
+        return scope
+    return replace(
+        scope,
+        protected_relation_keys=scope.protected_relation_keys | frozenset(pending),
+        protected_relation_reasons={**(scope.protected_relation_reasons or {}), **pending},
+    )

tests/e2e/src/sqlbuild/cli/commands/main/refactor/_test_types.py

@@ -97,3 +97,10 @@ class HiddenOriginE2ETestCase:
     description: str
     target_settings: str
     expected_output: str
+
+
+@dataclass(frozen=True)
+class PendingOriginJanitorE2ETestCase:
+    description: str
+    janitor_args: tuple[str, ...]
+    expected_order_ids: tuple[int, ...]

tests/e2e/src/sqlbuild/cli/commands/main/refactor/test_janitor_pending_origin.py

@@ -0,0 +1,72 @@
+"""CLI e2e coverage for the janitor keeping the origin of a pending model migration."""
+
+from __future__ import annotations
+
+import subprocess
+from pathlib import Path
+
+import pytest
+
+from tests.e2e.src.sqlbuild.cli.commands.main.refactor._test_types import (
+    PendingOriginJanitorE2ETestCase,
+)
+from tests.e2e.src.sqlbuild.cli.commands.main.refactor.helpers import (
+    load_raw_orders,
+    order_history_files,
+    order_ids,
+    relation_type,
+    sqb,
+    write_orders_project,
+)
+
+_PENDING_REASON: str = "analytics.order_history  pending migration origin for order_ledger"
+
+
+@pytest.mark.parametrize(
+    "test_case",
+    [
+        PendingOriginJanitorE2ETestCase(
+            description="janitor keeps the origin until the rename is migrated",
+            janitor_args=("janitor", "--retention-days", "0", "--auto-approve"),
+            expected_order_ids=(1, 2, 3, 4),
+        ),
+    ],
+    ids=lambda case: case.description,
+)
+def test_given_pending_rename_when_janitor_runs_then_origin_survives_until_migrated(
+    tmp_path: Path, test_case: PendingOriginJanitorE2ETestCase
+) -> None:
+    project_dir: Path = write_orders_project(
+        tmp_path=tmp_path, files=order_history_files(materialized="incremental", udf=False)
+    )
+    first: subprocess.CompletedProcess[str] = sqb(project_dir, "build")
+    renamed: subprocess.CompletedProcess[str] = sqb(
+        project_dir, "rename", "model:order_history", "order_ledger"
+    )
+
+    pending: subprocess.CompletedProcess[str] = sqb(project_dir, *test_case.janitor_args)
+    origin_after_pending: str | None = relation_type(project_dir=project_dir, name="order_history")
+    load_raw_orders(project_dir=project_dir, order_ids=(3, 4))
+    built: subprocess.CompletedProcess[str] = sqb(project_dir, "build")
+    migrated: tuple[int, ...] = order_ids(project_dir=project_dir, name="order_ledger")
+    settled: subprocess.CompletedProcess[str] = sqb(project_dir, *test_case.janitor_args)
+    old_name_after_settled: str | None = relation_type(
+        project_dir=project_dir, name="order_history"
+    )
+    dropped: subprocess.CompletedProcess[str] = sqb(
+        project_dir, *test_case.janitor_args, "--drop-old-name-view", "order_history"
+    )
+
+    assert first.returncode == 0, first.stdout + first.stderr
+    assert renamed.returncode == 0, renamed.stdout + renamed.stderr
+    assert pending.returncode == 0, pending.stdout + pending.stderr
+    assert _PENDING_REASON in pending.stdout
+    assert origin_after_pending == "BASE TABLE"
+    assert built.returncode == 0, built.stdout + built.stderr
+    assert migrated == test_case.expected_order_ids
+    assert settled.returncode == 0, settled.stdout + settled.stderr
+    assert "pending migration origin" not in settled.stdout
+    assert "analytics.order_history  -> model:order_ledger" in settled.stdout
+    assert old_name_after_settled == "VIEW"
+    assert dropped.returncode == 0, dropped.stdout + dropped.stderr
+    assert relation_type(project_dir=project_dir, name="order_history") is None

tests/unit/src/sqlbuild/executor/janitor/main/_test_types.py

@@ -22,6 +22,7 @@ class JanitorPlanTestCase:
     tracked_relations: tuple[tuple[str | None, str | None, str], ...] = field(default_factory=tuple)
     exclude_patterns: tuple[str, ...] = field(default_factory=tuple)
     protected_relation_keys: frozenset[JanitorRelationKey] = frozenset()
+    migrate_from: str | None = None
     expected_candidate_names: tuple[str, ...] = field(default_factory=tuple)
     expected_direct_state_table_names: tuple[str, ...] = field(default_factory=tuple)
     expected_skipped_relation_reasons: tuple[str, ...] = field(default_factory=tuple)

tests/unit/src/sqlbuild/executor/janitor/main/helpers.py

@@ -540,6 +540,7 @@ def build_project(
     source_database: str | None = None,
     destination_database: str | None = None,
     destination_schema: str = "analytics",
+    migrate_from: str | None = None,
 ) -> CompiledProject:
     source: CompiledSource = CompiledSource(
         key=CompiledObjectKey(
@@ -569,7 +570,7 @@ def build_project(
                 name="orders",
                 relative_path=Path("models/orders.sql"),
                 query_sql="select 1",
-                config=CompileModelConfig(),
+                config=CompileModelConfig(values={"migrate_from": migrate_from}),
                 destination=CompiledRelationLocation(
                     database=destination_database,
                     schema=destination_schema,

tests/unit/src/sqlbuild/executor/janitor/main/test_plan.py

@@ -164,6 +164,29 @@
             expected_candidate_names=(),
             expected_skipped_relation_reasons=("relation is not tracked by SQLBuild",),
         ),
+        JanitorPlanTestCase(
+            description="origin of a pending migration is kept",
+            relation_infos=(relation_info("old_orders", created_at=OLD_TIME),),
+            migrate_from="old_orders",
+            expected_candidate_names=(),
+            expected_skipped_relation_reasons=("pending migration origin for orders",),
+        ),
+        JanitorPlanTestCase(
+            description="qualified origin of a pending migration is kept",
+            relation_infos=(relation_info("OLD_ORDERS", created_at=OLD_TIME),),
+            migrate_from="analytics.old_orders",
+            expected_candidate_names=(),
+            expected_skipped_relation_reasons=("pending migration origin for orders",),
+        ),
+        JanitorPlanTestCase(
+            description="origin is cleaned normally once the destination exists",
+            relation_infos=(
+                relation_info("orders", created_at=OLD_TIME),
+                relation_info("old_orders", created_at=OLD_TIME),
+            ),
+            migrate_from="old_orders",
+            expected_candidate_names=("old_orders",),
+        ),
         JanitorPlanTestCase(
             description="existing direct state tables are eligible for history pruning",
             relation_infos=(
@@ -204,7 +227,9 @@ def test_given_project_and_warehouse_when_building_janitor_plan_then_returns_exp
     )
 
     plan: JanitorPlan = build_janitor_plan(
-        project=build_project(source_schema=test_case.source_schema),
+        project=build_project(
+            source_schema=test_case.source_schema, migrate_from=test_case.migrate_from
+        ),
         adapter=adapter,
         connection=object(),
         retention_days=test_case.retention_days,

website/src/content/docs/docs/cli/janitor.mdx

@@ -59,6 +59,10 @@ A compatibility view is never archived as a stale relation. Once it expires, or
 
 Run the same SQLBuild version for `sqb janitor` as for builds. An older janitor does not know compatibility views and archives them as stale relations.
 
+## Pending migration origins
+
+A relation that a project model still has to migrate from is never archived. While a model declares `migrate_from` and its own relation does not exist yet in the target, the janitor skips the origin and lists it as `pending migration origin for <model>`. A bare `migrate_from` name is matched in each schema the project builds into; a qualified name only at that location. Once the model has built, the origin is handled as usual: after a migration, the old name holds a compatibility view, which expires as described above.
+
 ## Relation age
 
 `retention_days` is measured from the last time a relation was written, as reported by the warehouse. For a relation the project no longer builds, that is roughly when it stopped being built.