kaxil commented on code in PR #74041:
URL: https://github.com/apache/airflow/pull/74041#discussion_r4167621343
##########
airflow-core/src/airflow/serialization/serialized_objects.py:
##########
@@ -2131,6 +2135,74 @@ def from_dict(cls, serialized_obj: dict) ->
SerializedDAG:
# Pass client_defaults directly to deserialize_dag
return cls.deserialize_dag(serialized_obj["dag"], client_defaults)
+ @classmethod
+ def fill_config_defaults(cls, serialized_obj: dict[str, Any]) -> None:
+ """
+ Fill in the Dag settings a serialized Dag leaves unset from the
Airflow config, as a Python Dag does.
+
+ A Lang-SDK runtime cannot read the Airflow config, so it leaves
``max_active_tasks``,
+ ``max_active_runs``, ``max_consecutive_failed_dag_runs``, ``catchup``
and
+ ``disable_bundle_versioning`` out unless the Dag sets them. A value
the Dag sets is kept.
+ *serialized_obj* is changed in place.
+ """
+ dag = serialized_obj.get("dag")
+ if not isinstance(dag, dict):
+ # validate_serialized_dag rejects it.
+ return
+ for key, get, section, option in (
+ ("max_active_tasks", conf.getint, "core",
"max_active_tasks_per_dag"),
+ ("max_active_runs", conf.getint, "core",
"max_active_runs_per_dag"),
+ (
+ "max_consecutive_failed_dag_runs",
+ conf.getint,
+ "core",
+ "max_consecutive_failed_dag_runs_per_dag",
+ ),
+ ("catchup", conf.getboolean, "scheduler", "catchup_by_default"),
+ ("disable_bundle_versioning", conf.getboolean, "dag_processor",
"disable_bundle_versioning"),
+ ):
+ if key not in dag:
+ dag[key] = get(section, option)
+
+ @classmethod
+ def validate_serialized_dag(cls, serialized_obj: dict[str, Any]) -> None:
+ """
+ Check that a serialized Dag, such as one a Lang-SDK runtime produced,
can be stored and loaded.
+
+ It must match the JSON schema, have unique task ids, deserialize, and
have no cycle in its task graph.
+ *serialized_obj* is not changed.
+
+ :raises DeserializationError: if it does not.
+ """
+ dag = serialized_obj.get("dag")
+ dag_id = dag.get("dag_id") if isinstance(dag, dict) else None
+ try:
+ cls.validate_schema(serialized_obj)
+ except ValidationError as e:
+ raise DeserializationError(
+ dag_id, f"Dag {dag_id!r} does not match the schema at
{e.json_path}: {e.message}"
+ ) from e
+ task_ids = collections.Counter(
+ task[Encoding.VAR]["task_id"] for task in
serialized_obj["dag"]["tasks"]
+ )
+ if duplicates := sorted(task_id for task_id, count in task_ids.items()
if count > 1):
+ raise DeserializationError(
+ dag_id, f"Dag {dag_id!r} has more than one task with id {',
'.join(map(repr, duplicates))}"
+ )
+ try:
+ cls.from_dict(copy.deepcopy(serialized_obj))
Review Comment:
`from_dict` skips the checks the SDK `DAG` runs when it is built, so a
filled Dag can pass here with a combination Python refuses. A `@continuous` Dag
that leaves out `max_active_runs` gets 16 and passes, though Python raises
`ContinuousTimetable requires max_active_runs <= 1` and the scheduler only
gates on `max_active_runs`. With `catchup_by_default = True`, a `@daily` Dag
with no `start_date` gets `catchup=True` and passes, where Python raises
`start_date is required when catchup=True`. #74035 lists the continuous case as
a known limitation, but since this method fills the value and is the validator,
could it run the `_validate_max_active_runs` and `_validate_catchup` rules
after `from_dict`?
##########
airflow-core/src/airflow/serialization/serialized_objects.py:
##########
@@ -2131,6 +2135,74 @@ def from_dict(cls, serialized_obj: dict) ->
SerializedDAG:
# Pass client_defaults directly to deserialize_dag
return cls.deserialize_dag(serialized_obj["dag"], client_defaults)
+ @classmethod
+ def fill_config_defaults(cls, serialized_obj: dict[str, Any]) -> None:
+ """
+ Fill in the Dag settings a serialized Dag leaves unset from the
Airflow config, as a Python Dag does.
+
+ A Lang-SDK runtime cannot read the Airflow config, so it leaves
``max_active_tasks``,
+ ``max_active_runs``, ``max_consecutive_failed_dag_runs``, ``catchup``
and
+ ``disable_bundle_versioning`` out unless the Dag sets them. A value
the Dag sets is kept.
+ *serialized_obj* is changed in place.
+ """
+ dag = serialized_obj.get("dag")
+ if not isinstance(dag, dict):
+ # validate_serialized_dag rejects it.
+ return
+ for key, get, section, option in (
+ ("max_active_tasks", conf.getint, "core",
"max_active_tasks_per_dag"),
+ ("max_active_runs", conf.getint, "core",
"max_active_runs_per_dag"),
+ (
+ "max_consecutive_failed_dag_runs",
+ conf.getint,
+ "core",
+ "max_consecutive_failed_dag_runs_per_dag",
+ ),
+ ("catchup", conf.getboolean, "scheduler", "catchup_by_default"),
+ ("disable_bundle_versioning", conf.getboolean, "dag_processor",
"disable_bundle_versioning"),
+ ):
+ if key not in dag:
+ dag[key] = get(section, option)
+
+ @classmethod
+ def validate_serialized_dag(cls, serialized_obj: dict[str, Any]) -> None:
+ """
+ Check that a serialized Dag, such as one a Lang-SDK runtime produced,
can be stored and loaded.
+
+ It must match the JSON schema, have unique task ids, deserialize, and
have no cycle in its task graph.
+ *serialized_obj* is not changed.
+
+ :raises DeserializationError: if it does not.
+ """
+ dag = serialized_obj.get("dag")
+ dag_id = dag.get("dag_id") if isinstance(dag, dict) else None
+ try:
+ cls.validate_schema(serialized_obj)
+ except ValidationError as e:
+ raise DeserializationError(
+ dag_id, f"Dag {dag_id!r} does not match the schema at
{e.json_path}: {e.message}"
+ ) from e
+ task_ids = collections.Counter(
+ task[Encoding.VAR]["task_id"] for task in
serialized_obj["dag"]["tasks"]
+ )
+ if duplicates := sorted(task_id for task_id, count in task_ids.items()
if count > 1):
+ raise DeserializationError(
+ dag_id, f"Dag {dag_id!r} has more than one task with id {',
'.join(map(repr, duplicates))}"
+ )
+ try:
+ cls.from_dict(copy.deepcopy(serialized_obj))
+ except Exception as e:
+ cause = e.__cause__ if isinstance(e, DeserializationError) and
e.__cause__ else e
+ raise DeserializationError(
+ dag_id, f"Dag {dag_id!r} cannot be deserialized:
{type(cause).__name__}: {cause}"
+ ) from e
+ downstream = {
Review Comment:
This walks only the raw `downstream_task_ids`, but `from_dict` also takes
edges from the legacy `_downstream_task_ids` key and from
`client_defaults["tasks"]`, so a payload with `a -> b -> a` in
`_downstream_task_ids` loads with that cycle and passes. Walking
`task_dict[...].downstream_task_ids` on the Dag `from_dict` returns would check
the edges the scheduler actually gets, and returning that Dag would also save
#74035's `_apply_team_rules` a second `deepcopy` + `from_dict`.
##########
airflow-core/tests/unit/serialization/test_dag_serialization.py:
##########
@@ -5155,3 +5155,139 @@ def get_weight(self, ti):
op = BaseOperator(task_id="empty_task",
weight_rule=NotRegisteredPriorityWeightStrategy())
with pytest.raises(ValueError, match="Unknown priority strategy"):
OperatorSerialization.serialize(op)
+
+
+class TestValidateSerializedDag:
+ @staticmethod
+ def _serialize() -> dict:
+ with DAG(dag_id="checked_dag", schedule=None) as dag:
+ BaseOperator(task_id="extract") >> BaseOperator(task_id="load")
+ return DagSerialization.to_dict(dag)
+
+ def test_accepts_a_dag_that_loads(self):
+ data = self._serialize()
+ data["__version"] = 2
+ before = copy.deepcopy(data)
+
+ DagSerialization.validate_serialized_dag(data)
+
+ assert data == before
+
+ @pytest.mark.parametrize(
+ ("change", "error"),
+ [
+ pytest.param(
+ {"max_active_runs": "many"},
+ "Dag 'checked_dag' does not match the schema at
$.dag.max_active_runs: "
+ "'many' is not of type 'number'",
+ id="schema",
+ ),
+ pytest.param(
+ {"timetable": {"__type": "no.such.Timetable", "__var": {}}},
+ "Dag 'checked_dag' cannot be deserialized:
TimetableNotRegistered: ",
+ id="deserialize",
+ ),
+ ],
+ )
+ def test_rejects_a_dag_that_does_not_load(self, change, error):
+ data = self._serialize()
+ data["dag"].update(change)
+
+ with pytest.raises(DeserializationError, match=f"^{re.escape(error)}"):
+ DagSerialization.validate_serialized_dag(data)
+
+ def test_rejects_a_task_entry_that_is_not_an_operator(self):
Review Comment:
`{}` is the only entry that exercises the new `tasks` schema, so dropping
the `__type` const, the `task_id` string type or `__var.required` would leave
every test here passing. Without the const, `{"__type": "dag", "__var":
{"task_id": "x"}}` passes and `from_dict` drops that task silently. Could this
be parametrized over `"x"`, a wrong `__type`, a `__var` with no `task_id`, and
`"task_id": 5`, each with its expected `json_path`?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
