Fix OpenAPI schema for get_log API (#50547) (#51357)

* Fix openapi schema for get_log API

* Fix test_log

(cherry picked from commit 08cc57d)

Co-authored-by: LIU ZHE YOU <[email protected]>

Author: 2 people

airflow-core/src/airflow/api_fastapi/common/headers.py

@@ -47,3 +47,30 @@ def header_accept_json_or_text_depends(
 
 
 HeaderAcceptJsonOrText = Annotated[Mimetype, Depends(header_accept_json_or_text_depends)]
+
+
+def header_accept_json_or_ndjson_depends(
+    accept: Annotated[
+        str,
+        Header(
+            json_schema_extra={
+                "type": "string",
+                "enum": [Mimetype.JSON, Mimetype.NDJSON, Mimetype.ANY],
+            }
+        ),
+    ] = Mimetype.ANY,
+) -> Mimetype:
+    if accept.startswith(Mimetype.ANY):
+        return Mimetype.ANY
+    if accept.startswith(Mimetype.JSON):
+        return Mimetype.JSON
+    if accept.startswith(Mimetype.NDJSON) or accept.startswith(Mimetype.ANY):
+        return Mimetype.NDJSON
+
+    raise HTTPException(
+        status_code=status.HTTP_406_NOT_ACCEPTABLE,
+        detail="Only application/json or application/x-ndjson is supported",
+    )
+
+
+HeaderAcceptJsonOrNdjson = Annotated[Mimetype, Depends(header_accept_json_or_ndjson_depends)]

airflow-core/src/airflow/api_fastapi/common/types.py

@@ -72,6 +72,7 @@ class Mimetype(str, Enum):
 
     TEXT = "text/plain"
     JSON = "application/json"
+    NDJSON = "application/x-ndjson"
     ANY = "*/*"
 
 

airflow-core/src/airflow/api_fastapi/core_api/openapi/v1-rest-api-generated.yaml

@@ -6354,7 +6354,7 @@ paths:
           type: string
           enum:
           - application/json
-          - text/plain
+          - application/x-ndjson
           - '*/*'
           default: '*/*'
           title: Accept
@@ -6365,10 +6365,12 @@ paths:
             application/json:
               schema:
                 $ref: '#/components/schemas/TaskInstancesLogResponse'
-            text/plain:
+            application/x-ndjson:
               schema:
                 type: string
-                example: 'content
+                example: '{"content": "content"}
+
+                  {"content": "content"}
 
                   '
         '401':

airflow-core/src/airflow/api_fastapi/core_api/routes/public/log.py

@@ -28,7 +28,7 @@
 
 from airflow.api_fastapi.common.dagbag import DagBagDep
 from airflow.api_fastapi.common.db.common import SessionDep
-from airflow.api_fastapi.common.headers import HeaderAcceptJsonOrText
+from airflow.api_fastapi.common.headers import HeaderAcceptJsonOrNdjson
 from airflow.api_fastapi.common.router import AirflowRouter
 from airflow.api_fastapi.common.types import Mimetype
 from airflow.api_fastapi.core_api.datamodels.log import TaskInstancesLogResponse
@@ -43,13 +43,14 @@
     tags=["Task Instance"], prefix="/dags/{dag_id}/dagRuns/{dag_run_id}/taskInstances"
 )
 
-text_example_response_for_get_log = {
-    Mimetype.TEXT: {
+ndjson_example_response_for_get_log = {
+    Mimetype.NDJSON: {
         "schema": {
             "type": "string",
             "example": textwrap.dedent(
                 """\
-    content
+    {"content": "content"}
+    {"content": "content"}
     """
             ),
         }
@@ -63,7 +64,7 @@
         **create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND]),
         status.HTTP_200_OK: {
             "description": "Successful Response",
-            "content": text_example_response_for_get_log,
+            "content": ndjson_example_response_for_get_log,
         },
     },
     dependencies=[Depends(requires_access_dag("GET", DagAccessEntity.TASK_LOGS))],
@@ -75,7 +76,7 @@ def get_log(
     dag_run_id: str,
     task_id: str,
     try_number: PositiveInt,
-    accept: HeaderAcceptJsonOrText,
+    accept: HeaderAcceptJsonOrNdjson,
     request: Request,
     dag_bag: DagBagDep,
     session: SessionDep,

airflow-core/src/airflow/ui/openapi-gen/queries/common.ts

@@ -1267,7 +1267,7 @@ export const UseTaskInstanceServiceGetLogKeyFn = (
     token,
     tryNumber,
   }: {
-    accept?: "application/json" | "text/plain" | "*/*";
+    accept?: "application/json" | "*/*" | "application/x-ndjson";
     dagId: string;
     dagRunId: string;
     fullContent?: boolean;

airflow-core/src/airflow/ui/openapi-gen/queries/ensureQueryData.ts

@@ -1744,7 +1744,7 @@ export const ensureUseTaskInstanceServiceGetLogData = (
     token,
     tryNumber,
   }: {
-    accept?: "application/json" | "text/plain" | "*/*";
+    accept?: "application/json" | "*/*" | "application/x-ndjson";
     dagId: string;
     dagRunId: string;
     fullContent?: boolean;

airflow-core/src/airflow/ui/openapi-gen/queries/prefetch.ts

@@ -1744,7 +1744,7 @@ export const prefetchUseTaskInstanceServiceGetLog = (
     token,
     tryNumber,
   }: {
-    accept?: "application/json" | "text/plain" | "*/*";
+    accept?: "application/json" | "*/*" | "application/x-ndjson";
     dagId: string;
     dagRunId: string;
     fullContent?: boolean;

airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts

@@ -2081,7 +2081,7 @@ export const useTaskInstanceServiceGetLog = <
     token,
     tryNumber,
   }: {
-    accept?: "application/json" | "text/plain" | "*/*";
+    accept?: "application/json" | "*/*" | "application/x-ndjson";
     dagId: string;
     dagRunId: string;
     fullContent?: boolean;

airflow-core/src/airflow/ui/openapi-gen/queries/suspense.ts

@@ -2058,7 +2058,7 @@ export const useTaskInstanceServiceGetLogSuspense = <
     token,
     tryNumber,
   }: {
-    accept?: "application/json" | "text/plain" | "*/*";
+    accept?: "application/json" | "*/*" | "application/x-ndjson";
     dagId: string;
     dagRunId: string;
     fullContent?: boolean;

airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts

@@ -2386,7 +2386,7 @@ export type PatchTaskInstanceDryRunData = {
 export type PatchTaskInstanceDryRunResponse = TaskInstanceCollectionResponse;
 
 export type GetLogData = {
-  accept?: "application/json" | "text/plain" | "*/*";
+  accept?: "application/json" | "application/x-ndjson" | "*/*";
   dagId: string;
   dagRunId: string;
   fullContent?: boolean;