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)