clean up the API

apache · Mar 8, 2024 · 9167fbc · 9167fbc
1 parent a2e1101
commit 9167fbc
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 18 deletions.
diff --git a/airflow/api_connexion/endpoints/dag_endpoint.py b/airflow/api_connexion/endpoints/dag_endpoint.py
@@ -85,7 +85,22 @@ def get_dag_details(
         dag_detail_schema = DAGDetailSchema(only=fields) if fields else DAGDetailSchema()
     except ValueError as e:
         raise BadRequest("DAGDetailSchema init error", detail=str(e))
-    return dag_detail_schema.dump(dag_model)
+    dag_schema_details = dag_detail_schema.dump(dag_model)
+    dag_schema_details["dataset_expression"] = simplify_dataset_expression(
+        dag_schema_details.get("dataset_expression")
+    )
+    return dag_schema_details
+
+
+def simplify_dataset_expression(dataset_expression: dict | None) -> dict | None:
+    """Simplifies a nested dataset expression into a 'any' or 'all' format with URIs."""
+    if dataset_expression is None:
+        return None
+    if dataset_expression.get("__type") == "dataset":
+        return dataset_expression["__var"]["uri"]
+
+    new_key = "any" if dataset_expression["__type"] == "dataset_any" else "all"
+    return {new_key: [simplify_dataset_expression(item) for item in dataset_expression["__var"]]}
 
 
 @security.requires_access_dag("GET")

diff --git a/tests/api_connexion/endpoints/test_dag_endpoint.py b/tests/api_connexion/endpoints/test_dag_endpoint.py
@@ -400,23 +400,10 @@ def test_should_respond_200_with_dataset_expression(self, url_safe_serializer):
             "dag_id": "test_dag",
             "dag_run_timeout": None,
             "dataset_expression": {
-                "__type": "dataset_any",
-                "__var": [
-                    {"__type": "dataset", "__var": {"extra": {"hi": "bye"}, "uri": "s3://dag1/output_1.txt"}},
-                    {
-                        "__type": "dataset_all",
-                        "__var": [
-                            {
-                                "__type": "dataset",
-                                "__var": {"extra": {"hi": "bye"}, "uri": "s3://dag2/output_1.txt"},
-                            },
-                            {
-                                "__type": "dataset",
-                                "__var": {"extra": {"hi": "bye"}, "uri": "s3://dag3/output_3.txt"},
-                            },
-                        ],
-                    },
-                ],
+                "any": [
+                    "s3://dag1/output_1.txt",
+                    {"all": ["s3://dag2/output_1.txt", "s3://dag3/output_3.txt"]},
+                ]
             },
             "default_view": None,
             "description": None,