uranusjr commented on code in PR #22975:
URL: https://github.com/apache/airflow/pull/22975#discussion_r849632316
##########
airflow/models/mappedoperator.py:
##########
@@ -686,33 +693,40 @@ def render_template_fields(
"""
if not jinja_env:
jinja_env = self.get_template_env()
- unmapped_task = self.unmap()
+ # Before we unmap we have to resolve the mapped arguments, otherwise
the real operator constructor
+ # could be called with an XComArg, rather than the value it resolves
to.
+ #
+ # We also need to resolve _all_ mapped arguments, even if they aren't
marked as templated
+ kwargs = self._get_unmap_kwargs()
+
+ template_fields = set(self.template_fields)
+
+ self._resolve_expansion_kwargs(kwargs, template_fields, context,
session)
+
+ unmapped_task = self.unmap(rendered_kwargs=kwargs)
self._do_render_template_fields(
parent=unmapped_task,
- template_fields=unmapped_task.template_fields,
+ template_fields=template_fields,
Review Comment:
I hope `template_fields` being unordered won’t cause bugs in user code.
##########
airflow/models/mappedoperator.py:
##########
@@ -670,10 +675,12 @@ def prepare_for_execution(self) -> "MappedOperator":
# we don't need to create a copy of the MappedOperator here.
return self
+ @provide_session
def render_template_fields(
self,
context: Context,
jinja_env: Optional["jinja2.Environment"] = None,
+ session: Session = NEW_SESSION,
Review Comment:
We might not be able to do this because `render_template_fields` is public
API that can be overriden. I think `@provide_session` can be put on
`_resolve_expansion_kwargs` to avoid the problem.
##########
airflow/models/mappedoperator.py:
##########
@@ -686,33 +693,40 @@ def render_template_fields(
"""
if not jinja_env:
jinja_env = self.get_template_env()
- unmapped_task = self.unmap()
+ # Before we unmap we have to resolve the mapped arguments, otherwise
the real operator constructor
+ # could be called with an XComArg, rather than the value it resolves
to.
+ #
+ # We also need to resolve _all_ mapped arguments, even if they aren't
marked as templated
+ kwargs = self._get_unmap_kwargs()
+
+ template_fields = set(self.template_fields)
+
+ self._resolve_expansion_kwargs(kwargs, template_fields, context,
session)
+
+ unmapped_task = self.unmap(rendered_kwargs=kwargs)
self._do_render_template_fields(
parent=unmapped_task,
- template_fields=unmapped_task.template_fields,
+ template_fields=template_fields,
context=context,
jinja_env=jinja_env,
seen_oids=set(),
)
return unmapped_task
- def _render_template_field(
Review Comment:
I think we can remove this in AbstractOperator now; the sole reason of this
function’s existence if for MappedOperator to hook into.
##########
airflow/models/mappedoperator.py:
##########
@@ -569,6 +573,7 @@ def _get_map_lengths(self, run_id: str, *, session:
Session) -> Dict[str, int]:
map_lengths[mapped_arg_name] += length
return map_lengths
+ @cache
def _resolve_map_lengths(self, run_id: str, *, session: Session) ->
Dict[str, int]:
Review Comment:
It’ll be difficult to actually hit this cache due to `session` :(
##########
airflow/models/mappedoperator.py:
##########
@@ -484,14 +486,16 @@ def _get_unmap_kwargs(self) -> Dict[str, Any]:
**self.mapped_kwargs,
}
- def unmap(self) -> "BaseOperator":
+ def unmap(self, rendered_kwargs: Optional[Dict[str, Any]] = None) ->
"BaseOperator":
Review Comment:
Probably should document `rendered_kwargs` (and maybe it’s better to call it
`unmap_kwargs` for consistency?)
##########
airflow/models/xcom_arg.py:
##########
@@ -136,12 +139,15 @@ def set_downstream(
"""Proxy to underlying operator set_downstream method. Required by
TaskMixin."""
self.operator.set_downstream(task_or_task_list, edge_modifier)
- def resolve(self, context: Context) -> Any:
+ @provide_session
+ def resolve(self, context: Context, session: "Session" = NEW_SESSION) ->
Any:
Review Comment:
IIRC I wondered if we could do this previously, but eventually did not due
to interface compatibility issues.
--
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]
