From 3ed63ee24c2e1afc10a12497215eb40ded3c5a80 Mon Sep 17 00:00:00 2001 From: vedantlodha Date: Mon, 30 Jan 2023 11:20:26 +0530 Subject: [PATCH 1/4] Fix dag run trigger with a note. Currently, triggering dag run with note gives 400. This PR fixes it. Closes: #28825 --- airflow/api_connexion/endpoints/dag_run_endpoint.py | 9 +++++++++ airflow/api_connexion/schemas/dag_run_schema.py | 2 +- 2 files changed, 10 insertions(+), 1 deletion(-) diff --git a/airflow/api_connexion/endpoints/dag_run_endpoint.py b/airflow/api_connexion/endpoints/dag_run_endpoint.py index c854761b58ebb..1ca3e258999cc 100644 --- a/airflow/api_connexion/endpoints/dag_run_endpoint.py +++ b/airflow/api_connexion/endpoints/dag_run_endpoint.py @@ -319,6 +319,15 @@ def post_dag_run(*, dag_id: str, session: Session = NEW_SESSION) -> APIResponse: dag_hash=get_airflow_app().dag_bag.dags_hash.get(dag_id), session=session, ) + dag_run_instance = ( + session.query(DagRun).filter(DagRun.dag_id == dag_id, DagRun.run_id == run_id).one_or_none() + ) + dag_run_note = post_body.get("note") + if dag_run_note: + from flask_login import current_user + + current_user_id = getattr(current_user, "id", None) + dag_run.note = (dag_run_note, current_user_id) return dagrun_schema.dump(dag_run) except ValueError as ve: raise BadRequest(detail=str(ve)) diff --git a/airflow/api_connexion/schemas/dag_run_schema.py b/airflow/api_connexion/schemas/dag_run_schema.py index 7ca857951b9fe..999b533728e0f 100644 --- a/airflow/api_connexion/schemas/dag_run_schema.py +++ b/airflow/api_connexion/schemas/dag_run_schema.py @@ -73,7 +73,7 @@ class Meta: data_interval_end = auto_field(dump_only=True) last_scheduling_decision = auto_field(dump_only=True) run_type = auto_field(dump_only=True) - note = auto_field(dump_only=True) + note = auto_field(dump_only=False) @pre_load def autogenerate(self, data, **kwargs): From fff361b4c09af57ac504438353c0eaaef3575a67 Mon Sep 17 00:00:00 2001 From: vedantlodha Date: Mon, 30 Jan 2023 13:20:34 +0530 Subject: [PATCH 2/4] make flask_login.current_user a global import --- airflow/api_connexion/endpoints/dag_run_endpoint.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/airflow/api_connexion/endpoints/dag_run_endpoint.py b/airflow/api_connexion/endpoints/dag_run_endpoint.py index 1ca3e258999cc..66c052e16237e 100644 --- a/airflow/api_connexion/endpoints/dag_run_endpoint.py +++ b/airflow/api_connexion/endpoints/dag_run_endpoint.py @@ -21,6 +21,7 @@ import pendulum from connexion import NoContent from flask import g +from flask_login import current_user from marshmallow import ValidationError from sqlalchemy import or_ from sqlalchemy.orm import Query, Session @@ -324,8 +325,6 @@ def post_dag_run(*, dag_id: str, session: Session = NEW_SESSION) -> APIResponse: ) dag_run_note = post_body.get("note") if dag_run_note: - from flask_login import current_user - current_user_id = getattr(current_user, "id", None) dag_run.note = (dag_run_note, current_user_id) return dagrun_schema.dump(dag_run) @@ -447,8 +446,6 @@ def set_dag_run_note(*, dag_id: str, dag_run_id: str, session: Session = NEW_SES except ValidationError as err: raise BadRequest(detail=str(err)) - from flask_login import current_user - current_user_id = getattr(current_user, "id", None) if dag_run.dag_run_note is None: dag_run.note = (new_note, current_user_id) From b53d4a30678e3cac673e5585006814dcf2bafae9 Mon Sep 17 00:00:00 2001 From: vedantlodha Date: Mon, 30 Jan 2023 20:31:55 +0530 Subject: [PATCH 3/4] precommit error fixes. --- airflow/api_connexion/endpoints/dag_run_endpoint.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/airflow/api_connexion/endpoints/dag_run_endpoint.py b/airflow/api_connexion/endpoints/dag_run_endpoint.py index 66c052e16237e..a703aabc55b85 100644 --- a/airflow/api_connexion/endpoints/dag_run_endpoint.py +++ b/airflow/api_connexion/endpoints/dag_run_endpoint.py @@ -320,9 +320,6 @@ def post_dag_run(*, dag_id: str, session: Session = NEW_SESSION) -> APIResponse: dag_hash=get_airflow_app().dag_bag.dags_hash.get(dag_id), session=session, ) - dag_run_instance = ( - session.query(DagRun).filter(DagRun.dag_id == dag_id, DagRun.run_id == run_id).one_or_none() - ) dag_run_note = post_body.get("note") if dag_run_note: current_user_id = getattr(current_user, "id", None) From 62034b85c808b2881e6088168b1c7514cb3c038c Mon Sep 17 00:00:00 2001 From: vedantlodha Date: Wed, 1 Feb 2023 23:21:44 +0530 Subject: [PATCH 4/4] add test coverage for the fix --- .../endpoints/test_dag_run_endpoint.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/tests/api_connexion/endpoints/test_dag_run_endpoint.py b/tests/api_connexion/endpoints/test_dag_run_endpoint.py index 9e113f2b13b88..59d413866fceb 100644 --- a/tests/api_connexion/endpoints/test_dag_run_endpoint.py +++ b/tests/api_connexion/endpoints/test_dag_run_endpoint.py @@ -1025,14 +1025,14 @@ def test_end_date_gte_lte(self, payload, expected_dag_run_ids): class TestPostDagRun(TestDagRunEndpoint): @pytest.mark.parametrize("logical_date_field_name", ["execution_date", "logical_date"]) @pytest.mark.parametrize( - "dag_run_id, logical_date", + "dag_run_id, logical_date, note", [ - pytest.param("TEST_DAG_RUN", "2020-06-11T18:00:00+00:00", id="both-present"), - pytest.param(None, "2020-06-11T18:00:00+00:00", id="only-date"), - pytest.param(None, None, id="both-missing"), + pytest.param("TEST_DAG_RUN", "2020-06-11T18:00:00+00:00", "test-note", id="all-present"), + pytest.param(None, "2020-06-11T18:00:00+00:00", None, id="only-date"), + pytest.param(None, None, None, id="all-missing"), ], ) - def test_should_respond_200(self, logical_date_field_name, dag_run_id, logical_date): + def test_should_respond_200(self, logical_date_field_name, dag_run_id, logical_date, note): self._create_dag("TEST_DAG_ID") # We'll patch airflow.utils.timezone.utcnow to always return this so we @@ -1044,7 +1044,7 @@ def test_should_respond_200(self, logical_date_field_name, dag_run_id, logical_d request_json[logical_date_field_name] = logical_date if dag_run_id is not None: request_json["dag_run_id"] = dag_run_id - + request_json["note"] = note with mock.patch("airflow.utils.timezone.utcnow", lambda: fixed_now): response = self.client.post( "api/v1/dags/TEST_DAG_ID/dagRuns", @@ -1075,7 +1075,7 @@ def test_should_respond_200(self, logical_date_field_name, dag_run_id, logical_d "data_interval_start": expected_logical_date, "last_scheduling_decision": None, "run_type": "manual", - "note": None, + "note": note, } def test_should_respond_400_if_a_dag_has_import_errors(self, session):