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.
  • Loading branch information
jedcunningham committed May 18, 2022
1 parent 5e3f652 commit 3849ebb
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 @@ -2835,3 +2835,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 3849ebb

Please sign in to comment.