Skip to content

Commit 962acf2

Browse files
jason810496claude
andcommitted
Check a serialized Dag as the SDK Dag is checked and return it
validate_serialized_dag now applies the max_active_runs and catchup rules the SDK Dag runs when it is built, to the Dag it deserializes. The cycle check walks that Dag's tasks, so edges from the legacy _downstream_task_ids key and from client_defaults count too. It returns the deserialized Dag. Co-Authored-By: Claude <noreply@anthropic.com>
1 parent c3113a9 commit 962acf2

2 files changed

Lines changed: 104 additions & 11 deletions

File tree

‎airflow-core/src/airflow/serialization/serialized_objects.py‎

Lines changed: 22 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -2165,13 +2165,15 @@ def fill_config_defaults(cls, serialized_obj: dict[str, Any]) -> None:
21652165
dag[key] = get(section, option)
21662166

21672167
@classmethod
2168-
def validate_serialized_dag(cls, serialized_obj: dict[str, Any]) -> None:
2168+
def validate_serialized_dag(cls, serialized_obj: dict[str, Any]) -> SerializedDAG:
21692169
"""
21702170
Check that a serialized Dag, such as one a Lang-SDK runtime produced, can be stored and loaded.
21712171
2172-
It must match the JSON schema, have unique task ids, deserialize, and have no cycle in its task graph.
2173-
*serialized_obj* is not changed.
2172+
It must match the JSON schema, have unique task ids and deserialize. Like a Dag built with the SDK,
2173+
it must not set a ``max_active_runs`` its timetable forbids, nor ``catchup`` without a
2174+
``start_date``, and its task graph must have no cycle. *serialized_obj* is not changed.
21742175
2176+
:return: The deserialized Dag.
21752177
:raises DeserializationError: if it does not.
21762178
"""
21772179
dag = serialized_obj.get("dag")
@@ -2190,18 +2192,30 @@ def validate_serialized_dag(cls, serialized_obj: dict[str, Any]) -> None:
21902192
dag_id, f"Dag {dag_id!r} has more than one task with id {', '.join(map(repr, duplicates))}"
21912193
)
21922194
try:
2193-
cls.from_dict(copy.deepcopy(serialized_obj))
2195+
deserialized = cls.from_dict(copy.deepcopy(serialized_obj))
21942196
except Exception as e:
21952197
cause = e.__cause__ if isinstance(e, DeserializationError) and e.__cause__ else e
21962198
raise DeserializationError(
21972199
dag_id, f"Dag {dag_id!r} cannot be deserialized: {type(cause).__name__}: {cause}"
21982200
) from e
2199-
downstream = {
2200-
task[Encoding.VAR]["task_id"]: task[Encoding.VAR].get("downstream_task_ids") or ()
2201-
for task in serialized_obj["dag"]["tasks"]
2202-
}
2201+
if (
2202+
limit := deserialized.timetable.active_runs_limit
2203+
) is not None and deserialized.max_active_runs > limit:
2204+
raise DeserializationError(
2205+
dag_id,
2206+
f"Dag {dag_id!r} sets max_active_runs {deserialized.max_active_runs}, "
2207+
f"but {type(deserialized.timetable).__name__} allows at most {limit}",
2208+
)
2209+
if (
2210+
deserialized.catchup
2211+
and deserialized.timetable.can_be_scheduled
2212+
and not (deserialized.start_date or "start_date" in deserialized.default_args)
2213+
):
2214+
raise DeserializationError(dag_id, f"Dag {dag_id!r} sets catchup but no start_date")
2215+
downstream = {task_id: task.downstream_task_ids for task_id, task in deserialized.task_dict.items()}
22032216
if (task_id := detect_cycle(downstream, downstream.__getitem__)) is not None:
22042217
raise DeserializationError(dag_id, f"Dag {dag_id!r} has a cycle through task {task_id!r}")
2218+
return deserialized
22052219

22062220

22072221
class TaskGroupSerialization(BaseSerialization):

‎airflow-core/tests/unit/serialization/test_dag_serialization.py‎

Lines changed: 82 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -5196,13 +5196,36 @@ def test_rejects_a_dag_that_does_not_load(self, change, error):
51965196
with pytest.raises(DeserializationError, match=f"^{re.escape(error)}"):
51975197
DagSerialization.validate_serialized_dag(data)
51985198

5199-
def test_rejects_a_task_entry_that_is_not_an_operator(self):
5199+
def test_returns_the_deserialized_dag(self):
5200+
dag = DagSerialization.validate_serialized_dag(self._serialize())
5201+
5202+
assert isinstance(dag, SerializedDAG)
5203+
assert dag.dag_id == "checked_dag"
5204+
assert dag.task_dict["extract"].downstream_task_ids == {"load"}
5205+
5206+
@pytest.mark.parametrize(
5207+
("entry", "json_path"),
5208+
[
5209+
pytest.param({}, "$.dag.tasks[2]", id="empty"),
5210+
pytest.param("x", "$.dag.tasks[2]", id="not-an-object"),
5211+
pytest.param(
5212+
{"__type": "dag", "__var": {"task_id": "t"}}, "$.dag.tasks[2]['__type']", id="wrong-type"
5213+
),
5214+
pytest.param({"__type": "operator", "__var": {}}, "$.dag.tasks[2]['__var']", id="no-task-id"),
5215+
pytest.param(
5216+
{"__type": "operator", "__var": {"task_id": 5}},
5217+
"$.dag.tasks[2]['__var'].task_id",
5218+
id="task-id-not-a-string",
5219+
),
5220+
],
5221+
)
5222+
def test_rejects_a_task_entry_that_is_not_an_operator(self, entry, json_path):
52005223
data = self._serialize()
5201-
data["dag"]["tasks"].append({})
5224+
data["dag"]["tasks"].append(entry)
52025225

52035226
with pytest.raises(
52045227
DeserializationError,
5205-
match=r"^Dag 'checked_dag' does not match the schema at \$\.dag\.tasks\[2\]: ",
5228+
match=rf"^Dag 'checked_dag' does not match the schema at {re.escape(json_path)}: ",
52065229
):
52075230
DagSerialization.validate_serialized_dag(data)
52085231

@@ -5232,6 +5255,62 @@ def test_rejects_a_dag_with_a_cycle(self):
52325255
with pytest.raises(DeserializationError, match="^Dag 'checked_dag' has a cycle through task 'load'$"):
52335256
DagSerialization.validate_serialized_dag(data)
52345257

5258+
def test_rejects_a_cycle_in_the_legacy_downstream_key(self):
5259+
data = self._serialize()
5260+
load = next(task for task in data["dag"]["tasks"] if task["__var"]["task_id"] == "load")
5261+
load["__var"]["_downstream_task_ids"] = ["extract"]
5262+
5263+
with pytest.raises(DeserializationError, match="^Dag 'checked_dag' has a cycle through task 'load'$"):
5264+
DagSerialization.validate_serialized_dag(data)
5265+
5266+
@pytest.mark.parametrize(
5267+
("change", "error"),
5268+
[
5269+
pytest.param(
5270+
{},
5271+
"Dag 'checked_dag' sets max_active_runs 16, but ContinuousTimetable allows at most 1",
5272+
id="unset",
5273+
),
5274+
pytest.param(
5275+
{"max_active_runs": 2},
5276+
"Dag 'checked_dag' sets max_active_runs 2, but ContinuousTimetable allows at most 1",
5277+
id="above-the-limit",
5278+
),
5279+
],
5280+
)
5281+
def test_rejects_max_active_runs_above_the_timetable_limit(self, change, error):
5282+
data = DagSerialization.to_dict(DAG(dag_id="checked_dag", schedule="@continuous", max_active_runs=1))
5283+
del data["dag"]["max_active_runs"]
5284+
data["dag"].update(change)
5285+
5286+
with pytest.raises(DeserializationError, match=f"^{re.escape(error)}$"):
5287+
DagSerialization.validate_serialized_dag(data)
5288+
5289+
def test_rejects_catchup_without_a_start_date(self):
5290+
data = DagSerialization.to_dict(DAG(dag_id="checked_dag", schedule="@daily"))
5291+
data["dag"]["catchup"] = True
5292+
5293+
with pytest.raises(DeserializationError, match="^Dag 'checked_dag' sets catchup but no start_date$"):
5294+
DagSerialization.validate_serialized_dag(data)
5295+
5296+
@pytest.mark.parametrize(
5297+
"kwargs",
5298+
[
5299+
pytest.param(
5300+
{
5301+
"schedule": "@daily",
5302+
"default_args": {"start_date": datetime(2024, 1, 1, tzinfo=dt_timezone.utc)},
5303+
},
5304+
id="start-date-in-default-args",
5305+
),
5306+
pytest.param({"schedule": None}, id="unscheduled"),
5307+
],
5308+
)
5309+
def test_accepts_catchup_the_sdk_dag_accepts(self, kwargs):
5310+
dag = DAG(dag_id="checked_dag", catchup=True, **kwargs)
5311+
5312+
assert DagSerialization.validate_serialized_dag(DagSerialization.to_dict(dag)).catchup is True
5313+
52355314

52365315
class TestFillConfigDefaults:
52375316
CONFIG = {

0 commit comments

Comments
 (0)