potiuk commented on code in PR #72242:
URL: https://github.com/apache/airflow/pull/72242#discussion_r3887356677
##########
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:
This is the actual fix and it looks right. Serializing `default` recursively
rather than passing it through is the detail that matters: `NOTSET` round-trips
as `arg_not_set` instead of the literal string `"NOTSET"`. Identity holds too -
`task-sdk/.../_internal/types.py:38` re-exports core's `NOTSET` singleton when
core is installed, so the `is not NOTSET` check compares the same object
`deserialize()` returns.
##########
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:
Suggest dropping this hunk (and its pair in `populate_operator`). The issue
is about `partial_kwargs`; template fields don't have the version-inflation bug.
On `main` today a directly-assigned DagParam already serializes stably,
because `serialize_template_field` reaches `DagParam.serialize` through its
`callable(inspect.getattr_static(obj, "serialize", None))` branch
(`helpers.py:80`). I checked:
```
main: MockOperator(task_id="t", arg1=dag.param("p","d"))
-> "arg1": {"dag_id": "...", "default": "d", "name": "p"} # no
memory address
```
So there's no stability gain here - but there is a cost. With the hunk, the
deserialized scheduler-side attribute is an object rather than JSON:
```python
json.dumps(restored.task_dict["t"].arg1)
# TypeError: Object of type SerializedDagParam is not JSON serializable
```
Making template fields JSON-encodable is `serialize_template_field`'s stated
first responsibility (`helpers.py:43`). I traced the consumers and did *not*
find a live break - `get_serialized_template_fields()`
(`renderedtifields.py:85`) takes a `SerializedBaseOperator` and re-enters
`serialize_template_field`, which handles the new object through that same
`.serialize()` branch. So this isn't a bug today; it just widens the blast
radius of a bugfix PR.
It also only tags a *directly* assigned DagParam, leaving an asymmetry:
```
arg1=dag.param("p","d") -> {"__type": "dag_param", ...} ->
SerializedDagParam
arg1=[dag.param("p","d")] -> [{"dag_id": ..., ...}] -> plain dict
arg1={"k": dag.param(...)} -> {"k": {...}} -> plain dict
```
If template fields should be tagged, it needs to recurse uniformly and the
encodability contract needs addressing explicitly - worth its own PR either way.
Minor: `SerializedDagParam` in this isinstance tuple is unreachable -
`_serialize_node` takes an `SdkOperator`.
##########
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):
Review Comment:
Both of these methods are unreachable in airflow-core; suggest dropping them
and keeping `SerializedDagParam` as a plain data holder.
- `resolve()`: nothing resolves scheduler-side objects - the worker
re-parses the Dag file with the SDK, so the real `DagParam.resolve` runs there.
The precedent is right next door: the whole `SchedulerXComArg` family in
`definitions/xcom_arg.py` deliberately has no `resolve()`, only
`iter_references()` and map-length helpers. I also confirmed scheduler-side
`partial_kwargs` is read only for scheduling attributes (`owner`, `retries`,
`pool`, ...) in `definitions/mappedoperator.py:146-262`; user kwargs are never
touched.
- `iter_references()`: `SchedulerXComArg.iter_xcom_references`
(`xcom_arg.py:74`) only dispatches to `arg.iter_references()` for
`ReferenceMixin` instances, and `SerializedDagParam` isn't one.
As written it's ~12 lines hand-copied from `DagParam.resolve` with nothing
keeping the two in sync, and the three `test_serialized_dagparam_resolve_*`
tests exercise only themselves. If the resolve semantics are deliberately kept
for something planned, a comment saying so would help - otherwise the next
reader will reasonably assume the scheduler resolves params.
Also: `iter_references` is missing a return annotation, unlike the rest of
the file.
##########
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:
Non-blocking, just so it's a conscious call: a new `__type` value means an
older Airflow reading a newly written blob hits `TypeError: Invalid type
dag_param in deserialization`, and `SERIALIZER_VERSION` stays at 3. Consistent
with how other `DAT` members were added, and Airflow expects components on one
version - flagging only because rolling upgrades touch this path.
##########
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:
Pair of the comment above - suggest dropping this too. It also only unwraps
a DagParam at the top level of a template field, so `arg1=[dag.param(...)]`
stays a plain dict while `arg1=dag.param(...)` becomes a `SerializedDagParam`.
##########
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:
This test passes on unmodified `main` - I ran it. `do(dag.param(...))` is a
non-mapped taskflow task, so `op_kwargs` goes through
`serialize_template_field`, which was already address-stable. It pins
pre-existing behaviour rather than this change.
Per AGENTS.md ("every test must fail without the PR's change") this should
either go or be reshaped into something that does fail without the fix - e.g.
asserting the `dag_param` type tag.
Same applies more mildly to
`test_dagparam_jinja_string_in_partial_stays_string` (already noted in the PR
body); that one reads as a deliberate negative guard, so it's more defensible.
--
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]
