Return Response instead of pydantic model

airflow/api_fastapi/core_api/openapi/v1-generated.yaml

@@ -774,8 +774,7 @@ paths:
           description: Successful Response
           content:
             application/json:
-              schema:
-                $ref: '#/components/schemas/DAGSourceModel'
+              schema: {}
         '400':
           content:
             application/json:
@@ -1629,18 +1628,6 @@ components:
       - dataset_triggered
       title: DAGRunTypes
       description: DAG Run Types for responses.
-    DAGSourceModel:
-      properties:
-        content:
-          anyOf:
-          - type: string
-          - type: 'null'
-          title: Content
-      type: object
-      required:
-      - content
-      title: DAGSourceModel
-      description: DAG Source Model class.
     DagProcessorInfoSchema:
       properties:
         status:

airflow/api_fastapi/core_api/routes/public/dag_sources.py

@@ -16,7 +16,7 @@
 # under the License.
 from __future__ import annotations
 
-from fastapi import Depends, HTTPException, Request
+from fastapi import Depends, HTTPException, Request, Response
 from itsdangerous import BadSignature, URLSafeSerializer
 from sqlalchemy.orm import Session
 from typing_extensions import Annotated
@@ -28,7 +28,9 @@
 from airflow.models.dagcode import DagCode
 
 dag_sources_router = AirflowRouter(tags=["DagSource"], prefix="/dagSources")
-supported_mime_types = ["text/plain", "application/json"]
+mime_type_text = "text/plain"
+mime_type_json = "application/json"
+supported_mime_types = [mime_type_text, mime_type_json]
 
 
 def _get_matching_mime_type(request: Request) -> str | None:
@@ -48,7 +50,7 @@ async def get_dag_source(
     file_token: str,
     session: Annotated[Session, Depends(get_session)],
     request: Request,
-) -> DAGSourceModel:
+) -> Response:
     """Get source code using file token."""
     auth_s = URLSafeSerializer(request.app.state.secret_key)
 
@@ -59,7 +61,12 @@ async def get_dag_source(
         raise HTTPException(404, "DAG source not found")
 
     return_type = _get_matching_mime_type(request)
-    if return_type is not None:
-        return DAGSourceModel.model_validate(dag_source)
+
+    if return_type == mime_type_text:
+        return Response(dag_source, media_type=return_type)
+
+    if return_type == mime_type_json:
+        content = DAGSourceModel(content=dag_source).model_dump_json()
+        return Response(content, media_type=return_type)
 
     raise HTTPException(406, "Content not available for Accept header")

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

@@ -246,7 +246,7 @@ export const prefetchUseDagRunServiceGetDagRun = (
  * Get source code using file token.
  * @param data The data for the request.
  * @param data.fileToken
- * @returns DAGSourceModel Successful Response
+ * @returns unknown Successful Response
  * @throws ApiError
  */
 export const prefetchUseDagSourceServiceGetDagSource = (

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

@@ -309,7 +309,7 @@ export const useDagRunServiceGetDagRun = <
  * Get source code using file token.
  * @param data The data for the request.
  * @param data.fileToken
- * @returns DAGSourceModel Successful Response
+ * @returns unknown Successful Response
  * @throws ApiError
  */
 export const useDagSourceServiceGetDagSource = <

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

@@ -304,7 +304,7 @@ export const useDagRunServiceGetDagRunSuspense = <
  * Get source code using file token.
  * @param data The data for the request.
  * @param data.fileToken
- * @returns DAGSourceModel Successful Response
+ * @returns unknown Successful Response
  * @throws ApiError
  */
 export const useDagSourceServiceGetDagSourceSuspense = <

airflow/ui/openapi-gen/requests/schemas.gen.ts

@@ -1013,26 +1013,6 @@ export const $DAGRunTypes = {
   description: "DAG Run Types for responses.",
 } as const;
 
-export const $DAGSourceModel = {
-  properties: {
-    content: {
-      anyOf: [
-        {
-          type: "string",
-        },
-        {
-          type: "null",
-        },
-      ],
-      title: "Content",
-    },
-  },
-  type: "object",
-  required: ["content"],
-  title: "DAGSourceModel",
-  description: "DAG Source Model class.",
-} as const;
-
 export const $DagProcessorInfoSchema = {
   properties: {
     status: {

airflow/ui/openapi-gen/requests/services.gen.ts

@@ -444,7 +444,7 @@ export class DagSourceService {
    * Get source code using file token.
    * @param data The data for the request.
    * @param data.fileToken
-   * @returns DAGSourceModel Successful Response
+   * @returns unknown Successful Response
    * @throws ApiError
    */
   public static getDagSource(

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

@@ -177,13 +177,6 @@ export type DAGRunTypes = {
   dataset_triggered: number;
 };
 
-/**
- * DAG Source Model class.
- */
-export type DAGSourceModel = {
-  content: string | null;
-};
-
 /**
  * Schema for DagProcessor info.
  */
@@ -433,7 +426,7 @@ export type GetDagSourceData = {
   fileToken: string;
 };
 
-export type GetDagSourceResponse = DAGSourceModel;
+export type GetDagSourceResponse = unknown;
 
 export type GetHealthResponse = HealthInfoSchema;
 
@@ -806,7 +799,7 @@ export type $OpenApiTs = {
         /**
          * Successful Response
          */
-        200: DAGSourceModel;
+        200: unknown;
         /**
          * Bad Request
          */

tests/api_fastapi/core_api/routes/public/test_dag_sources.py

@@ -21,6 +21,7 @@
 import os
 
 import pytest
+from httpx import Response
 
 from airflow.models.dag import DAG
 from airflow.models.dagbag import DagBag
@@ -29,6 +30,7 @@
 
 pytestmark = pytest.mark.db_test
 
+ROOT_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), os.pardir, os.pardir))
 EXAMPLE_DAG_FILE = os.path.join("airflow", "example_dags", "example_bash_operator.py")
 TEST_DAG_ID = "latest_only"
 API_PREFIX = "/public/dagSources"
@@ -37,13 +39,13 @@
 class TestGetDAGSource:
     @pytest.fixture(autouse=True)
     def setup(self, url_safe_serializer) -> None:
-        clear_db_dag_code()
+        self.clear_db()
         self.test_dag, self.dag_docstring = self.create_dag_source()
         fileloc = url_safe_serializer.dumps(self.test_dag.fileloc)
         self.dag_sources_url = f"{API_PREFIX}/{fileloc}"
 
     def teardown_method(self) -> None:
-        clear_db_dag_code()
+        self.clear_db()
 
     @staticmethod
     def _get_dag_file_docstring(fileloc: str) -> str | None:
@@ -59,27 +61,40 @@ def create_dag_source(self) -> tuple[DAG, str | None]:
         test_dag: DAG = dagbag.dags[TEST_DAG_ID]
         return test_dag, self._get_dag_file_docstring(test_dag.fileloc)
 
-    @staticmethod
-    def clear_db():
+    def clear_db(self):
         clear_db_dags()
         clear_db_serialized_dags()
         clear_db_dag_code()
 
     def test_should_respond_200_text(self, test_client):
-        response = test_client.get(self.dag_sources_url, headers={"Accept": "text/plain"})
+        response: Response = test_client.get(self.dag_sources_url, headers={"Accept": "text/plain"})
+
+        assert isinstance(response, Response)
         assert 200 == response.status_code
         assert len(self.dag_docstring) > 0
-        assert self.dag_docstring == response.data.decode()
-        assert "text/plain" == response.headers["Content-Type"]
+        assert self.dag_docstring in response.content.decode()
+        assert response.headers["Content-Type"].startswith("text/plain")
 
     def test_should_respond_200_json(self, test_client):
-        response = test_client.get(self.dag_sources_url, headers={"Accept": "application/json"})
+        response: Response = test_client.get(
+            self.dag_sources_url,
+            headers={"Accept": "application/json"},
+        )
+        assert isinstance(response, Response)
         assert 200 == response.status_code
         assert len(self.dag_docstring) > 0
-        assert isinstance(response.json, dict)
-        assert len(response.json["content"]) > 0
-        assert self.dag_docstring == response.json["content"]
-        assert "application/json" == response.headers["Content-Type"]
+        res_json = response.json()
+        assert isinstance(res_json, dict)
+        assert len(res_json["content"]) > 0
+        assert self.dag_docstring in res_json["content"]
+        assert response.headers["Content-Type"].startswith("application/json")
+
+    def test_should_respond_406_unsupport_mime_type(self, test_client):
+        response = test_client.get(
+            self.dag_sources_url,
+            headers={"Accept": "text/html"},
+        )
+        assert 406 == response.status_code
 
     def test_should_respond_404(self, test_client):
         wrong_fileloc = "abcd1234"