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_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 test_xcom_disable_pickle_type(self):
json_obj = {"key": "value"}
execution_date = timezone.utcnow()
key = "xcom_test1"
dag_id = "test_dag1"
task_id = "test_task1"
XCom.set(key=key,
value=json_obj,
dag_id=dag_id,
task_id=task_id,
execution_date=execution_date)
ret_value = XCom.get_many(key=key,
dag_ids=dag_id,
task_ids=task_id,
execution_date=execution_date).first().value
self.assertEqual(ret_value, json_obj)
session = settings.Session()
ret_value = session.query(XCom).filter(XCom.key == key, XCom.dag_id == dag_id,
XCom.task_id == task_id,
XCom.execution_date == execution_date
).first().value
self.assertEqual(ret_value, json_obj)
def test_xcom_get_many(self):
json_obj = {"key": "value"}
execution_date = timezone.utcnow()
key = "xcom_test4"
dag_id1 = "test_dag4"
task_id1 = "test_task4"
dag_id2 = "test_dag5"
task_id2 = "test_task5"
XCom.set(key=key,
value=json_obj,
dag_id=dag_id1,
task_id=task_id1,
execution_date=execution_date)
XCom.set(key=key,
value=json_obj,
dag_id=dag_id2,
task_id=task_id2,
execution_date=execution_date)
results = XCom.get_many(key=key, execution_date=execution_date)
for result in results:
self.assertEqual(result.value, json_obj)
def test_xcom_enable_pickle_type(self):
json_obj = {"key": "value"}
execution_date = timezone.utcnow()
key = "xcom_test2"
dag_id = "test_dag2"
task_id = "test_task2"
XCom.set(key=key,
value=json_obj,
dag_id=dag_id,
task_id=task_id,
execution_date=execution_date)
ret_value = (XCom.get_many(
key=key,
dag_ids=dag_id,
task_ids=task_id,
execution_date=execution_date).first().value)
assert ret_value == json_obj
session = settings.Session()
ret_value = (session.query(XCom).filter(
XCom.key == key,
XCom.dag_id == dag_id,
XCom.task_id == task_id,
XCom.execution_date == execution_date,
).first().value)
assert ret_value == json_obj
def get_xcoms(task_instance):
from airflow.models.xcom import XCom
execution_date = task_instance.execution_date
task_id = task_instance.task_id
dag_id = task_instance.dag_id
results = XCom.get_many(execution_date, task_ids=task_id, dag_ids=dag_id)
return [(xcom.key, str(xcom.value)) for xcom in results]
