def get_xcom_entries(dag_id: str, dag_run_id: str, task_id: str, session: Session, limit: Optional[int], offset: Optional[int] = None) -> XComCollectionSchema: """ Get all XCom values """ query = session.query(XCom) if dag_id != '~': query = query.filter(XCom.dag_id == dag_id) query.join( DR, and_(XCom.dag_id == DR.dag_id, XCom.execution_date == DR.execution_date)) else: query.join(DR, XCom.execution_date == DR.execution_date) if task_id != '~': query = query.filter(XCom.task_id == task_id) if dag_run_id != '~': query = query.filter(DR.run_id == dag_run_id) query = query.order_by(XCom.execution_date, XCom.task_id, XCom.dag_id, XCom.key) total_entries = session.query(func.count(XCom.key)).scalar() query = query.offset(offset).limit(limit) return xcom_collection_schema.dump( XComCollection(xcom_entries=query.all(), total_entries=total_entries))
def get_xcom_entries(dag_id: str, dag_run_id: str, task_id: str, session: Session) -> XComCollectionSchema: """ Get all XCom values """ offset = request.args.get(parameters.page_offset, 0) limit = min(int(request.args.get(parameters.page_limit, 100)), 100) query = session.query(XCom) if dag_id != '~': query = query.filter(XCom.dag_id == dag_id) query.join( DR, and_(XCom.dag_id == DR.dag_id, XCom.execution_date == DR.execution_date)) else: query.join(DR, XCom.execution_date == DR.execution_date) if task_id != '~': query = query.filter(XCom.task_id == task_id) if dag_run_id != '~': query = query.filter(DR.run_id == dag_run_id) query = query.order_by(XCom.execution_date, XCom.task_id, XCom.dag_id, XCom.key) total_entries = session.query(func.count(XCom.key)).scalar() query = query.offset(offset).limit(limit) return xcom_collection_schema.dump( XComCollection(xcom_entries=query.all(), total_entries=total_entries))
def get_xcom_entries( dag_id: str, dag_run_id: str, task_id: str, session: Session, limit: Optional[int], offset: Optional[int] = None, ) -> XComCollectionSchema: """Get all XCom values""" query = session.query(XCom) if dag_id == '~': appbuilder = current_app.appbuilder readable_dag_ids = appbuilder.sm.get_readable_dag_ids(g.user) query = query.filter(XCom.dag_id.in_(readable_dag_ids)) query = query.join( DR, and_(XCom.dag_id == DR.dag_id, XCom.execution_date == DR.execution_date)) else: query = query.filter(XCom.dag_id == dag_id) query = query.join( DR, and_(XCom.dag_id == DR.dag_id, XCom.execution_date == DR.execution_date)) if task_id != '~': query = query.filter(XCom.task_id == task_id) if dag_run_id != '~': query = query.filter(DR.run_id == dag_run_id) query = query.order_by(XCom.execution_date, XCom.task_id, XCom.dag_id, XCom.key) total_entries = query.count() query = query.offset(offset).limit(limit) return xcom_collection_schema.dump( XComCollection(xcom_entries=query.all(), total_entries=total_entries))
def test_serialize(self, session): xcom_model_1 = XCom( key='test_key_1', timestamp=self.time_1, execution_date=self.time_1, task_id='test_task_id_1', dag_id='test_dag_1', ) xcom_model_2 = XCom( key='test_key_2', timestamp=self.time_2, execution_date=self.time_2, task_id='test_task_id_2', dag_id='test_dag_2', ) xcom_models = [xcom_model_1, xcom_model_2] session.add_all(xcom_models) session.commit() xcom_models_query = session.query(XCom).filter( or_(XCom.execution_date == self.time_1, XCom.execution_date == self.time_2)) xcom_models_queried = xcom_models_query.all() deserialized_xcoms = xcom_collection_schema.dump( XComCollection( xcom_entries=xcom_models_queried, total_entries=xcom_models_query.count(), )) self.assertEqual( deserialized_xcoms, { 'xcom_entries': [ { 'key': 'test_key_1', 'timestamp': self.default_time_1, 'execution_date': self.default_time_1, 'task_id': 'test_task_id_1', 'dag_id': 'test_dag_1', }, { 'key': 'test_key_2', 'timestamp': self.default_time_2, 'execution_date': self.default_time_2, 'task_id': 'test_task_id_2', 'dag_id': 'test_dag_2', }, ], 'total_entries': len(xcom_models), }, )