def get_task_instances_batch(session=None):
"""
Get list of task instances.
"""
body = request.get_json()
try:
data = task_instance_batch_form.load(body)
except ValidationError as err:
raise BadRequest(detail=str(err.messages))
base_query = session.query(TI)
base_query = _apply_array_filter(base_query,
key=TI.dag_id,
values=data["dag_ids"])
base_query = _apply_range_filter(
base_query,
key=TI.execution_date,
value_range=(data["execution_date_gte"], data["execution_date_lte"]),
)
base_query = _apply_range_filter(
base_query,
key=TI.start_date,
value_range=(data["start_date_gte"], data["start_date_lte"]),
)
base_query = _apply_range_filter(base_query,
key=TI.end_date,
value_range=(data["end_date_gte"],
data["end_date_lte"]))
base_query = _apply_range_filter(base_query,
key=TI.duration,
value_range=(data["duration_gte"],
data["duration_lte"]))
base_query = _apply_array_filter(base_query,
key=TI.state,
values=data["state"])
base_query = _apply_array_filter(base_query,
key=TI.pool,
values=data["pool"])
base_query = _apply_array_filter(base_query,
key=TI.queue,
values=data["queue"])
# Count elements before joining extra columns
total_entries = base_query.with_entities(func.count('*')).scalar()
# Add join
base_query = base_query.join(
SlaMiss,
and_(
SlaMiss.dag_id == TI.dag_id,
SlaMiss.task_id == TI.task_id,
SlaMiss.execution_date == TI.execution_date,
),
isouter=True,
)
ti_query = base_query.add_entity(SlaMiss)
task_instances = ti_query.all()
return task_instance_collection_schema.dump(
TaskInstanceCollection(task_instances=task_instances,
total_entries=total_entries))
def get_task_instances(
limit: int,
dag_id: Optional[str] = None,
dag_run_id: Optional[str] = None,
execution_date_gte: Optional[str] = None,
execution_date_lte: Optional[str] = None,
start_date_gte: Optional[str] = None,
start_date_lte: Optional[str] = None,
end_date_gte: Optional[str] = None,
end_date_lte: Optional[str] = None,
duration_gte: Optional[float] = None,
duration_lte: Optional[float] = None,
state: Optional[List[str]] = None,
pool: Optional[List[str]] = None,
queue: Optional[List[str]] = None,
offset: Optional[int] = None,
session=None,
):
"""Get list of task instances."""
# Because state can be 'none'
states = _convert_state(state)
base_query = session.query(TI).join(TI.dag_run)
if dag_id != "~":
base_query = base_query.filter(TI.dag_id == dag_id)
if dag_run_id != "~":
base_query = base_query.filter(TI.run_id == dag_run_id)
base_query = _apply_range_filter(
base_query,
key=DR.execution_date,
value_range=(execution_date_gte, execution_date_lte),
)
base_query = _apply_range_filter(
base_query, key=TI.start_date, value_range=(start_date_gte, start_date_lte)
)
base_query = _apply_range_filter(base_query, key=TI.end_date, value_range=(end_date_gte, end_date_lte))
base_query = _apply_range_filter(base_query, key=TI.duration, value_range=(duration_gte, duration_lte))
base_query = _apply_array_filter(base_query, key=TI.state, values=states)
base_query = _apply_array_filter(base_query, key=TI.pool, values=pool)
base_query = _apply_array_filter(base_query, key=TI.queue, values=queue)
# Count elements before joining extra columns
total_entries = base_query.with_entities(func.count('*')).scalar()
# Add join
base_query = base_query.join(
SlaMiss,
and_(
SlaMiss.dag_id == TI.dag_id,
SlaMiss.task_id == TI.task_id,
SlaMiss.execution_date == DR.execution_date,
),
isouter=True,
)
ti_query = base_query.add_entity(SlaMiss)
task_instances = ti_query.offset(offset).limit(limit).all()
return task_instance_collection_schema.dump(
TaskInstanceCollection(task_instances=task_instances, total_entries=total_entries)
)
def get_mapped_task_instances(
*,
dag_id: str,
dag_run_id: str,
task_id: str,
execution_date_gte: Optional[str] = None,
execution_date_lte: Optional[str] = None,
start_date_gte: Optional[str] = None,
start_date_lte: Optional[str] = None,
end_date_gte: Optional[str] = None,
end_date_lte: Optional[str] = None,
duration_gte: Optional[float] = None,
duration_lte: Optional[float] = None,
state: Optional[List[str]] = None,
pool: Optional[List[str]] = None,
queue: Optional[List[str]] = None,
limit: Optional[int] = None,
offset: Optional[int] = None,
order_by: Optional[str] = None,
session: Session = NEW_SESSION,
) -> APIResponse:
"""Get list of task instances."""
# Because state can be 'none'
states = _convert_state(state)
base_query = (session.query(TI).filter(TI.dag_id == dag_id,
TI.run_id == dag_run_id,
TI.task_id == task_id,
TI.map_index >= 0).join(TI.dag_run))
# 0 can mean a mapped TI that expanded to an empty list, so it is not an automatic 404
if base_query.with_entities(func.count('*')).scalar() == 0:
dag = current_app.dag_bag.get_dag(dag_id)
if not dag:
error_message = f"DAG {dag_id} not found"
raise NotFound(error_message)
task = dag.get_task(task_id)
if not task:
error_message = f"Task id {task_id} not found"
raise NotFound(error_message)
if not task.is_mapped:
error_message = f"Task id {task_id} is not mapped"
raise NotFound(error_message)
# Other search criteria
query = _apply_range_filter(
base_query,
key=DR.execution_date,
value_range=(execution_date_gte, execution_date_lte),
)
query = _apply_range_filter(query,
key=TI.start_date,
value_range=(start_date_gte, start_date_lte))
query = _apply_range_filter(query,
key=TI.end_date,
value_range=(end_date_gte, end_date_lte))
query = _apply_range_filter(query,
key=TI.duration,
value_range=(duration_gte, duration_lte))
query = _apply_array_filter(query, key=TI.state, values=states)
query = _apply_array_filter(query, key=TI.pool, values=pool)
query = _apply_array_filter(query, key=TI.queue, values=queue)
# Count elements before joining extra columns
total_entries = query.with_entities(func.count('*')).scalar()
# Add SLA miss
query = (query.join(
SlaMiss,
and_(
SlaMiss.dag_id == TI.dag_id,
SlaMiss.task_id == TI.task_id,
SlaMiss.execution_date == DR.execution_date,
),
isouter=True,
).add_entity(SlaMiss).options(joinedload(
TI.rendered_task_instance_fields)))
if order_by:
if order_by == 'state':
query = query.order_by(TI.state.asc(), TI.map_index.asc())
elif order_by == '-state':
query = query.order_by(TI.state.desc(), TI.map_index.asc())
elif order_by == '-map_index':
query = query.order_by(TI.map_index.desc())
else:
raise BadRequest(
detail=f"Ordering with '{order_by}' is not supported")
else:
query = query.order_by(TI.map_index.asc())
task_instances = query.offset(offset).limit(limit).all()
return task_instance_collection_schema.dump(
TaskInstanceCollection(task_instances=task_instances,
total_entries=total_entries))
