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",
}