Vamsi-klu commented on code in PR #72242:
URL: https://github.com/apache/airflow/pull/72242#discussion_r4002503605
##########
airflow-core/src/airflow/serialization/enums.py:
##########
@@ -79,6 +79,7 @@ class DagAttributeTypes(str, Enum):
TASK_GROUP = "taskgroup"
EDGE_INFO = "edgeinfo"
PARAM = "param"
+ DAG_PARAM = "dag_param"
Review Comment:
Thanks for flagging this. I am keeping `SERIALIZER_VERSION` at 3 because
this follows how other `DAT` values have been added, and Airflow components are
expected to run the same version. The focused round-trip tests cover the new
type on the current version.
---
Drafted-by: Codex (GPT-5); reviewed by @Vamsi-klu before posting
##########
airflow-core/tests/unit/serialization/test_dag_serialization.py:
##########
@@ -5128,3 +5128,172 @@ 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)
+
+
+def _encoded_dag_params(obj):
+ found = []
+ if isinstance(obj, dict):
+ type_ = obj.get(Encoding.TYPE, obj.get("__type"))
+ if type_ in (DagAttributeTypes.DAG_PARAM, "dag_param"):
+ found.append(obj)
+ for value in obj.values():
+ found.extend(_encoded_dag_params(value))
+ elif isinstance(obj, list):
+ for value in obj:
+ found.extend(_encoded_dag_params(value))
+ return found
+
+
+def _taskflow_mapped_dag_with_param(*, default="p_default_val"):
+ from airflow.sdk import task
+
+ with DAG("test-dagparam-mapped", schedule=None, start_date=datetime(2020,
1, 1)) as dag:
+
+ @task
+ def add(value):
+ return value
+
+ add.partial(value=dag.param("p", default)).expand(value=[1, 2, 3])
+ return dag
+
+
+def test_dagparam_in_taskflow_partial_is_serialized_stably():
+ first = DagSerialization.to_dict(_taskflow_mapped_dag_with_param())
+ second = DagSerialization.to_dict(_taskflow_mapped_dag_with_param())
+ blob = json.dumps(first)
+ assert "object at 0x" not in blob.lower()
+ encoded = _encoded_dag_params(first)
+ assert encoded
+ for item in encoded:
+ var = item[Encoding.VAR]
+ assert item[Encoding.TYPE] == DagAttributeTypes.DAG_PARAM
+ assert var["name"] == "p"
+ assert var["dag_id"] == "test-dagparam-mapped"
+ assert var["default"] == "p_default_val"
+ assert first == second
+
+
+def test_dagparam_in_taskflow_partial_roundtrip():
+ serialized = DagSerialization.to_dict(_taskflow_mapped_dag_with_param())
+ restored = DagSerialization.from_dict(serialized)
+ value = restored.task_dict["add"].partial_kwargs["op_kwargs"]["value"]
+ assert isinstance(value, SerializedDagParam)
+ assert value.name == "p"
+ assert value.default == "p_default_val"
+ assert value.dag_id == "test-dagparam-mapped"
+
+
+def test_dagparam_in_mapped_operator_partial():
+ with DAG("test-dagparam-mapped-op", schedule=None,
start_date=datetime(2020, 1, 1)) as dag:
+ MockOperator.partial(task_id="t", arg1=dag.param("p",
"from_partial")).expand(arg2=["a", "b"])
+
+ serialized = DagSerialization.to_dict(dag)
+ assert "object at 0x" not in json.dumps(serialized).lower()
+ restored = DagSerialization.from_dict(serialized)
+ arg1 = restored.task_dict["t"].partial_kwargs["arg1"]
+ assert isinstance(arg1, SerializedDagParam)
+ assert arg1.name == "p"
+ assert arg1.default == "from_partial"
+
+
+def test_dagparam_in_non_mapped_operator_field():
+ with DAG("test-dagparam-plain", schedule=None, start_date=datetime(2020,
1, 1)) as dag:
+ MockOperator(task_id="t", arg1=dag.param("subject", "Hi from
Airflow!"))
+
+ serialized = DagSerialization.to_dict(dag)
+ assert "object at 0x" not in json.dumps(serialized).lower()
+ restored = DagSerialization.from_dict(serialized)
+ arg1 = restored.task_dict["t"].arg1
+ assert isinstance(arg1, SerializedDagParam)
+ assert arg1.name == "subject"
+ assert arg1.default == "Hi from Airflow!"
+
+
+def test_dagparam_notset_default_is_not_stringified():
+ with DAG("test-dagparam-notset", schedule=None, start_date=datetime(2020,
1, 1)) as dag:
+ param = dag.param("p")
+ MockOperator.partial(task_id="t", arg1=param).expand(arg2=[1])
+
+ encoded = BaseSerialization.serialize(param, strict=True)
+ assert encoded[Encoding.TYPE] == DagAttributeTypes.DAG_PARAM
+ assert encoded[Encoding.VAR]["default"][Encoding.TYPE] ==
DagAttributeTypes.ARG_NOT_SET
+
+ restored = DagSerialization.from_dict(DagSerialization.to_dict(dag))
+ arg1 = restored.task_dict["t"].partial_kwargs["arg1"]
+ assert isinstance(arg1, SerializedDagParam)
+ assert arg1.default is NOTSET
+ assert arg1.default != "NOTSET"
+
+
+def test_dagparam_jinja_string_in_partial_stays_string():
+ with DAG("test-dagparam-jinja", schedule=None, start_date=datetime(2020,
1, 1)) as dag:
+ MockOperator.partial(task_id="t", arg1="{{ params.p
}}").expand(arg2=[1])
+
+ serialized = DagSerialization.to_dict(dag)
+ assert _encoded_dag_params(serialized) == []
+ restored = DagSerialization.from_dict(serialized)
+ assert restored.task_dict["t"].partial_kwargs["arg1"] == "{{ params.p }}"
+
+
+def test_two_dagparams_in_one_partial():
+ from airflow.sdk import task
+
+ with DAG("test-dagparam-two", schedule=None, start_date=datetime(2020, 1,
1)) as dag:
+
+ @task
+ def add(left, right, extra):
+ return left, right, extra
+
+ add.partial(left=dag.param("left", "L"), right=dag.param("right",
"R")).expand(extra=[1, 2])
+
+ restored = DagSerialization.from_dict(DagSerialization.to_dict(dag))
+ op_kwargs = restored.task_dict["add"].partial_kwargs["op_kwargs"]
+ assert isinstance(op_kwargs["left"], SerializedDagParam)
+ assert isinstance(op_kwargs["right"], SerializedDagParam)
+ assert op_kwargs["left"].name == "left"
+ assert op_kwargs["right"].name == "right"
+
+
+def test_serialized_dagparam_resolve_prefers_dag_run_conf():
+ param = SerializedDagParam(dag_id="d", name="p", default="from_default")
+ context = {
+ "dag_run": type("DR", (), {"conf": {"p": "from_conf"}})(),
+ "params": {"p": "from_params"},
+ }
+ assert param.resolve(context) == "from_conf"
+ context["dag_run"].conf = {}
+ assert param.resolve(context) == "from_default"
+ param_notset = SerializedDagParam(dag_id="d", name="p")
+ assert param_notset.resolve(
+ {"dag_run": type("DR", (), {"conf": {}})(), "params": {"p":
"from_params"}}
+ ) == ("from_params")
+
+
+def test_serialized_dagparam_resolve_skips_conf_when_name_missing():
+ param = SerializedDagParam(dag_id="d", name="p", default="from_default")
+ context = {
+ "dag_run": type("DR", (), {"conf": {"other": "x"}})(),
+ "params": {"p": "from_params"},
+ }
+ assert param.resolve(context) == "from_default"
+
+
+def test_serialized_dagparam_resolve_raises_when_unresolved():
+ param = SerializedDagParam(dag_id="d", name="p")
+ with pytest.raises(RuntimeError, match="No value could be resolved for
parameter p"):
+ param.resolve({"dag_run": type("DR", (), {"conf": {}})(), "params":
{}})
+
+
+def test_dagparam_nested_in_taskflow_call_is_address_stable():
Review Comment:
Good catch. I removed the non-mapped TaskFlow test in `f752022d15`. The
remaining positive regression cases exercise mapped `.partial()` and the new
`dag_param` encoding.
---
Drafted-by: Codex (GPT-5); reviewed by @Vamsi-klu before posting
##########
airflow-core/src/airflow/serialization/definitions/param.py:
##########
@@ -82,6 +83,44 @@ def dump(self) -> dict[str, Any]:
}
+class SerializedDagParam:
+ """
+ Scheduler-side DagParam: a late-bound name, not a schema Param.
+
+ ``resolve()`` matches SDK ``DagParam.resolve``: ``dag_run.conf``, then the
+ serialized default, then ``context["params"]``.
+ """
+
+ def __init__(self, *, dag_id: str, name: str, default: Any = NOTSET):
+ self.dag_id = dag_id
+ self.name = name
+ self.default = default
+
+ def iter_references(self):
+ return ()
+
+ def resolve(self, context: Mapping[str, Any]) -> Any:
Review Comment:
Agreed. I removed both methods and the three direct resolve tests in
`f752022d15`. `SerializedDagParam` is now only a scheduler-side data holder, so
there is no duplicate runtime resolution logic to drift from the Task SDK.
---
Drafted-by: Codex (GPT-5); reviewed by @Vamsi-klu before posting
##########
airflow-core/src/airflow/serialization/definitions/param.py:
##########
@@ -82,6 +83,44 @@ def dump(self) -> dict[str, Any]:
}
+class SerializedDagParam:
Review Comment:
Good suggestion. `SerializedDagParam` now uses `@attrs.define(kw_only=True)`
in `f752022d15`, which gives it structural equality and a useful representation
while keeping it a small data holder.
---
Drafted-by: Codex (GPT-5); reviewed by @Vamsi-klu before posting
##########
airflow-core/src/airflow/serialization/serialized_objects.py:
##########
@@ -1100,7 +1110,10 @@ def _serialize_node(cls, op: SdkOperator) -> dict[str,
Any]:
)
value = getattr(op, template_field, None)
if not cls._is_excluded(value, template_field, op):
- serialize_op[template_field] =
serialize_template_field(value, template_field)
+ if isinstance(value, (DagParam, SerializedDagParam)):
Review Comment:
I agreed with keeping this fix narrow and removed the non-mapped
template-field branch in `f752022d15`. Existing template fields keep their
current serialized shape, and the PR now only changes mapped `.partial()`
values.
---
Drafted-by: Codex (GPT-5); reviewed by @Vamsi-klu before posting
##########
airflow-core/src/airflow/serialization/serialized_objects.py:
##########
@@ -1160,7 +1173,8 @@ def populate_operator(
v = v_in # surpass PLW2901
# Use centralized field deserialization logic
if k in encoded_op.get("template_fields", []):
- pass # Template fields are handled separately
+ if isinstance(v, dict) and v.get(Encoding.TYPE) ==
DAT.DAG_PARAM and Encoding.VAR in v:
Review Comment:
Good point. I removed the paired template-field deserialization change in
`f752022d15`, so this PR no longer introduces the top-level versus nested
asymmetry for non-mapped template fields.
---
Drafted-by: Codex (GPT-5); reviewed by @Vamsi-klu before posting
##########
airflow-core/src/airflow/serialization/serialized_objects.py:
##########
@@ -611,6 +611,10 @@ def serialize(
return cls._encode(cls._serialize_param(var), type_=DAT.PARAM)
elif isinstance(var, XComArg):
return cls._encode(serialize_xcom_arg(var), type_=DAT.XCOM_REF)
+ elif isinstance(var, (DagParam, SerializedDagParam)):
Review Comment:
You were right. I narrowed the serialization branch to `DagParam` only in
`f752022d15`, so an already-deserialized `SerializedDagParam` is no longer
treated as a supported serialization input.
---
Drafted-by: Codex (GPT-5); reviewed by @Vamsi-klu before posting
--
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]
