Add DAG Source endpoint (apache#9322)

airflow/api_connexion/endpoints/dag_source_endpoint.py

@@ -14,13 +14,34 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
+import logging
 
-# TODO(mik-laj): We have to implement it.
-#     Do you want to help? Please look at: https://github.com/apache/airflow/issues/8137
+from flask import Response, current_app, request
+from itsdangerous import BadSignature, URLSafeSerializer
 
+from airflow.api_connexion.exceptions import NotFound
+from airflow.api_connexion.schemas.dag_source_schema import dag_source_schema
+from airflow.models.dagcode import DagCode
 
-def get_dag_source():
+log = logging.getLogger(__name__)
+
+
+def get_dag_source(file_token: str):
     """
     Get source code using file token
     """
-    raise NotImplementedError("Not implemented yet.")
+    secret_key = current_app.config["SECRET_KEY"]
+    auth_s = URLSafeSerializer(secret_key)
+    try:
+        path = auth_s.loads(file_token)
+        dag_source = DagCode.code(path)
+    except (BadSignature, FileNotFoundError):
+        raise NotFound("Dag source not found")
+
+    return_type = request.accept_mimetypes.best_match(['text/plain', 'application/json'])
+    if return_type == 'text/plain':
+        return Response(dag_source, headers={'Content-Type': return_type})
+    if return_type == 'application/json':
+        content = dag_source_schema.dumps(dict(content=dag_source))
+        return Response(content, headers={'Content-Type': return_type})
+    return Response("Not Allowed Accept Header", status=406)

airflow/api_connexion/openapi/v1.yaml

@@ -1025,12 +1025,18 @@ paths:
                 properties:
                   content:
                     type: string
+            plain/text:
+              schema:
+                type: string
+
         '401':
           $ref: '#/components/responses/Unauthenticated'
         '403':
           $ref: '#/components/responses/PermissionDenied'
         '404':
           $ref: '#/components/responses/NotFound'
+        '406':
+          $ref: '#/components/responses/NotAcceptable'
 
   /config:
     get:
@@ -2405,6 +2411,13 @@ components:
         application/json:
           schema:
             $ref: '#/components/schemas/Error'
+    # 406
+    'NotAcceptable':
+      description: A specified Accept header is not allowed.
+      content:
+        application/json:
+          schema:
+            $ref: '#/components/schemas/Error'
     # 409
     'AlreadyExists':
       description: The resource that a client tried to create already exists.

airflow/api_connexion/schemas/dag_source_schema.py

@@ -0,0 +1,26 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from marshmallow import Schema, fields
+
+
+class DagSourceSchema(Schema):
+    """Dag Source schema"""
+    content = fields.String(dump_only=True)
+
+
+dag_source_schema = DagSourceSchema()

tests/api_connexion/endpoints/test_dag_source_endpoint.py

@@ -14,33 +14,120 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
+import ast
+import os
 import unittest
+from unittest import mock
 
-import pytest
+from itsdangerous import URLSafeSerializer
+from parameterized import parameterized
 
+from airflow import DAG
+from airflow.configuration import conf
+from airflow.models import DagBag
 from airflow.www import app
+from tests.test_utils.db import clear_db_dag_code, clear_db_dags, clear_db_serialized_dags
 
+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")
 
-class TestDagSourceEndpoint(unittest.TestCase):
+
+class TestGetSource(unittest.TestCase):
     @classmethod
     def setUpClass(cls) -> None:
         super().setUpClass()
         cls.app = app.create_app(testing=True)  # type:ignore
 
     def setUp(self) -> None:
         self.client = self.app.test_client()  # type:ignore
+        self.clear_db()
 
+    def tearDown(self) -> None:
+        self.clear_db()
 
-class TestGetSource(unittest.TestCase):
-    @classmethod
-    def setUpClass(cls) -> None:
-        super().setUpClass()
-        cls.app = app.create_app(testing=True)  # type:ignore
+    @staticmethod
+    def clear_db():
+        clear_db_dags()
+        clear_db_serialized_dags()
+        clear_db_dag_code()
 
-    def setUp(self) -> None:
-        self.client = self.app.test_client()  # type:ignore
+    @staticmethod
+    def _get_dag_file_docstring(fileloc: str) -> str:
+        with open(fileloc) as f:
+            file_contents = f.read()
+        module = ast.parse(file_contents)
+        docstring = ast.get_docstring(module)
+        return docstring
+
+    @parameterized.expand([(True,), (False,)])
+    def test_should_response_200_text(self, store_dag_code):
+        serializer = URLSafeSerializer(conf.get('webserver', 'SECRET_KEY'))
+        with mock.patch(
+            "airflow.models.dag.settings.STORE_DAG_CODE", store_dag_code
+        ), mock.patch("airflow.models.dagcode.STORE_DAG_CODE", store_dag_code):
+            dagbag = DagBag(dag_folder=EXAMPLE_DAG_FILE)
+            dagbag.sync_to_db()
+            first_dag: DAG = next(iter(dagbag.dags.values()))
+            dag_docstring = self._get_dag_file_docstring(first_dag.fileloc)
+
+            url = f"/api/v1/dagSources/{serializer.dumps(first_dag.fileloc)}"
+            response = self.client.get(url, headers={
+                "Accept": "text/plain"
+            })
+
+            self.assertEqual(200, response.status_code)
+            self.assertIn(dag_docstring, response.data.decode())
+            self.assertEqual('text/plain', response.headers['Content-Type'])
+
+    @parameterized.expand([(True,), (False,)])
+    def test_should_response_200_json(self, store_dag_code):
+        serializer = URLSafeSerializer(conf.get('webserver', 'SECRET_KEY'))
+        with mock.patch(
+            "airflow.models.dag.settings.STORE_DAG_CODE", store_dag_code
+        ), mock.patch("airflow.models.dagcode.STORE_DAG_CODE", store_dag_code):
+            dagbag = DagBag(dag_folder=EXAMPLE_DAG_FILE)
+            dagbag.sync_to_db()
+            first_dag: DAG = next(iter(dagbag.dags.values()))
+            dag_docstring = self._get_dag_file_docstring(first_dag.fileloc)
+
+            url = f"/api/v1/dagSources/{serializer.dumps(first_dag.fileloc)}"
+            response = self.client.get(url, headers={
+                "Accept": 'application/json'
+            })
+
+            self.assertEqual(200, response.status_code)
+            self.assertIn(
+                dag_docstring,
+                response.json['content']
+            )
+            self.assertEqual('application/json', response.headers['Content-Type'])
+
+    @parameterized.expand([(True,), (False,)])
+    def test_should_response_406(self, store_dag_code):
+        serializer = URLSafeSerializer(conf.get('webserver', 'SECRET_KEY'))
+        with mock.patch(
+            "airflow.models.dag.settings.STORE_DAG_CODE", store_dag_code
+        ), mock.patch("airflow.models.dagcode.STORE_DAG_CODE", store_dag_code):
+            dagbag = DagBag(dag_folder=EXAMPLE_DAG_FILE)
+            dagbag.sync_to_db()
+            first_dag: DAG = next(iter(dagbag.dags.values()))
+
+            url = f"/api/v1/dagSources/{serializer.dumps(first_dag.fileloc)}"
+            response = self.client.get(url, headers={
+                "Accept": 'image/webp'
+            })
+
+            self.assertEqual(406, response.status_code)
+
+    @parameterized.expand([(True,), (False,)])
+    def test_should_response_404(self, store_dag_code):
+        with mock.patch(
+            "airflow.models.dag.settings.STORE_DAG_CODE", store_dag_code
+        ), mock.patch("airflow.models.dagcode.STORE_DAG_CODE", store_dag_code):
+            wrong_fileloc = "abcd1234"
+            url = f"/api/v1/dagSources/{wrong_fileloc}"
+            response = self.client.get(url, headers={
+                "Accept": 'application/json'
+            })
 
-    @pytest.mark.skip(reason="Not implemented yet")
-    def test_should_response_200(self):
-        response = self.client.get("/api/v1/health")
-        assert response.status_code == 200
+            self.assertEqual(404, response.status_code)