def get_lineage(dag_id: str,
execution_date: datetime.datetime,
session=None) -> Dict[str, Dict[str, Any]]:
"""
Gets the lineage information for dag specified
"""
dag = check_and_get_dag(dag_id)
check_and_get_dagrun(dag, execution_date)
inlets: List[XCom] = XCom.get_many(dag_ids=dag_id,
execution_date=execution_date,
key=PIPELINE_INLETS,
session=session).all()
outlets: List[XCom] = XCom.get_many(dag_ids=dag_id,
execution_date=execution_date,
key=PIPELINE_OUTLETS,
session=session).all()
lineage: Dict[str, Dict[str, Any]] = {}
for meta in inlets:
lineage[meta.task_id] = {'inlets': meta.value}
for meta in outlets:
lineage[meta.task_id]['outlets'] = meta.value
return {'task_ids': lineage}
def get_task_instance(dag_id: str, task_id: str, execution_date: datetime) -> TaskInstance:
"""Return the task instance identified by the given dag_id, task_id and execution_date."""
dag = check_and_get_dag(dag_id, task_id)
dagrun = check_and_get_dagrun(dag=dag, execution_date=execution_date)
# Get task instance object and check that it exists
task_instance = dagrun.get_task_instance(task_id)
if not task_instance:
error_message = f'Task {task_id} instance for date {execution_date} not found'
raise TaskInstanceNotFound(error_message)
return task_instance
def get_dag_run_state(dag_id: str, execution_date: datetime) -> Dict[str, str]:
"""Return the Dag Run state identified by the given dag_id and execution_date.
:param dag_id: DAG id
:param execution_date: execution date
:return: Dictionary storing state of the object
"""
dag = check_and_get_dag(dag_id=dag_id)
dagrun = check_and_get_dagrun(dag, execution_date)
return {'state': dagrun.get_state()}
def get_dag_run_state(
dag_id, execution_date): # type: (str, datetime) -> Dict[str, str]
"""Return the task object identified by the given dag_id and task_id.
:param dag_id: DAG id
:param execution_date: execution date
:return: Dictionary storing state of the object
"""
dag = check_and_get_dag(dag_id=dag_id)
dagrun = check_and_get_dagrun(dag, execution_date)
return {'state': dagrun.get_state()}
def get_lineage(dag_id: str, execution_date: datetime.datetime, *, session) -> Dict[str, Dict[str, Any]]:
"""Gets the lineage information for dag specified."""
dag = check_and_get_dag(dag_id)
dagrun = check_and_get_dagrun(dag, execution_date)
inlets = XCom.get_many(dag_ids=dag_id, run_id=dagrun.run_id, key=PIPELINE_INLETS, session=session)
outlets = XCom.get_many(dag_ids=dag_id, run_id=dagrun.run_id, key=PIPELINE_OUTLETS, session=session)
lineage: Dict[str, Dict[str, Any]] = collections.defaultdict(dict)
for meta in inlets:
lineage[meta.task_id]["inlets"] = meta.value
for meta in outlets:
lineage[meta.task_id]["outlets"] = meta.value
return {"task_ids": {k: v for k, v in lineage.items()}}
def get_task_instance(
dag_id, task_id,
execution_date): # type: (str, str, datetime) -> TaskInstance
"""Return the task object identified by the given dag_id and task_id."""
dag = check_and_get_dag(dag_id, task_id)
dagrun = check_and_get_dagrun(dag=dag, execution_date=execution_date)
# 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