From c20e54d95b1aef26b4c29cc6c4ac5a31ba0cbaf3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tobiasz=20K=C4=99dzierski?= Date: Tue, 16 Jun 2020 05:26:23 +0200 Subject: [PATCH 01/11] DAG Source endpoint --- .../endpoints/dag_source_endpoint.py | 23 +++++++++++++++---- airflow/api_connexion/openapi/v1.yaml | 5 +--- 2 files changed, 20 insertions(+), 8 deletions(-) diff --git a/airflow/api_connexion/endpoints/dag_source_endpoint.py b/airflow/api_connexion/endpoints/dag_source_endpoint.py index 71982b80f48a4..70f430dde37e1 100644 --- a/airflow/api_connexion/endpoints/dag_source_endpoint.py +++ b/airflow/api_connexion/endpoints/dag_source_endpoint.py @@ -15,12 +15,27 @@ # specific language governing permissions and limitations # under the License. -# TODO(mik-laj): We have to implement it. -# Do you want to help? Please look at: https://github.com/apache/airflow/issues/8137 +import logging +from flask import current_app +from itsdangerous import BadSignature, URLSafeSerializer -def get_dag_source(): +from airflow.api_connexion.exceptions import NotFound +from airflow.models.dagcode import DagCode + +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): + NotFound("Dag Source not found") + else: + return dag_source diff --git a/airflow/api_connexion/openapi/v1.yaml b/airflow/api_connexion/openapi/v1.yaml index 8e341d62f0f10..ff5e6e679a59f 100644 --- a/airflow/api_connexion/openapi/v1.yaml +++ b/airflow/api_connexion/openapi/v1.yaml @@ -1021,10 +1021,7 @@ paths: content: application/json: schema: - type: object - properties: - content: - type: string + type: string '401': $ref: '#/components/responses/Unauthenticated' '403': From 8eb8808747c4737580cc7c45d40b8e1e5dfd28a0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Kamil=20Bregu=C5=82a?= Date: Thu, 18 Jun 2020 05:29:00 +0200 Subject: [PATCH 02/11] fixup! DAG Source endpoint --- .../endpoints/dag_source_endpoint.py | 17 ++++-- airflow/api_connexion/openapi/v1.yaml | 7 +++ .../endpoints/test_dag_source_endpoint.py | 55 +++++++++++++++++-- 3 files changed, 69 insertions(+), 10 deletions(-) diff --git a/airflow/api_connexion/endpoints/dag_source_endpoint.py b/airflow/api_connexion/endpoints/dag_source_endpoint.py index 70f430dde37e1..b8af430943425 100644 --- a/airflow/api_connexion/endpoints/dag_source_endpoint.py +++ b/airflow/api_connexion/endpoints/dag_source_endpoint.py @@ -14,10 +14,10 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. - +import json import logging -from flask import current_app +from flask import Response, current_app, request from itsdangerous import BadSignature, URLSafeSerializer from airflow.api_connexion.exceptions import NotFound @@ -36,6 +36,13 @@ def get_dag_source(file_token: str): path = auth_s.loads(file_token) dag_source = DagCode.code(path) except (BadSignature, FileNotFoundError): - NotFound("Dag Source not found") - else: - return dag_source + raise NotFound("Dag Source not found") + + from accept_types import get_best_match + return_type = get_best_match(request.headers.get('Accept'), ['text/plain', 'application/json']) + if return_type == 'text/plain': + return Response(dag_source, headers={'Content-Type': return_type}) + if return_type == 'application/json': + content = json.dumps(dict(content=dag_source)) + return Response(content, headers={'Content-Type': return_type}) + return '', 406 diff --git a/airflow/api_connexion/openapi/v1.yaml b/airflow/api_connexion/openapi/v1.yaml index ff5e6e679a59f..fa99742c984b7 100644 --- a/airflow/api_connexion/openapi/v1.yaml +++ b/airflow/api_connexion/openapi/v1.yaml @@ -1020,8 +1020,15 @@ paths: description: Successful response. content: application/json: + schema: + type: object + properties: + content: + type: string + plain/text: schema: type: string + '401': $ref: '#/components/responses/Unauthenticated' '403': diff --git a/tests/api_connexion/endpoints/test_dag_source_endpoint.py b/tests/api_connexion/endpoints/test_dag_source_endpoint.py index 9b07955eb0c6c..6b577bdf0d849 100644 --- a/tests/api_connexion/endpoints/test_dag_source_endpoint.py +++ b/tests/api_connexion/endpoints/test_dag_source_endpoint.py @@ -14,11 +14,20 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. +import os import unittest -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.config import conf_vars + +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): @@ -40,7 +49,43 @@ def setUpClass(cls) -> None: def setUp(self) -> None: self.client = self.app.test_client() # type:ignore - @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 + @parameterized.expand([("True",), ("False",)]) + def test_should_response_200_text(self, store_serialized_dags): + serializer = URLSafeSerializer(conf.get('webserver', 'SECRET_KEY')) + with conf_vars( + {("core", "store_serialized_dags"): store_serialized_dags} + ): + 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": "text/plain" + }) + + self.assertEqual(200, response.status_code) + self.assertIn("Example DAG demonstrating the usage of the BashOperator.", response.data.decode()) + self.assertEqual('text/plain', response.headers['Content-Type']) + + @parameterized.expand([("True",), ("False",)]) + def test_should_response_200_json(self, store_serialized_dags): + serializer = URLSafeSerializer(conf.get('webserver', 'SECRET_KEY')) + with conf_vars( + {("core", "store_serialized_dags"): store_serialized_dags} + ): + 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": 'application/json' + }) + + self.assertEqual(200, response.status_code) + self.assertIn( + "Example DAG demonstrating the usage of the BashOperator.", + response.json['content'] + ) + self.assertEqual('application/json', response.headers['Content-Type']) From b88121910de7316c4756f7f1252511f5b6596213 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tobiasz=20K=C4=99dzierski?= Date: Sun, 28 Jun 2020 17:56:42 +0200 Subject: [PATCH 03/11] fixup! fixup! DAG Source endpoint --- .../api_connexion/endpoints/dag_source_endpoint.py | 3 +-- .../endpoints/test_dag_source_endpoint.py | 14 ++++++++++++-- 2 files changed, 13 insertions(+), 4 deletions(-) diff --git a/airflow/api_connexion/endpoints/dag_source_endpoint.py b/airflow/api_connexion/endpoints/dag_source_endpoint.py index b8af430943425..f3aeba9d23517 100644 --- a/airflow/api_connexion/endpoints/dag_source_endpoint.py +++ b/airflow/api_connexion/endpoints/dag_source_endpoint.py @@ -38,8 +38,7 @@ def get_dag_source(file_token: str): except (BadSignature, FileNotFoundError): raise NotFound("Dag Source not found") - from accept_types import get_best_match - return_type = get_best_match(request.headers.get('Accept'), ['text/plain', 'application/json']) + 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': diff --git a/tests/api_connexion/endpoints/test_dag_source_endpoint.py b/tests/api_connexion/endpoints/test_dag_source_endpoint.py index 6b577bdf0d849..9867932432754 100644 --- a/tests/api_connexion/endpoints/test_dag_source_endpoint.py +++ b/tests/api_connexion/endpoints/test_dag_source_endpoint.py @@ -14,6 +14,7 @@ # 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 @@ -49,6 +50,13 @@ def setUpClass(cls) -> None: def setUp(self) -> None: self.client = self.app.test_client() # type:ignore + def _get_dag_file_docstring(self, fileloc: str) -> str: + with open(fileloc) as fd: + file_contents = fd.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_serialized_dags): serializer = URLSafeSerializer(conf.get('webserver', 'SECRET_KEY')) @@ -58,6 +66,7 @@ def test_should_response_200_text(self, store_serialized_dags): 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={ @@ -65,7 +74,7 @@ def test_should_response_200_text(self, store_serialized_dags): }) self.assertEqual(200, response.status_code) - self.assertIn("Example DAG demonstrating the usage of the BashOperator.", response.data.decode()) + self.assertIn(dag_docstring, response.data.decode()) self.assertEqual('text/plain', response.headers['Content-Type']) @parameterized.expand([("True",), ("False",)]) @@ -77,6 +86,7 @@ def test_should_response_200_json(self, store_serialized_dags): 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={ @@ -85,7 +95,7 @@ def test_should_response_200_json(self, store_serialized_dags): self.assertEqual(200, response.status_code) self.assertIn( - "Example DAG demonstrating the usage of the BashOperator.", + dag_docstring, response.json['content'] ) self.assertEqual('application/json', response.headers['Content-Type']) From c3a52369e8cfa779cc4e6aacb2e218d08d295450 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tobiasz=20K=C4=99dzierski?= Date: Sun, 5 Jul 2020 12:56:28 +0200 Subject: [PATCH 04/11] fixup! fixup! fixup! DAG Source endpoint --- tests/api_connexion/endpoints/test_dag_source_endpoint.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/tests/api_connexion/endpoints/test_dag_source_endpoint.py b/tests/api_connexion/endpoints/test_dag_source_endpoint.py index 9867932432754..7619924e650b6 100644 --- a/tests/api_connexion/endpoints/test_dag_source_endpoint.py +++ b/tests/api_connexion/endpoints/test_dag_source_endpoint.py @@ -50,9 +50,10 @@ def setUpClass(cls) -> None: def setUp(self) -> None: self.client = self.app.test_client() # type:ignore - def _get_dag_file_docstring(self, fileloc: str) -> str: - with open(fileloc) as fd: - file_contents = fd.read() + @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 From bd06c239cf94d573636d8ff9c91b56dd174e5b29 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tobiasz=20K=C4=99dzierski?= Date: Sun, 5 Jul 2020 13:42:45 +0200 Subject: [PATCH 05/11] Add 406 to spec --- airflow/api_connexion/endpoints/dag_source_endpoint.py | 2 +- airflow/api_connexion/openapi/v1.yaml | 9 +++++++++ 2 files changed, 10 insertions(+), 1 deletion(-) diff --git a/airflow/api_connexion/endpoints/dag_source_endpoint.py b/airflow/api_connexion/endpoints/dag_source_endpoint.py index f3aeba9d23517..e31179e5e4325 100644 --- a/airflow/api_connexion/endpoints/dag_source_endpoint.py +++ b/airflow/api_connexion/endpoints/dag_source_endpoint.py @@ -44,4 +44,4 @@ def get_dag_source(file_token: str): if return_type == 'application/json': content = json.dumps(dict(content=dag_source)) return Response(content, headers={'Content-Type': return_type}) - return '', 406 + 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 fa99742c984b7..9da8f73672b18 100644 --- a/airflow/api_connexion/openapi/v1.yaml +++ b/airflow/api_connexion/openapi/v1.yaml @@ -1035,6 +1035,8 @@ paths: $ref: '#/components/responses/PermissionDenied' '404': $ref: '#/components/responses/NotFound' + '406': + $ref: '#/components/responses/NotAcceptable' /config: get: @@ -2409,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. From 814768de22b01717f478265f09f3adf149a05b5e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tobiasz=20K=C4=99dzierski?= Date: Sun, 5 Jul 2020 13:51:56 +0200 Subject: [PATCH 06/11] Add tests for 404 and 406 --- .../endpoints/test_dag_source_endpoint.py | 38 +++++++++++++++++-- 1 file changed, 34 insertions(+), 4 deletions(-) diff --git a/tests/api_connexion/endpoints/test_dag_source_endpoint.py b/tests/api_connexion/endpoints/test_dag_source_endpoint.py index 7619924e650b6..6490af07355bb 100644 --- a/tests/api_connexion/endpoints/test_dag_source_endpoint.py +++ b/tests/api_connexion/endpoints/test_dag_source_endpoint.py @@ -59,10 +59,10 @@ def _get_dag_file_docstring(fileloc: str) -> str: return docstring @parameterized.expand([("True",), ("False",)]) - def test_should_response_200_text(self, store_serialized_dags): + def test_should_response_200_text(self, store_dag_code): serializer = URLSafeSerializer(conf.get('webserver', 'SECRET_KEY')) with conf_vars( - {("core", "store_serialized_dags"): store_serialized_dags} + {("core", "store_serialized_dags"): store_dag_code} ): dagbag = DagBag(dag_folder=EXAMPLE_DAG_FILE) dagbag.sync_to_db() @@ -79,10 +79,10 @@ def test_should_response_200_text(self, store_serialized_dags): self.assertEqual('text/plain', response.headers['Content-Type']) @parameterized.expand([("True",), ("False",)]) - def test_should_response_200_json(self, store_serialized_dags): + def test_should_response_200_json(self, store_dag_code): serializer = URLSafeSerializer(conf.get('webserver', 'SECRET_KEY')) with conf_vars( - {("core", "store_serialized_dags"): store_serialized_dags} + {("core", "store_serialized_dags"): store_dag_code} ): dagbag = DagBag(dag_folder=EXAMPLE_DAG_FILE) dagbag.sync_to_db() @@ -100,3 +100,33 @@ def test_should_response_200_json(self, store_serialized_dags): 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 conf_vars( + {("core", "store_serialized_dags"): 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 conf_vars( + {("core", "store_serialized_dags"): store_dag_code} + ): + wrong_fileloc = "abcd1234" + url = f"/api/v1/dagSources/{wrong_fileloc}" + response = self.client.get(url, headers={ + "Accept": 'application/json' + }) + + self.assertEqual(404, response.status_code) From 8a72df0a0fff77c04f78cb7860c2f1af3b2f9978 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tobiasz=20K=C4=99dzierski?= Date: Wed, 8 Jul 2020 11:42:06 +0200 Subject: [PATCH 07/11] Add marshmallow schema --- .../endpoints/dag_source_endpoint.py | 6 ++--- .../schemas/dag_source_schema.py | 26 +++++++++++++++++++ 2 files changed, 29 insertions(+), 3 deletions(-) create mode 100644 airflow/api_connexion/schemas/dag_source_schema.py diff --git a/airflow/api_connexion/endpoints/dag_source_endpoint.py b/airflow/api_connexion/endpoints/dag_source_endpoint.py index e31179e5e4325..ba2ae65b12ffc 100644 --- a/airflow/api_connexion/endpoints/dag_source_endpoint.py +++ b/airflow/api_connexion/endpoints/dag_source_endpoint.py @@ -14,13 +14,13 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -import json import logging 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 log = logging.getLogger(__name__) @@ -42,6 +42,6 @@ def get_dag_source(file_token: str): if return_type == 'text/plain': return Response(dag_source, headers={'Content-Type': return_type}) if return_type == 'application/json': - content = json.dumps(dict(content=dag_source)) - return Response(content, headers={'Content-Type': return_type}) + content = dag_source_schema.dumps(dict(content=dag_source)) + return Response(content.data, headers={'Content-Type': return_type}) return Response("Not Allowed Accept Header", status=406) 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() From 6c5b49c61bcf3b56fb405bf2c09a8ca7d106494f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tobiasz=20K=C4=99dzierski?= Date: Sun, 12 Jul 2020 18:28:27 +0200 Subject: [PATCH 08/11] Remove unused testing class --- .../endpoints/test_dag_source_endpoint.py | 50 +++++++++---------- 1 file changed, 23 insertions(+), 27 deletions(-) diff --git a/tests/api_connexion/endpoints/test_dag_source_endpoint.py b/tests/api_connexion/endpoints/test_dag_source_endpoint.py index 6490af07355bb..5fda1d2427858 100644 --- a/tests/api_connexion/endpoints/test_dag_source_endpoint.py +++ b/tests/api_connexion/endpoints/test_dag_source_endpoint.py @@ -17,6 +17,7 @@ import ast import os import unittest +from unittest import mock from itsdangerous import URLSafeSerializer from parameterized import parameterized @@ -25,22 +26,12 @@ from airflow.configuration import conf from airflow.models import DagBag from airflow.www import app -from tests.test_utils.config import conf_vars +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): - @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 - - class TestGetSource(unittest.TestCase): @classmethod def setUpClass(cls) -> None: @@ -50,6 +41,11 @@ def setUpClass(cls) -> None: def setUp(self) -> None: self.client = self.app.test_client() # type:ignore + def tearDown(self) -> None: + clear_db_dags() + clear_db_serialized_dags() + clear_db_dag_code() + @staticmethod def _get_dag_file_docstring(fileloc: str) -> str: with open(fileloc) as f: @@ -58,12 +54,12 @@ def _get_dag_file_docstring(fileloc: str) -> str: docstring = ast.get_docstring(module) return docstring - @parameterized.expand([("True",), ("False",)]) + @parameterized.expand([(True,), (False,)]) def test_should_response_200_text(self, store_dag_code): serializer = URLSafeSerializer(conf.get('webserver', 'SECRET_KEY')) - with conf_vars( - {("core", "store_serialized_dags"): 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): dagbag = DagBag(dag_folder=EXAMPLE_DAG_FILE) dagbag.sync_to_db() first_dag: DAG = next(iter(dagbag.dags.values())) @@ -78,12 +74,12 @@ def test_should_response_200_text(self, store_dag_code): self.assertIn(dag_docstring, response.data.decode()) self.assertEqual('text/plain', response.headers['Content-Type']) - @parameterized.expand([("True",), ("False",)]) + @parameterized.expand([(True,), (False,)]) def test_should_response_200_json(self, store_dag_code): serializer = URLSafeSerializer(conf.get('webserver', 'SECRET_KEY')) - with conf_vars( - {("core", "store_serialized_dags"): 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): dagbag = DagBag(dag_folder=EXAMPLE_DAG_FILE) dagbag.sync_to_db() first_dag: DAG = next(iter(dagbag.dags.values())) @@ -101,12 +97,12 @@ def test_should_response_200_json(self, store_dag_code): ) self.assertEqual('application/json', response.headers['Content-Type']) - @parameterized.expand([("True",), ("False",)]) + @parameterized.expand([(True,), (False,)]) def test_should_response_406(self, store_dag_code): serializer = URLSafeSerializer(conf.get('webserver', 'SECRET_KEY')) - with conf_vars( - {("core", "store_serialized_dags"): 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): dagbag = DagBag(dag_folder=EXAMPLE_DAG_FILE) dagbag.sync_to_db() first_dag: DAG = next(iter(dagbag.dags.values())) @@ -118,11 +114,11 @@ def test_should_response_406(self, store_dag_code): self.assertEqual(406, response.status_code) - @parameterized.expand([("True",), ("False",)]) + @parameterized.expand([(True,), (False,)]) def test_should_response_404(self, store_dag_code): - with conf_vars( - {("core", "store_serialized_dags"): 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={ From cd66b790424b7f6e9a9de85c1bc0c14d9fb78ab7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tobiasz=20K=C4=99dzierski?= Date: Sun, 12 Jul 2020 18:30:51 +0200 Subject: [PATCH 09/11] Clear db during setUp and tearDown to avoid side effects --- tests/api_connexion/endpoints/test_dag_source_endpoint.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/tests/api_connexion/endpoints/test_dag_source_endpoint.py b/tests/api_connexion/endpoints/test_dag_source_endpoint.py index 5fda1d2427858..49a7bac3bd1e6 100644 --- a/tests/api_connexion/endpoints/test_dag_source_endpoint.py +++ b/tests/api_connexion/endpoints/test_dag_source_endpoint.py @@ -40,8 +40,13 @@ 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() + + @staticmethod + def clear_db(): clear_db_dags() clear_db_serialized_dags() clear_db_dag_code() From 3874a4592b3f6fd53e122fd9ffdbcf330b5bc9ac Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tobiasz=20K=C4=99dzierski?= Date: Mon, 13 Jul 2020 07:49:08 +0200 Subject: [PATCH 10/11] Apply suggestions from code review Co-authored-by: Tomek Urbaszek --- airflow/api_connexion/endpoints/dag_source_endpoint.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/airflow/api_connexion/endpoints/dag_source_endpoint.py b/airflow/api_connexion/endpoints/dag_source_endpoint.py index ba2ae65b12ffc..40ff202b732f6 100644 --- a/airflow/api_connexion/endpoints/dag_source_endpoint.py +++ b/airflow/api_connexion/endpoints/dag_source_endpoint.py @@ -36,7 +36,7 @@ def get_dag_source(file_token: str): path = auth_s.loads(file_token) dag_source = DagCode.code(path) except (BadSignature, FileNotFoundError): - raise NotFound("Dag Source not found") + raise NotFound("Dag source not found") return_type = request.accept_mimetypes.best_match(['text/plain', 'application/json']) if return_type == 'text/plain': From e7ff213a8d8f153c708e89fbac05b5dbcd1ff287 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tobiasz=20K=C4=99dzierski?= Date: Mon, 13 Jul 2020 08:54:31 +0200 Subject: [PATCH 11/11] Update endpoint for new version of marshmallow --- airflow/api_connexion/endpoints/dag_source_endpoint.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/airflow/api_connexion/endpoints/dag_source_endpoint.py b/airflow/api_connexion/endpoints/dag_source_endpoint.py index 40ff202b732f6..32fce382226e1 100644 --- a/airflow/api_connexion/endpoints/dag_source_endpoint.py +++ b/airflow/api_connexion/endpoints/dag_source_endpoint.py @@ -43,5 +43,5 @@ def get_dag_source(file_token: str): 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.data, headers={'Content-Type': return_type}) + return Response(content, headers={'Content-Type': return_type}) return Response("Not Allowed Accept Header", status=406)