def create_ti(task: "Operator", indexes: Tuple[int,
...]) -> Generator:
for map_index in indexes:
ti = TI(task, run_id=self.run_id, map_index=map_index)
task_instance_mutation_hook(ti)
created_counts[ti.operator] += 1
yield ti
def verify_integrity(self, session: Session = None):
"""
Verifies the DagRun by checking for removed tasks or tasks that are not in the
database yet. It will set state to removed or add the task if required.
:param session: Sqlalchemy ORM Session
:type session: Session
"""
dag = self.get_dag()
tis = self.get_task_instances(session=session)
# check for removed or restored tasks
task_ids = set()
for ti in tis:
task_instance_mutation_hook(ti)
task_ids.add(ti.task_id)
task = None
try:
task = dag.get_task(ti.task_id)
except AirflowException:
if ti.state == State.REMOVED:
pass # ti has already been removed, just ignore it
elif self.state is not State.RUNNING and not dag.partial:
self.log.warning(
"Failed to get task '%s' for dag '%s'. "
"Marking it as removed.", ti, dag)
Stats.incr("task_removed_from_dag.{}".format(dag.dag_id),
1, 1)
ti.state = State.REMOVED
should_restore_task = (task
is not None) and ti.state == State.REMOVED
if should_restore_task:
self.log.info(
"Restoring task '%s' which was previously "
"removed from DAG '%s'", ti, dag)
Stats.incr("task_restored_to_dag.{}".format(dag.dag_id), 1, 1)
ti.state = State.NONE
session.merge(ti)
# check for missing tasks
for task in dag.task_dict.values():
if task.start_date > self.execution_date and not self.is_backfill:
continue
if task.task_id not in task_ids:
Stats.incr("task_instance_created-{}".format(task.task_type),
1, 1)
ti = TI(task, self.execution_date)
task_instance_mutation_hook(ti)
session.add(ti)
try:
session.commit()
except IntegrityError as err:
self.log.info(str(err))
self.log.info('Hit IntegrityError while creating the TIs for '
f'{dag.dag_id} - {self.execution_date}.')
self.log.info('Doing session rollback.')
session.rollback()
def verify_integrity(self, session=None):
"""
Verifies the DagRun by checking for removed tasks or tasks that are not in the
database yet. It will set state to removed or add the task if required.
"""
from airflow.models.taskinstance import TaskInstance # Avoid circular import
dag = self.get_dag()
tis = self.get_task_instances(session=session)
# check for removed or restored tasks
task_ids = set()
for ti in tis:
task_instance_mutation_hook(ti)
task_ids.add(ti.task_id)
task = None
try:
task = dag.get_task(ti.task_id)
except AirflowException:
if ti.state == State.REMOVED:
pass # ti has already been removed, just ignore it
elif self.state is not State.RUNNING and not dag.partial:
self.log.warning("Failed to get task '{}' for dag '{}'. "
"Marking it as removed.".format(ti, dag))
Stats.incr(
"task_removed_from_dag.{}".format(dag.dag_id), 1, 1)
ti.state = State.REMOVED
is_task_in_dag = task is not None
should_restore_task = is_task_in_dag and ti.state == State.REMOVED
if should_restore_task:
self.log.info("Restoring task '{}' which was previously "
"removed from DAG '{}'".format(ti, dag))
Stats.incr("task_restored_to_dag.{}".format(dag.dag_id), 1, 1)
ti.state = State.NONE
session.merge(ti)
# check for missing tasks
for task in six.itervalues(dag.task_dict):
if task.start_date > self.execution_date and not self.is_backfill:
continue
if task.task_id not in task_ids:
Stats.incr(
"task_instance_created-{}".format(task.__class__.__name__),
1, 1)
ti = TaskInstance(task, self.execution_date)
task_instance_mutation_hook(ti)
session.add(ti)
try:
session.commit()
except IntegrityError as err:
self.log.info(str(err))
self.log.info(
'Hit IntegrityError while creating the TIs for %s - %s',
dag.dag_id, self.execution_date
)
self.log.info('Doing session rollback.')
session.rollback()
def verify_integrity(self, session: Session = NEW_SESSION):
"""
Verifies the DagRun by checking for removed tasks or tasks that are not in the
database yet. It will set state to removed or add the task if required.
:param session: Sqlalchemy ORM Session
:type session: Session
"""
from airflow.settings import task_instance_mutation_hook
dag = self.get_dag()
tis = self.get_task_instances(session=session)
# check for removed or restored tasks
task_ids = set()
for ti in tis:
task_instance_mutation_hook(ti)
task_ids.add(ti.task_id)
task = None
try:
task = dag.get_task(ti.task_id)
except AirflowException:
if ti.state == State.REMOVED:
pass # ti has already been removed, just ignore it
elif self.state != State.RUNNING and not dag.partial:
self.log.warning("Failed to get task '%s' for dag '%s'. Marking it as removed.", ti, dag)
Stats.incr(f"task_removed_from_dag.{dag.dag_id}", 1, 1)
ti.state = State.REMOVED
should_restore_task = (task is not None) and ti.state == State.REMOVED
if should_restore_task:
self.log.info("Restoring task '%s' which was previously removed from DAG '%s'", ti, dag)
Stats.incr(f"task_restored_to_dag.{dag.dag_id}", 1, 1)
ti.state = State.NONE
session.merge(ti)
# check for missing tasks
for task in dag.task_dict.values():
if task.start_date > self.execution_date and not self.is_backfill:
continue
if task.task_id not in task_ids:
Stats.incr(f"task_instance_created-{task.task_type}", 1, 1)
ti = TI(task, execution_date=None, run_id=self.run_id)
task_instance_mutation_hook(ti)
session.add(ti)
try:
session.flush()
except IntegrityError as err:
self.log.info(str(err))
self.log.info('Hit IntegrityError while creating the TIs for %s- %s', dag.dag_id, self.run_id)
self.log.info('Doing session rollback.')
# TODO[HA]: We probably need to savepoint this so we can keep the transaction alive.
session.rollback()
def verify_integrity(self, session=None):
"""
Verifies the DagRun by checking for removed tasks or tasks that are not in the
database yet. It will set state to removed or add the task if required.
"""
dag = self.get_dag()
tis = self.get_task_instances(session=session)
# check for removed or restored tasks
task_ids = []
for ti in tis:
task_instance_mutation_hook(ti)
task_ids.append(ti.task_id)
task = None
try:
task = dag.get_task(ti.task_id)
except AirflowException:
if ti.state == State.REMOVED:
pass # ti has already been removed, just ignore it
elif self.state is not State.RUNNING and not dag.partial:
self.log.warning("Failed to get task '{}' for dag '{}'. "
"Marking it as removed.".format(ti, dag))
Stats.incr("task_removed_from_dag.{}".format(dag.dag_id),
1, 1)
ti.state = State.REMOVED
should_restore_task = (task
is not None) and ti.state == State.REMOVED
if should_restore_task:
self.log.info("Restoring task '{}' which was previously "
"removed from DAG '{}'".format(ti, dag))
Stats.incr("task_restored_to_dag.{}".format(dag.dag_id), 1, 1)
ti.state = State.NONE
session.merge(ti)
# check for missing tasks
for task in dag.task_dict.values():
if task.start_date > self.execution_date and not self.is_backfill:
continue
if task.task_id not in task_ids:
Stats.incr(
"task_instance_created-{}".format(task.__class__.__name__),
1, 1)
ti = TI(task, self.execution_date)
task_instance_mutation_hook(ti)
session.add(ti)
session.commit()
def expand_mapped_task(self, run_id: str, *,
session: Session) -> Sequence["TaskInstance"]:
"""Create the mapped task instances for mapped task.
:return: The mapped task instances, in ascending order by map index.
"""
from airflow.models.taskinstance import TaskInstance
from airflow.settings import task_instance_mutation_hook
total_length = functools.reduce(
operator.mul,
self._get_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())
ret: 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
session.flush()
return ret
# Otherwise convert this into the first mapped index, and create
# TaskInstance for other indexes.
unmapped_ti.map_index = 0
state = unmapped_ti.state
self.log.debug("Updated in place to become %s", unmapped_ti)
ret.append(unmapped_ti)
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.
# TODO: Change `TaskInstance` ctor to take Operator, not BaseOperator
ti = TaskInstance(self,
run_id=run_id,
map_index=index,
state=state) # type: ignore
self.log.debug("Expanding TIs upserted %s", ti)
task_instance_mutation_hook(ti)
ti = session.merge(ti)
ti.task = self
ret.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 ret
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 create_ti(task: "BaseOperator") -> TI:
ti = TI(task, run_id=self.run_id)
task_instance_mutation_hook(ti)
created_counts[ti.operator] += 1
return ti
def verify_integrity(self, session: Session = NEW_SESSION):
"""
Verifies the DagRun by checking for removed tasks or tasks that are not in the
database yet. It will set state to removed or add the task if required.
:param session: Sqlalchemy ORM Session
:type session: Session
"""
from airflow.settings import task_instance_mutation_hook
dag = self.get_dag()
tis = self.get_task_instances(session=session)
# check for removed or restored tasks
task_ids = set()
for ti in tis:
task_instance_mutation_hook(ti)
task_ids.add(ti.task_id)
task = None
try:
task = dag.get_task(ti.task_id)
except AirflowException:
if ti.state == State.REMOVED:
pass # ti has already been removed, just ignore it
elif self.state != State.RUNNING and not dag.partial:
self.log.warning(
"Failed to get task '%s' for dag '%s'. Marking it as removed.",
ti, dag)
Stats.incr(f"task_removed_from_dag.{dag.dag_id}", 1, 1)
ti.state = State.REMOVED
should_restore_task = (task
is not None) and ti.state == State.REMOVED
if should_restore_task:
self.log.info(
"Restoring task '%s' which was previously removed from DAG '%s'",
ti, dag)
Stats.incr(f"task_restored_to_dag.{dag.dag_id}", 1, 1)
ti.state = State.NONE
session.merge(ti)
def task_filter(task: "BaseOperator"):
return task.task_id not in task_ids and (
self.is_backfill or task.start_date <= self.execution_date)
created_counts: Dict[str, int] = defaultdict(int)
# Set for the empty default in airflow.settings -- if it's not set this means it has been changed
hook_is_noop = getattr(task_instance_mutation_hook, 'is_noop', False)
if hook_is_noop:
def create_ti_mapping(task: "BaseOperator"):
created_counts[task.task_type] += 1
return TI.insert_mapping(self.run_id, task)
else:
def create_ti(task: "BaseOperator") -> TI:
ti = TI(task, run_id=self.run_id)
task_instance_mutation_hook(ti)
created_counts[ti.operator] += 1
return ti
# Create missing tasks
tasks = list(filter(task_filter, dag.task_dict.values()))
try:
if hook_is_noop:
session.bulk_insert_mappings(TI, map(create_ti_mapping, tasks))
else:
session.bulk_save_objects(map(create_ti, tasks))
for task_type, count in created_counts.items():
Stats.incr(f"task_instance_created-{task_type}", count)
session.flush()
except IntegrityError as err:
self.log.info(str(err))
self.log.info(
'Hit IntegrityError while creating the TIs for %s- %s',
dag.dag_id, self.run_id)
self.log.info('Doing session rollback.')
# TODO[HA]: We probably need to savepoint this so we can keep the transaction alive.
session.rollback()
def verify_integrity(self, session: Session = NEW_SESSION):
"""
Verifies the DagRun by checking for removed tasks or tasks that are not in the
database yet. It will set state to removed or add the task if required.
:param session: Sqlalchemy ORM Session
"""
from airflow.settings import task_instance_mutation_hook
dag = self.get_dag()
tis = self.get_task_instances(session=session)
# check for removed or restored tasks
task_ids = set()
for ti in tis:
task_instance_mutation_hook(ti)
task_ids.add(ti.task_id)
task = None
try:
task = dag.get_task(ti.task_id)
should_restore_task = (
task is not None) and ti.state == State.REMOVED
if should_restore_task:
self.log.info(
"Restoring task '%s' which was previously removed from DAG '%s'",
ti, dag)
Stats.incr(f"task_restored_to_dag.{dag.dag_id}", 1, 1)
ti.state = State.NONE
except AirflowException:
if ti.state == State.REMOVED:
pass # ti has already been removed, just ignore it
elif self.state != State.RUNNING and not dag.partial:
self.log.warning(
"Failed to get task '%s' for dag '%s'. Marking it as removed.",
ti, dag)
Stats.incr(f"task_removed_from_dag.{dag.dag_id}", 1, 1)
ti.state = State.REMOVED
continue
if not task.is_mapped:
continue
task = cast("MappedOperator", task)
num_mapped_tis = task.parse_time_mapped_ti_count
# Check if the number of mapped literals has changed and we need to mark this TI as removed
if num_mapped_tis is not None:
if ti.map_index >= num_mapped_tis:
self.log.debug(
"Removing task '%s' as the map_index is longer than the literal mapping list (%s)",
ti,
num_mapped_tis,
)
ti.state = State.REMOVED
elif ti.map_index < 0:
self.log.debug(
"Removing the unmapped TI '%s' as the mapping can now be performed",
ti)
ti.state = State.REMOVED
else:
self.log.info("Restoring mapped task '%s'", ti)
Stats.incr(f"task_restored_to_dag.{dag.dag_id}", 1, 1)
ti.state = State.NONE
else:
# What if it is _now_ dynamically mapped, but wasn't before?
total_length = task.run_time_mapped_ti_count(self.run_id,
session=session)
if total_length is None:
# Not all upstreams finished, so we can't tell what should be here. Remove everything.
if ti.map_index >= 0:
self.log.debug(
"Removing the unmapped TI '%s' as the mapping can't be resolved yet",
ti)
ti.state = State.REMOVED
continue
# Upstreams finished, check there aren't any extras
if ti.map_index >= total_length:
self.log.debug(
"Removing task '%s' as the map_index is longer than the resolved mapping list (%d)",
ti,
total_length,
)
ti.state = State.REMOVED
...
def task_filter(task: "Operator") -> bool:
return task.task_id not in task_ids and (
self.is_backfill or task.start_date <= self.execution_date and
(task.end_date is None
or self.execution_date <= task.end_date))
created_counts: Dict[str, int] = defaultdict(int)
# Set for the empty default in airflow.settings -- if it's not set this means it has been changed
hook_is_noop = getattr(task_instance_mutation_hook, 'is_noop', False)
if hook_is_noop:
def create_ti_mapping(task: "Operator",
indexes: Tuple[int, ...]) -> Generator:
created_counts[task.task_type] += 1
for map_index in indexes:
yield TI.insert_mapping(self.run_id,
task,
map_index=map_index)
creator = create_ti_mapping
else:
def create_ti(task: "Operator", indexes: Tuple[int,
...]) -> Generator:
for map_index in indexes:
ti = TI(task, run_id=self.run_id, map_index=map_index)
task_instance_mutation_hook(ti)
created_counts[ti.operator] += 1
yield ti
creator = create_ti
# Create missing tasks -- and expand any MappedOperator that _only_ have literals as input
def expand_mapped_literals(
task: "Operator") -> Tuple["Operator", Sequence[int]]:
if not task.is_mapped:
return (task, (-1, ))
task = cast("MappedOperator", task)
count = task.parse_time_mapped_ti_count or task.run_time_mapped_ti_count(
self.run_id, session=session)
if not count:
return (task, (-1, ))
return (task, range(count))
tasks_and_map_idxs = map(expand_mapped_literals,
filter(task_filter, dag.task_dict.values()))
tasks = itertools.chain.from_iterable(
itertools.starmap(creator, tasks_and_map_idxs))
try:
if hook_is_noop:
session.bulk_insert_mappings(TI, tasks)
else:
session.bulk_save_objects(tasks)
for task_type, count in created_counts.items():
Stats.incr(f"task_instance_created-{task_type}", count)
session.flush()
except IntegrityError:
self.log.info(
'Hit IntegrityError while creating the TIs for %s- %s',
dag.dag_id,
self.run_id,
exc_info=True,
)
self.log.info('Doing session rollback.')
# TODO[HA]: We probably need to savepoint this so we can keep the transaction alive.
session.rollback()