Skip to content

Commit b2998c4

Browse files
committed
Remove callbacks from DAG default_args when serializating it (#57397)
(cherry picked from commit 3b72458)
1 parent faeb1bf commit b2998c4

2 files changed

Lines changed: 79 additions & 0 deletions

File tree

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

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2512,6 +2512,23 @@ def serialize_dag(cls, dag: DAG) -> dict:
25122512
serialized_dag["has_on_success_callback"] = True
25132513
if dag.has_on_failure_callback:
25142514
serialized_dag["has_on_failure_callback"] = True
2515+
2516+
# TODO: Move this logic to a better place -- ideally before serializing contents of default_args.
2517+
# There is some duplication with this and SerializedBaseOperator.partial_kwargs serialization.
2518+
# Ideally default_args goes through same logic as fields of SerializedBaseOperator.
2519+
if serialized_dag.get("default_args", {}):
2520+
default_args_dict = serialized_dag["default_args"][Encoding.VAR]
2521+
callbacks_to_remove = []
2522+
for k, v in list(default_args_dict.items()):
2523+
if k in [
2524+
f"on_{x}_callback" for x in ("execute", "failure", "success", "retry", "skipped")
2525+
]:
2526+
if bool(v):
2527+
default_args_dict[f"has_{k}"] = True
2528+
callbacks_to_remove.append(k)
2529+
for k in callbacks_to_remove:
2530+
del default_args_dict[k]
2531+
25152532
return serialized_dag
25162533
except SerializationError:
25172534
raise

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

Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4283,3 +4283,65 @@ def test_partial_kwargs_end_to_end_deserialization(self):
42834283
assert "owner" in deserialized_task.partial_kwargs
42844284
assert deserialized_task.partial_kwargs["retry_delay"] == timedelta(seconds=600)
42854285
assert deserialized_task.partial_kwargs["owner"] == "custom_owner"
4286+
4287+
4288+
@pytest.mark.parametrize(
4289+
["callbacks", "expected_has_flags", "absent_keys"],
4290+
[
4291+
pytest.param(
4292+
{
4293+
"on_failure_callback": lambda ctx: None,
4294+
"on_success_callback": lambda ctx: None,
4295+
"on_retry_callback": lambda ctx: None,
4296+
},
4297+
["has_on_failure_callback", "has_on_success_callback", "has_on_retry_callback"],
4298+
["on_failure_callback", "on_success_callback", "on_retry_callback"],
4299+
id="multiple_callbacks",
4300+
),
4301+
pytest.param(
4302+
{"on_failure_callback": lambda ctx: None},
4303+
["has_on_failure_callback"],
4304+
["on_failure_callback", "has_on_success_callback", "on_success_callback"],
4305+
id="single_callback",
4306+
),
4307+
pytest.param(
4308+
{"on_failure_callback": lambda ctx: None, "on_execute_callback": None},
4309+
["has_on_failure_callback"],
4310+
["on_failure_callback", "has_on_execute_callback", "on_execute_callback"],
4311+
id="callback_with_none",
4312+
),
4313+
pytest.param(
4314+
{},
4315+
[],
4316+
[
4317+
"has_on_execute_callback",
4318+
"has_on_failure_callback",
4319+
"has_on_success_callback",
4320+
"has_on_retry_callback",
4321+
"has_on_skipped_callback",
4322+
],
4323+
id="no_callbacks",
4324+
),
4325+
],
4326+
)
4327+
def test_dag_default_args_callbacks_serialization(callbacks, expected_has_flags, absent_keys):
4328+
"""Test callbacks in DAG default_args are serialized as boolean flags."""
4329+
default_args = {"owner": "test_owner", "retries": 2, **callbacks}
4330+
4331+
with DAG(dag_id="test_default_args_callbacks", default_args=default_args) as dag:
4332+
BashOperator(task_id="task1", bash_command="echo 1", dag=dag)
4333+
4334+
serialized_dag_dict = SerializedDAG.serialize_dag(dag)
4335+
default_args_dict = serialized_dag_dict["default_args"][Encoding.VAR]
4336+
4337+
for flag in expected_has_flags:
4338+
assert default_args_dict.get(flag) is True
4339+
4340+
for key in absent_keys:
4341+
assert key not in default_args_dict
4342+
4343+
assert default_args_dict["owner"] == "test_owner"
4344+
assert default_args_dict["retries"] == 2
4345+
4346+
deserialized_dag = SerializedDAG.deserialize_dag(serialized_dag_dict)
4347+
assert deserialized_dag.dag_id == "test_default_args_callbacks"

0 commit comments

Comments
 (0)