def expand_mapped_task(
self, run_id: str, *,
session: Session) -> Tuple[Sequence["TaskInstance"], int]:
"""Create the mapped task instances for mapped task.
:return: The newly created mapped TaskInstances (if any) in ascending order by map index, and the
maximum map_index.
"""
from airflow.models.taskinstance import TaskInstance
from airflow.settings import task_instance_mutation_hook
total_length: Optional[int]
try:
total_length = self._get_specified_expand_input(
).get_total_map_length(run_id, session=session)
except NotFullyPopulated as e:
self.log.info(
"Cannot expand %r for run %s; missing upstream values: %s",
self,
run_id,
sorted(e.missing),
)
total_length = None
state: Optional[TaskInstanceState] = None
unmapped_ti: Optional[TaskInstance] = (
session.query(TaskInstance).filter(
TaskInstance.dag_id == self.dag_id,
TaskInstance.task_id == self.task_id,
TaskInstance.run_id == run_id,
TaskInstance.map_index == -1,
or_(TaskInstance.state.in_(State.unfinished),
TaskInstance.state.is_(None)),
).one_or_none())
all_expanded_tis: List[TaskInstance] = []
if unmapped_ti:
# The unmapped task instance still exists and is unfinished, i.e. we
# haven't tried to run it before.
if total_length is None:
# If the map length cannot be calculated (due to unavailable
# upstream sources), fail the unmapped task.
unmapped_ti.state = TaskInstanceState.UPSTREAM_FAILED
indexes_to_map: Iterable[int] = ()
elif total_length < 1:
# If the upstream maps this to a zero-length value, simply mark
# the unmapped task instance as SKIPPED (if needed).
self.log.info(
"Marking %s as SKIPPED since the map has %d values to expand",
unmapped_ti,
total_length,
)
unmapped_ti.state = TaskInstanceState.SKIPPED
indexes_to_map = ()
else:
# Otherwise convert this into the first mapped index, and create
# TaskInstance for other indexes.
unmapped_ti.map_index = 0
self.log.debug("Updated in place to become %s", unmapped_ti)
all_expanded_tis.append(unmapped_ti)
indexes_to_map = range(1, total_length)
state = unmapped_ti.state
elif not total_length:
# Nothing to fixup.
indexes_to_map = ()
else:
# Only create "missing" ones.
current_max_mapping = (session.query(
func.max(TaskInstance.map_index)).filter(
TaskInstance.dag_id == self.dag_id,
TaskInstance.task_id == self.task_id,
TaskInstance.run_id == run_id,
).scalar())
indexes_to_map = range(current_max_mapping + 1, total_length)
for index in indexes_to_map:
# TODO: Make more efficient with bulk_insert_mappings/bulk_save_mappings.
ti = TaskInstance(self,
run_id=run_id,
map_index=index,
state=state)
self.log.debug("Expanding TIs upserted %s", ti)
task_instance_mutation_hook(ti)
ti = session.merge(ti)
ti.refresh_from_task(
self) # session.merge() loses task information.
all_expanded_tis.append(ti)
# Coerce the None case to 0 -- these two are almost treated identically,
# except the unmapped ti (if exists) is marked to different states.
total_expanded_ti_count = total_length or 0
# Set to "REMOVED" any (old) TaskInstances with map indices greater
# than the current map value
session.query(TaskInstance).filter(
TaskInstance.dag_id == self.dag_id,
TaskInstance.task_id == self.task_id,
TaskInstance.run_id == run_id,
TaskInstance.map_index >= total_expanded_ti_count,
).update({TaskInstance.state: TaskInstanceState.REMOVED})
session.flush()
return all_expanded_tis, total_expanded_ti_count - 1
def expand_mapped_task(
self, run_id: str, *,
session: Session) -> Tuple[Sequence["TaskInstance"], int]:
"""Create the mapped task instances for mapped task.
:return: The newly created mapped TaskInstances (if any) in ascending order by map index, and the
maximum map_index.
"""
from airflow.models.taskinstance import TaskInstance
from airflow.settings import task_instance_mutation_hook
total_length = functools.reduce(
operator.mul,
self._resolve_map_lengths(run_id, session=session).values())
state: Optional[TaskInstanceState] = None
unmapped_ti: Optional[TaskInstance] = (
session.query(TaskInstance).filter(
TaskInstance.dag_id == self.dag_id,
TaskInstance.task_id == self.task_id,
TaskInstance.run_id == run_id,
TaskInstance.map_index == -1,
or_(TaskInstance.state.in_(State.unfinished),
TaskInstance.state.is_(None)),
).one_or_none())
all_expanded_tis: List[TaskInstance] = []
if unmapped_ti:
# The unmapped task instance still exists and is unfinished, i.e. we
# haven't tried to run it before.
if total_length < 1:
# If the upstream maps this to a zero-length value, simply marked the
# unmapped task instance as SKIPPED (if needed).
self.log.info(
"Marking %s as SKIPPED since the map has %d values to expand",
unmapped_ti,
total_length,
)
unmapped_ti.state = TaskInstanceState.SKIPPED
else:
# Otherwise convert this into the first mapped index, and create
# TaskInstance for other indexes.
unmapped_ti.map_index = 0
self.log.debug("Updated in place to become %s", unmapped_ti)
all_expanded_tis.append(unmapped_ti)
state = unmapped_ti.state
indexes_to_map = range(1, total_length)
else:
# Only create "missing" ones.
current_max_mapping = (session.query(
func.max(TaskInstance.map_index)).filter(
TaskInstance.dag_id == self.dag_id,
TaskInstance.task_id == self.task_id,
TaskInstance.run_id == run_id,
).scalar())
indexes_to_map = range(current_max_mapping + 1, total_length)
for index in indexes_to_map:
# TODO: Make more efficient with bulk_insert_mappings/bulk_save_mappings.
ti = TaskInstance(self,
run_id=run_id,
map_index=index,
state=state)
self.log.debug("Expanding TIs upserted %s", ti)
task_instance_mutation_hook(ti)
ti = session.merge(ti)
ti.refresh_from_task(
self) # session.merge() loses task information.
all_expanded_tis.append(ti)
# Set to "REMOVED" any (old) TaskInstances with map indices greater
# than the current map value
session.query(TaskInstance).filter(
TaskInstance.dag_id == self.dag_id,
TaskInstance.task_id == self.task_id,
TaskInstance.run_id == run_id,
TaskInstance.map_index >= total_length,
).update({TaskInstance.state: TaskInstanceState.REMOVED})
session.flush()
return all_expanded_tis, total_length
