jason810496 commented on code in PR #74041:
URL: https://github.com/apache/airflow/pull/74041#discussion_r4176794545
##########
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:
Done in 2bcf050883. `validate_serialized_dag` now applies both rules to the
Dag `from_dict` returns, under the same conditions as the SDK `DAG`. So a
`@continuous` Dag left at the default `max_active_runs` of 16 and a `@daily`
Dag with `catchup` but no `start_date` now raise `DeserializationError`.
##########
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:
Done in 2bcf050883. The cycle check now walks `downstream_task_ids` on the
tasks of the Dag `from_dict` returns, so edges from `_downstream_task_ids` and
`client_defaults["tasks"]` count too. `validate_serialized_dag` returns that
Dag, and in #74035 edca324650 passes it to `_apply_team_rules`, which no longer
runs its own `deepcopy` and `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:
Done in 2bcf050883. The test now runs over `{}`, `"x"`, a wrong `__type`, a
`__var` with no `task_id` and `"task_id": 5`, each with its expected
`json_path`. Since `e.json_path` renders `$.dag.tasks[2]['__type']` under
jsonschema 4.26 and `$.dag.tasks[2].__type` under 4.23, a7efc7c851 builds the
path from the error's `absolute_path`, which gives the latter under both.
--
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]
