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_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_user():
"""Create a new user"""
try:
data = user_schema.load(request.json)
except ValidationError as e:
raise BadRequest(detail=str(e.messages))
security_manager = current_app.appbuilder.sm
username = data["username"]
email = data["email"]
if security_manager.find_user(username=username):
detail = f"Username `{username}` already exists. Use PATCH to update."
raise AlreadyExists(detail=detail)
if security_manager.find_user(email=email):
detail = f"The email `{email}` is already taken."
raise AlreadyExists(detail=detail)
roles_to_add = []
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_add.append(role)
if missing_role_names:
detail = f"Unknown roles: {', '.join(repr(n) for n in missing_role_names)}"
raise BadRequest(detail=detail)
if roles_to_add:
default_role = roles_to_add.pop()
else: # No roles provided, use the F.A.B's default registered user role.
default_role = security_manager.find_role(
security_manager.auth_user_registration_role)
user = security_manager.add_user(role=default_role, **data)
if not user:
detail = f"Failed to add user `{username}`."
return Unknown(detail=detail)
if roles_to_add:
user.roles.extend(roles_to_add)
security_manager.update_user(user)
return user_schema.dump(user)
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))
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 = current_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=State.QUEUED,
conf=post_body.get("conf"),
external_trigger=True,
dag_hash=current_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 delete_dag(dag_id: str, session: Session = NEW_SESSION) -> APIResponse:
"""Delete the specific DAG."""
from airflow.api.common import delete_dag as delete_dag_module
try:
delete_dag_module.delete_dag(dag_id, session=session)
except DagNotFound:
raise NotFound(f"Dag with id: '{dag_id}' not found")
except AirflowException:
raise AlreadyExists(detail=f"Task instances of dag with id: '{dag_id}' are still running")
return NoContent, 204
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))
execution_date = 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 == execution_date),
).first())
if not dagrun_instance:
dag_run = current_app.dag_bag.get_dag(dag_id).create_dagrun(
run_type=DagRunType.MANUAL,
run_id=run_id,
execution_date=execution_date,
state=State.QUEUED,
conf=post_body.get("conf"),
external_trigger=True,
dag_hash=current_app.dag_bag.dags_hash.get(dag_id),
)
return dagrun_schema.dump(dag_run)
if dagrun_instance.execution_date == 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_connection(*, session: Session = NEW_SESSION) -> APIResponse:
"""Create connection entry"""
body = request.json
try:
data = connection_schema.load(body)
except ValidationError as err:
raise BadRequest(detail=str(err.messages))
conn_id = data['conn_id']
query = session.query(Connection)
connection = query.filter_by(conn_id=conn_id).first()
if not connection:
connection = Connection(**data)
session.add(connection)
session.commit()
return connection_schema.dump(connection)
raise AlreadyExists(detail=f"Connection already exist. ID: {conn_id}")
def post_role() -> APIResponse:
"""Create a new 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=data['name'])
if not role:
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": data["name"], "perms": perms}])
return role_schema.dump(role)
detail = f"Role with name {role.name!r} already exists; please update with the PATCH endpoint"
raise AlreadyExists(detail=detail)
def delete_dag(dag_id: str, session: Session):
"""Delete the specific DAG."""
# TODO: This function is shared with the /delete endpoint used by the web
# UI, so we're reusing it to simplify maintenance. Refactor the function to
# another place when the experimental/legacy API is removed.
from airflow.api.common.experimental import delete_dag
try:
delete_dag.delete_dag(dag_id, session=session)
except DagNotFound:
raise NotFound(f"Dag with id: '{dag_id}' not found")
except AirflowException:
raise AlreadyExists(
detail=
f"Task instances of dag with id: '{dag_id}' are still running")
return NoContent, 204
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")
post_body = dagrun_schema.load(request.json, session=session)
dagrun_instance = (
session.query(DagRun).filter(DagRun.dag_id == dag_id, DagRun.run_id == post_body["run_id"]).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)
raise AlreadyExists(
detail=f"DAGRun with DAG ID: '{dag_id}' and DAGRun ID: '{post_body['run_id']}' already exists"
)
def post_role():
"""Create a new 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=data['name'])
if not role:
perms = [(item['action']['name'], item['resource']['name']) for item in data['permissions'] if item]
_check_action_and_resource(security_manager, perms)
security_manager.init_role(role_name=data['name'], perms=perms)
return role_schema.dump(role)
raise AlreadyExists(
detail=f"Role with name `{role.name}` already exist. Please update with patch endpoint"
)
def post_connection(session):
"""
Create connection entry
"""
body = request.json
try:
result = connection_schema.load(body)
except ValidationError as err:
raise BadRequest(detail=str(err.messages))
data = result.data
conn_id = data['conn_id']
query = session.query(Connection)
connection = query.filter_by(conn_id=conn_id).first()
if not connection:
connection = Connection(**data)
session.add(connection)
session.commit()
return connection_schema.dump(connection)
raise AlreadyExists("Connection already exist. ID: %s" % conn_id)
def post_pool(session):
"""Create a pool"""
required_fields = ["name", "slots"
] # Pool would require both fields in the post request
for field in required_fields:
if field not in request.json.keys():
raise BadRequest(detail=f"'{field}' is a required property")
try:
post_body = pool_schema.load(request.json, 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 post_pool(session):
"""Create a pool"""
required_fields = {"name", "slots"
} # Pool would require both fields in the post request
fields_diff = required_fields - set(request.json.keys())
if fields_diff:
raise BadRequest(
detail=f"Missing required property(ies): {sorted(fields_diff)}")
try:
post_body = pool_schema.load(request.json, 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_user(*,
username: str,
update_mask: UpdateMask = None) -> APIResponse:
"""Update a user"""
try:
data = user_schema.load(request.json)
except ValidationError as e:
raise BadRequest(detail=str(e.messages))
security_manager = get_airflow_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)
# Check unique username
new_username = data.get('username')
if new_username and new_username != username:
if security_manager.find_user(username=new_username):
raise AlreadyExists(
detail=f"The username `{new_username}` already exists")
# Check unique email
email = data.get('email')
if email and email != user.email:
if security_manager.find_user(email=email):
raise AlreadyExists(detail=f"The email `{email}` already exists")
# Get fields to update.
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)
data = masked_data
roles_to_update: Optional[List[Role]]
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)