muditagrawal2007 commented on PR #74920:
URL: https://github.com/apache/airflow/pull/74920#issuecomment-6238181300
Thanks for the feedback — both review comments (the mapped-task and
serialization coverage) along with @abhijeets25012-tech's note have been
addressed in the just-pushed commit. A summary of what changed and why:
**Mapped-task test** (`task-sdk/tests/task_sdk/bases/test_operator.py`)
- Added `test_template_fields_extra_propagates_through_mapped_operator`. It
builds an operator with `template_fields_extra`, runs `.partial().expand()`,
asserts `MappedOperator.template_fields_extra` is preserved and
`template_fields` is the merged `[arg1, arg2]` order, and asserts
`unmap(...).render_template_fields(...)` substitutes Jinja in both the original
and the extras fields. This covers the `@task`-decorated mapped-task case
indirectly via the same `MappedOperator` codepath and directly via
`partial().expand()`.
- The four pre-existing tests were edited to declare an `__init__` on the
test operator classes — the `BaseOperatorMeta._apply_defaults` validator
rejects unknown kwargs, so passing `arg1`/`arg2`/`arg3` directly was tripping
`TypeError: Invalid arguments were passed … **kwargs: {'arg1': …}`. The
original PR had been raising this same exception in CI; my fix unblocks the
tests.
**Serialization round-trip tests**
(`airflow-core/tests/unit/serialization/test_dag_serialization.py`)
- `test_serialize_template_fields_extra_roundtrip` exercises
`SerializedBaseOperator` directly: it asserts `template_fields_extra` appears
in the serialized blob, that the extras show up in the merged
`template_fields`, that deserialize yields a `SerializedBaseOperator` whose
`template_fields_extra` equals the original, and that the merged
`template_fields` contains the extras.
- `test_serialize_template_fields_extra_roundtrip_mapped` covers the
`SerializedMappedOperator` path — it asserts the merged `template_fields`
survives through `serialize_mapped_operator` ⇄ `deserialize_operator`.
**Incidental fixes the round-trip tests surfaced** (included in the same
commit)
1. `BaseSerialization.serialize` now handles `frozenset` (the type of
`template_fields_extra`) the same way as `set`. Previously a frozenset was
str()-ified to `'frozenset({'arg3'})'` and the round-trip was unrecoverable.
2. `DecoratedMappedOperator.__init__` now receives `template_fields_extra`
alongside the merged `template_fields` — without this, every `@task.mapped`
call errored with `TypeError: DecoratedMappedOperator.__init__() missing 1
required keyword-only argument: 'template_fields_extra'`. This makes the
taskflow-decorated mapped path symmetric with the classic `OperatorPartial`
path.
3. `SerializedMappedOperator` gained `template_fields_extra: Collection[str]
= ()`, was added to `get_serialized_fields()`, and `deserialize_operator`
propagates it. The PR description claimed mapped and serialized tasks
round-trip the override; that wasn't true for `SerializedMappedOperator` until
now.
4. `dag_version_diff._DIFF_V1_PUBLIC_TASK_FIELDS`,
`test_dag_serialization.validate_deserialized_task`, and
`test_no_new_fields_added_to_base_operator` were updated to acknowledge the new
field — adding a serialized field otherwise fails the schema-tracking and
no-new-fields tripwires (CI otherwise fails on a different test, but it would
have caught the next maintainer).
5. `schema.json` now declares `template_fields_extra` next to
`template_fields` so the operator schema validator sees it.
`additionalProperties: true` means it isn't strictly required, but it's the
right thing to do for documentation and downstream consumers.
**On the two follow-up review comments specifically:**
- *Ordering and deduplication*: `OperatorPartial.__call__` and
`DecoratedMappedOperator.__init__` use `dict.fromkeys((*template_fields,
*template_fields_extra))`, which both preserves first-seen order and removes
duplicates — so a class with `template_fields = ("a", "b")` and
`template_fields_extra = frozenset({"b", "c"})` ends up with `["a", "b", "c"]`.
Happy to add an explicit test for that case if you'd like; the merge order in
`test_template_fields_extra_class_level_merges_into_template_fields` already
covers the case where `template_fields` comes first.
- *`__init__` on the test helpers*: added to all four (and the new fifth)
classes so `BaseOperatorMeta._apply_defaults` doesn't reject the constructor
kwargs. Agreed module-level would be cleaner; for now the helpers stay inline
because each test pins a slightly different
`template_fields`/`template_fields_extra` combination.
All new and updated tests pass under `uv run --project task-sdk pytest` and
`uv run --project airflow-core pytest airflow-core/tests/unit/serialization/`.
Thanks again for the review!
--
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]
