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