AIP-84 Migrate private graph_data endpoint (#44394)

* Working basic version

* Add tests

airflow/api_fastapi/core_api/datamodels/ui/graph.py

@@ -0,0 +1,60 @@
+# 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 __future__ import annotations
+
+from typing import Literal
+
+from airflow.api_fastapi.core_api.base import BaseModel
+
+
+class EdgeResponse(BaseModel):
+    """Edge serializer for responses."""
+
+    is_setup_teardown: bool | None = None
+    label: str | None = None
+    source_id: str
+    target_id: str
+
+
+class NodeValueResponse(BaseModel):
+    """Graph Node Value responses."""
+
+    isMapped: bool | None = None
+    label: str | None = None
+    labelStyle: str | None = None
+    style: str | None = None
+    tooltip: str | None = None
+    rx: int
+    ry: int
+    clusterLabelPos: str | None = None
+    setupTeardownType: Literal["setup", "teardown"] | None = None
+
+
+class NodeResponse(BaseModel):
+    """Node serializer for responses."""
+
+    children: list[NodeResponse] | None = None
+    id: str | None
+    value: NodeValueResponse
+
+
+class GraphDataResponse(BaseModel):
+    """Graph Data serializer for responses."""
+
+    edges: list[EdgeResponse]
+    nodes: NodeResponse
+    arrange: Literal["BT", "LR", "RL", "TB"]

airflow/api_fastapi/core_api/openapi/v1-generated.yaml

@@ -34,47 +34,26 @@ paths:
             application/json:
               schema:
                 $ref: '#/components/schemas/HTTPValidationError'
-  /ui/dashboard/historical_metrics_data:
+  /ui/config:
     get:
       tags:
-      - Dashboard
-      summary: Historical Metrics
-      description: Return cluster activity historical metrics.
-      operationId: historical_metrics
-      parameters:
-      - name: start_date
-        in: query
-        required: true
-        schema:
-          type: string
-          title: Start Date
-      - name: end_date
-        in: query
-        required: false
-        schema:
-          anyOf:
-          - type: string
-          - type: 'null'
-          title: End Date
+      - Config
+      summary: Get Configs
+      description: Get configs for UI.
+      operationId: get_configs
       responses:
         '200':
           description: Successful Response
           content:
             application/json:
               schema:
-                $ref: '#/components/schemas/HistoricalMetricDataResponse'
-        '400':
+                $ref: '#/components/schemas/ConfigResponse'
+        '404':
+          description: Not Found
           content:
             application/json:
               schema:
                 $ref: '#/components/schemas/HTTPExceptionResponse'
-          description: Bad Request
-        '422':
-          description: Validation Error
-          content:
-            application/json:
-              schema:
-                $ref: '#/components/schemas/HTTPValidationError'
   /ui/dags/recent_dag_runs:
     get:
       tags:
@@ -174,26 +153,102 @@ paths:
             application/json:
               schema:
                 $ref: '#/components/schemas/HTTPValidationError'
-  /ui/config:
+  /ui/dashboard/historical_metrics_data:
     get:
       tags:
-      - Config
-      summary: Get Configs
-      description: Get configs for UI.
-      operationId: get_configs
+      - Dashboard
+      summary: Historical Metrics
+      description: Return cluster activity historical metrics.
+      operationId: historical_metrics
+      parameters:
+      - name: start_date
+        in: query
+        required: true
+        schema:
+          type: string
+          title: Start Date
+      - name: end_date
+        in: query
+        required: false
+        schema:
+          anyOf:
+          - type: string
+          - type: 'null'
+          title: End Date
       responses:
         '200':
           description: Successful Response
           content:
             application/json:
               schema:
-                $ref: '#/components/schemas/ConfigResponse'
-        '404':
-          description: Not Found
+                $ref: '#/components/schemas/HistoricalMetricDataResponse'
+        '400':
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/HTTPExceptionResponse'
+          description: Bad Request
+        '422':
+          description: Validation Error
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/HTTPValidationError'
+  /ui/graph/graph_data:
+    get:
+      tags:
+      - Graph
+      summary: Graph Data
+      description: Get Graph Data.
+      operationId: graph_data
+      parameters:
+      - name: dag_id
+        in: query
+        required: true
+        schema:
+          type: string
+          title: Dag Id
+      - name: root
+        in: query
+        required: false
+        schema:
+          anyOf:
+          - type: string
+          - type: 'null'
+          title: Root
+      - name: include_upstream
+        in: query
+        required: false
+        schema:
+          type: boolean
+          default: false
+          title: Include Upstream
+      - name: include_downstream
+        in: query
+        required: false
+        schema:
+          type: boolean
+          default: false
+          title: Include Downstream
+      responses:
+        '200':
+          description: Successful Response
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/GraphDataResponse'
+        '400':
           content:
             application/json:
               schema:
                 $ref: '#/components/schemas/HTTPExceptionResponse'
+          description: Bad Request
+        '422':
+          description: Validation Error
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/HTTPValidationError'
   /public/assets:
     get:
       tags:
@@ -7296,6 +7351,30 @@ components:
         This is the set of allowable values for the ``warning_type`` field
 
         in the DagWarning model.'
+    EdgeResponse:
+      properties:
+        is_setup_teardown:
+          anyOf:
+          - type: boolean
+          - type: 'null'
+          title: Is Setup Teardown
+        label:
+          anyOf:
+          - type: string
+          - type: 'null'
+          title: Label
+        source_id:
+          type: string
+          title: Source Id
+        target_id:
+          type: string
+          title: Target Id
+      type: object
+      required:
+      - source_id
+      - target_id
+      title: EdgeResponse
+      description: Edge serializer for responses.
     EventLogCollectionResponse:
       properties:
         event_logs:
@@ -7407,6 +7486,30 @@ components:
       - name
       title: FastAPIAppResponse
       description: Serializer for Plugin FastAPI App responses.
+    GraphDataResponse:
+      properties:
+        edges:
+          items:
+            $ref: '#/components/schemas/EdgeResponse'
+          type: array
+          title: Edges
+        nodes:
+          $ref: '#/components/schemas/NodeResponse'
+        arrange:
+          type: string
+          enum:
+          - BT
+          - LR
+          - RL
+          - TB
+          title: Arrange
+      type: object
+      required:
+      - edges
+      - nodes
+      - arrange
+      title: GraphDataResponse
+      description: Graph Data serializer for responses.
     HTTPExceptionResponse:
       properties:
         detail:
@@ -7584,6 +7687,80 @@ components:
       - unixname
       title: JobResponse
       description: Job serializer for responses.
+    NodeResponse:
+      properties:
+        children:
+          anyOf:
+          - items:
+              $ref: '#/components/schemas/NodeResponse'
+            type: array
+          - type: 'null'
+          title: Children
+        id:
+          anyOf:
+          - type: string
+          - type: 'null'
+          title: Id
+        value:
+          $ref: '#/components/schemas/NodeValueResponse'
+      type: object
+      required:
+      - id
+      - value
+      title: NodeResponse
+      description: Node serializer for responses.
+    NodeValueResponse:
+      properties:
+        isMapped:
+          anyOf:
+          - type: boolean
+          - type: 'null'
+          title: Ismapped
+        label:
+          anyOf:
+          - type: string
+          - type: 'null'
+          title: Label
+        labelStyle:
+          anyOf:
+          - type: string
+          - type: 'null'
+          title: Labelstyle
+        style:
+          anyOf:
+          - type: string
+          - type: 'null'
+          title: Style
+        tooltip:
+          anyOf:
+          - type: string
+          - type: 'null'
+          title: Tooltip
+        rx:
+          type: integer
+          title: Rx
+        ry:
+          type: integer
+          title: Ry
+        clusterLabelPos:
+          anyOf:
+          - type: string
+          - type: 'null'
+          title: Clusterlabelpos
+        setupTeardownType:
+          anyOf:
+          - type: string
+            enum:
+            - setup
+            - teardown
+          - type: 'null'
+          title: Setupteardowntype
+      type: object
+      required:
+      - rx
+      - ry
+      title: NodeValueResponse
+      description: Graph Node Value responses.
     PatchTaskInstanceBody:
       properties:
         dry_run:

airflow/api_fastapi/core_api/routes/ui/__init__.py

@@ -21,10 +21,12 @@
 from airflow.api_fastapi.core_api.routes.ui.config import config_router
 from airflow.api_fastapi.core_api.routes.ui.dags import dags_router
 from airflow.api_fastapi.core_api.routes.ui.dashboard import dashboard_router
+from airflow.api_fastapi.core_api.routes.ui.graph import graph_data_router
 
 ui_router = AirflowRouter(prefix="/ui")
 
 ui_router.include_router(assets_router)
-ui_router.include_router(dashboard_router)
-ui_router.include_router(dags_router)
 ui_router.include_router(config_router)
+ui_router.include_router(dags_router)
+ui_router.include_router(dashboard_router)
+ui_router.include_router(graph_data_router)

airflow/api_fastapi/core_api/routes/ui/dashboard.py

@@ -35,11 +35,11 @@
 from airflow.api_fastapi.common.router import AirflowRouter
 from airflow.utils import timezone
 
-dashboard_router = AirflowRouter(tags=["Dashboard"])
+dashboard_router = AirflowRouter(tags=["Dashboard"], prefix="/dashboard")
 
 
 @dashboard_router.get(
-    "/dashboard/historical_metrics_data",
+    "/historical_metrics_data",
     include_in_schema=False,
     responses=create_openapi_http_exception_doc([status.HTTP_400_BAD_REQUEST]),
 )