diff --git a/airflow/api_connexion/endpoints/dag_source_endpoint.py b/airflow/api_connexion/endpoints/dag_source_endpoint.py index 71982b80f48a4..32fce382226e1 100644 --- a/airflow/api_connexion/endpoints/dag_source_endpoint.py +++ b/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) diff --git a/airflow/api_connexion/openapi/v1.yaml b/airflow/api_connexion/openapi/v1.yaml index 8e341d62f0f10..9da8f73672b18 100644 --- a/airflow/api_connexion/openapi/v1.yaml +++ b/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. diff --git a/airflow/api_connexion/schemas/dag_source_schema.py b/airflow/api_connexion/schemas/dag_source_schema.py new file mode 100644 index 0000000000000..6ce65c80ac642 --- /dev/null +++ b/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() diff --git a/tests/api_connexion/endpoints/test_dag_source_endpoint.py b/tests/api_connexion/endpoints/test_dag_source_endpoint.py index 9b07955eb0c6c..49a7bac3bd1e6 100644 --- a/tests/api_connexion/endpoints/test_dag_source_endpoint.py +++ b/tests/api_connexion/endpoints/test_dag_source_endpoint.py @@ -14,14 +14,25 @@ # 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() @@ -29,18 +40,94 @@ def setUpClass(cls) -> None: 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)