def get_dag_run_state(self):
if self.run_id is None:
raise DagRunNotFound(f"dag_id {self.dag_name}; {self.run_id}")
dag_runs = DagRun.find(dag_id=self.dag_name, run_id=self.run_id)
return dag_runs[0].state if dag_runs else None
def get_task_instance(dag_id, task_id, execution_date):
"""Return the task object identified by the given dag_id and task_id."""
dagbag = DagBag()
# Check DAG exists.
if dag_id not in dagbag.dags:
error_message = "Dag id {} not found".format(dag_id)
raise DagNotFound(error_message)
# Get DAG object and check Task Exists
dag = dagbag.get_dag(dag_id)
if not dag.has_task(task_id):
error_message = 'Task {} not found in dag {}'.format(task_id, dag_id)
raise TaskNotFound(error_message)
# Get DagRun object and check that it exists
dagrun = dag.get_dagrun(execution_date=execution_date)
if not dagrun:
error_message = ('Dag Run for date {} not found in dag {}'
.format(execution_date, dag_id))
raise DagRunNotFound(error_message)
# Get task instance object and check that it exists
task_instance = dagrun.get_task_instance(task_id)
if not task_instance:
error_message = ('Task {} instance for date {} not found'
.format(task_id, execution_date))
raise TaskInstanceNotFound(error_message)
return task_instance
def task_states_for_dag_run(args, session=None):
"""Get the status of all task instances in a DagRun"""
dag_run = (session.query(DagRun).filter(
DagRun.run_id == args.execution_date_or_run_id,
DagRun.dag_id == args.dag_id).one_or_none())
if not dag_run:
try:
execution_date = timezone.parse(args.execution_date_or_run_id)
dag_run = (session.query(DagRun).filter(
DagRun.execution_date == execution_date,
DagRun.dag_id == args.dag_id).one_or_none())
except (ParserError, TypeError) as err:
raise AirflowException(
f"Error parsing the supplied execution_date. Error: {str(err)}"
)
if dag_run is None:
raise DagRunNotFound(
f"DagRun for {args.dag_id} with run_id or execution_date of {args.execution_date_or_run_id!r} "
"not found")
AirflowConsole().print_as(
data=dag_run.task_instances,
output=args.output,
mapper=lambda ti: {
"dag_id": ti.dag_id,
"execution_date": dag_run.execution_date.isoformat(),
"task_id": ti.task_id,
"state": ti.state,
"start_date": ti.start_date.isoformat() if ti.start_date else "",
"end_date": ti.end_date.isoformat() if ti.end_date else "",
},
)
def _get_dag_run(dag, exec_date_or_run_id, create_if_necessary, session):
dag_run = dag.get_dagrun(run_id=exec_date_or_run_id, session=session)
if dag_run:
return dag_run
execution_date = None
with suppress(ParserError, TypeError):
execution_date = timezone.parse(exec_date_or_run_id)
if create_if_necessary and not execution_date:
return DagRun(dag_id=dag.dag_id, run_id=exec_date_or_run_id)
try:
return (
session.query(DagRun)
.filter(
DagRun.dag_id == dag.dag_id,
DagRun.execution_date == execution_date,
)
.one()
)
except NoResultFound:
if create_if_necessary:
return DagRun(dag.dag_id, execution_date=execution_date)
raise DagRunNotFound(
f"DagRun for {dag.dag_id} with run_id or execution_date of {exec_date_or_run_id!r} not found"
) from None
def check_and_get_dagrun(dag: DagModel, execution_date: datetime) -> DagRun:
"""Get DagRun object and check that it exists"""
dagrun = dag.get_dagrun(execution_date=execution_date)
if not dagrun:
error_message = f'Dag Run for date {execution_date} not found in dag {dag.dag_id}'
raise DagRunNotFound(error_message)
return dagrun
def _get_dag_run(
*,
dag: DAG,
exec_date_or_run_id: str,
create_if_necessary: CreateIfNecessary,
session: Session,
) -> Tuple[DagRun, bool]:
"""Try to retrieve a DAG run from a string representing either a run ID or logical date.
This checks DAG runs like this:
1. If the input ``exec_date_or_run_id`` matches a DAG run ID, return the run.
2. Try to parse the input as a date. If that works, and the resulting
date matches a DAG run's logical date, return the run.
3. If ``create_if_necessary`` is *False* and the input works for neither of
the above, raise ``DagRunNotFound``.
4. Try to create a new DAG run. If the input looks like a date, use it as
the logical date; otherwise use it as a run ID and set the logical date
to the current time.
"""
dag_run = dag.get_dagrun(run_id=exec_date_or_run_id, session=session)
if dag_run:
return dag_run, False
try:
execution_date: Optional[datetime.datetime] = timezone.parse(exec_date_or_run_id)
except (ParserError, TypeError):
execution_date = None
try:
dag_run = (
session.query(DagRun)
.filter(DagRun.dag_id == dag.dag_id, DagRun.execution_date == execution_date)
.one()
)
except NoResultFound:
if not create_if_necessary:
raise DagRunNotFound(
f"DagRun for {dag.dag_id} with run_id or execution_date of {exec_date_or_run_id!r} not found"
) from None
else:
return dag_run, False
if execution_date is not None:
dag_run_execution_date = execution_date
else:
dag_run_execution_date = timezone.utcnow()
if create_if_necessary == "memory":
dag_run = DagRun(dag.dag_id, run_id=exec_date_or_run_id, execution_date=dag_run_execution_date)
return dag_run, True
elif create_if_necessary == "db":
dag_run = dag.create_dagrun(
state=DagRunState.QUEUED,
execution_date=dag_run_execution_date,
run_id=_generate_temporary_run_id(),
session=session,
)
return dag_run, True
raise ValueError(f"unknown create_if_necessary value: {create_if_necessary!r}")
def check_and_get_dagrun(
dag, execution_date): # type: (DagModel, datetime) -> DagRun
"""Get DagRun object and check that it exists"""
dagrun = dag.get_dagrun(execution_date=execution_date)
if not dagrun:
error_message = ('Dag Run for date {} not found in dag {}'.format(
execution_date, dag.dag_id))
raise DagRunNotFound(error_message)
return dagrun
def _get_dag_run(
*,
dag: DAG,
exec_date_or_run_id: str,
create_if_necessary: bool,
session: Session,
) -> DagRun:
"""Try to retrieve a DAG run from a string representing either a run ID or logical date.
This checks DAG runs like this:
1. If the input ``exec_date_or_run_id`` matches a DAG run ID, return the run.
2. Try to parse the input as a date. If that works, and the resulting
date matches a DAG run's logical date, return the run.
3. If ``create_if_necessary`` is *False* and the input works for neither of
the above, raise ``DagRunNotFound``.
4. Try to create a new DAG run. If the input looks like a date, use it as
the logical date; otherwise use it as a run ID and set the logical date
to the current time.
"""
dag_run = dag.get_dagrun(run_id=exec_date_or_run_id, session=session)
if dag_run:
return dag_run
try:
execution_date: Optional[datetime.datetime] = timezone.parse(
exec_date_or_run_id)
except (ParserError, TypeError):
execution_date = None
try:
return (session.query(DagRun).filter(
DagRun.dag_id == dag.dag_id,
DagRun.execution_date == execution_date).one())
except NoResultFound:
if not create_if_necessary:
raise DagRunNotFound(
f"DagRun for {dag.dag_id} with run_id or execution_date of {exec_date_or_run_id!r} not found"
) from None
if execution_date is not None:
return DagRun(dag.dag_id,
run_id=IN_MEMORY_RUN_ID,
execution_date=execution_date)
return DagRun(dag.dag_id,
run_id=exec_date_or_run_id,
execution_date=timezone.utcnow())
def get_dag_run_state(dag_id, execution_date):
"""Return the task object identified by the given dag_id and task_id."""
dagbag = DagBag()
# Check DAG exists.
if dag_id not in dagbag.dags:
error_message = "Dag id {} not found".format(dag_id)
raise DagNotFound(error_message)
# Get DAG object and check Task Exists
dag = dagbag.get_dag(dag_id)
# Get DagRun object and check that it exists
dagrun = dag.get_dagrun(execution_date=execution_date)
if not dagrun:
error_message = ('Dag Run for date {} not found in dag {}'.format(
execution_date, dag_id))
raise DagRunNotFound(error_message)
return {'state': dagrun.get_state()}
def task_states_for_dag_run(args, session=None):
"""Get the status of all task instances in a DagRun"""
dag_run = (session.query(DagRun).filter(
DagRun.run_id == args.execution_date_or_run_id,
DagRun.dag_id == args.dag_id).one_or_none())
if not dag_run:
try:
execution_date = timezone.parse(args.execution_date_or_run_id)
dag_run = (session.query(DagRun).filter(
DagRun.execution_date == execution_date,
DagRun.dag_id == args.dag_id).one_or_none())
except (ParserError, TypeError) as err:
raise AirflowException(
f"Error parsing the supplied execution_date. Error: {str(err)}"
)
if dag_run is None:
raise DagRunNotFound(
f"DagRun for {args.dag_id} with run_id or execution_date of {args.execution_date_or_run_id!r} "
"not found")
has_mapped_instances = any(ti.map_index >= 0
for ti in dag_run.task_instances)
def format_task_instance(ti: TaskInstance) -> Dict[str, str]:
data = {
"dag_id": ti.dag_id,
"execution_date": dag_run.execution_date.isoformat(),
"task_id": ti.task_id,
"state": ti.state,
"start_date": ti.start_date.isoformat() if ti.start_date else "",
"end_date": ti.end_date.isoformat() if ti.end_date else "",
}
if has_mapped_instances:
data["map_index"] = str(ti.map_index) if ti.map_index >= 0 else ""
return data
AirflowConsole().print_as(data=dag_run.task_instances,
output=args.output,
mapper=format_task_instance)
