def get_task_instance( *, dag_id: str, dag_run_id: str, task_id: str, session: Session = NEW_SESSION, ) -> APIResponse: """Get task instance""" query = (session.query(TI).filter( TI.dag_id == dag_id, DR.run_id == dag_run_id, TI.task_id == task_id).join(TI.dag_run).outerjoin( SlaMiss, and_( SlaMiss.dag_id == TI.dag_id, SlaMiss.execution_date == DR.execution_date, SlaMiss.task_id == TI.task_id, ), ).add_entity(SlaMiss)) query = query.outerjoin( RTIF, and_( RTIF.dag_id == TI.dag_id, RTIF.execution_date == DR.execution_date, RTIF.task_id == TI.task_id, ), ).add_entity(RTIF) task_instance = query.one_or_none() if task_instance is None: raise NotFound("Task instance not found") return task_instance_schema.dump(task_instance)
def test_task_instance_schema_without_sla(self, session): ti = TI(task=self.task, **self.default_ti_init) for key, value in self.default_ti_extras.items(): setattr(ti, key, value) session.add(ti) session.commit() serialized_ti = task_instance_schema.dump((ti, None)) expected_json = { "dag_id": "TEST_DAG_ID", "duration": 10000.0, "end_date": "2020-01-03T00:00:00+00:00", "execution_date": "2020-01-01T00:00:00+00:00", "executor_config": "{}", "hostname": "", "max_tries": 0, "operator": "DummyOperator", "pid": 100, "pool": "default_pool", "pool_slots": 1, "priority_weight": 1, "queue": "default_queue", "queued_when": None, "sla_miss": None, "start_date": "2020-01-02T00:00:00+00:00", "state": "running", "task_id": "TEST_TASK_ID", "try_number": 0, "unixname": getuser(), } assert serialized_ti == expected_json
def get_mapped_task_instance( *, dag_id: str, dag_run_id: str, task_id: str, map_index: int, session: Session = NEW_SESSION, ) -> APIResponse: """Get task instance""" query = (session.query(TI).filter( TI.dag_id == dag_id, TI.run_id == dag_run_id, TI.task_id == task_id, TI.map_index == map_index).join(TI.dag_run).outerjoin( SlaMiss, and_( SlaMiss.dag_id == TI.dag_id, SlaMiss.execution_date == DR.execution_date, SlaMiss.task_id == TI.task_id, ), ).add_entity(SlaMiss).options( joinedload(TI.rendered_task_instance_fields))) task_instance = query.one_or_none() if task_instance is None: raise NotFound("Task instance not found") return task_instance_schema.dump(task_instance)
def test_task_instance_schema_with_sla(self, session): ti = TI(task=self.task, **self.default_ti_init) for key, value in self.default_ti_extras.items(): setattr(ti, key, value) sla_miss = SlaMiss( task_id="TEST_TASK_ID", dag_id="TEST_DAG_ID", execution_date=self.default_time, ) session.add(ti) session.add(sla_miss) session.commit() serialized_ti = task_instance_schema.dump((ti, sla_miss)) expected_json = { "dag_id": "TEST_DAG_ID", "duration": 10000.0, "end_date": "2020-01-03T00:00:00+00:00", "execution_date": "2020-01-01T00:00:00+00:00", "executor_config": "{}", "hostname": "", "max_tries": 0, "operator": "DummyOperator", "pid": 100, "pool": "default_pool", "pool_slots": 1, "priority_weight": 1, "queue": "default_queue", "queued_when": None, "sla_miss": { "dag_id": "TEST_DAG_ID", "description": None, "email_sent": False, "execution_date": "2020-01-01T00:00:00+00:00", "notification_sent": False, "task_id": "TEST_TASK_ID", "timestamp": None, }, "start_date": "2020-01-02T00:00:00+00:00", "state": "running", "task_id": "TEST_TASK_ID", "try_number": 0, "unixname": getpass.getuser(), } self.assertDictEqual(serialized_ti, expected_json)
def get_task_instance(dag_id: str, dag_run_id: str, task_id: str, session=None): """Get task instance""" query = (session.query(TI).filter( TI.dag_id == dag_id, DR.run_id == dag_run_id, TI.task_id == task_id).join(TI.dag_run).options(eagerload( TI.dag_run)).outerjoin( SlaMiss, and_( SlaMiss.dag_id == TI.dag_id, SlaMiss.execution_date == DR.execution_date, SlaMiss.task_id == TI.task_id, ), ).add_entity(SlaMiss)) task_instance = query.one_or_none() if task_instance is None: raise NotFound("Task instance not found") return task_instance_schema.dump(task_instance)
def get_task_instance( *, dag_id: str, dag_run_id: str, task_id: str, session: Session = NEW_SESSION, ) -> APIResponse: """Get task instance""" query = ( session.query(TI) .filter(TI.dag_id == dag_id, TI.run_id == dag_run_id, TI.task_id == task_id) .join(TI.dag_run) .outerjoin( SlaMiss, and_( SlaMiss.dag_id == TI.dag_id, SlaMiss.execution_date == DR.execution_date, SlaMiss.task_id == TI.task_id, ), ) .add_entity(SlaMiss) .options(joinedload(TI.rendered_task_instance_fields)) ) try: task_instance = query.one_or_none() except MultipleResultsFound: raise NotFound( "Task instance not found", detail="Task instance is mapped, add the map_index value to the URL" ) if task_instance is None: raise NotFound("Task instance not found") if task_instance[0].map_index != -1: raise NotFound( "Task instance not found", detail="Task instance is mapped, add the map_index value to the URL" ) return task_instance_schema.dump(task_instance)