def patch_pool(
*,
pool_name: str,
update_mask: UpdateMask = None,
session: Session = NEW_SESSION,
) -> APIResponse:
"""Update a pool"""
request_dict = get_json_request_dict()
# Only slots can be modified in 'default_pool'
try:
if pool_name == Pool.DEFAULT_POOL_NAME and request_dict[
"name"] != Pool.DEFAULT_POOL_NAME:
if update_mask and len(
update_mask) == 1 and update_mask[0].strip() == "slots":
pass
else:
raise BadRequest(
detail="Default Pool's name can't be modified")
except KeyError:
pass
pool = session.query(Pool).filter(Pool.pool == pool_name).first()
if not pool:
raise NotFound(detail=f"Pool with name:'{pool_name}' not found")
try:
patch_body = pool_schema.load(request_dict)
except ValidationError as err:
raise BadRequest(detail=str(err.messages))
if update_mask:
update_mask = [i.strip() for i in update_mask]
_patch_body = {}
try:
update_mask = [
pool_schema.declared_fields[field].attribute
if pool_schema.declared_fields[field].attribute else field
for field in update_mask
]
except KeyError as err:
raise BadRequest(
detail=f"Invalid field: {err.args[0]} in update mask")
_patch_body = {field: patch_body[field] for field in update_mask}
patch_body = _patch_body
else:
required_fields = {"name", "slots"}
fields_diff = required_fields - set(get_json_request_dict().keys())
if fields_diff:
raise BadRequest(
detail=f"Missing required property(ies): {sorted(fields_diff)}"
)
for key, value in patch_body.items():
setattr(pool, key, value)
session.commit()
return pool_schema.dump(pool)
def update_dag_run_state(*,
dag_id: str,
dag_run_id: str,
session: Session = NEW_SESSION) -> APIResponse:
"""Set a state of a dag run."""
dag_run: Optional[DagRun] = (session.query(DagRun).filter(
DagRun.dag_id == dag_id, DagRun.run_id == dag_run_id).one_or_none())
if dag_run is None:
error_message = f'Dag Run id {dag_run_id} not found in dag {dag_id}'
raise NotFound(error_message)
try:
post_body = set_dagrun_state_form_schema.load(get_json_request_dict())
except ValidationError as err:
raise BadRequest(detail=str(err))
state = post_body['state']
dag = get_airflow_app().dag_bag.get_dag(dag_id)
if state == DagRunState.SUCCESS:
set_dag_run_state_to_success(dag=dag,
run_id=dag_run.run_id,
commit=True)
elif state == DagRunState.QUEUED:
set_dag_run_state_to_queued(dag=dag,
run_id=dag_run.run_id,
commit=True)
else:
set_dag_run_state_to_failed(dag=dag,
run_id=dag_run.run_id,
commit=True)
dag_run = session.query(DagRun).get(dag_run.id)
return dagrun_schema.dump(dag_run)
def post_clear_task_instances(*,
dag_id: str,
session: Session = NEW_SESSION) -> APIResponse:
"""Clear task instances."""
body = get_json_request_dict()
try:
data = clear_task_instance_form.load(body)
except ValidationError as err:
raise BadRequest(detail=str(err.messages))
dag = get_airflow_app().dag_bag.get_dag(dag_id)
if not dag:
error_message = f"Dag id {dag_id} not found"
raise NotFound(error_message)
reset_dag_runs = data.pop('reset_dag_runs')
dry_run = data.pop('dry_run')
# We always pass dry_run here, otherwise this would try to confirm on the terminal!
task_instances = dag.clear(dry_run=True,
dag_bag=get_airflow_app().dag_bag,
**data)
if not dry_run:
clear_task_instances(
task_instances.all(),
session,
dag=dag,
dag_run_state=DagRunState.QUEUED if reset_dag_runs else False,
)
return task_instance_reference_collection_schema.dump(
TaskInstanceReferenceCollection(task_instances=task_instances.all()))
def post_variables() -> Response:
"""Create a variable"""
try:
data = variable_schema.load(get_json_request_dict())
except ValidationError as err:
raise BadRequest("Invalid Variable schema", detail=str(err.messages))
Variable.set(data["key"], data["val"])
return variable_schema.dump(data)
def post_clear_task_instances(*, dag_id: str, session: Session = NEW_SESSION) -> APIResponse:
"""Clear task instances."""
body = get_json_request_dict()
try:
data = clear_task_instance_form.load(body)
except ValidationError as err:
raise BadRequest(detail=str(err.messages))
dag = get_airflow_app().dag_bag.get_dag(dag_id)
if not dag:
error_message = f"Dag id {dag_id} not found"
raise NotFound(error_message)
reset_dag_runs = data.pop('reset_dag_runs')
dry_run = data.pop('dry_run')
# We always pass dry_run here, otherwise this would try to confirm on the terminal!
dag_run_id = data.pop('dag_run_id', None)
future = data.pop('include_future', False)
past = data.pop('include_past', False)
downstream = data.pop('include_downstream', False)
upstream = data.pop('include_upstream', False)
if dag_run_id is not None:
dag_run: Optional[DR] = (
session.query(DR).filter(DR.dag_id == dag_id, DR.run_id == dag_run_id).one_or_none()
)
if dag_run is None:
error_message = f'Dag Run id {dag_run_id} not found in dag {dag_id}'
raise NotFound(error_message)
data['start_date'] = dag_run.logical_date
data['end_date'] = dag_run.logical_date
if past:
data['start_date'] = None
if future:
data['end_date'] = None
task_ids = data.pop('task_ids', None)
if task_ids is not None:
task_id = [task[0] if isinstance(task, tuple) else task for task in task_ids]
dag = dag.partial_subset(
task_ids_or_regex=task_id,
include_downstream=downstream,
include_upstream=upstream,
)
if len(dag.task_dict) > 1:
# If we had upstream/downstream etc then also include those!
task_ids.extend(tid for tid in dag.task_dict if tid != task_id)
task_instances = dag.clear(dry_run=True, dag_bag=get_airflow_app().dag_bag, task_ids=task_ids, **data)
if not dry_run:
clear_task_instances(
task_instances.all(),
session,
dag=dag,
dag_run_state=DagRunState.QUEUED if reset_dag_runs else False,
)
return task_instance_reference_collection_schema.dump(
TaskInstanceReferenceCollection(task_instances=task_instances.all())
)
def get_task_instances_batch(session: Session = NEW_SESSION) -> APIResponse:
"""Get list of task instances."""
body = get_json_request_dict()
try:
data = task_instance_batch_form.load(body)
except ValidationError as err:
raise BadRequest(detail=str(err.messages))
states = _convert_state(data['state'])
base_query = session.query(TI).join(TI.dag_run)
base_query = _apply_array_filter(base_query,
key=TI.dag_id,
values=data["dag_ids"])
base_query = _apply_range_filter(
base_query,
key=DR.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=states)
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 == DR.execution_date,
),
isouter=True,
).add_entity(SlaMiss)
ti_query = base_query.options(joinedload(TI.rendered_task_instance_fields))
task_instances = ti_query.all()
return task_instance_collection_schema.dump(
TaskInstanceCollection(task_instances=task_instances,
total_entries=total_entries))
def post_set_task_instances_state(*,
dag_id: str,
session: Session = NEW_SESSION
) -> APIResponse:
"""Set a state of task instances."""
body = get_json_request_dict()
try:
data = set_task_instance_state_form.load(body)
except ValidationError as err:
raise BadRequest(detail=str(err.messages))
error_message = f"Dag ID {dag_id} not found"
dag = get_airflow_app().dag_bag.get_dag(dag_id)
if not dag:
raise NotFound(error_message)
task_id = data['task_id']
task = dag.task_dict.get(task_id)
if not task:
error_message = f"Task ID {task_id} not found"
raise NotFound(error_message)
execution_date = data.get('execution_date')
run_id = data.get('dag_run_id')
if (execution_date and (session.query(TI).filter(
TI.task_id == task_id, TI.dag_id == dag_id, TI.execution_date
== execution_date).one_or_none()) is None):
raise NotFound(
detail=
f"Task instance not found for task {task_id!r} on execution_date {execution_date}"
)
if run_id and not session.query(TI).get({
'task_id': task_id,
'dag_id': dag_id,
'run_id': run_id,
'map_index': -1
}):
error_message = f"Task instance not found for task {task_id!r} on DAG run with ID {run_id!r}"
raise NotFound(detail=error_message)
tis = dag.set_task_instance_state(
task_id=task_id,
run_id=run_id,
execution_date=execution_date,
state=data["new_state"],
upstream=data["include_upstream"],
downstream=data["include_downstream"],
future=data["include_future"],
past=data["include_past"],
commit=not data["dry_run"],
session=session,
)
return task_instance_reference_collection_schema.dump(
TaskInstanceReferenceCollection(task_instances=tis))
def post_dag_run(*,
dag_id: str,
session: Session = NEW_SESSION) -> APIResponse:
"""Trigger a DAG."""
dm = session.query(DagModel).filter(DagModel.dag_id == dag_id).first()
if not dm:
raise NotFound(title="DAG not found",
detail=f"DAG with dag_id: '{dag_id}' not found")
if dm.has_import_errors:
raise BadRequest(
title="DAG cannot be triggered",
detail=f"DAG with dag_id: '{dag_id}' has import errors",
)
try:
post_body = dagrun_schema.load(get_json_request_dict(),
session=session)
except ValidationError as err:
raise BadRequest(detail=str(err))
logical_date = pendulum.instance(post_body["execution_date"])
run_id = post_body["run_id"]
dagrun_instance = (session.query(DagRun).filter(
DagRun.dag_id == dag_id,
or_(DagRun.run_id == run_id, DagRun.execution_date == logical_date),
).first())
if not dagrun_instance:
try:
dag = get_airflow_app().dag_bag.get_dag(dag_id)
dag_run = dag.create_dagrun(
run_type=DagRunType.MANUAL,
run_id=run_id,
execution_date=logical_date,
data_interval=dag.timetable.infer_manual_data_interval(
run_after=logical_date),
state=DagRunState.QUEUED,
conf=post_body.get("conf"),
external_trigger=True,
dag_hash=get_airflow_app().dag_bag.dags_hash.get(dag_id),
)
return dagrun_schema.dump(dag_run)
except ValueError as ve:
raise BadRequest(detail=str(ve))
if dagrun_instance.execution_date == logical_date:
raise AlreadyExists(detail=(
f"DAGRun with DAG ID: '{dag_id}' and "
f"DAGRun logical date: '{logical_date.isoformat(sep=' ')}' already exists"
), )
raise AlreadyExists(
detail=
f"DAGRun with DAG ID: '{dag_id}' and DAGRun ID: '{run_id}' already exists"
)
def post_pool(*, session: Session = NEW_SESSION) -> APIResponse:
"""Create a pool"""
required_fields = {"name", "slots"
} # Pool would require both fields in the post request
fields_diff = required_fields - set(get_json_request_dict().keys())
if fields_diff:
raise BadRequest(
detail=f"Missing required property(ies): {sorted(fields_diff)}")
try:
post_body = pool_schema.load(get_json_request_dict(), session=session)
except ValidationError as err:
raise BadRequest(detail=str(err.messages))
pool = Pool(**post_body)
try:
session.add(pool)
session.commit()
return pool_schema.dump(pool)
except IntegrityError:
raise AlreadyExists(detail=f"Pool: {post_body['pool']} already exists")
def patch_variable(*, variable_key: str, update_mask: UpdateMask = None) -> Response:
"""Update a variable by key"""
try:
data = variable_schema.load(get_json_request_dict())
except ValidationError as err:
raise BadRequest("Invalid Variable schema", detail=str(err.messages))
if data["key"] != variable_key:
raise BadRequest("Invalid post body", detail="key from request body doesn't match uri parameter")
if update_mask:
if "key" in update_mask:
raise BadRequest("key is a ready only field")
if "value" not in update_mask:
raise BadRequest("No field to update")
Variable.set(data["key"], data["val"])
return variable_schema.dump(data)
def clear_dag_run(*,
dag_id: str,
dag_run_id: str,
session: Session = NEW_SESSION) -> APIResponse:
"""Clear a dag run."""
dag_run: Optional[DagRun] = (session.query(DagRun).filter(
DagRun.dag_id == dag_id, DagRun.run_id == dag_run_id).one_or_none())
if dag_run is None:
error_message = f'Dag Run id {dag_run_id} not found in dag {dag_id}'
raise NotFound(error_message)
try:
post_body = clear_dagrun_form_schema.load(get_json_request_dict())
except ValidationError as err:
raise BadRequest(detail=str(err))
dry_run = post_body.get('dry_run', False)
dag = get_airflow_app().dag_bag.get_dag(dag_id)
start_date = dag_run.logical_date
end_date = dag_run.logical_date
if dry_run:
task_instances = dag.clear(
start_date=start_date,
end_date=end_date,
task_ids=None,
include_subdags=True,
include_parentdag=True,
only_failed=False,
dry_run=True,
)
return task_instance_reference_collection_schema.dump(
TaskInstanceReferenceCollection(task_instances=task_instances))
else:
dag.clear(
start_date=start_date,
end_date=end_date,
task_ids=None,
include_subdags=True,
include_parentdag=True,
only_failed=False,
)
dag_run.refresh_from_db()
return dagrun_schema.dump(dag_run)
def get_dag_runs_batch(*, session: Session = NEW_SESSION) -> APIResponse:
"""Get list of DAG Runs"""
body = get_json_request_dict()
try:
data = dagruns_batch_form_schema.load(body)
except ValidationError as err:
raise BadRequest(detail=str(err.messages))
appbuilder = get_airflow_app().appbuilder
readable_dag_ids = appbuilder.sm.get_readable_dag_ids(g.user)
query = session.query(DagRun)
if data.get("dag_ids"):
dag_ids = set(data["dag_ids"]) & set(readable_dag_ids)
query = query.filter(DagRun.dag_id.in_(dag_ids))
else:
query = query.filter(DagRun.dag_id.in_(readable_dag_ids))
states = data.get("states")
if states:
query = query.filter(DagRun.state.in_(states))
dag_runs, total_entries = _fetch_dag_runs(
query,
end_date_gte=data["end_date_gte"],
end_date_lte=data["end_date_lte"],
execution_date_gte=data["execution_date_gte"],
execution_date_lte=data["execution_date_lte"],
start_date_gte=data["start_date_gte"],
start_date_lte=data["start_date_lte"],
limit=data["page_limit"],
offset=data["page_offset"],
order_by=data.get("order_by", "id"),
)
return dagrun_collection_schema.dump(
DAGRunCollection(dag_runs=dag_runs, total_entries=total_entries))