def _get_ti(
task: BaseOperator,
exec_date_or_run_id: str,
map_index: int,
*,
create_if_necessary: bool = False,
session: Session = NEW_SESSION,
) -> TaskInstance:
"""Get the task instance through DagRun.run_id, if that fails, get the TI the old way"""
if task.is_mapped:
if map_index < 0:
raise RuntimeError("No map_index passed to mapped task")
elif map_index >= 0:
raise RuntimeError("map_index passed to non-mapped task")
dag_run = _get_dag_run(
dag=task.dag,
exec_date_or_run_id=exec_date_or_run_id,
create_if_necessary=create_if_necessary,
session=session,
)
ti_or_none = dag_run.get_task_instance(task.task_id, map_index=map_index, session=session)
if ti_or_none is None:
if not create_if_necessary:
raise TaskInstanceNotFound(
f"TaskInstance for {task.dag.dag_id}, {task.task_id}, map={map_index} with "
f"run_id or execution_date of {exec_date_or_run_id!r} not found"
)
# TODO: Validate map_index is in range?
ti = TaskInstance(task, run_id=dag_run.run_id, map_index=map_index)
ti.dag_run = dag_run
else:
ti = ti_or_none
ti.refresh_from_task(task)
return ti
def _get_ti(
task: BaseOperator,
exec_date_or_run_id: str,
*,
create_if_necessary: bool = False,
session: Session = NEW_SESSION,
) -> TaskInstance:
"""Get the task instance through DagRun.run_id, if that fails, get the TI the old way"""
dag_run = _get_dag_run(
dag=task.dag,
exec_date_or_run_id=exec_date_or_run_id,
create_if_necessary=create_if_necessary,
session=session,
)
ti_or_none = dag_run.get_task_instance(task.task_id)
if ti_or_none is None:
if not create_if_necessary:
raise TaskInstanceNotFound(
f"TaskInstance for {task.dag.dag_id}, {task.task_id} with "
f"run_id or execution_date of {exec_date_or_run_id!r} not found"
)
ti = TaskInstance(task, run_id=dag_run.run_id)
ti.dag_run = dag_run
else:
ti = ti_or_none
ti.refresh_from_task(task)
return ti
def _get_ti(task, exec_date_or_run_id, create_if_necessary=False, session=None):
"""Get the task instance through DagRun.run_id, if that fails, get the TI the old way"""
dag_run = _get_dag_run(task.dag, exec_date_or_run_id, create_if_necessary, session)
ti = dag_run.get_task_instance(task.task_id)
if not ti and create_if_necessary:
ti = TaskInstance(task, run_id=None)
ti.dag_run = dag_run
ti.refresh_from_task(task)
return ti
def _get_ti(task, exec_date_or_run_id):
"""Get the task instance through DagRun.run_id, if that fails, get the TI the old way"""
dag_run = task.dag.get_dagrun(run_id=exec_date_or_run_id)
if not dag_run:
try:
execution_date = timezone.parse(exec_date_or_run_id)
ti = TaskInstance(task, execution_date)
ti.refresh_from_db()
return ti
except (ParserError, TypeError):
raise AirflowException(
f"DagRun with run_id: {exec_date_or_run_id} not found")
ti = dag_run.get_task_instance(task.task_id)
ti.refresh_from_task(task)
return ti
