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)
