def get_xcom_entry(dag_id: str, task_id: str, dag_run_id: str, xcom_key: str,
session: Session) -> XComCollectionItemSchema:
"""Get an XCom entry"""
query = session.query(XCom).filter(XCom.dag_id == dag_id,
XCom.task_id == task_id,
XCom.key == xcom_key)
query = query.join(
DR,
and_(XCom.dag_id == DR.dag_id,
XCom.execution_date == DR.execution_date))
query = query.filter(DR.run_id == dag_run_id)
query_object = query.one_or_none()
if not query_object:
raise NotFound("XCom entry not found")
return xcom_collection_item_schema.dump(query_object)
def test_serialize(self, session):
xcom_model = XCom(
key='test_key',
timestamp=self.default_time_parsed,
execution_date=self.default_time_parsed,
task_id='test_task_id',
dag_id='test_dag',
)
session.add(xcom_model)
session.commit()
xcom_model = session.query(XCom).first()
deserialized_xcom = xcom_collection_item_schema.dump(xcom_model)
assert deserialized_xcom == {
'key': 'test_key',
'timestamp': self.default_time,
'execution_date': self.default_time,
'task_id': 'test_task_id',
'dag_id': 'test_dag',
}