This is an automated email from the ASF dual-hosted git repository.

pierrejeambrun pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git


The following commit(s) were added to refs/heads/main by this push:
     new 8382712ec5d AIP-84 Add operator name field to structure data endpoint 
response (#44651)
8382712ec5d is described below

commit 8382712ec5d129b0bc4cb65cca81d05c59d5ef28
Author: Pierre Jeambrun <[email protected]>
AuthorDate: Thu Dec 5 01:58:53 2024 +0800

    AIP-84 Add operator name field to structure data endpoint response (#44651)
---
 .../api_fastapi/core_api/datamodels/ui/structure.py  |  1 +
 .../api_fastapi/core_api/openapi/v1-generated.yaml   |  5 +++++
 airflow/ui/openapi-gen/requests/schemas.gen.ts       | 11 +++++++++++
 airflow/ui/openapi-gen/requests/types.gen.ts         |  1 +
 airflow/utils/task_group.py                          |  6 ++++++
 .../api_fastapi/core_api/routes/ui/test_structure.py |  3 +++
 tests/utils/test_task_group.py                       | 20 +++++++++++++++-----
 7 files changed, 42 insertions(+), 5 deletions(-)

diff --git a/airflow/api_fastapi/core_api/datamodels/ui/structure.py 
b/airflow/api_fastapi/core_api/datamodels/ui/structure.py
index cc234500eb7..9461ba85a76 100644
--- a/airflow/api_fastapi/core_api/datamodels/ui/structure.py
+++ b/airflow/api_fastapi/core_api/datamodels/ui/structure.py
@@ -40,6 +40,7 @@ class NodeResponse(BaseModel):
     tooltip: str | None = None
     setup_teardown_type: Literal["setup", "teardown"] | None = None
     type: Literal["join", "sensor", "task", "asset_condition"]
+    operator: str | None = None
 
 
 class StructureDataResponse(BaseModel):
diff --git a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml 
b/airflow/api_fastapi/core_api/openapi/v1-generated.yaml
index 010a4141e06..9973ee7c6bb 100644
--- a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml
+++ b/airflow/api_fastapi/core_api/openapi/v1-generated.yaml
@@ -7813,6 +7813,11 @@ components:
           - task
           - asset_condition
           title: Type
+        operator:
+          anyOf:
+          - type: string
+          - type: 'null'
+          title: Operator
       type: object
       required:
       - id
diff --git a/airflow/ui/openapi-gen/requests/schemas.gen.ts 
b/airflow/ui/openapi-gen/requests/schemas.gen.ts
index cb8e2e7396e..566452ee11e 100644
--- a/airflow/ui/openapi-gen/requests/schemas.gen.ts
+++ b/airflow/ui/openapi-gen/requests/schemas.gen.ts
@@ -3174,6 +3174,17 @@ export const $NodeResponse = {
       enum: ["join", "sensor", "task", "asset_condition"],
       title: "Type",
     },
+    operator: {
+      anyOf: [
+        {
+          type: "string",
+        },
+        {
+          type: "null",
+        },
+      ],
+      title: "Operator",
+    },
   },
   type: "object",
   required: ["id", "label", "type"],
diff --git a/airflow/ui/openapi-gen/requests/types.gen.ts 
b/airflow/ui/openapi-gen/requests/types.gen.ts
index 13321ad4424..78fbc609499 100644
--- a/airflow/ui/openapi-gen/requests/types.gen.ts
+++ b/airflow/ui/openapi-gen/requests/types.gen.ts
@@ -767,6 +767,7 @@ export type NodeResponse = {
   tooltip?: string | null;
   setup_teardown_type?: "setup" | "teardown" | null;
   type: "join" | "sensor" | "task" | "asset_condition";
+  operator?: string | null;
 };
 
 export type type = "join" | "sensor" | "task" | "asset_condition";
diff --git a/airflow/utils/task_group.py b/airflow/utils/task_group.py
index 2f850826ca9..c8ee4923acb 100644
--- a/airflow/utils/task_group.py
+++ b/airflow/utils/task_group.py
@@ -80,6 +80,7 @@ class 
MappedTaskGroup(airflow.sdk.definitions.taskgroup.MappedTaskGroup):
 def task_group_to_dict(task_item_or_group):
     """Create a nested dict representation of this TaskGroup and its children 
used to construct the Graph."""
     from airflow.models.abstractoperator import AbstractOperator
+    from airflow.models.baseoperator import BaseOperator
     from airflow.models.mappedoperator import MappedOperator
     from airflow.sensors.base import BaseSensorOperator
 
@@ -87,6 +88,7 @@ def task_group_to_dict(task_item_or_group):
         setup_teardown_type = {}
         is_mapped = {}
         node_type = {"type": "task"}
+        node_operator = {}
         if task.is_setup is True:
             setup_teardown_type["setup_teardown_type"] = "setup"
         elif task.is_teardown is True:
@@ -95,12 +97,16 @@ def task_group_to_dict(task_item_or_group):
             is_mapped["is_mapped"] = True
         if isinstance(task, BaseSensorOperator):
             node_type["type"] = "sensor"
+        if isinstance(task, BaseOperator) or isinstance(task, MappedOperator):
+            node_operator["operator"] = task.operator_name
+
         return {
             "id": task.task_id,
             "label": task.label,
             **is_mapped,
             **setup_teardown_type,
             **node_type,
+            **node_operator,
         }
 
     task_group = task_item_or_group
diff --git a/tests/api_fastapi/core_api/routes/ui/test_structure.py 
b/tests/api_fastapi/core_api/routes/ui/test_structure.py
index 22cf1262e91..e270a49c6e8 100644
--- a/tests/api_fastapi/core_api/routes/ui/test_structure.py
+++ b/tests/api_fastapi/core_api/routes/ui/test_structure.py
@@ -79,6 +79,7 @@ class TestStructureDataEndpoint:
                             "id": "task_1",
                             "is_mapped": None,
                             "label": "task_1",
+                            "operator": "EmptyOperator",
                             "setup_teardown_type": None,
                             "tooltip": None,
                             "type": "task",
@@ -88,6 +89,7 @@ class TestStructureDataEndpoint:
                             "id": "task_2",
                             "is_mapped": None,
                             "label": "task_2",
+                            "operator": "EmptyOperator",
                             "setup_teardown_type": None,
                             "tooltip": None,
                             "type": "task",
@@ -118,6 +120,7 @@ class TestStructureDataEndpoint:
                             "id": "task_1",
                             "is_mapped": None,
                             "label": "task_1",
+                            "operator": "EmptyOperator",
                             "setup_teardown_type": None,
                             "tooltip": None,
                             "type": "task",
diff --git a/tests/utils/test_task_group.py b/tests/utils/test_task_group.py
index bd336fa4c94..71d6748b44a 100644
--- a/tests/utils/test_task_group.py
+++ b/tests/utils/test_task_group.py
@@ -186,19 +186,29 @@ EXPECTED_JSON = {
                     "tooltip": "",
                     "is_mapped": False,
                     "children": [
-                        {"id": "group234.group34.task3", "label": "task3", 
"type": "task"},
-                        {"id": "group234.group34.task4", "label": "task4", 
"type": "task"},
+                        {
+                            "id": "group234.group34.task3",
+                            "label": "task3",
+                            "operator": "EmptyOperator",
+                            "type": "task",
+                        },
+                        {
+                            "id": "group234.group34.task4",
+                            "label": "task4",
+                            "operator": "EmptyOperator",
+                            "type": "task",
+                        },
                         {"id": "group234.group34.downstream_join_id", "label": 
"", "type": "join"},
                     ],
                     "type": "task",
                 },
-                {"id": "group234.task2", "label": "task2", "type": "task"},
+                {"id": "group234.task2", "label": "task2", "operator": 
"EmptyOperator", "type": "task"},
                 {"id": "group234.upstream_join_id", "label": "", "type": 
"join"},
             ],
             "type": "task",
         },
-        {"id": "task1", "label": "task1", "type": "task"},
-        {"id": "task5", "label": "task5", "type": "task"},
+        {"id": "task1", "label": "task1", "operator": "EmptyOperator", "type": 
"task"},
+        {"id": "task5", "label": "task5", "operator": "EmptyOperator", "type": 
"task"},
     ],
     "type": "task",
 }

Reply via email to