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)