def test_xcom_get_one_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_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_one_enable_pickle_type(self): json_obj = {"key": "value"} execution_date = timezone.utcnow() key = "xcom_test3" dag_id = "test_dag" task_id = "test_task3" 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) 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_link(self, operator: BaseOperator, dttm: datetime): pipelines_conf = XCom.get_one( dag_id=operator.dag.dag_id, task_id=operator.task_id, execution_date=dttm, key=DataFusionPipelinesLink.key, ) return (DATAFUSION_PIPELINES_LINK.format(uri=pipelines_conf["uri"], ) if pipelines_conf else "")
def get_link(self, operator, dttm): # Fetch the correct execution date for the triggerED dag which is # stored in xcom during execution of the triggerING task. trigger_execution_date_iso = XCom.get_one( execution_date=dttm, key=XCOM_EXECUTION_DATE_ISO, task_id=operator.task_id, dag_id=operator.dag_id ) query = {"dag_id": operator.trigger_dag_id, "base_date": trigger_execution_date_iso} return build_airflow_url_with_query(query)
def get_link(self, operator: BaseOperator, dttm: datetime): storage_conf = XCom.get_one( dag_id=operator.dag.dag_id, task_id=operator.task_id, execution_date=dttm, key=StorageLink.key, ) return (GCS_STORAGE_LINK.format( uri=storage_conf["uri"], project_id=storage_conf["project_id"], ) if storage_conf else "")
def get_link(self, operator: BaseOperator, dttm: datetime): instance_conf = XCom.get_one( dag_id=operator.dag.dag_id, task_id=operator.task_id, execution_date=dttm, key=DataFusionInstanceLink.key, ) return (DATAFUSION_INSTANCE_LINK.format( region=instance_conf["region"], instance_name=instance_conf["instance_name"], project_id=instance_conf["project_id"], ) if instance_conf else "")
def get_link( self, operator, dttm: Optional[datetime] = None, ti_key: Optional["TaskInstanceKey"] = None, ) -> str: if ti_key: conf = XCom.get_one(key=self.key, ti_key=ti_key) else: assert dttm conf = XCom.get_one( dag_id=operator.dag.dag_id, task_id=operator.task_id, execution_date=dttm, key=self.key, ) return (conf["url"].format( region=conf["region"], service_id=conf["service_id"], project_id=conf["project_id"], ) if conf else "")
def get_link( self, operator: "AbstractOperator", *, ti_key: "TaskInstanceKey", ) -> str: # Fetch the correct execution date for the triggerED dag which is # stored in xcom during execution of the triggerING task. when = XCom.get_one(ti_key=ti_key, key=XCOM_EXECUTION_DATE_ISO) query = { "dag_id": cast(TriggerDagRunOperator, operator).trigger_dag_id, "base_date": when } return build_airflow_url_with_query(query)
def get_link(self, operator: BaseOperator, dttm: datetime): pipelines_conf = XCom.get_one( key=VertexAITrainingPipelinesLink.key, dag_id=operator.dag.dag_id, task_id=operator.task_id, execution_date=dttm, ) return ( VERTEX_AI_TRAINING_PIPELINES_LINK.format( project_id=pipelines_conf["project_id"], ) if pipelines_conf else "" )
def get_link(self, operator: BaseOperator, dttm: datetime): datasets_conf = XCom.get_one( key=VertexAIDatasetListLink.key, dag_id=operator.dag.dag_id, task_id=operator.task_id, execution_date=dttm, ) return ( VERTEX_AI_DATASET_LIST_LINK.format( project_id=datasets_conf["project_id"], ) if datasets_conf else "" )
def get_link(self, operator: BaseOperator, dttm: datetime): conf = XCom.get_one( dag_id=operator.dag.dag_id, task_id=operator.task_id, execution_date=dttm, key=DataprocMetastoreLink.key, ) return ( conf["url"].format( region=conf["region"], service_id=conf["service_id"], project_id=conf["project_id"], ) if conf else "" )
def get_link(self, operator: BaseOperator, dttm: datetime): model_conf = XCom.get_one( key=VertexAIModelLink.key, dag_id=operator.dag.dag_id, task_id=operator.task_id, execution_date=dttm, ) return ( VERTEX_AI_MODEL_LINK.format( region=model_conf["region"], model_id=model_conf["model_id"], project_id=model_conf["project_id"], ) if model_conf else "" )
def test_xcom_deserialize_with_pickle_to_json_switch(self): json_obj = {"key": "value"} execution_date = timezone.utcnow() key = "xcom_test3" dag_id = "test_dag" task_id = "test_task3" with conf_vars({("core", "enable_xcom_pickling"): "True"}): XCom.set(key=key, value=json_obj, dag_id=dag_id, task_id=task_id, execution_date=execution_date) with conf_vars({("core", "enable_xcom_pickling"): "False"}): ret_value = XCom.get_one(key=key, dag_id=dag_id, task_id=task_id, execution_date=execution_date) assert ret_value == json_obj