From 049b9e7bffacf5f2fefceb29cd548ec841d95722 Mon Sep 17 00:00:00 2001 From: EphraimBuddy Date: Tue, 23 Jun 2020 01:52:46 +0100 Subject: [PATCH 1/8] add crud endpoint for xcom --- .../api_connexion/endpoints/xcom_endpoint.py | 97 +++++++++++++++++-- airflow/api_connexion/openapi/v1.yaml | 1 + .../endpoints/test_xcom_endpoint.py | 93 ++++++++++++++++-- 3 files changed, 173 insertions(+), 18 deletions(-) diff --git a/airflow/api_connexion/endpoints/xcom_endpoint.py b/airflow/api_connexion/endpoints/xcom_endpoint.py index cd317ad819979..3b0cf95b50e10 100644 --- a/airflow/api_connexion/endpoints/xcom_endpoint.py +++ b/airflow/api_connexion/endpoints/xcom_endpoint.py @@ -14,25 +14,51 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -from flask import request from sqlalchemy import and_, func -from sqlalchemy.orm.session import Session +from typing import List, Optional from airflow.api_connexion import parameters -from airflow.api_connexion.exceptions import NotFound from airflow.api_connexion.schemas.xcom_schema import ( XComCollection, XComCollectionItemSchema, XComCollectionSchema, xcom_collection_item_schema, xcom_collection_schema, ) +from flask import Response, request +from marshmallow import ValidationError +from sqlalchemy.orm.session import Session + +from airflow.api_connexion.exceptions import AlreadyExists, BadRequest, NotFound +from airflow.api_connexion.schemas.xcom_schema import xcom_schema from airflow.models import DagRun as DR, XCom from airflow.utils.session import provide_session -def delete_xcom_entry(): +@provide_session +def delete_xcom_entry( + dag_id: str, + dag_run_id: str, + task_id: str, + xcom_key: str, + session: Session +) -> Response: """ Delete an XCom entry """ - raise NotImplementedError("Not implemented yet.") + + dag_run = session.query(DR).filter(DR.run_id == dag_run_id, + DR.dag_id == dag_id).first() + if not dag_run: + raise NotFound(f'DAGRun with dag_id:{dag_id} and run_id:{dag_run_id} not found') + + query = session.query(XCom) + query = query.filter(XCom.dag_id == dag_id, + XCom.task_id == task_id, + XCom.key == xcom_key) + entry = query.delete() + + if not entry: + raise NotFound("XCom entry not found") + + return Response('', 204) @provide_session @@ -89,15 +115,68 @@ def get_xcom_entry( return xcom_collection_item_schema.dump(query_object) -def patch_xcom_entry(): +@provide_session +def patch_xcom_entry( + dag_id: str, + dag_run_id: str, + task_id: str, + xcom_key: str, + update_mask: Optional[List[str]], + session: Session +): """ Update an XCom entry """ - raise NotImplementedError("Not implemented yet.") + + query = session.query(XCom).filter(XCom.dag_id == dag_id, + XCom.key == xcom_key, + XCom.task_id == task_id) + xcom_entry = query.join(DR).filter(DR.run_id == dag_run_id).first() + if not xcom_entry: + raise NotFound(f"Xcom not found") + try: + body = xcom_schema.load(request.json) + except ValidationError as err: + raise BadRequest() + data = body.data + if update_mask: + update_mask = [i.strip() for i in update_mask] + data_ = {} + for field in update_mask: + if field in data: + data_[field] = data[field] + else: + raise BadRequest(f"Unknown field {field} in update_mask") + data = data_ + for key in data: + setattr(xcom_entry, key, data[key]) + session.add(xcom_entry) + session.commit() + return xcom_schema.dumps(xcom_entry) -def post_xcom_entries(): +@provide_session +def post_xcom_entries(dag_id: str, + dag_run_id: str, + task_id: str, + session: Session + ): """ Create an XCom entry """ - raise NotImplementedError("Not implemented yet.") + dag_run = session.query(DR).filter(DR.run_id == dag_run_id, + DR.dag_id == dag_id).first() + if not dag_run: + raise NotFound(f'DAGRun with dag_id:{dag_id} and run_id:{dag_run_id} not found') + + xcom_entry = xcom_schema.load(request.json) + + xcom = session.query(XCom).filter(XCom.dag_id == dag_id, + XCom.task_id == task_id).first() + if not xcom: + data = xcom_entry.data + xcom = XCom(**data) + session.add(xcom) + session.flush() + return xcom_schema.dump(xcom) + raise AlreadyExists("XCom entry already exists") diff --git a/airflow/api_connexion/openapi/v1.yaml b/airflow/api_connexion/openapi/v1.yaml index 1794b65264c09..21f7cc8a28cae 100644 --- a/airflow/api_connexion/openapi/v1.yaml +++ b/airflow/api_connexion/openapi/v1.yaml @@ -1616,6 +1616,7 @@ components: properties: value: type: string + nullable: true # Python objects # Based on diff --git a/tests/api_connexion/endpoints/test_xcom_endpoint.py b/tests/api_connexion/endpoints/test_xcom_endpoint.py index 3f36e602d8534..1a4cab21231a1 100644 --- a/tests/api_connexion/endpoints/test_xcom_endpoint.py +++ b/tests/api_connexion/endpoints/test_xcom_endpoint.py @@ -25,6 +25,11 @@ from airflow.utils.types import DagRunType from airflow.www import app +DAG_ID = 'test-dag-id' +TASK_ID = 'test-task-id' +EXECUTION_DATE = '2005-04-02T00:00:00+00:00' +XCOM_KEY = 'test-xcom-key' + class TestXComEndpoint(unittest.TestCase): @classmethod @@ -37,6 +42,7 @@ def setUp(self) -> None: Setup For XCom endpoint TC """ self.client = self.app.test_client() # type:ignore + self.execution_date_parsed = parse_execution_date(EXECUTION_DATE) # clear existing xcoms with create_session() as session: session.query(XCom).delete() @@ -50,14 +56,64 @@ def tearDown(self) -> None: session.query(XCom).delete() session.query(DR).delete() + def _create_xcom(self, session): + + xcom_model = XCom(key=XCOM_KEY, + execution_date=self.execution_date_parsed, + task_id=TASK_ID, + dag_id=DAG_ID, + timestamp=self.execution_date_parsed) + session.add(xcom_model) + session.commit() + + def _create_dag_run(self, session): + dag_run_id = DR.generate_run_id(DagRunType.MANUAL, self.execution_date_parsed) + dagrun = DR(dag_id=DAG_ID, + run_id=dag_run_id, + execution_date=self.execution_date_parsed, + start_date=self.execution_date_parsed, + run_type=DagRunType.MANUAL.value) + session.add(dagrun) + session.commit() + return dag_run_id + class TestDeleteXComEntry(TestXComEndpoint): - @pytest.mark.skip(reason="Not implemented yet") - def test_should_response_200(self): - response = self.client.delete( - "/dags/TEST_DAG_ID/taskInstances/TEST_TASK_ID/2005-04-02T00:00:00Z/xcomEntries/XCOM_KEY" + + @provide_session + def test_delete_should_response_204(self, session): + self._create_xcom(session) + dag_run_id = self._create_dag_run(session) + url = (f"/api/v1/dags/{DAG_ID}/dagRuns/{dag_run_id}/taskInstances" + f"/{TASK_ID}/xcomEntries/{XCOM_KEY}") + response = self.client.delete(url) + self.assertEqual(response.status_code, 204) + + @provide_session + def test_delete_should_response_404_for_invalid_dagrun(self, session): + self._create_xcom(session) + dag_run_id = self._create_dag_run(session) + url = (f"/api/v1/dags/INVALID_DAG_ID/dagRuns/{dag_run_id}/taskInstances" + f"/{TASK_ID}/xcomEntries/{XCOM_KEY}") + response = self.client.delete(url) + self.assertEqual(response.status_code, 404) + self.assertEqual( + response.json['title'], + f'DAGRun with dag_id:INVALID_DAG_ID and run_id:{dag_run_id} not found' + ) + + @provide_session + def test_delete_should_response_404_for_invalid_xcom(self, session): + self._create_xcom(session) + dag_run_id = self._create_dag_run(session) + url = (f"/api/v1/dags/{DAG_ID}/dagRuns/{dag_run_id}/taskInstances" + f"/{TASK_ID}/xcomEntries/INVALID_XCOM_KEY") + response = self.client.delete(url) + self.assertEqual(response.status_code, 404) + self.assertEqual( + response.json['title'], + "XCom entry not found" ) - assert response.status_code == 204 class TestGetXComEntry(TestXComEndpoint): @@ -100,6 +156,7 @@ def test_should_response_200(self, session): class TestGetXComEntries(TestXComEndpoint): + @provide_session def test_should_response_200(self, session): dag_id = 'test-dag-id' @@ -255,9 +312,27 @@ def test_should_response_200(self): class TestPostXComEntry(TestXComEndpoint): - @pytest.mark.skip(reason="Not implemented yet") - def test_should_response_200(self): + + @provide_session + def test_post_should_response_200(self, session): + payload = { + 'dag_id': DAG_ID, + 'execution_date': EXECUTION_DATE, + 'key': XCOM_KEY, + 'task_id': TASK_ID, + 'timestamp': EXECUTION_DATE, + 'value': "{}" + } + + dag_run_id = self._create_dag_run(session) + + url = (f"/api/v1/dags/{DAG_ID}/dagRuns/{dag_run_id}/taskInstances" + f"/{TASK_ID}/xcomEntries") response = self.client.post( - "/dags/TEST_DAG_ID/taskInstances/TEST_TASK_ID/2005-04-02T00:00:00Z/xcomEntries/XCOM_KEY" + url, + json=payload ) - assert response.status_code == 200 + self.assertEqual(response.status_code, 200) + xcom_entry = session.query(XCom).all() + self.assertEqual(len(xcom_entry), 1) + self.assertEqual(xcom_entry[0].key, XCOM_KEY) From 18297c63ebf7c9b4615fbfa529dc15bb69d1c8d9 Mon Sep 17 00:00:00 2001 From: EphraimBuddy Date: Tue, 23 Jun 2020 04:41:44 +0100 Subject: [PATCH 2/8] add some tests on patch and post --- .../api_connexion/endpoints/xcom_endpoint.py | 44 ++++--- airflow/api_connexion/schemas/xcom_schema.py | 9 +- .../endpoints/test_xcom_endpoint.py | 122 ++++++++++++++++-- 3 files changed, 146 insertions(+), 29 deletions(-) diff --git a/airflow/api_connexion/endpoints/xcom_endpoint.py b/airflow/api_connexion/endpoints/xcom_endpoint.py index 3b0cf95b50e10..feb5094d2332d 100644 --- a/airflow/api_connexion/endpoints/xcom_endpoint.py +++ b/airflow/api_connexion/endpoints/xcom_endpoint.py @@ -49,12 +49,9 @@ def delete_xcom_entry( if not dag_run: raise NotFound(f'DAGRun with dag_id:{dag_id} and run_id:{dag_run_id} not found') - query = session.query(XCom) - query = query.filter(XCom.dag_id == dag_id, - XCom.task_id == task_id, - XCom.key == xcom_key) - entry = query.delete() - + entry = session.query(XCom).filter(XCom.dag_id == dag_id, + XCom.task_id == task_id, + XCom.key == xcom_key).delete() if not entry: raise NotFound("XCom entry not found") @@ -117,27 +114,31 @@ def get_xcom_entry( @provide_session def patch_xcom_entry( + session: Session, dag_id: str, dag_run_id: str, task_id: str, xcom_key: str, - update_mask: Optional[List[str]], - session: Session + update_mask: Optional[List[str]] = None ): """ Update an XCom entry """ + dag_run = session.query(DR).filter(DR.run_id == dag_run_id, + DR.dag_id == dag_id).first() + if not dag_run: + raise NotFound(f'DAGRun with dag_id:{dag_id} and run_id:{dag_run_id} not found') + + xcom_entry = session.query(XCom).filter(XCom.dag_id == dag_id, + XCom.task_id == task_id, + XCom.key == xcom_key).first() - query = session.query(XCom).filter(XCom.dag_id == dag_id, - XCom.key == xcom_key, - XCom.task_id == task_id) - xcom_entry = query.join(DR).filter(DR.run_id == dag_run_id).first() if not xcom_entry: - raise NotFound(f"Xcom not found") + raise NotFound("XCom not found") try: body = xcom_schema.load(request.json) except ValidationError as err: - raise BadRequest() + raise BadRequest(detail=str(err.messages.get('_schema', err.messages))) data = body.data if update_mask: update_mask = [i.strip() for i in update_mask] @@ -152,7 +153,7 @@ def patch_xcom_entry( setattr(xcom_entry, key, data[key]) session.add(xcom_entry) session.commit() - return xcom_schema.dumps(xcom_entry) + return xcom_schema.dump(xcom_entry) @provide_session @@ -164,12 +165,15 @@ def post_xcom_entries(dag_id: str, """ Create an XCom entry """ - dag_run = session.query(DR).filter(DR.run_id == dag_run_id, - DR.dag_id == dag_id).first() + # TODO add 409 to spec? + + dag_run = session.query(DR).filter(DR.run_id == dag_run_id, DR.dag_id == dag_id).first() if not dag_run: raise NotFound(f'DAGRun with dag_id:{dag_id} and run_id:{dag_run_id} not found') - - xcom_entry = xcom_schema.load(request.json) + try: + xcom_entry = xcom_schema.load(request.json) + except ValidationError as err: + raise BadRequest(detail=str(err.messages.get('_schema', err.messages))) xcom = session.query(XCom).filter(XCom.dag_id == dag_id, XCom.task_id == task_id).first() @@ -177,6 +181,6 @@ def post_xcom_entries(dag_id: str, data = xcom_entry.data xcom = XCom(**data) session.add(xcom) - session.flush() + session.commit() return xcom_schema.dump(xcom) raise AlreadyExists("XCom entry already exists") diff --git a/airflow/api_connexion/schemas/xcom_schema.py b/airflow/api_connexion/schemas/xcom_schema.py index 5adc36da34da6..96e3e88a6e2c6 100644 --- a/airflow/api_connexion/schemas/xcom_schema.py +++ b/airflow/api_connexion/schemas/xcom_schema.py @@ -16,7 +16,7 @@ # under the License. from typing import List, NamedTuple -from marshmallow import Schema, fields +from marshmallow import Schema, ValidationError, fields, validates_schema from marshmallow_sqlalchemy import SQLAlchemySchema, auto_field from airflow.models import XCom @@ -37,6 +37,13 @@ class Meta: task_id = auto_field() dag_id = auto_field() + @validates_schema(pass_original=True) + def check_unknown_fields(self, data, original_data): + """ Validates unknown field """ + unknown = set(original_data) - set(self.fields) + if unknown: + raise ValidationError(f'Extra arguments passed: {list(unknown)}') + class XComSchema(XComCollectionItemSchema): """ diff --git a/tests/api_connexion/endpoints/test_xcom_endpoint.py b/tests/api_connexion/endpoints/test_xcom_endpoint.py index 1a4cab21231a1..7354fa66241ae 100644 --- a/tests/api_connexion/endpoints/test_xcom_endpoint.py +++ b/tests/api_connexion/endpoints/test_xcom_endpoint.py @@ -16,7 +16,6 @@ # under the License. import unittest -import pytest from parameterized import parameterized from airflow.models import DagRun as DR, XCom @@ -28,6 +27,7 @@ DAG_ID = 'test-dag-id' TASK_ID = 'test-task-id' EXECUTION_DATE = '2005-04-02T00:00:00+00:00' +EXECUTION_DATE_2 = '2020-04-02T00:00:00+00:00' XCOM_KEY = 'test-xcom-key' @@ -57,7 +57,6 @@ def tearDown(self) -> None: session.query(DR).delete() def _create_xcom(self, session): - xcom_model = XCom(key=XCOM_KEY, execution_date=self.execution_date_parsed, task_id=TASK_ID, @@ -303,12 +302,40 @@ def _create_xcoms(self, count): class TestPatchXComEntry(TestXComEndpoint): - @pytest.mark.skip(reason="Not implemented yet") - def test_should_response_200(self): + + @parameterized.expand( + [ + ( + { + 'dag_id': DAG_ID, + 'execution_date': EXECUTION_DATE_2, + 'key': XCOM_KEY, + 'task_id': TASK_ID, + }, + ), + ( + { + 'dag_id': "CHANGED_DAG_ID", + 'execution_date': EXECUTION_DATE, + 'key': XCOM_KEY, + 'task_id': TASK_ID, + }, + ), + ] + ) + @provide_session + def test_patch_should_response_200(self, payload, session): + self._create_xcom(session) + dag_run_id = self._create_dag_run(session) + + url = (f"/api/v1/dags/{DAG_ID}/dagRuns/{dag_run_id}/taskInstances" + f"/{TASK_ID}/xcomEntries/{XCOM_KEY}") + response = self.client.patch( - "/dags/TEST_DAG_ID/taskInstances/TEST_TASK_ID/2005-04-02T00:00:00Z/xcomEntries" + url, + json=payload ) - assert response.status_code == 200 + self.assertEqual(response.status_code, 200) class TestPostXComEntry(TestXComEndpoint): @@ -320,8 +347,7 @@ def test_post_should_response_200(self, session): 'execution_date': EXECUTION_DATE, 'key': XCOM_KEY, 'task_id': TASK_ID, - 'timestamp': EXECUTION_DATE, - 'value': "{}" + 'timestamp': EXECUTION_DATE } dag_run_id = self._create_dag_run(session) @@ -336,3 +362,83 @@ def test_post_should_response_200(self, session): xcom_entry = session.query(XCom).all() self.assertEqual(len(xcom_entry), 1) self.assertEqual(xcom_entry[0].key, XCOM_KEY) + + @parameterized.expand( + [ + ( + { + 'dag_id': DAG_ID, + 'execution_date': EXECUTION_DATE, + 'key': XCOM_KEY, + 'task_id': TASK_ID, + 'var': 'Unknown' # unknown field + }, + '["Extra arguments passed: [\'var\']"]' + ), + ( + { + 'dag_id': DAG_ID, + 'execution_date': EXECUTION_DATE, + 'key': XCOM_KEY, + # task_id missing + }, + "{'task_id': ['Missing data for required field.']}" + ) + ] + ) + @provide_session + def test_post_should_response_400_for_invalid_payload(self, payload, error_message, session): + dag_run_id = self._create_dag_run(session) + + url = (f"/api/v1/dags/{DAG_ID}/dagRuns/{dag_run_id}/taskInstances" + f"/{TASK_ID}/xcomEntries") + + response = self.client.post(url, json=payload) + self.assertEqual(response.status_code, 400) + self.assertEqual(response.json, + {'detail': error_message, + 'status': 400, + 'title': 'Bad request', + 'type': 'about:blank'} + ) + + def test_post_should_response_404_no_dag_run(self): + url = (f"/api/v1/dags/{DAG_ID}/dagRuns/INVALID_DAG_RUN_ID/taskInstances" + f"/{TASK_ID}/xcomEntries") + payload = { + 'dag_id': DAG_ID, + 'execution_date': EXECUTION_DATE, + 'key': XCOM_KEY, + 'task_id': TASK_ID, + } + response = self.client.post(url, json=payload) + self.assertEqual(response.status_code, 404) + self.assertEqual(response.json, + {'detail': None, + 'status': 404, + 'title': f'DAGRun with dag_id:{DAG_ID} and' + ' run_id:INVALID_DAG_RUN_ID not found', + 'type': 'about:blank'} + ) + + @provide_session + def test_post_should_response_409_for_already_exist(self, session): + dag_run_id = self._create_dag_run(session) + url = (f"/api/v1/dags/{DAG_ID}/dagRuns/{dag_run_id}/taskInstances" + f"/{TASK_ID}/xcomEntries") + payload = { + 'dag_id': DAG_ID, + 'execution_date': EXECUTION_DATE, + 'key': XCOM_KEY, + 'task_id': TASK_ID, + } + response = self.client.post(url, json=payload) + self.assertEqual(response.status_code, 200) + # Another requesst + response = self.client.post(url, json=payload) + self.assertEqual(response.json, + {'detail': None, + 'status': 409, + 'title': "XCom entry already exists", + 'type': 'about:blank'} + ) From 2a5ecdc7a410564e9ea9a2f880253f11cd2b3f7e Mon Sep 17 00:00:00 2001 From: EphraimBuddy Date: Tue, 23 Jun 2020 20:40:26 +0100 Subject: [PATCH 3/8] fix value serialization --- airflow/api_connexion/endpoints/xcom_endpoint.py | 8 +++++--- airflow/api_connexion/schemas/xcom_schema.py | 2 +- tests/api_connexion/endpoints/test_xcom_endpoint.py | 8 ++++++-- 3 files changed, 12 insertions(+), 6 deletions(-) diff --git a/airflow/api_connexion/endpoints/xcom_endpoint.py b/airflow/api_connexion/endpoints/xcom_endpoint.py index feb5094d2332d..07ce964dcf739 100644 --- a/airflow/api_connexion/endpoints/xcom_endpoint.py +++ b/airflow/api_connexion/endpoints/xcom_endpoint.py @@ -138,7 +138,7 @@ def patch_xcom_entry( try: body = xcom_schema.load(request.json) except ValidationError as err: - raise BadRequest(detail=str(err.messages.get('_schema', err.messages))) + raise BadRequest(detail=err.messages.get('_schema', [str(err.messages)])[0]) data = body.data if update_mask: update_mask = [i.strip() for i in update_mask] @@ -149,6 +149,7 @@ def patch_xcom_entry( else: raise BadRequest(f"Unknown field {field} in update_mask") data = data_ + data['value'] = XCom.serialize_value(data['value']) for key in data: setattr(xcom_entry, key, data[key]) session.add(xcom_entry) @@ -173,13 +174,14 @@ def post_xcom_entries(dag_id: str, try: xcom_entry = xcom_schema.load(request.json) except ValidationError as err: - raise BadRequest(detail=str(err.messages.get('_schema', err.messages))) + raise BadRequest(detail=err.messages.get('_schema', [str(err.messages)])[0]) xcom = session.query(XCom).filter(XCom.dag_id == dag_id, XCom.task_id == task_id).first() if not xcom: data = xcom_entry.data - xcom = XCom(**data) + data['value'] = XCom.serialize_value(data['value']) + xcom=XCom(**data) session.add(xcom) session.commit() return xcom_schema.dump(xcom) diff --git a/airflow/api_connexion/schemas/xcom_schema.py b/airflow/api_connexion/schemas/xcom_schema.py index 96e3e88a6e2c6..2dbf790696ced 100644 --- a/airflow/api_connexion/schemas/xcom_schema.py +++ b/airflow/api_connexion/schemas/xcom_schema.py @@ -50,7 +50,7 @@ class XComSchema(XComCollectionItemSchema): XCom schema """ - value = auto_field() + value = auto_field(required=True) class XComCollection(NamedTuple): diff --git a/tests/api_connexion/endpoints/test_xcom_endpoint.py b/tests/api_connexion/endpoints/test_xcom_endpoint.py index 7354fa66241ae..e2b4a9738e8f7 100644 --- a/tests/api_connexion/endpoints/test_xcom_endpoint.py +++ b/tests/api_connexion/endpoints/test_xcom_endpoint.py @@ -347,7 +347,8 @@ def test_post_should_response_200(self, session): 'execution_date': EXECUTION_DATE, 'key': XCOM_KEY, 'task_id': TASK_ID, - 'timestamp': EXECUTION_DATE + 'timestamp': EXECUTION_DATE, + 'value': "{'key':'val'}" } dag_run_id = self._create_dag_run(session) @@ -371,15 +372,17 @@ def test_post_should_response_200(self, session): 'execution_date': EXECUTION_DATE, 'key': XCOM_KEY, 'task_id': TASK_ID, + 'value': 'someval', 'var': 'Unknown' # unknown field }, - '["Extra arguments passed: [\'var\']"]' + "Extra arguments passed: ['var']" ), ( { 'dag_id': DAG_ID, 'execution_date': EXECUTION_DATE, 'key': XCOM_KEY, + 'value': 'someval' # task_id missing }, "{'task_id': ['Missing data for required field.']}" @@ -431,6 +434,7 @@ def test_post_should_response_409_for_already_exist(self, session): 'execution_date': EXECUTION_DATE, 'key': XCOM_KEY, 'task_id': TASK_ID, + 'value': "value" } response = self.client.post(url, json=payload) self.assertEqual(response.status_code, 200) From 4357e267c9040066272dbff4ff5f28790532aa1e Mon Sep 17 00:00:00 2001 From: EphraimBuddy Date: Wed, 24 Jun 2020 23:59:09 +0100 Subject: [PATCH 4/8] add more tests --- .../api_connexion/endpoints/xcom_endpoint.py | 37 ++-- .../endpoints/test_xcom_endpoint.py | 175 +++++++++++++++--- 2 files changed, 165 insertions(+), 47 deletions(-) diff --git a/airflow/api_connexion/endpoints/xcom_endpoint.py b/airflow/api_connexion/endpoints/xcom_endpoint.py index 07ce964dcf739..4917770263296 100644 --- a/airflow/api_connexion/endpoints/xcom_endpoint.py +++ b/airflow/api_connexion/endpoints/xcom_endpoint.py @@ -26,7 +26,7 @@ from marshmallow import ValidationError from sqlalchemy.orm.session import Session -from airflow.api_connexion.exceptions import AlreadyExists, BadRequest, NotFound +from airflow.api_connexion.exceptions import BadRequest, NotFound from airflow.api_connexion.schemas.xcom_schema import xcom_schema from airflow.models import DagRun as DR, XCom from airflow.utils.session import provide_session @@ -136,22 +136,28 @@ def patch_xcom_entry( if not xcom_entry: raise NotFound("XCom not found") try: - body = xcom_schema.load(request.json) + body = xcom_schema.load(request.json, partial=("dag_id", "task_id", + "key", "execution_date")) except ValidationError as err: raise BadRequest(detail=err.messages.get('_schema', [str(err.messages)])[0]) data = body.data + # Check that other attributes are not being updated + for field in data: + if field != 'value' and data[field] != getattr(xcom_entry, field): + raise BadRequest(detail=f"'{field}' cannot be updated") if update_mask: update_mask = [i.strip() for i in update_mask] data_ = {} for field in update_mask: - if field in data: + if field in data and field != 'value': + raise BadRequest(detail=f"{field} cannot be updated") + elif field in data and field == 'value': data_[field] = data[field] else: - raise BadRequest(f"Unknown field {field} in update_mask") + raise BadRequest(detail=f"Unknown field '{field}' in update_mask") data = data_ data['value'] = XCom.serialize_value(data['value']) - for key in data: - setattr(xcom_entry, key, data[key]) + xcom_entry.value = data['value'] session.add(xcom_entry) session.commit() return xcom_schema.dump(xcom_entry) @@ -166,7 +172,6 @@ def post_xcom_entries(dag_id: str, """ Create an XCom entry """ - # TODO add 409 to spec? dag_run = session.query(DR).filter(DR.run_id == dag_run_id, DR.dag_id == dag_id).first() if not dag_run: @@ -175,14 +180,10 @@ def post_xcom_entries(dag_id: str, xcom_entry = xcom_schema.load(request.json) except ValidationError as err: raise BadRequest(detail=err.messages.get('_schema', [str(err.messages)])[0]) - - xcom = session.query(XCom).filter(XCom.dag_id == dag_id, - XCom.task_id == task_id).first() - if not xcom: - data = xcom_entry.data - data['value'] = XCom.serialize_value(data['value']) - xcom=XCom(**data) - session.add(xcom) - session.commit() - return xcom_schema.dump(xcom) - raise AlreadyExists("XCom entry already exists") + session.query(XCom).filter(XCom.dag_id == dag_id, XCom.task_id == task_id).delete() + data = xcom_entry.data + data['value'] = XCom.serialize_value(data['value']) + xcom = XCom(**data) + session.add(xcom) + session.commit() + return xcom_schema.dump(xcom) diff --git a/tests/api_connexion/endpoints/test_xcom_endpoint.py b/tests/api_connexion/endpoints/test_xcom_endpoint.py index e2b4a9738e8f7..01ef24829867f 100644 --- a/tests/api_connexion/endpoints/test_xcom_endpoint.py +++ b/tests/api_connexion/endpoints/test_xcom_endpoint.py @@ -58,6 +58,7 @@ def tearDown(self) -> None: def _create_xcom(self, session): xcom_model = XCom(key=XCOM_KEY, + value=XCom.serialize_value("value"), execution_date=self.execution_date_parsed, task_id=TASK_ID, dag_id=DAG_ID, @@ -308,35 +309,173 @@ class TestPatchXComEntry(TestXComEndpoint): ( { 'dag_id': DAG_ID, - 'execution_date': EXECUTION_DATE_2, + 'execution_date': EXECUTION_DATE, 'key': XCOM_KEY, 'task_id': TASK_ID, + 'value': "{'key': 'value'}" }, + '"{\'key\': \'value\'}"' ), ( { - 'dag_id': "CHANGED_DAG_ID", - 'execution_date': EXECUTION_DATE, - 'key': XCOM_KEY, - 'task_id': TASK_ID, + 'value': "new_value" + }, + '"new_value"' + ), + ] + ) + @provide_session + def test_patch_should_response_200(self, payload, expected_value, session): + self._create_xcom(session) + dag_run_id = self._create_dag_run(session) + + url = (f"/api/v1/dags/{DAG_ID}/dagRuns/{dag_run_id}/taskInstances" + f"/{TASK_ID}/xcomEntries/{XCOM_KEY}") + + response = self.client.patch( + url, + json=payload + ) + + self.assertEqual(response.status_code, 200) + + self.assertEqual(response.json['value'], expected_value) + + @parameterized.expand( + [ + ( + { + 'dag_id': 'DAG_ID', + "value": 'new_value' + }, + "'dag_id' cannot be updated" + ), + ( + { + 'task_id': 'TASK_ID', + "value": 'new_value' }, + "'task_id' cannot be updated" + ), + ( + { + 'key': 'XCOM_KEY', + 'value': 'new_value', + }, + "'key' cannot be updated" + ), + ( + { + 'execution_date': EXECUTION_DATE_2, # trying to change the execution date + 'value': "new_value" + }, + "'execution_date' cannot be updated" + ), + ( + { + 'timestamp': EXECUTION_DATE_2, # trying to change the timestamp + 'value': "new_value" + }, + "'timestamp' cannot be updated" ), ] ) @provide_session - def test_patch_should_response_200(self, payload, session): + def test_patch_should_response_400_for_invalid_payload(self, payload, expected_message, session): self._create_xcom(session) dag_run_id = self._create_dag_run(session) url = (f"/api/v1/dags/{DAG_ID}/dagRuns/{dag_run_id}/taskInstances" f"/{TASK_ID}/xcomEntries/{XCOM_KEY}") + response = self.client.patch( + url, + json=payload + ) + self.assertEqual(response.status_code, 400) + self.assertEqual(response.json['detail'], expected_message) + + @provide_session + def test_patch_should_response_400_for_missing_value_field(self, session): + self._create_xcom(session) + dag_run_id = self._create_dag_run(session) + + payload = { + 'dag_id': DAG_ID, + 'execution_date': EXECUTION_DATE, + 'key': XCOM_KEY, + 'task_id': TASK_ID + } + + url = (f"/api/v1/dags/{DAG_ID}/dagRuns/{dag_run_id}/taskInstances" + f"/{TASK_ID}/xcomEntries/{XCOM_KEY}") + + response = self.client.patch( + url, + json=payload + ) + self.assertEqual(response.status_code, 400) + self.assertEqual(response.json['detail'], + "{'value': ['Missing data for required field.']}" + ) + + @provide_session + def test_patch_should_response_200_with_update_mask(self, session): + self._create_xcom(session) + dag_run_id = self._create_dag_run(session) + + payload = { + "value": "{'key':'var'}" + } + + url = (f"/api/v1/dags/{DAG_ID}/dagRuns/{dag_run_id}/taskInstances" + f"/{TASK_ID}/xcomEntries/{XCOM_KEY}?update_mask=value") + response = self.client.patch( url, json=payload ) self.assertEqual(response.status_code, 200) + @parameterized.expand( + [ + ( + { + "value": 'new_value' + }, + "update_mask=dags, value", + "Unknown field 'dags' in update_mask" + ), + ( + + { + "execution_date": EXECUTION_DATE, + "value": 'new_value' + }, + "update_mask=value, execution_date", + "execution_date cannot be updated" + ) + ] + ) + @provide_session + def test_patch_should_response_400_for_invalid_fields_in_update_mask( + self, payload, update_mask, error_message, session + ): + self._create_xcom(session) + dag_run_id = self._create_dag_run(session) + + url = (f"/api/v1/dags/{DAG_ID}/dagRuns/{dag_run_id}/taskInstances" + f"/{TASK_ID}/xcomEntries/{XCOM_KEY}?{update_mask}") + + response = self.client.patch( + url, + json=payload + ) + self.assertEqual(response.status_code, 400) + self.assertEqual(response.json['detail'], + error_message + ) + class TestPostXComEntry(TestXComEndpoint): @@ -413,6 +552,7 @@ def test_post_should_response_404_no_dag_run(self): 'execution_date': EXECUTION_DATE, 'key': XCOM_KEY, 'task_id': TASK_ID, + 'value': "value" } response = self.client.post(url, json=payload) self.assertEqual(response.status_code, 404) @@ -423,26 +563,3 @@ def test_post_should_response_404_no_dag_run(self): ' run_id:INVALID_DAG_RUN_ID not found', 'type': 'about:blank'} ) - - @provide_session - def test_post_should_response_409_for_already_exist(self, session): - dag_run_id = self._create_dag_run(session) - url = (f"/api/v1/dags/{DAG_ID}/dagRuns/{dag_run_id}/taskInstances" - f"/{TASK_ID}/xcomEntries") - payload = { - 'dag_id': DAG_ID, - 'execution_date': EXECUTION_DATE, - 'key': XCOM_KEY, - 'task_id': TASK_ID, - 'value': "value" - } - response = self.client.post(url, json=payload) - self.assertEqual(response.status_code, 200) - # Another requesst - response = self.client.post(url, json=payload) - self.assertEqual(response.json, - {'detail': None, - 'status': 409, - 'title': "XCom entry already exists", - 'type': 'about:blank'} - ) From 17cd801cdd3b9cf5c3191470123729f845b218d6 Mon Sep 17 00:00:00 2001 From: EphraimBuddy Date: Thu, 25 Jun 2020 00:10:41 +0100 Subject: [PATCH 5/8] fixup add more tests --- .../endpoints/test_xcom_endpoint.py | 29 +++++++++++++++++++ 1 file changed, 29 insertions(+) diff --git a/tests/api_connexion/endpoints/test_xcom_endpoint.py b/tests/api_connexion/endpoints/test_xcom_endpoint.py index 01ef24829867f..9d0ec95aa41a6 100644 --- a/tests/api_connexion/endpoints/test_xcom_endpoint.py +++ b/tests/api_connexion/endpoints/test_xcom_endpoint.py @@ -341,6 +341,35 @@ def test_patch_should_response_200(self, payload, expected_value, session): self.assertEqual(response.json['value'], expected_value) + @provide_session + def test_patch_should_response_404(self, session): + self._create_xcom(session) + dag_run_id = self._create_dag_run(session) + payload = {"value": "new-val"} + url = (f"/api/v1/dags/{DAG_ID}/dagRuns/{dag_run_id}/taskInstances" + f"/{TASK_ID}/xcomEntries/INVALID_XCOMKEY") + + response = self.client.patch( + url, + json=payload + ) + + self.assertEqual(response.status_code, 404) + + @provide_session + def test_patch_should_response_404_invalid_dagrun(self, session): + self._create_xcom(session) + payload = {"value": "new-val"} + url = (f"/api/v1/dags/{DAG_ID}/dagRuns/INVALID_DAGRUN_ID/taskInstances" + f"/{TASK_ID}/xcomEntries/{XCOM_KEY}") + + response = self.client.patch( + url, + json=payload + ) + + self.assertEqual(response.status_code, 404) + @parameterized.expand( [ ( From c4e665cce8a79e806f7d92becf1a9a2c32d47fce Mon Sep 17 00:00:00 2001 From: EphraimBuddy Date: Thu, 25 Jun 2020 00:17:08 +0100 Subject: [PATCH 6/8] reverse value set in spec --- airflow/api_connexion/openapi/v1.yaml | 1 - 1 file changed, 1 deletion(-) diff --git a/airflow/api_connexion/openapi/v1.yaml b/airflow/api_connexion/openapi/v1.yaml index 21f7cc8a28cae..1794b65264c09 100644 --- a/airflow/api_connexion/openapi/v1.yaml +++ b/airflow/api_connexion/openapi/v1.yaml @@ -1616,7 +1616,6 @@ components: properties: value: type: string - nullable: true # Python objects # Based on From a6f990a79bd384b079b332db42a726f5ef813023 Mon Sep 17 00:00:00 2001 From: EphraimBuddy Date: Thu, 25 Jun 2020 12:57:05 +0100 Subject: [PATCH 7/8] isort fix --- airflow/api_connexion/endpoints/xcom_endpoint.py | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/airflow/api_connexion/endpoints/xcom_endpoint.py b/airflow/api_connexion/endpoints/xcom_endpoint.py index 4917770263296..32c5b191410a5 100644 --- a/airflow/api_connexion/endpoints/xcom_endpoint.py +++ b/airflow/api_connexion/endpoints/xcom_endpoint.py @@ -14,20 +14,19 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -from sqlalchemy import and_, func from typing import List, Optional -from airflow.api_connexion import parameters -from airflow.api_connexion.schemas.xcom_schema import ( - XComCollection, XComCollectionItemSchema, XComCollectionSchema, xcom_collection_item_schema, - xcom_collection_schema, -) from flask import Response, request from marshmallow import ValidationError +from sqlalchemy import and_, func from sqlalchemy.orm.session import Session +from airflow.api_connexion import parameters from airflow.api_connexion.exceptions import BadRequest, NotFound -from airflow.api_connexion.schemas.xcom_schema import xcom_schema +from airflow.api_connexion.schemas.xcom_schema import ( + XComCollection, XComCollectionItemSchema, XComCollectionSchema, xcom_collection_item_schema, + xcom_collection_schema, xcom_schema, +) from airflow.models import DagRun as DR, XCom from airflow.utils.session import provide_session From d8094208c2ebd5366ddde8dfc490ad50f769737d Mon Sep 17 00:00:00 2001 From: Ephraim Anierobi Date: Sun, 28 Jun 2020 23:54:08 +0100 Subject: [PATCH 8/8] Update airflow/api_connexion/endpoints/xcom_endpoint.py Co-authored-by: Omair Khan --- airflow/api_connexion/endpoints/xcom_endpoint.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/airflow/api_connexion/endpoints/xcom_endpoint.py b/airflow/api_connexion/endpoints/xcom_endpoint.py index 32c5b191410a5..05065995de772 100644 --- a/airflow/api_connexion/endpoints/xcom_endpoint.py +++ b/airflow/api_connexion/endpoints/xcom_endpoint.py @@ -38,7 +38,7 @@ def delete_xcom_entry( task_id: str, xcom_key: str, session: Session -) -> Response: +): """ Delete an XCom entry """