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
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 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 pull_xcom_call(**kwargs):
# xcom will get all values, that was written before this date with using Xcom directly (without context object)
get_mane_xcoms_values__with_xcom_class = XCom.get_many(
execution_date=make_aware(datetime(2020, 27, 4, 14, 50, 00, 00)),
dag_ids=["write_to_xcom"],
include_prior_dates=True)
print('XCom.get_many ')
print(get_mane_xcoms_values__with_xcom_class)
get_xcom_with_ti = kwargs['ti'].xcom_pull(dag_id="write_to_xcom",
include_prior_dates=True)
print('ti.xcom_pull with include_prior_dates')
print(get_xcom_with_ti)
def test_pull_parameters_by_task(env, create_dag, create_df_task, push_xcom):
dag = create_dag
# df_task = create_df_task(dag)
push_task = PythonOperator(task_id='three',
python_callable=push_xcom,
op_args=['three_value'],
dag=dag,
provide_context=True)
push_ti = TaskInstance(task=push_task, execution_date=datetime.now())
push_task.execute(push_ti.get_template_context())
xcom = XCom.get_many(execution_date=push_ti.execution_date,
dag_ids=[dag.dag_id],
task_ids=[push_task.task_id])
pprint(xcom[0])
assert len(xcom) == 1
assert xcom[0].value == 'three_value'
assert False
def test_pull_parameters_by_specific_key(env, create_dag, create_df_task,
push_xcom_key):
dag = create_dag
df_task = create_df_task(dag)
push_task = PythonOperator(task_id='two',
python_callable=push_xcom_key,
op_args=['two', 'two_value'],
dag=dag,
provide_context=True)
push_ti = TaskInstance(task=push_task, execution_date=datetime.now())
push_task.execute(push_ti.get_template_context())
xcom = XCom.get_many(execution_date=push_ti.execution_date,
dag_ids=[dag.dag_id])
assert len(xcom) == 1
assert xcom[0].value == 'two_value'
merged = df_task.merge_parameters(push_ti)
assert merged['two-specific-key'] == 'two_value'
def _get_task_instance_xcom_dict(dag_id, task_id, execution_date):
try:
results = XCom.get_many(dag_ids=dag_id,
task_ids=task_id,
execution_date=execution_date)
if not results:
return {}
xcom_dict = {xcom.key: str(xcom.value) for xcom in results}
sliced_xcom = (dict(
itertools.islice(xcom_dict.items(), MAX_XCOM_LENGTH))
if len(xcom_dict) > MAX_XCOM_LENGTH else xcom_dict)
for key, value in six.iteritems(sliced_xcom):
sliced_xcom[key] = shorten_xcom_value(value)
return sliced_xcom
except Exception as e:
logging.info("Failed to get xcom dict. Exception: {}".format(e))
return {}
