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, )