def get_dag(*, dag_id: str, session: Session = NEW_SESSION) -> APIResponse: """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 get_connection(connection_id, session): """ Get a connection entry """ query = session.query(Connection) query = query.filter(Connection.conn_id == connection_id) connection = query.one_or_none() if connection is None: raise NotFound("Connection not found") return connection_collection_item_schema.dump(connection)
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_import_error(import_error_id, session): """ Get an import error """ query = session.query(ImportError) error = query.filter(ImportError.id == import_error_id).one_or_none() if error is None: raise NotFound("Import error not found") return import_error_schema.dump(error)
def post_set_task_instances_state(dag_id, session): """Set a state of task instances.""" body = request.get_json() 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" try: dag = current_app.dag_bag.get_dag(dag_id) if not dag: raise NotFound(error_message) except SerializedDagNotFound: # If DAG is not found in serialized_dag table 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) tis = set_state( tasks=[task], execution_date=data["execution_date"], upstream=data["include_upstream"], downstream=data["include_downstream"], future=data["include_future"], past=data["include_past"], state=data["new_state"], commit=not data["dry_run"], ) execution_dates = {ti.execution_date for ti in tis} execution_date_to_run_id_map = dict( session.query(DR.execution_date, DR.run_id).filter( DR.dag_id == dag_id, DR.execution_date.in_(execution_dates))) tis_with_run_id = [ (ti, execution_date_to_run_id_map.get(ti.execution_date)) for ti in tis ] return task_instance_reference_collection_schema.dump( TaskInstanceReferenceCollection(task_instances=tis_with_run_id))
def get_import_error(import_error_id, session): """Get an import error""" error = session.query(ImportErrorModel).filter(ImportErrorModel.id == import_error_id).one_or_none() if error is None: raise NotFound( "Import error not found", detail=f"The ImportError with import_error_id: `{import_error_id}` was not found", ) return import_error_schema.dump(error)
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 delete_pool(*, pool_name: str, session: Session = NEW_SESSION) -> APIResponse: """Delete a pool""" if pool_name == "default_pool": raise BadRequest(detail="Default Pool can't be deleted") affected_count = session.query(Pool).filter( Pool.pool == pool_name).delete() if affected_count == 0: raise NotFound(detail=f"Pool with name:'{pool_name}' not found") return Response(status=HTTPStatus.NO_CONTENT)
def get_dag_run(dag_id, dag_run_id, session): """ Get a DAG Run. """ query = session.query(DagRun) query = query.filter(DagRun.dag_id == dag_id) query = query.filter(DagRun.run_id == dag_run_id) dag_run = query.one_or_none() if dag_run is None: raise NotFound("DAGRun not found") return dagrun_schema.dump(dag_run)
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", detail= f"DAGRun with DAG ID: '{dag_id}' and DagRun ID: '{dag_run_id}' not found", ) return dagrun_schema.dump(dag_run)
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", detail= f"The Connection with connection_id: `{connection_id}` was not found", ) return connection_schema.dump(connection)
def get_log(session, dag_id, dag_run_id, task_id, task_try_number, full_content=False, token=None): """Get logs for specific task instance""" key = current_app.config["SECRET_KEY"] if not token: metadata = {} else: try: metadata = URLSafeSerializer(key).loads(token) except BadSignature: raise BadRequest("Bad Signature. Please use only the tokens provided by the API.") if metadata.get('download_logs') and metadata['download_logs']: full_content = True if full_content: metadata['download_logs'] = True else: metadata['download_logs'] = False task_log_reader = TaskLogReader() if not task_log_reader.supports_read: raise BadRequest("Task log handler does not support read logs.") ti = ( session.query(TaskInstance) .filter(TaskInstance.task_id == task_id, TaskInstance.run_id == dag_run_id) .join(TaskInstance.dag_run) .options(eagerload(TaskInstance.dag_run)) .one_or_none() ) if ti is None: metadata['end_of_log'] = True raise NotFound(title="TaskInstance not found") dag = current_app.dag_bag.get_dag(dag_id) if dag: try: ti.task = dag.get_task(ti.task_id) except TaskNotFound: pass return_type = request.accept_mimetypes.best_match(['text/plain', 'application/json']) # return_type would be either the above two or None if return_type == 'application/json' or return_type is None: # default logs, metadata = task_log_reader.read_log_chunks(ti, task_try_number, metadata) logs = logs[0] if task_try_number is not None else logs token = URLSafeSerializer(key).dumps(metadata) return logs_schema.dump(LogResponseObject(continuation_token=token, content=logs)) # text/plain. Stream logs = task_log_reader.read_log_stream(ti, task_try_number, metadata) return Response(logs, headers={"Content-Type": return_type})
def get_pool(pool_name, session): """ Get a pool """ pool_id = pool_name query = session.query(Pool) obj = query.filter(Pool.pool == pool_id).one_or_none() if obj is None: raise NotFound("Pool not found") return pool_schema.dump(obj)
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 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") return dag_schema.dump(dag)
def delete_user(username): """Delete a user""" 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) user.roles = [] # Clear foreign keys on this user first. security_manager.get_session.delete(user) security_manager.get_session.commit()
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 get_dataset(uri: str, session: Session = NEW_SESSION) -> APIResponse: """Get a Dataset""" dataset = (session.query(DatasetModel).filter( DatasetModel.uri == uri).options( joinedload(DatasetModel.consuming_dags), joinedload(DatasetModel.producing_tasks)).one_or_none()) if not dataset: raise NotFound( "Dataset not found", detail=f"The Dataset with uri: `{uri}` was not found", ) return dataset_schema.dump(dataset)
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 delete_dag_run(*, dag_id: str, dag_run_id: str, session: Session = NEW_SESSION) -> APIResponse: """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, HTTPStatus.NO_CONTENT
def delete_connection(connection_id, session): """Delete a connection entry""" connection = session.query(Connection).filter_by( conn_id=connection_id).one_or_none() if connection is None: raise NotFound( 'Connection not found', detail= f"The Connection with connection_id: `{connection_id}` was not found", ) session.delete(connection) return NoContent, 204
def get_extra_links( *, dag_id: str, dag_run_id: str, task_id: str, session: Session = NEW_SESSION, ) -> APIResponse: """Get extra links for task instance""" from airflow.models.taskinstance import TaskInstance dagbag: DagBag = current_app.dag_bag dag: DAG = dagbag.get_dag(dag_id) if not dag: raise NotFound("DAG not found", detail=f'DAG with ID = "{dag_id}" not found') try: task = dag.get_task(task_id) except TaskNotFound: raise NotFound("Task not found", detail=f'Task with ID = "{task_id}" not found') ti = ( session.query(TaskInstance) .filter( TaskInstance.dag_id == dag_id, TaskInstance.run_id == dag_run_id, TaskInstance.task_id == task_id, ) .one_or_none() ) if not ti: raise NotFound("DAG Run not found", detail=f'DAG Run with ID = "{dag_run_id}" not found') all_extra_link_pairs = ( (link_name, task.get_extra_links(ti, link_name)) for link_name in task.extra_links ) all_extra_links = { link_name: link_url if link_url else None for link_name, link_url in all_extra_link_pairs } return all_extra_links
def get_import_error(*, import_error_id: int, session: Session = NEW_SESSION) -> APIResponse: """Get an import error""" error = session.query(ImportErrorModel).get(import_error_id) if error is None: raise NotFound( "Import error not found", detail= f"The ImportError with import_error_id: `{import_error_id}` was not found", ) return import_error_schema.dump(error)
def get_connection(*, connection_id: str, session: Session = NEW_SESSION) -> APIResponse: """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", detail= f"The Connection with connection_id: `{connection_id}` was not found", ) return connection_schema.dump(connection)
def patch_pool(pool_name, session, update_mask=None): """Update a pool""" # Only slots can be modified in 'default_pool' try: if pool_name == Pool.DEFAULT_POOL_NAME and request.json[ "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.json) 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(request.json.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 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, TI.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) .options(joinedload(TI.rendered_task_instance_fields)) ) try: task_instance = query.one_or_none() except MultipleResultsFound: raise NotFound( "Task instance not found", detail="Task instance is mapped, add the map_index value to the URL" ) if task_instance is None: raise NotFound("Task instance not found") if task_instance[0].map_index != -1: raise NotFound( "Task instance not found", detail="Task instance is mapped, add the map_index value to the URL" ) return task_instance_schema.dump(task_instance)
def post_set_task_instances_state(dag_id, session): """Set a state of task instances.""" body = request.get_json() 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 = current_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['execution_date'] try: session.query(TI).filter_by(execution_date=execution_date, task_id=task_id, dag_id=dag_id).one() except NoResultFound: raise NotFound(f"Task instance not found for task {task_id} on execution_date {execution_date}") tis = dag.set_task_instance_state( task_id=task_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 delete_connection(*, connection_id: str, session: Session = NEW_SESSION) -> APIResponse: """Delete a connection entry""" connection = session.query(Connection).filter_by( conn_id=connection_id).one_or_none() if connection is None: raise NotFound( 'Connection not found', detail= f"The Connection with connection_id: `{connection_id}` was not found", ) session.delete(connection) return NoContent, HTTPStatus.NO_CONTENT
def get_dag_run(*, dag_id: str, dag_run_id: str, session: Session = NEW_SESSION) -> APIResponse: """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", detail= f"DAGRun with DAG ID: '{dag_id}' and DagRun ID: '{dag_run_id}' not found", ) return dagrun_schema.dump(dag_run)
def delete_user(*, username: str) -> APIResponse: """Delete a user""" 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) user.roles = [] # Clear foreign keys on this user first. security_manager.get_session.delete(user) security_manager.get_session.commit() return NoContent, HTTPStatus.NO_CONTENT