Skip to content

Commit

Permalink
Fix auto upstream dep when expanding non-templated field (#23771)
Browse files Browse the repository at this point in the history
If you tried to expand via xcom into a non-templated field without
explicitly setting the upstream task dependency, the scheduler would
crash because the upstream task dependency wasn't being set
automatically. It was being set only for templated fields, but now we do
it for both.

(cherry picked from commit 3849ebb)
  • Loading branch information
jedcunningham authored and ephraimbuddy committed May 19, 2022
1 parent 9920202 commit 27762af
Show file tree
Hide file tree
Showing 2 changed files with 21 additions and 2 deletions.
3 changes: 1 addition & 2 deletions airflow/models/mappedoperator.py
Original file line number Diff line number Diff line change
Expand Up @@ -300,8 +300,7 @@ def __attrs_post_init__(self):
if self.dag:
self.dag.add_task(self)
for k, v in self.mapped_kwargs.items():
if k in self.template_fields:
XComArg.apply_upstream_relationship(self, v)
XComArg.apply_upstream_relationship(self, v)
for k, v in self.partial_kwargs.items():
if k in self.template_fields:
XComArg.apply_upstream_relationship(self, v)
Expand Down
20 changes: 20 additions & 0 deletions tests/models/test_taskinstance.py
Original file line number Diff line number Diff line change
Expand Up @@ -2855,3 +2855,23 @@ def add_one(x):
ti.refresh_from_task(dag.get_task("add_one"))
with pytest.raises(XComForMappingNotPushed):
ti.run()


def test_expand_non_templated_field(dag_maker, session):
"""Test expand on non-templated fields sets upstream deps properly."""

class SimpleBashOperator(BashOperator):
template_fields = ()

with dag_maker(dag_id="product_same_types", session=session) as dag:

@dag.task
def get_extra_env():
return [{"foo": "bar"}, {"foo": "biz"}]

SimpleBashOperator.partial(task_id="echo", bash_command="echo $FOO").expand(env=get_extra_env())

dag_maker.create_dagrun()

echo_task = dag.get_task("echo")
assert "get_extra_env" in echo_task.upstream_task_ids

0 comments on commit 27762af

Please sign in to comment.