Skip to content

fix: resolve dag_run.log_template_id to target's active template during AF3 migrations - #183

Open
sahas-kamsani wants to merge 1 commit into
mainfrom
fix/af3-resolve-log-template-id
Open

sahas-kamsani wants to merge 1 commit into
mainfrom
fix/af3-resolve-log-template-id

Conversation

@sahas-kamsani

Copy link
Copy Markdown

Fixes: #182
Related: #179, #181


What this changes

#181 stopped the log_template_id/backfill_id/trigger_id foreign-key violations from #179 by
removing those fields from the migration insert payload entirely. That's correct for
backfill_id/trigger_id (safe to leave NULL — no active backfill/deferral is a valid state), but
not for log_template_id: Airflow has no fallback for a NULL log_template_id when resolving a
task's log path, so every migrated dag_run row was left unable to serve logs for any task, on any
try, even when the worker wrote the log successfully. #181 traded a loud, migration-blocking error
for a silent, deferred one — worse UX, just delayed.

This PR resolves log_template_id to the target deployment's own currently-active log_template
row instead of leaving it unset, mirroring how Airflow itself resolves the value via
LogTemplate.latest_id() when creating a new DagRun.

Only astronomer_starship/_af3/starship_compatability.py needs changing. _af2 doesn't send
log_template_id in the first place — its get_dag_runs projects through the attrs-filtered
dag_run_attrs(), which never lists log_template_id, whereas AF3's get_dag_runs does a raw
select(table) column dump that picks up every reflected column. So this bug — and this fix — is
AF3-source → AF3-target specific.

Patch

diff --git a/astronomer_starship/_af3/starship_compatability.py b/astronomer_starship/_af3/starship_compatability.py
index 7e1e8cd..323b3c4 100644
--- a/astronomer_starship/_af3/starship_compatability.py
+++ b/astronomer_starship/_af3/starship_compatability.py
@@ -29,6 +29,10 @@ class StarshipAirflow(BaseStarshipAirflow):
 class StarshipAirflow30(StarshipAirflow):
     """Airflow 3.0 compatibility layer."""
 
+    _cached_target_log_template_id = None
+    """Per-instance cache for `_get_target_log_template_id`. A migration is short-lived, so
+    the target's active `log_template` row can't change mid-run."""
+
     @classmethod
     def pool_attrs(cls) -> "Dict[str, AttrDesc]":
         return {
@@ -983,6 +987,11 @@ class StarshipAirflow30(StarshipAirflow):
         if not items:
             return []
 
+        # Resolved once per batch (and cached across batches) rather than per-row --
+        # the target's active log_template can't change mid-migration.
+        target_log_template_id = self._get_target_log_template_id() if table_name == "dag_run" else None
+
+        remapped_log_template_count = 0
         for item in items:
             if "id" in item and table_name not in ["task_instance", "task_instance_history"]:
                 del item["id"]
@@ -991,8 +1000,15 @@ class StarshipAirflow30(StarshipAirflow):
             item.pop("dag_version_id", None)
             item.pop("created_dag_version_id", None)
             if table_name == "dag_run":
-                item.pop("log_template_id", None)
                 item.pop("backfill_id", None)
+                if "log_template_id" in item:
+                    # log_template_id is a source-local FK, just like backfill_id/trigger_id --
+                    # but unlike those, Airflow has no fallback for a NULL log_template_id when
+                    # resolving a task's log path, so dropping it (as we do for backfill_id)
+                    # would silently break log retrieval for every task in this dag_run, forever.
+                    # Resolve it to the target's own active template instead.
+                    item["log_template_id"] = target_log_template_id
+                    remapped_log_template_count += 1
             if table_name == "task_instance":
                 item.pop("trigger_id", None)
 
@@ -1000,6 +1016,24 @@ class StarshipAirflow30(StarshipAirflow):
                 # Drop executor_config, because its original type may have gotten lost
                 # and pickling it will not recover it
                 item["executor_config"] = pickle.dumps({})
+
+        if remapped_log_template_count:
+            if target_log_template_id is None:
+                logger.warning(
+                    "Migrated %d dag_run row(s) with log_template_id remapped to NULL: "
+                    "the target deployment has no rows in its log_template table. "
+                    "Log retrieval will remain unavailable for these rows until Airflow "
+                    "creates one (e.g. on next scheduler start).",
+                    remapped_log_template_count,
+                )
+            else:
+                logger.warning(
+                    "Migrated %d dag_run row(s) with log_template_id remapped to the "
+                    "target's active log_template (id=%s). Source log_template ids are "
+                    "not portable across deployments.",
+                    remapped_log_template_count,
+                    target_log_template_id,
+                )
         try:
             engine = self.session.get_bind()
             metadata = MetaData(bind=engine)
@@ -1038,6 +1072,31 @@ class StarshipAirflow30(StarshipAirflow):
             self.session.rollback()
             raise e
 
+    def _get_target_log_template_id(self):
+        """Resolve the target deployment's currently-active log_template id.
+
+        Mirrors how Airflow itself resolves log_template_id when creating a new DagRun --
+        see ``LogTemplate.latest_id()`` in ``airflow.models.log_template``, which is simply
+        the highest ``id`` in the local ``log_template`` table. Source log_template ids are
+        never portable across deployments (the id spaces are independent), so this always
+        queries the target's own table rather than trusting any value from the source.
+
+        Returns None (and lets the caller fall back to NULL) if the target's log_template
+        table has no rows, which shouldn't happen in practice -- Airflow seeds one on first
+        scheduler start -- but is handled rather than raising, so a single edge case doesn't
+        fail an entire migration batch.
+
+        Cached per-instance since a migration run is short-lived and this can't change
+        mid-run.
+        """
+        from sqlalchemy import text
+
+        if self._cached_target_log_template_id is None:
+            self._cached_target_log_template_id = self.session.execute(
+                text("SELECT id FROM log_template ORDER BY id DESC LIMIT 1")
+            ).scalar()
+        return self._cached_target_log_template_id
+
     def get_latest_dag_version_id(self, dag_id: str):
         from sqlalchemy import MetaData, desc, select
 
diff --git a/tests/af3_insert_sanitization_test.py b/tests/af3_insert_sanitization_test.py
index 054e576..5b4079d 100644
--- a/tests/af3_insert_sanitization_test.py
+++ b/tests/af3_insert_sanitization_test.py
@@ -43,16 +43,33 @@ class FakeInsert:
         return self
 
 
+class FakeScalarResult:
+    def __init__(self, value):
+        self._value = value
+
+    def scalar(self):
+        return self._value
+
+
 class FakeSession:
-    def __init__(self):
+    def __init__(self, log_template_id=2):
         self.statement = None
         self.committed = False
+        self.log_template_id = log_template_id
+        self.log_template_lookups = 0
 
     def get_bind(self):
         return object()
 
     def execute(self, statement):
-        self.statement = statement
+        # The real code only ever hands this fake two kinds of statement: the final
+        # FakeInsert (recorded for assertions), or a raw `text()` select used to resolve
+        # the target's active log_template id (answered from `self.log_template_id`).
+        if isinstance(statement, FakeInsert):
+            self.statement = statement
+            return None
+        self.log_template_lookups += 1
+        return FakeScalarResult(self.log_template_id)
 
     def commit(self):
         self.committed = True
@@ -61,11 +78,11 @@ class FakeSession:
         raise AssertionError("rollback should not be called")
 
 
-def test_dag_run_direct_insert_strips_source_log_template_id(monkeypatch):
+def test_dag_run_direct_insert_resolves_log_template_id_to_target_active_template(monkeypatch):
     import sqlalchemy
     import sqlalchemy.dialects.postgresql
 
-    fake_session = FakeSession()
+    fake_session = FakeSession(log_template_id=2)
     starship = StarshipAirflow31()
     starship._session = fake_session
 
@@ -76,7 +93,7 @@ def test_dag_run_direct_insert_strips_source_log_template_id(monkeypatch):
         {
             "dag_id": "example_dag",
             "run_id": "scheduled__2026-08-01T00:00:00+00:00",
-            "log_template_id": 3,
+            "log_template_id": 3,  # source-local id -- must NOT survive onto the target
             "dag_version_id": "source-task-version-id",
             "created_dag_version_id": "source-created-version-id",
         }
@@ -85,12 +102,67 @@ def test_dag_run_direct_insert_strips_source_log_template_id(monkeypatch):
     result = starship.insert_directly("dag_run", items)
 
     inserted = fake_session.statement.items[0]
-    assert "log_template_id" not in inserted
+    # resolved to the target's own active template, not the source's id (3)
+    assert inserted["log_template_id"] == 2
     assert "dag_version_id" not in inserted
     assert "created_dag_version_id" not in inserted
     assert fake_session.statement.conflict_target == ["dag_id", "run_id"]
     assert fake_session.committed is True
-    assert result == [{"dag_id": "example_dag", "run_id": "scheduled__2026-08-01T00:00:00+00:00"}]
+    assert result == [
+        {
+            "dag_id": "example_dag",
+            "run_id": "scheduled__2026-08-01T00:00:00+00:00",
+            "log_template_id": 2,
+        }
+    ]
+
+
+def test_dag_run_direct_insert_resolves_log_template_id_once_per_batch(monkeypatch):
+    """The target's active log_template id shouldn't be looked up once per row."""
+    import sqlalchemy
+    import sqlalchemy.dialects.postgresql
+
+    fake_session = FakeSession(log_template_id=2)
+    starship = StarshipAirflow31()
+    starship._session = fake_session
+
+    monkeypatch.setattr(sqlalchemy, "MetaData", FakeMetaData)
+    monkeypatch.setattr(sqlalchemy.dialects.postgresql, "insert", FakeInsert)
+
+    items = [{"dag_id": "example_dag", "run_id": f"scheduled__{i}", "log_template_id": 3} for i in range(5)]
+
+    starship.insert_directly("dag_run", items)
+
+    assert all(item["log_template_id"] == 2 for item in fake_session.statement.items)
+    assert fake_session.log_template_lookups == 1
+
+
+def test_dag_run_direct_insert_leaves_log_template_id_null_when_target_has_none(monkeypatch, caplog):
+    """Target log_template table with zero rows shouldn't crash the whole batch -- and should
+    say so loudly instead of silently repeating the #181 gap."""
+    import sqlalchemy
+    import sqlalchemy.dialects.postgresql
+
+    fake_session = FakeSession(log_template_id=None)
+    starship = StarshipAirflow31()
+    starship._session = fake_session
+
+    monkeypatch.setattr(sqlalchemy, "MetaData", FakeMetaData)
+    monkeypatch.setattr(sqlalchemy.dialects.postgresql, "insert", FakeInsert)
+
+    items = [
+        {
+            "dag_id": "example_dag",
+            "run_id": "scheduled__2026-08-01T00:00:00+00:00",
+            "log_template_id": 3,
+        }
+    ]
+
+    starship.insert_directly("dag_run", items)
+
+    inserted = fake_session.statement.items[0]
+    assert inserted["log_template_id"] is None
+    assert "no rows in its log_template table" in caplog.text
 
 
 def test_dag_run_direct_insert_strips_source_backfill_id(monkeypatch):

Design decisions

  • Resolution strategy: SELECT id FROM log_template ORDER BY id DESC LIMIT 1 against the
    target's own session (self.session) — this is exactly MAX(id), matching Airflow's own
    LogTemplate.latest_id(). Never trusts the source's id — id spaces aren't portable across
    deployments (a real customer migration hit exactly this pitfall during manual backfill).
  • Resolved once per batch, not per row — computed before the loop and cached on the instance,
    since the target's active template can't change mid-migration. Covered by
    test_dag_run_direct_insert_resolves_log_template_id_once_per_batch.
  • Target log_template table has zero rows (shouldn't happen in practice — Airflow seeds one
    on first scheduler start): resolves to None/NULL rather than raising, so this one edge case
    doesn't abort an entire migration batch — but logs a logger.warning so it's visible in
    migration output rather than silent, which is the exact complaint this PR exists to fix. Open to
    feedback if maintainers would rather this fail loudly instead — flagging it explicitly since it
    was called out as an open question in earlier drafts of this fix.

Testing

  • tests/af3_insert_sanitization_test.py::test_dag_run_direct_insert_resolves_log_template_id_to_target_active_template
    — replaces the old "strips" regression test added in Strip dag_run log_template_id during AF3 migrations #181; asserts the source's id (3) does
    not survive and the target's active id (2) is what actually gets inserted.
  • ..._resolves_log_template_id_once_per_batch — 5-row batch, asserts exactly 1 lookup query.
  • ..._leaves_log_template_id_null_when_target_has_none — empty-table edge case: asserts no crash
    and asserts the warning is logged (via caplog).
  • Existing backfill_id/trigger_id stripping tests (also from Strip dag_run log_template_id during AF3 migrations #181) are unchanged and still
    pass, confirming no regression there.

Verified locally: pytest tests/af3_insert_sanitization_test.py -v5/5 passed, against a real
apache-airflow==3.0.6 install (not mocked out) and the repo's pinned ruff==0.14.5 (ruff check

  • ruff format --check both clean). Full pytest -c pyproject.toml run with no other regressions
    (the one unrelated failure, a just build-backend shell-out in validation_test.py, is a sandbox
    limitation — no just binary available there — not related to this change).

Not yet covered: a live end-to-end migration through the UI/API against two real Airflow 3
instances. See the companion doc starship-local-testing-guide.md for a walkthrough using this
repo's own dev3/ tooling — worth running once before merge.

…ng AF3 migrations

#181 stopped the log_template_id/backfill_id/trigger_id foreign-key violations from #179 by
dropping those fields from the migration insert payload entirely. That's correct for
backfill_id/trigger_id (NULL is a valid state), but not for log_template_id: Airflow has no
fallback for a NULL log_template_id when resolving a task's log path, so every migrated
dag_run row was left unable to serve logs for any task, on any try -- even when the worker
wrote the log successfully. This traded a loud, migration-blocking FK error for a silent,
deferred log-visibility gap.

Resolve log_template_id to the target deployment's own currently-active log_template row
instead of dropping it, mirroring how Airflow itself resolves the value via
LogTemplate.latest_id() when creating a new DagRun. Resolved once per insert_directly batch
and cached per-instance. Falls back to NULL (with a logged warning) if the target's
log_template table is empty, rather than failing the whole batch.

Updates the existing regression test from #181 (which asserted the field was stripped) to
assert it's resolved to the target's id instead, and adds coverage for the once-per-batch
caching and the empty-target-table fallback.

Fixes #<issue-number>
Related: #179, #181

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
@codecov-commenter

Copy link
Copy Markdown

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 24.89%. Comparing base (792a6ed) to head (3a9867f).

Additional details and impacted files
@@            Coverage Diff             @@
##             main     #183      +/-   ##
==========================================
+ Coverage   23.07%   24.89%   +1.82%     
==========================================
  Files          22       22              
  Lines        2102     2153      +51     
==========================================
+ Hits          485      536      +51     
  Misses       1617     1617              

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
  • 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

@sahas-kamsani

Copy link
Copy Markdown
Author

Tested against real Astro Cloud deployments (before/after)

Verified this fix live against two Airflow-3.1.8-matched Astro Cloud deployments (source/target),
switching the target between the published pre-fix astronomer-starship==2.9.1 wheel and this
branch (pip install git+https://github.com/astronomer/starship.git@fix/af3-resolve-log-template-id;
confirmed via direct image inspection that the built image actually contains
_get_target_log_template_id and no longer contains the item.pop("log_template_id", ...) line).
Note: this branch still self-reports as pip version 2.9.1 since it doesn't bump __version__
don't rely on the version string alone when testing this, check the actual installed code.

Before (pre-fix 2.9.1): migrated dag_run.log_template_id = NULL on target (independently
verified via a fresh GET, not just the insert's echoed response).

After (this branch): migrated dag_run.log_template_id resolves to target's real active
log_template.id. To confirm this is genuine active resolution and not just "no longer stripped"
(which would also make real values pass through unchanged and look identical when source/target
happen to share the same Airflow version), I sent a deliberately impossible log_template_id
(424242) from source — target's DB ended up with its own real active id, not 424242 and not
NULL. Matches _get_target_log_template_id()'s
SELECT id FROM log_template ORDER BY id DESC LIMIT 1 exactly.

No regression observed in backfill_id/trigger_id stripping behavior during this testing.

One limitation worth flagging for reviewers: I could not get an end-to-end "log now renders in the
UI" confirmation in this test setup, because source and target are separate Astro deployments with
separate log storage buckets — Starship migrates DB rows, not the physical log files, so the log
file itself never existed in target's bucket regardless of log_template_id. A valid
log_template_id is necessary but not sufficient for a log to render; this fix resolves the
DB-level correctness (a valid FK, consistent with how Airflow itself resolves the field), which is
what's in scope here. Full render-through-the-UI verification would need a same-storage or
storage-copied test setup.

Also hit and worked around a separate, pre-existing issue while setting this up: the
DAG/operator-based migration path (StarshipAirflowMigrationDAG) can't run against an Airflow-3
source today (RuntimeError: Direct database access via the ORM is not allowed in Airflow 3.0
unrelated to this PR's change, StarshipLocalHook does direct ORM calls that AF3's Task SDK
forbids from inside a task). Drove the migration via direct calls to
/api/starship/dag_runs//api/starship/task_instances instead to get this evidence.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants