def patch_connection(connection_id, session, update_mask=None):
"""
Update a connection entry
"""
try:
data = connection_schema.load(request.json, partial=True)
except ValidationError as err:
# If validation get to here, it is extra field validation.
raise BadRequest(detail=str(err.messages))
non_update_fields = ['connection_id', 'conn_id']
connection = session.query(Connection).filter_by(
conn_id=connection_id).first()
if connection is None:
raise NotFound(
"Connection not found",
detail=
f"The Connection with connection_id: `{connection_id}` was not found",
)
if data.get('conn_id', None) and connection.conn_id != data['conn_id']:
raise BadRequest(detail="The connection_id cannot be updated.")
if update_mask:
update_mask = [i.strip() for i in update_mask]
data_ = {}
for field in update_mask:
if field in data and field not in non_update_fields:
data_[field] = data[field]
else:
raise BadRequest(
detail=f"'{field}' is unknown or cannot be updated.")
data = data_
for key in data:
setattr(connection, key, data[key])
session.add(connection)
session.commit()
return connection_schema.dump(connection)
def patch_dag(session, dag_id, update_mask=None):
"""Update the specific DAG"""
dag = session.query(DagModel).filter(
DagModel.dag_id == dag_id).one_or_none()
if not dag:
raise NotFound(f"Dag with id: '{dag_id}' not found")
try:
patch_body = dag_schema.load(request.json, session=session)
except ValidationError as err:
raise BadRequest("Invalid Dag schema", detail=str(err.messages))
if update_mask:
patch_body_ = {}
if len(update_mask) > 1:
raise BadRequest(
detail=
"Only `is_paused` field can be updated through the REST API")
update_mask = update_mask[0]
if update_mask != 'is_paused':
raise BadRequest(
detail=
"Only `is_paused` field can be updated through the REST API")
patch_body_[update_mask] = patch_body[update_mask]
patch_body = patch_body_
setattr(dag, 'is_paused', patch_body['is_paused'])
session.commit()
return dag_schema.dump(dag)
def delete_dag_run(dag_id, dag_run_id, session):
"""
Delete a DAG Run
"""
if session.query(DagRun).filter(DagRun.dag_id == dag_id, DagRun.run_id == dag_run_id).delete() == 0:
raise NotFound(detail=f"DAGRun with DAG ID: '{dag_id}' and DagRun ID: '{dag_run_id}' not found")
return NoContent, 204
def patch_role(role_name, update_mask=None):
"""Update a role"""
appbuilder = current_app.appbuilder
security_manager = appbuilder.sm
body = request.json
try:
data = role_schema.load(body)
except ValidationError as err:
raise BadRequest(detail=str(err.messages))
role = security_manager.find_role(name=role_name)
if not role:
raise NotFound(title="Role not found",
detail=f"Role with name: `{role_name} was not found")
if update_mask:
update_mask = [i.strip() for i in update_mask]
data_ = {}
for field in update_mask:
if field in data and not field == "permissions":
data_[field] = data[field]
elif field == "actions":
data_["permissions"] = data['permissions']
else:
raise BadRequest(detail=f"'{field}' in update_mask is unknown")
data = data_
perms = data.get("permissions", [])
if perms:
perms = [(item['permission']['name'], item['view_menu']['name'])
for item in data['permissions'] if item]
_check_action_and_resource(security_manager, perms)
security_manager.update_role(pk=role.id, name=data['name'])
security_manager.init_role(role_name=data['name'],
perms=perms or role.permissions)
return role_schema.dump(role)
def get_user(*, username: str) -> APIResponse:
"""Get a user"""
ab_security_manager = current_app.appbuilder.sm
user = ab_security_manager.find_user(username=username)
if not user:
raise NotFound(title="User not found", detail=f"The User with username `{username}` was not found")
return user_collection_item_schema.dump(user)
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 post_dag_run(dag_id, session):
"""Trigger a DAG."""
if not session.query(DagModel).filter(DagModel.dag_id == dag_id).first():
raise NotFound(title="DAG not found",
detail=f"DAG with dag_id: '{dag_id}' not found")
try:
post_body = dagrun_schema.load(request.json, session=session)
except ValidationError as err:
raise BadRequest(detail=str(err))
dagrun_instance = (session.query(DagRun).filter(
DagRun.dag_id == dag_id,
or_(DagRun.run_id == post_body["run_id"],
DagRun.execution_date == post_body["execution_date"]),
).first())
if not dagrun_instance:
dag_run = DagRun(dag_id=dag_id,
run_type=DagRunType.MANUAL,
**post_body)
session.add(dag_run)
session.commit()
return dagrun_schema.dump(dag_run)
if dagrun_instance.execution_date == post_body["execution_date"]:
raise AlreadyExists(
detail=f"DAGRun with DAG ID: '{dag_id}' and "
f"DAGRun ExecutionDate: '{post_body['execution_date']}' already exists"
)
raise AlreadyExists(
detail=
f"DAGRun with DAG ID: '{dag_id}' and DAGRun ID: '{post_body['run_id']}' already exists"
)
def post_clear_task_instances(dag_id: str, session=None):
"""Clear task instances."""
body = request.get_json()
try:
data = clear_task_instance_form.load(body)
except ValidationError as err:
raise BadRequest(detail=str(err.messages))
dag = current_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=current_app.dag_bag,
**data)
if not dry_run:
clear_task_instances(
task_instances,
session,
dag=dag,
dag_run_state=State.RUNNING if reset_dag_runs else False)
task_instances = task_instances.join(
DR, and_(DR.dag_id == TI.dag_id,
DR.execution_date == TI.execution_date)).add_column(DR.run_id)
return task_instance_reference_collection_schema.dump(
TaskInstanceReferenceCollection(task_instances=task_instances.all()))
def get_dag_details(*, dag_id: str) -> APIResponse:
"""Get details of DAG."""
dag: DAG = current_app.dag_bag.get_dag(dag_id)
if not dag:
raise NotFound("DAG not found",
detail=f"The DAG with dag_id: {dag_id} was not found")
return dag_detail_schema.dump(dag)
def patch_role(*, role_name: str, update_mask: UpdateMask = None) -> APIResponse:
"""Update a role"""
appbuilder = current_app.appbuilder
security_manager = appbuilder.sm
body = request.json
try:
data = role_schema.load(body)
except ValidationError as err:
raise BadRequest(detail=str(err.messages))
role = security_manager.find_role(name=role_name)
if not role:
raise NotFound(title="Role not found", detail=f"Role with name {role_name!r} was not found")
if update_mask:
update_mask = [i.strip() for i in update_mask]
data_ = {}
for field in update_mask:
if field in data and not field == "permissions":
data_[field] = data[field]
elif field == "actions":
data_["permissions"] = data['permissions']
else:
raise BadRequest(detail=f"'{field}' in update_mask is unknown")
data = data_
if "permissions" in data:
perms = [(item["action"]["name"], item["resource"]["name"]) for item in data["permissions"] if item]
_check_action_and_resource(security_manager, perms)
security_manager.bulk_sync_roles([{"role": role_name, "perms": perms}])
new_name = data.get("name")
if new_name is not None and new_name != role.name:
security_manager.update_role(role_id=role.id, name=new_name)
return role_schema.dump(role)
def get_variable(*, variable_key: str) -> Response:
"""Get a variables by key"""
try:
var = Variable.get(variable_key)
except KeyError:
raise NotFound("Variable not found")
return variable_schema.dump({"key": variable_key, "val": var})
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 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 post_clear_task_instances(dag_id: str, session=None):
"""
Clear task instances.
"""
body = request.get_json()
try:
data = clear_task_instance_form.load(body)
except ValidationError as err:
raise BadRequest(detail=str(err.messages))
dag = current_app.dag_bag.get_dag(dag_id)
if not dag:
error_message = "Dag id {} not found".format(dag_id)
raise NotFound(error_message)
reset_dag_runs = data.pop('reset_dag_runs')
task_instances = dag.clear(get_tis=True, **data)
if not data["dry_run"]:
clear_task_instances(
task_instances,
session,
dag=dag,
activate_dag_runs=False, # We will set DagRun state later.
)
if reset_dag_runs:
dag.set_dag_runs_state(
session=session,
start_date=data["start_date"],
end_date=data["end_date"],
state=State.RUNNING,
)
task_instances = task_instances.join(
DR, and_(DR.dag_id == TI.dag_id,
DR.execution_date == TI.execution_date)).add_column(DR.run_id)
return task_instance_reference_collection_schema.dump(
TaskInstanceReferenceCollection(task_instances=task_instances.all()))
def post_clear_task_instances(*,
dag_id: str,
session: Session = NEW_SESSION) -> APIResponse:
"""Clear task instances."""
body = request.get_json()
try:
data = clear_task_instance_form.load(body)
except ValidationError as err:
raise BadRequest(detail=str(err.messages))
dag = current_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=current_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 get_role(*, role_name: str) -> APIResponse:
"""Get role"""
ab_security_manager = current_app.appbuilder.sm
role = ab_security_manager.find_role(name=role_name)
if not role:
raise NotFound(title="Role not found", detail=f"Role with name {role_name!r} was not found")
return role_schema.dump(role)
def delete_variable(variable_key: str) -> Response:
"""
Delete variable
"""
if Variable.delete(variable_key) == 0:
raise NotFound("Variable not found")
return Response(status=204)
def get_role(role_name):
"""Get role"""
ab_security_manager = current_app.appbuilder.sm
role = ab_security_manager.find_role(name=role_name)
if not role:
raise NotFound(title="Role not found", detail=f"The Role with name `{role_name}` was not found")
return role_schema.dump(role)
def patch_user(username, update_mask=None):
"""Update a role"""
try:
data = user_schema.load(request.json)
except ValidationError as e:
raise BadRequest(detail=str(e.messages))
security_manager = current_app.appbuilder.sm
user = security_manager.find_user(username=username)
if user is None:
detail = f"The User with username `{username}` was not found"
raise NotFound(title="User not found", detail=detail)
# Get fields to update. 'username' is always excluded (and it's an error to
# include it in update_maek).
if update_mask is not None:
masked_data = {}
missing_mask_names = []
for field in update_mask:
field = field.strip()
try:
masked_data[field] = data[field]
except KeyError:
missing_mask_names.append(field)
if missing_mask_names:
detail = f"Unknown update masks: {', '.join(repr(n) for n in missing_mask_names)}"
raise BadRequest(detail=detail)
if "username" in masked_data:
raise BadRequest("Cannot update fields: 'username'")
data = masked_data
else:
data.pop("username", None)
if "roles" in data:
roles_to_update = []
missing_role_names = []
for role_data in data.pop("roles", ()):
role_name = role_data["name"]
role = security_manager.find_role(role_name)
if role is None:
missing_role_names.append(role_name)
else:
roles_to_update.append(role)
if missing_role_names:
detail = f"Unknown roles: {', '.join(repr(n) for n in missing_role_names)}"
raise BadRequest(detail=detail)
else:
roles_to_update = None # Don't change existing value.
if "password" in data:
user.password = generate_password_hash(data.pop("password"))
if roles_to_update is not None:
user.roles = roles_to_update
for key, value in data.items():
setattr(user, key, value)
security_manager.update_user(user)
return user_schema.dump(user)
def get_connection(connection_id, session):
"""
Get a connection entry
"""
connection = session.query(Connection).filter(Connection.conn_id == connection_id).one_or_none()
if connection is None:
raise NotFound("Connection not found")
return connection_collection_item_schema.dump(connection)
def get_pool(pool_name, session):
"""
Get a pool
"""
obj = session.query(Pool).filter(Pool.pool == pool_name).one_or_none()
if obj is None:
raise NotFound(detail=f"Pool with name:'{pool_name}' not found")
return pool_schema.dump(obj)
def get_dag_details(dag_id):
"""
Get details of DAG.
"""
dag: DAG = current_app.dag_bag.get_dag(dag_id)
if not dag:
raise NotFound("DAG not found")
return dag_detail_schema.dump(dag)
def get_event_log(event_log_id, session):
"""
Get a log entry
"""
event_log = session.query(Log).filter(Log.id == event_log_id).one_or_none()
if event_log is None:
raise NotFound("Event Log not found")
return event_log_schema.dump(event_log)
def get_dag(dag_id, session):
"""Get basic information about a DAG."""
dag = session.query(DagModel).filter(DagModel.dag_id == dag_id).one_or_none()
if dag is None:
raise NotFound("DAG not found", detail=f"The DAG with dag_id: {dag_id} was not found")
return dag_schema.dump(dag)
def delete_role(role_name):
"""Delete a role"""
ab_security_manager = current_app.appbuilder.sm
role = ab_security_manager.find_role(name=role_name)
if not role:
raise NotFound(title="Role not found", detail=f"The Role with name `{role_name}` was not found")
ab_security_manager.delete_role(role_name=role_name)
return NoContent, 204
def get_event_log(*,
event_log_id: int,
session: Session = NEW_SESSION) -> APIResponse:
"""Get a log entry"""
event_log = session.query(Log).get(event_log_id)
if event_log is None:
raise NotFound("Event Log not found")
return event_log_schema.dump(event_log)
def delete_role(*, role_name: str) -> APIResponse:
"""Delete a role"""
ab_security_manager = current_app.appbuilder.sm
role = ab_security_manager.find_role(name=role_name)
if not role:
raise NotFound(title="Role not found", detail=f"Role with name {role_name!r} was not found")
ab_security_manager.delete_role(role_name=role_name)
return NoContent, 204
def delete_pool(pool_name: str, session):
"""Delete a pool"""
if pool_name == "default_pool":
raise BadRequest(detail="Default Pool can't be deleted")
elif session.query(Pool).filter(Pool.pool == pool_name).delete() == 0:
raise NotFound(detail=f"Pool with name:'{pool_name}' not found")
else:
return Response(status=204)
def get_tasks(dag_id):
"""Get tasks for DAG"""
dag: DAG = current_app.dag_bag.get_dag(dag_id)
if not dag:
raise NotFound("DAG not found")
tasks = dag.tasks
task_collection = TaskCollection(tasks=tasks, total_entries=len(tasks))
return task_collection_schema.dump(task_collection)
def get_dag_run(dag_id, dag_run_id, session):
"""
Get a DAG Run.
"""
dag_run = session.query(DagRun).filter(DagRun.dag_id == dag_id, DagRun.run_id == dag_run_id).one_or_none()
if dag_run is None:
raise NotFound("DAGRun not found")
return dagrun_schema.dump(dag_run)
