add schema-level validation

apache · Jan 22, 2024 · 54d9dcc · 54d9dcc
1 parent 672b16b
commit 54d9dcc
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 17 deletions.
diff --git a/airflow/api_connexion/endpoints/dag_run_endpoint.py b/airflow/api_connexion/endpoints/dag_run_endpoint.py
@@ -338,21 +338,8 @@ def post_dag_run(*, dag_id: str, session: Session = NEW_SESSION) -> APIResponse:
         try:
             dag = get_airflow_app().dag_bag.get_dag(dag_id)
 
-            data_interval_start_exists = post_body.get("data_interval_start") is not None
-            data_interval_end_exists = post_body.get("data_interval_end") is not None
-
-            if (
-                data_interval_start_exists
-                and not data_interval_end_exists
-                or data_interval_end_exists
-                and not data_interval_start_exists
-            ):
-                raise BadRequest(
-                    detail="Both 'data_interval_start' and 'data_interval_end' must be specified, you cannot specify only one",
-                )
-
             interval = None
-            if data_interval_start_exists and data_interval_end_exists:
+            if post_body.get("data_interval_start") and post_body.get("data_interval_end"):
                 interval = DataInterval(
                     start=pendulum.instance(post_body["data_interval_start"]),
                     end=pendulum.instance(post_body["data_interval_end"]),

diff --git a/airflow/api_connexion/schemas/dag_run_schema.py b/airflow/api_connexion/schemas/dag_run_schema.py
@@ -20,7 +20,7 @@
 import json
 from typing import NamedTuple
 
-from marshmallow import fields, post_dump, pre_load, validate
+from marshmallow import ValidationError, fields, post_dump, pre_load, validate, validates_schema
 from marshmallow.schema import Schema
 from marshmallow.validate import Range
 from marshmallow_sqlalchemy import SQLAlchemySchema, auto_field
@@ -121,6 +121,18 @@ def autofill(self, data, **kwargs):
 
         return ret_data
 
+    @validates_schema
+    def validate_data_interval_dates(self, data, **kwargs):
+        data_interval_start_exists = data.get("data_interval_start") is not None
+        data_interval_end_exists = data.get("data_interval_end") is not None
+
+        if (data_interval_start_exists and not data_interval_end_exists) or (
+            data_interval_end_exists and not data_interval_start_exists
+        ):
+            raise ValidationError(
+                "Both 'data_interval_start' and 'data_interval_end' must be specified, you cannot specify only one"
+            )
+
 
 class SetDagRunStateFormSchema(Schema):
     """Schema for handling the request of setting state of DAG run."""

diff --git a/tests/api_connexion/endpoints/test_dag_run_endpoint.py b/tests/api_connexion/endpoints/test_dag_run_endpoint.py
@@ -1373,12 +1373,12 @@ def test_should_response_400_for_conflicting_execution_date_logical_date(self):
             (
                 "2020-11-10T08:25:56.939143+00:00",
                 None,
-                "Both 'data_interval_start' and 'data_interval_end' must be specified, you cannot specify only one",
+                "{'_schema': [\"Both 'data_interval_start' and 'data_interval_end' must be specified, you cannot specify only one\"]}",
             ),
             (
                 None,
                 "2020-11-10T08:25:56.939143+00:00",
-                "Both 'data_interval_start' and 'data_interval_end' must be specified, you cannot specify only one",
+                "{'_schema': [\"Both 'data_interval_start' and 'data_interval_end' must be specified, you cannot specify only one\"]}",
             ),
         ],
     )