From 55b30b28eaeb29908f98dcdb483405a4d9ec3fd0 Mon Sep 17 00:00:00 2001 From: Kalyan R Date: Tue, 19 Nov 2024 15:25:40 +0530 Subject: [PATCH] use status (#44161) Co-authored-by: Amogh Desai --- airflow/api_fastapi/core_api/routes/public/assets.py | 12 ++++++------ .../core_api/routes/public/connections.py | 4 +++- .../api_fastapi/core_api/routes/public/dag_run.py | 3 ++- .../core_api/routes/public/import_error.py | 5 ++++- .../core_api/routes/public/task_instances.py | 6 +++--- 5 files changed, 18 insertions(+), 12 deletions(-) diff --git a/airflow/api_fastapi/core_api/routes/public/assets.py b/airflow/api_fastapi/core_api/routes/public/assets.py index b94b825f76294..fe578b519183b 100644 --- a/airflow/api_fastapi/core_api/routes/public/assets.py +++ b/airflow/api_fastapi/core_api/routes/public/assets.py @@ -80,7 +80,7 @@ def _generate_queued_event_where_clause( @assets_router.get( "/assets", - responses=create_openapi_http_exception_doc([401, 403, 404]), + responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND]), ) def get_assets( limit: QueryLimit, @@ -115,7 +115,7 @@ def get_assets( @assets_router.get( "/assets/events", - responses=create_openapi_http_exception_doc([404]), + responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND]), ) def get_asset_events( limit: QueryLimit, @@ -165,7 +165,7 @@ def get_asset_events( @assets_router.post( "/assets/events", - responses=create_openapi_http_exception_doc([404]), + responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND]), ) def create_asset_event( body: CreateAssetEventsBody, @@ -174,7 +174,7 @@ def create_asset_event( """Create asset events.""" asset = session.scalar(select(AssetModel).where(AssetModel.uri == body.uri).limit(1)) if not asset: - raise HTTPException(404, f"Asset with uri: `{body.uri}` was not found") + raise HTTPException(status.HTTP_404_NOT_FOUND, f"Asset with uri: `{body.uri}` was not found") timestamp = timezone.utcnow() assets_event = asset_manager.register_asset_change( @@ -185,13 +185,13 @@ def create_asset_event( ) if not assets_event: - raise HTTPException(404, f"Asset with uri: `{body.uri}` was not found") + raise HTTPException(status.HTTP_404_NOT_FOUND, f"Asset with uri: `{body.uri}` was not found") return AssetEventResponse.model_validate(assets_event, from_attributes=True) @assets_router.get( "/assets/{uri:path}", - responses=create_openapi_http_exception_doc([401, 403, 404]), + responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND]), ) def get_asset( uri: str, diff --git a/airflow/api_fastapi/core_api/routes/public/connections.py b/airflow/api_fastapi/core_api/routes/public/connections.py index 406d689d5e9e9..ff17dbb3b27b1 100644 --- a/airflow/api_fastapi/core_api/routes/public/connections.py +++ b/airflow/api_fastapi/core_api/routes/public/connections.py @@ -166,7 +166,9 @@ def patch_connection( connection = session.scalar(select(Connection).filter_by(conn_id=connection_id).limit(1)) if connection is None: - raise HTTPException(404, f"The Connection with connection_id: `{connection_id}` was not found") + raise HTTPException( + status.HTTP_404_NOT_FOUND, f"The Connection with connection_id: `{connection_id}` was not found" + ) if update_mask: data = patch_body.model_dump(include=set(update_mask) - non_update_fields) diff --git a/airflow/api_fastapi/core_api/routes/public/dag_run.py b/airflow/api_fastapi/core_api/routes/public/dag_run.py index decc7ff2b285c..ce9c4410aeb69 100644 --- a/airflow/api_fastapi/core_api/routes/public/dag_run.py +++ b/airflow/api_fastapi/core_api/routes/public/dag_run.py @@ -195,7 +195,8 @@ def clear_dag_run( dag_run = session.scalar(select(DagRun).filter_by(dag_id=dag_id, run_id=dag_run_id)) if dag_run is None: raise HTTPException( - 404, f"The DagRun with dag_id: `{dag_id}` and run_id: `{dag_run_id}` was not found" + status.HTTP_404_NOT_FOUND, + f"The DagRun with dag_id: `{dag_id}` and run_id: `{dag_run_id}` was not found", ) dag: DAG = request.app.state.dag_bag.get_dag(dag_id) diff --git a/airflow/api_fastapi/core_api/routes/public/import_error.py b/airflow/api_fastapi/core_api/routes/public/import_error.py index 13540c5a475a4..5312904fc7b35 100644 --- a/airflow/api_fastapi/core_api/routes/public/import_error.py +++ b/airflow/api_fastapi/core_api/routes/public/import_error.py @@ -53,7 +53,10 @@ def get_import_error( """Get an import error.""" error = session.scalar(select(ParseImportError).where(ParseImportError.id == import_error_id)) if error is None: - raise HTTPException(404, f"The ImportError with import_error_id: `{import_error_id}` was not found") + raise HTTPException( + status.HTTP_404_NOT_FOUND, + f"The ImportError with import_error_id: `{import_error_id}` was not found", + ) return ImportErrorResponse.model_validate( error, diff --git a/airflow/api_fastapi/core_api/routes/public/task_instances.py b/airflow/api_fastapi/core_api/routes/public/task_instances.py index c805095f58208..7baa4a4af0606 100644 --- a/airflow/api_fastapi/core_api/routes/public/task_instances.py +++ b/airflow/api_fastapi/core_api/routes/public/task_instances.py @@ -138,15 +138,15 @@ def get_mapped_task_instances( dag = request.app.state.dag_bag.get_dag(dag_id) if not dag: error_message = f"DAG {dag_id} not found" - raise HTTPException(404, error_message) + raise HTTPException(status.HTTP_404_NOT_FOUND, error_message) try: task = dag.get_task(task_id) except TaskNotFound: error_message = f"Task id {task_id} not found" - raise HTTPException(404, error_message) + raise HTTPException(status.HTTP_404_NOT_FOUND, error_message) if not task.get_needs_expansion(): error_message = f"Task id {task_id} is not mapped" - raise HTTPException(404, error_message) + raise HTTPException(status.HTTP_404_NOT_FOUND, error_message) task_instance_select, total_entries = paginated_select( base_query,