Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
44 changes: 21 additions & 23 deletions airflow-core/src/airflow/models/serialized_dag.py
Original file line number Diff line number Diff line change
Expand Up @@ -547,28 +547,6 @@ def _prefetch_dag_write_metadata(
if not dag_id_list:
return {}

# Fetch the serialized_dag (last_updated, dag_hash) of the latest DagVersion per dag_id,
# ordering by version_number so it stays consistent with the DagVersion picked by dv_subq.
sd_subq = (
select(
cls.dag_id.label("dag_id"),
cls.last_updated.label("last_updated"),
cls.dag_hash.label("dag_hash"),
func.row_number()
.over(partition_by=cls.dag_id, order_by=DagVersion.version_number.desc())
.label("rn"),
)
.join(DagVersion, cls.dag_version_id == DagVersion.id)
.where(cls.dag_id.in_(dag_id_list))
.subquery()
)
sd_rows = session.execute(
select(sd_subq.c.dag_id, sd_subq.c.last_updated, sd_subq.c.dag_hash).where(sd_subq.c.rn == 1)
).all()
sd_by_dag_id: dict[str, tuple[datetime, str]] = {
row.dag_id: (row.last_updated, row.dag_hash) for row in sd_rows
}

# Fetch latest DagVersion per dag_id, ordering by version_number to match write_dag.
dv_subq = (
select(
Expand All @@ -586,6 +564,23 @@ def _prefetch_dag_write_metadata(
).all()
dv_by_dag_id: dict[str, DagVersion] = {dv.dag_id: dv for dv in dag_versions}

# Fetch the serialized_dag (last_updated, dag_hash) of the latest DagVersion per dag_id,
# outer join with dv_subq so None is set when latest dag version has no serialized entry.
sd_subq = (
select(
dv_subq.c.dag_id,
cls.last_updated.label("last_updated"),
cls.dag_hash.label("dag_hash"),
)
.where(dv_subq.c.rn == 1)
.outerjoin(cls, cls.dag_version_id == dv_subq.c.id)
.subquery()
)
sd_rows = session.execute(select(sd_subq.c.dag_id, sd_subq.c.last_updated, sd_subq.c.dag_hash)).all()
sd_by_dag_id: dict[str, tuple[datetime, str]] = {
row.dag_id: (row.last_updated, row.dag_hash) for row in sd_rows
}

return {
dag_id: DagWriteMetadata(
last_updated=sd_by_dag_id[dag_id][0] if dag_id in sd_by_dag_id else None,
Expand Down Expand Up @@ -723,10 +718,13 @@ def write_dag(
)
)

if dag_version and not has_task_instances:
if dag_version and not has_task_instances and serialized_dag_hash is not None:

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

serialized_dag_hash is not None isn't the same predicate as "the latest dag_version has a serialized row". dag_hash is nullable=False and sd_subq in _prefetch_dag_write_metadata inner-joins serialized_dag to dag_version, so the hash comes from the highest version that still has a row, while dag_version here is the absolute latest. If only the latest version's row is missing, the hash is non-None, this guard passes, the UPDATE matches 0 rows and we return False at line 747 again.

Checked on this branch in breeze: v1 with a serialized row plus a bare latest v2 keeps returning False, with no row ever created for v2. If the incoming hash matches v1's, it returns even earlier, at line 710 on the unchanged-hash short-circuit. Since the only production creator of a dag_version row is line 769, in the same transaction as the insert at line 784, both shapes come out of the same accident class.

Pairing the prefetched hash with the version this branch actually updates covers both (outerjoin serialized_dag off the latest-version subquery so dag_hash is None whenever the latest version has no row). I tried that in a probe and both shapes then heal on the next parse. Letting rowcount == 0 fall through to DagVersion.write_dag is a good complement since it's the direct signal, but on its own it misses the unchanged-hash case, and it leaves min_update_interval reading the older row's last_updated. If the clause does stay here, a line of comment saying what it detects would help, since as written it reads as null-safety on a non-nullable column.

# This is for dynamic DAGs that the hashes changes often. We should update
# the serialized dag, the dag_version and the dag_code instead of a new version
# if the dag_version is not associated with any task instances
# exception is when the *latest* dag version has no corresponding serialized dag
# which is denoted by serialized_dag_hash == None, then fall through to write

new_serialized_dag = cls(dag)

# Use direct UPDATE to avoid loading the full serialized DAG
Expand Down
223 changes: 217 additions & 6 deletions airflow-core/tests/unit/models/test_serialized_dag.py
Original file line number Diff line number Diff line change
Expand Up @@ -34,7 +34,7 @@
from airflow.models.dag import DagModel
from airflow.models.dag_version import DagVersion
from airflow.models.deadline_alert import DeadlineAlert as DAM
from airflow.models.serialized_dag import SerializedDagModel as SDM
from airflow.models.serialized_dag import DagWriteMetadata, SerializedDagModel as SDM
from airflow.providers.standard.operators.bash import BashOperator
from airflow.providers.standard.operators.empty import EmptyOperator
from airflow.providers.standard.operators.python import PythonOperator
Expand Down Expand Up @@ -688,6 +688,36 @@ def test_prefetch_dag_write_metadata_returns_latest_version(self, dag_maker, ses
assert metadata.dag_version is not None
assert metadata.dag_version.version_number == 2

def test_prefetch_dag_write_metadata_latest_version_without_sdm_entry(self, dag_maker, session):
"""When the latest DagVersion has no SDM entry, last_updated and dag_hash are None."""
with dag_maker("prefetch_no_sdm_dag") as dag:
EmptyOperator(task_id="task1")

v1 = session.scalar(select(DagVersion).where(DagVersion.dag_id == dag.dag_id))
assert v1 is not None
assert v1.version_number == 1

# v1 has an SDM entry (written by dag_maker); create v2 directly without one.
v2 = DagVersion.write_dag(
dag_id=dag.dag_id,
bundle_name="dag_maker",
bundle_version="v2",
session=session,
)
session.flush()

assert v2.version_number == 2

result = SDM._prefetch_dag_write_metadata([dag.dag_id], session=session)
metadata = result[dag.dag_id]

assert metadata.dag_version is not None
assert metadata.dag_version.version_number == 2
assert metadata.dag_version.dag_id == dag.dag_id
assert metadata.dag_version.bundle_version == "v2"
assert metadata.last_updated is None
assert metadata.dag_hash is None

def test_new_dag_version_created_when_bundle_name_changes_and_hash_unchanged(self, dag_maker, session):
"""Test that new dag_version is created if bundle_name changes but DAG is unchanged."""
# Create and write initial DAG
Expand Down Expand Up @@ -875,10 +905,11 @@ def __init__(self, *, task_id: str, **kwargs):
# Hashes should be identical
assert hash_1 == hash_2, "Hashes should be identical when dicts are sorted consistently"

def test_dynamic_dag_update_preserves_null_check(self, dag_maker, session):
def test_new_dag_version_is_created_when_version_exists_but_serialized_dag_row_missing(
self, dag_maker, session
):
"""
Test that dynamic DAG update gracefully handles case where SerializedDagModel doesn't exist.
This preserves the null-check fix from PR #56422 and tests the direct UPDATE path.
Test that dynamic DAG update creates a SerializedDagModel if it doesn't exist.
"""
with dag_maker(dag_id="test_missing_serdag", serialized=True, session=session) as dag:
EmptyOperator(task_id="task1")
Expand Down Expand Up @@ -906,16 +937,196 @@ def test_dynamic_dag_update_preserves_null_check(self, dag_maker, session):
# Verify no SerializedDagModel exists
assert SDM.get("test_missing_serdag", session=session) is None

# Try to update - should return False gracefully (not crash)
# Try to update - should create a new serialized dag row under the latest DagVersion
result = SDM.write_dag(
dag=lazy_dag,
bundle_name="test_bundle",
bundle_version=None,
min_update_interval=None,
session=session,
)
session.commit()

assert result is True
latest_version = session.scalar(
select(DagVersion)
.where(DagVersion.dag_id == "test_missing_serdag")
.order_by(DagVersion.version_number.desc())
.limit(1)
)
assert latest_version is not None
serialized_dag = SDM.get("test_missing_serdag", session=session)
assert serialized_dag is not None

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This asserts that some row exists for the dag_id, but the property that makes the dag schedulable again is that the row belongs to the latest DagVersion. SDM.get("test_missing_serdag", session=session) already orders by version_number, so using it plus an assert on version_number would pin what the test name claims. As written the test still passes in the case where an older version holds the row and the latest one doesn't.

Worth keeping a case for the rowcount == 0 branch too. After this change nothing in the file exercises it, and if a later cleanup drops it as dead code that path would run the dag_version and dag_code updates and return True having written no serialized row.

assert serialized_dag.dag_version_id == latest_version.id

def test_write_dag_returns_false_when_update_finds_no_serialized_row(self, dag_maker, session):
"""
Test that write_dag returns False when the UPDATE path finds no serialized row.

This exercises the rowcount == 0 branch: _prefetched carries a non-None dag_hash
(so the dynamic-update UPDATE path is taken), but the SDM row has been deleted
between prefetch and execution, so the UPDATE affects 0 rows.
"""
with dag_maker(dag_id="test_rowcount_zero", serialized=True, session=session) as dag:
EmptyOperator(task_id="task1")

lazy_dag = LazyDeserializedDAG.from_dag(dag)
SDM.write_dag(
dag=lazy_dag,
bundle_name="test_bundle",
bundle_version=None,
session=session,
)
session.commit()

dag_version = session.scalar(
select(DagVersion)
.where(DagVersion.dag_id == "test_rowcount_zero")
.order_by(DagVersion.version_number.desc())
.limit(1)
)
assert dag_version is not None

assert result is False # Should return False when SerializedDagModel is missing
# Build prefetched metadata that looks like the SDM row exists (non-None hash),
# then delete the actual row so the UPDATE finds nothing.
stale_hash = session.scalar(select(SDM.dag_hash).where(SDM.dag_id == "test_rowcount_zero"))
assert stale_hash is not None
session.execute(delete(SDM).where(SDM.dag_id == "test_rowcount_zero"))
session.commit()

prefetched = DagWriteMetadata(
last_updated=None,
dag_hash=stale_hash,
dag_version=dag_version,
)
# Change the dag so the hash differs, forcing the dynamic-update branch.
from airflow.sdk.definitions.dag import DAG as SdkDAG

new_dag_obj = SdkDAG(dag_id="test_rowcount_zero", schedule=None)
with new_dag_obj:
EmptyOperator(task_id="task1")
EmptyOperator(task_id="task2")
new_lazy = LazyDeserializedDAG.from_dag(new_dag_obj)

result = SDM.write_dag(
dag=new_lazy,
bundle_name="test_bundle",
bundle_version=None,
min_update_interval=None,
session=session,
_prefetched=prefetched,
)

assert result is False
assert SDM.get("test_rowcount_zero", session=session) is None

def test_latest_dag_version_with_no_serialized_row_and_changed_hash_heals(self, dag_maker, session):
"""
v1 has a serialized row, v2 is a bare DagVersion (no serialized row), and
the incoming Dag hash differs from v1's. The v2 dag should be serialized successfully
"""
with dag_maker(dag_id="test_v2_bare_changed_hash", serialized=True, session=session) as dag:
EmptyOperator(task_id="task1")

lazy_dag = LazyDeserializedDAG.from_dag(dag)
SDM.write_dag(dag=lazy_dag, bundle_name="test_bundle", bundle_version=None, session=session)
session.commit()

# Create a bare v2 DagVersion with no corresponding serialized row.
DagVersion.write_dag(dag_id="test_v2_bare_changed_hash", bundle_name="test_bundle", session=session)
session.commit()

v2 = session.scalar(
select(DagVersion)
.where(DagVersion.dag_id == "test_v2_bare_changed_hash")
.order_by(DagVersion.version_number.desc())
.limit(1)
)
assert v2 is not None
assert v2.version_number == 2

# A changed dag produces a different hash — bypasses the unchanged-hash short-circuit.
from airflow.sdk.definitions.dag import DAG as SdkDAG

changed_dag = SdkDAG(dag_id="test_v2_bare_changed_hash", schedule=None)
with changed_dag:
EmptyOperator(task_id="task1")
EmptyOperator(task_id="task2")
changed_lazy = LazyDeserializedDAG.from_dag(changed_dag)

result = SDM.write_dag(
dag=changed_lazy,
bundle_name="test_bundle",
bundle_version=None,
min_update_interval=None,
session=session,
)
session.commit()

assert result is True
# The code falls through to DagVersion.write_dag, creating a new v3 with the serialized
# row attached to it — v2 stays bare. Assert the latest version now has a serialized row.
latest = session.scalar(
select(DagVersion)
.where(DagVersion.dag_id == "test_v2_bare_changed_hash")
.order_by(DagVersion.version_number.desc())
.limit(1)
)
assert latest is not None
assert latest.version_number == 3
assert session.scalar(select(SDM).where(SDM.dag_version_id == latest.id)) is not None

def test_latest_dag_version_with_no_serialized_row_and_unchanged_hash_heals(self, dag_maker, session):
"""
v1 has a serialized row, v2 is a bare DagVersion (no serialized row), and
the incoming Dag hash matches v1's. The v2 dag should be serialized successfully
"""
with dag_maker(dag_id="test_v2_bare_unchanged_hash", serialized=True, session=session) as dag:
EmptyOperator(task_id="task1")

lazy_dag = LazyDeserializedDAG.from_dag(dag)
SDM.write_dag(dag=lazy_dag, bundle_name="test_bundle", bundle_version=None, session=session)
session.commit()

v1_hash = session.scalar(select(SDM.dag_hash).where(SDM.dag_id == "test_v2_bare_unchanged_hash"))
assert v1_hash is not None

# Create a bare v2 DagVersion with no corresponding serialized row.
DagVersion.write_dag(dag_id="test_v2_bare_unchanged_hash", bundle_name="test_bundle", session=session)
session.commit()

v2 = session.scalar(
select(DagVersion)
.where(DagVersion.dag_id == "test_v2_bare_unchanged_hash")
.order_by(DagVersion.version_number.desc())
.limit(1)
)
assert v2 is not None
assert v2.version_number == 2

# Same dag — same hash as v1. Without the fix the prefetch returns v1's hash (inner
# join excludes v2), the unchanged-hash guard fires, and v2 never gets a serialized row.
result = SDM.write_dag(
dag=lazy_dag,
bundle_name="test_bundle",
bundle_version=None,
min_update_interval=None,
session=session,
)
session.commit()

assert result is True
# The code falls through to DagVersion.write_dag, creating a new v3 with the serialized
# row attached to it — v2 stays bare. Assert the latest version now has a serialized row.
latest = session.scalar(
select(DagVersion)
.where(DagVersion.dag_id == "test_v2_bare_unchanged_hash")
.order_by(DagVersion.version_number.desc())
.limit(1)
)
assert latest is not None
assert latest.version_number == 3
assert session.scalar(select(SDM).where(SDM.dag_version_id == latest.id)) is not None

def test_dynamic_dag_update_success(self, dag_maker, session):
"""
Expand Down