use status (apache#44161)

Co-authored-by: Amogh Desai <[email protected]>
astronomer · Nov 19, 2024 · 55b30b2 · 55b30b2
1 parent a60d105
commit 55b30b2
Show file tree

Hide file tree

Showing 5 changed files with 18 additions and 12 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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,