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
