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"
configuration.set("core", "xcom_enable_pickling", "True")
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_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_enable_pickle_type(self):
json_obj = {"key": "value"}
execution_date = timezone.utcnow()
key = "xcom_test2"
dag_id = "test_dag2"
task_id = "test_task2"
configuration.set("core", "enable_xcom_pickling", "True")
XCom.set(key=key,
value=json_obj,
dag_id=dag_id,
task_id=task_id,
execution_date=execution_date)
ret_value = XCom.get_one(key=key,
dag_id=dag_id,
task_id=task_id,
execution_date=execution_date)
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_enable_pickle_type(self):
json_obj = {"key": "value"}
execution_date = timezone.utcnow()
key = "xcom_test2"
dag_id = "test_dag2"
task_id = "test_task2"
configuration.set("core", "enable_xcom_pickling", "True")
XCom.set(key=key,
value=json_obj,
dag_id=dag_id,
task_id=task_id,
execution_date=execution_date)
ret_value = XCom.get_one(key=key,
dag_id=dag_id,
task_id=task_id,
execution_date=execution_date)
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"
configuration.set("core", "xcom_enable_pickling", "True")
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 _create_xcom_entries(self,
dag_id,
dag_run_id,
execution_date,
task_id,
session=None):
for i in [1, 2]:
XCom.set(
key=f'test-xcom-key-{i}',
value="TEST",
execution_date=execution_date,
task_id=task_id,
dag_id=dag_id,
)
dag = DagModel(dag_id=dag_id)
session.add(dag)
dagrun = DR(
dag_id=dag_id,
run_id=dag_run_id,
execution_date=execution_date,
start_date=execution_date,
run_type=DagRunType.MANUAL.value,
)
session.add(dagrun)
def _create_invalid_xcom_entries(self, execution_date, session=None):
"""
Invalid XCom entries to test join query
"""
for i in [1, 2]:
XCom.set(
key=f'invalid-xcom-key-{i}',
value="TEST",
execution_date=execution_date,
task_id="invalid_task",
dag_id="invalid_dag",
)
dag = DagModel(dag_id="invalid_dag")
session.add(dag)
dagrun = DR(
dag_id="invalid_dag",
run_id="invalid_run_id",
execution_date=execution_date + timedelta(days=1),
start_date=execution_date,
run_type=DagRunType.MANUAL,
)
session.add(dagrun)
dagrun = DR(
dag_id="invalid_dag_1",
run_id="invalid_run_id",
execution_date=execution_date,
start_date=execution_date,
run_type=DagRunType.MANUAL,
)
session.commit()
def _xcom_set(self, context, key, value, task_id):
XCom.set(
key=key,
value=value,
task_id=task_id,
dag_id=self.dag_id,
execution_date=context["ti"].execution_date,
)
def _xcom_set(self, context, key, value, task_id):
XCom.set(
key=key,
value=value,
task_id=task_id,
dag_id=self.dag_id,
run_id=context["run_id"],
)
def run_airflow_branch_task(dag_id: str, task_id: str, batch: Any, **context):
batch_num = 1
dag_context = make_typhoon_dag_context(context)
XCom.set(
key=f'{uuid4()}:{batch_num}',
value=(batch_num, batch),
execution_date=dag_context.execution_time,
task_id=task_id,
dag_id=dag_id,
)
def _create_xcom_entry(self, dag_id, dag_run_id, execution_date, task_id, xcom_key, session=None):
XCom.set(key=xcom_key,
value="TEST_VALUE",
execution_date=execution_date,
task_id=task_id,
dag_id=dag_id,)
dagrun = DR(dag_id=dag_id,
run_id=dag_run_id,
execution_date=execution_date,
start_date=execution_date,
run_type=DagRunType.MANUAL.value)
session.add(dagrun)
def send(
self,
dag_context: DagContext,
source_id: str,
destination: 'TaskInterface',
batch_num: int,
batch: Any,
batch_group_id: str,
):
XCom.set(
key=f'{batch_group_id}:{batch_num:04d}',
value=(batch_num, batch),
execution_date=dag_context.execution_time,
task_id=source_id,
dag_id=self.dag_id,
)
