From df30dac303d728ac3a9f82774a56996de2f71354 Mon Sep 17 00:00:00 2001 From: pierrejeambrun Date: Fri, 19 Jun 2026 17:54:41 +0200 Subject: [PATCH] Restore graph task and group coloring via Chakra palette tokens Custom operator and TaskGroup colors disappeared in Airflow 3: the ui_color/ui_fgcolor attributes were kept but the new UI ignored them, so teams that relied on color to parse large graphs at a glance lost that signal (see the demand on the linked issue). Bring the coloring back in a way that fits the new opinionated theming: the value must be a Chakra palette token (e.g. blue.500), which resolves through a theme-controlled CSS variable and therefore stays legible in both light and dark mode and under custom themes -- the dark-mode and accessibility concerns that drove the original removal. The graph paints the node fill from ui_color and the operator label from ui_fgcolor, as in Airflow 2.x, while the border keeps showing run state. Legacy hex/named values can no longer adapt to the theme, so they are ignored by the graph and emit a warning for user-authored operators and task groups. --- .../newsfragments/68760.significant.rst | 18 +++++ .../core_api/datamodels/ui/structure.py | 2 + .../core_api/openapi/_private_ui.yaml | 10 +++ .../core_api/services/ui/task_group.py | 17 ++++ .../ui/openapi-gen/requests/schemas.gen.ts | 22 +++++ .../ui/openapi-gen/requests/types.gen.ts | 2 + .../ui/src/components/Graph/TaskNode.test.tsx | 71 ++++++++++++++++ .../ui/src/components/Graph/TaskNode.tsx | 16 +++- .../components/Graph/elkGraphUtils.test.ts | 32 ++++++++ .../ui/src/components/Graph/elkGraphUtils.ts | 4 + .../ui/src/components/Graph/reactflowUtils.ts | 2 + .../core_api/routes/ui/test_structure.py | 77 ++++++++++++++++++ .../src/airflow_shared/dagnode/ui_color.py | 40 +++++++++ shared/dagnode/tests/dagnode/test_ui_color.py | 49 +++++++++++ task-sdk/src/airflow/sdk/bases/operator.py | 24 ++++++ .../src/airflow/sdk/definitions/taskgroup.py | 28 ++++++- .../task_sdk/definitions/test_taskgroup.py | 24 ++++++ .../task_sdk/definitions/test_ui_color.py | 81 +++++++++++++++++++ 18 files changed, 514 insertions(+), 5 deletions(-) create mode 100644 airflow-core/newsfragments/68760.significant.rst create mode 100644 airflow-core/src/airflow/ui/src/components/Graph/TaskNode.test.tsx create mode 100644 shared/dagnode/src/airflow_shared/dagnode/ui_color.py create mode 100644 shared/dagnode/tests/dagnode/test_ui_color.py create mode 100644 task-sdk/tests/task_sdk/definitions/test_ui_color.py diff --git a/airflow-core/newsfragments/68760.significant.rst b/airflow-core/newsfragments/68760.significant.rst new file mode 100644 index 0000000000000..83516685307d6 --- /dev/null +++ b/airflow-core/newsfragments/68760.significant.rst @@ -0,0 +1,18 @@ +Custom graph colors return, using Chakra palette tokens + +The ``ui_color`` and ``ui_fgcolor`` attributes on operators and TaskGroups color the graph view +again (the node fill from ``ui_color`` and the operator label from ``ui_fgcolor``), as they did in +Airflow 2. They were silently ignored since Airflow 3.0. + +**What changed:** + +- The value must now be a Chakra palette token such as ``blue.500`` (including tokens added through + custom UI theming). Tokens resolve through theme-controlled CSS variables, so the color stays + legible in both light and dark mode. +- Legacy values that cannot adapt to the theme (raw hex like ``#fff``, CSS names like + ``CornflowerBlue``, or a bare family like ``blue``) are ignored by the graph. +- Setting a non-token ``ui_color``/``ui_fgcolor`` on a user-authored operator or TaskGroup emits a + warning at Dag parse time. Built-in operators and providers do not warn. + +The node border keeps reflecting the task-instance run state; the custom color is applied to the +fill and label only. diff --git a/airflow-core/src/airflow/api_fastapi/core_api/datamodels/ui/structure.py b/airflow-core/src/airflow/api_fastapi/core_api/datamodels/ui/structure.py index 57f052dbc42a1..5d89b08ee98b6 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/datamodels/ui/structure.py +++ b/airflow-core/src/airflow/api_fastapi/core_api/datamodels/ui/structure.py @@ -42,6 +42,8 @@ class NodeResponse(BaseNodeResponse): setup_teardown_type: Literal["setup", "teardown"] | None = None operator: str | None = None asset_condition_type: Literal["or-gate", "and-gate"] | None = None + ui_color: str | None = None + ui_fgcolor: str | None = None class StructureDataResponse(BaseGraphResponse[EdgeResponse, NodeResponse]): diff --git a/airflow-core/src/airflow/api_fastapi/core_api/openapi/_private_ui.yaml b/airflow-core/src/airflow/api_fastapi/core_api/openapi/_private_ui.yaml index d786e259b4e77..d3f6bb0d4108b 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/openapi/_private_ui.yaml +++ b/airflow-core/src/airflow/api_fastapi/core_api/openapi/_private_ui.yaml @@ -3319,6 +3319,16 @@ components: - and-gate - type: 'null' title: Asset Condition Type + ui_color: + anyOf: + - type: string + - type: 'null' + title: Ui Color + ui_fgcolor: + anyOf: + - type: string + - type: 'null' + title: Ui Fgcolor type: object required: - id diff --git a/airflow-core/src/airflow/api_fastapi/core_api/services/ui/task_group.py b/airflow-core/src/airflow/api_fastapi/core_api/services/ui/task_group.py index 47d49757e69e4..4c58ca2ba2222 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/services/ui/task_group.py +++ b/airflow-core/src/airflow/api_fastapi/core_api/services/ui/task_group.py @@ -23,6 +23,7 @@ from functools import cache from operator import methodcaller +from airflow._shared.dagnode.ui_color import is_chakra_color_token from airflow.configuration import conf from airflow.serialization.definitions.baseoperator import SerializedBaseOperator from airflow.serialization.definitions.mappedoperator import SerializedMappedOperator, is_mapped @@ -37,6 +38,20 @@ def get_task_group_children_getter() -> Callable: return methodcaller("hierarchical_alphabetical_sort") +def _ui_colors(node) -> dict[str, str]: + """ + Return the node's ``ui_color``/``ui_fgcolor`` keys, keeping only valid Chakra color tokens. + + Legacy values (raw hex, CSS names, the unset defaults) are dropped so the graph only ever + receives colors the UI can resolve against the theme. + """ + return { + key: value + for key in ("ui_color", "ui_fgcolor") + if is_chakra_color_token(value := getattr(node, key, None)) + } + + def task_group_to_dict(task_item_or_group, parent_group_is_mapped=False): """Create a nested dict representation of this TaskGroup and its children used to construct the Graph.""" if isinstance(task := task_item_or_group, (SerializedBaseOperator, SerializedMappedOperator)): @@ -47,6 +62,7 @@ def task_group_to_dict(task_item_or_group, parent_group_is_mapped=False): "label": task_display_name, "operator": task.operator_name, "type": "task", + **_ui_colors(task), } if task.is_setup: node_operator["setup_teardown_type"] = "setup" @@ -78,6 +94,7 @@ def task_group_to_dict(task_item_or_group, parent_group_is_mapped=False): "is_mapped": mapped, "children": children, "type": "task", + **_ui_colors(task_group), } return node diff --git a/airflow-core/src/airflow/ui/openapi-gen/requests/schemas.gen.ts b/airflow-core/src/airflow/ui/openapi-gen/requests/schemas.gen.ts index 413dbef3d7e42..1815fcc8d029f 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/requests/schemas.gen.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/requests/schemas.gen.ts @@ -9849,6 +9849,28 @@ export const $NodeResponse = { } ], title: 'Asset Condition Type' + }, + ui_color: { + anyOf: [ + { + type: 'string' + }, + { + type: 'null' + } + ], + title: 'Ui Color' + }, + ui_fgcolor: { + anyOf: [ + { + type: 'string' + }, + { + type: 'null' + } + ], + title: 'Ui Fgcolor' } }, type: 'object', diff --git a/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts b/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts index 31f716d85f885..1c4158c93457f 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts @@ -2477,6 +2477,8 @@ export type NodeResponse = { setup_teardown_type?: 'setup' | 'teardown' | null; operator?: string | null; asset_condition_type?: 'or-gate' | 'and-gate' | null; + ui_color?: string | null; + ui_fgcolor?: string | null; }; export type OklchColor = string; diff --git a/airflow-core/src/airflow/ui/src/components/Graph/TaskNode.test.tsx b/airflow-core/src/airflow/ui/src/components/Graph/TaskNode.test.tsx new file mode 100644 index 0000000000000..cf7fe550ab78f --- /dev/null +++ b/airflow-core/src/airflow/ui/src/components/Graph/TaskNode.test.tsx @@ -0,0 +1,71 @@ +/*! + * 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. + */ +import { render } from "@testing-library/react"; +import { ReactFlowProvider } from "@xyflow/react"; +import type { ReactNode } from "react"; +import { describe, expect, it, vi } from "vitest"; + +import { Wrapper } from "src/utils/Wrapper"; + +import { TaskNode } from "./TaskNode"; +import type { CustomNodeProps } from "./reactflowUtils"; + +vi.mock("src/context/groups", () => ({ + useGroups: vi.fn(() => ({ toggleGroupId: vi.fn() })), +})); + +const TestWrapper = ({ children }: { readonly children: ReactNode }) => ( + + {children} + +); + +// Chakra/Panda hashes color props into atomic class names rather than inline styles, so the resolved +// colour cannot be read back in jsdom. Instead we render two otherwise-identical nodes that differ +// only in the prop under test: any markup difference is attributable to that prop (hashing is +// deterministic), and identical markup proves the prop had no effect. +const renderHtml = (data: Partial): string => { + const { container } = render( + // The xyflow NodeProps surface is large; the component only reads `data` and `id`. + , + { wrapper: TestWrapper }, + ); + + return container.innerHTML; +}; + +describe("TaskNode operator colors", () => { + it("tints a leaf task when ui_color is set", () => { + expect(renderHtml({ operator: "BashOperator", uiColor: "blue.500" })).not.toBe( + renderHtml({ operator: "BashOperator" }), + ); + }); + + it("colors the operator text when ui_fgcolor is set", () => { + expect(renderHtml({ operator: "BashOperator", uiFgcolor: "red.700" })).not.toBe( + renderHtml({ operator: "BashOperator" }), + ); + }); + + it("does not tint a group node (2.x parity: groups keep their own background)", () => { + expect(renderHtml({ isGroup: true, uiColor: "blue.500" })).toBe(renderHtml({ isGroup: true })); + }); +}); diff --git a/airflow-core/src/airflow/ui/src/components/Graph/TaskNode.tsx b/airflow-core/src/airflow/ui/src/components/Graph/TaskNode.tsx index 5d8affeca3059..1abe5fe3b9bf8 100644 --- a/airflow-core/src/airflow/ui/src/components/Graph/TaskNode.tsx +++ b/airflow-core/src/airflow/ui/src/components/Graph/TaskNode.tsx @@ -47,6 +47,8 @@ export const TaskNode = ({ taskInstance, team, tooltip, + uiColor, + uiFgcolor, width = 0, }, id, @@ -84,6 +86,14 @@ export const TaskNode = ({ .map(([_state, count]) => count) .reduce((sum, val) => sum + val, 0); + // Custom operator fill (ui_color). Like Airflow 2.x, only leaf tasks are tinted and the colour is + // blended toward the theme background so it stays subtle against the run-state border and works in + // both light and dark mode. The value is a Chakra token, resolved through its CSS variable. + const customBg = + uiColor !== undefined && uiColor !== null && !isGroup + ? `color-mix(in srgb, var(--chakra-colors-${uiColor.replaceAll(".", "-")}, transparent) 25%, var(--chakra-colors-bg))` + : undefined; + return ( @@ -96,8 +106,8 @@ export const TaskNode = ({ tooltip={isGroup ? tooltip : undefined} > { expect(rootEdgeIds).toEqual(new Set(["start-group_a", "group_a-final_task"])); }); }); + +describe("generateElkGraph — operator colors", () => { + it("maps ui_color/ui_fgcolor onto the formatted leaf node", () => { + const root = generateElkGraph({ + direction: "RIGHT", + edges: [], + font: "12px sans-serif", + nodes: [buildNode({ id: "t1", label: "t1", ui_color: "blue.500", ui_fgcolor: "red.700" })], + openGroupIds: [], + }); + + const node = root.children?.[0] as FormattedNode; + + expect(node.uiColor).toBe("blue.500"); + expect(node.uiFgcolor).toBe("red.700"); + }); + + it("leaves the colors undefined when the node has none", () => { + const root = generateElkGraph({ + direction: "RIGHT", + edges: [], + font: "12px sans-serif", + nodes: [buildNode({ id: "t1", label: "t1" })], + openGroupIds: [], + }); + + const node = root.children?.[0] as FormattedNode; + + expect(node.uiColor).toBeUndefined(); + expect(node.uiFgcolor).toBeUndefined(); + }); +}); diff --git a/airflow-core/src/airflow/ui/src/components/Graph/elkGraphUtils.ts b/airflow-core/src/airflow/ui/src/components/Graph/elkGraphUtils.ts index 593dec17c9a24..995e28323034f 100644 --- a/airflow-core/src/airflow/ui/src/components/Graph/elkGraphUtils.ts +++ b/airflow-core/src/airflow/ui/src/components/Graph/elkGraphUtils.ts @@ -49,6 +49,8 @@ export type FormattedNode = { isOpen?: boolean; setupTeardownType?: NodeResponse["setup_teardown_type"]; team?: string | null; + uiColor?: string | null; + uiFgcolor?: string | null; } & ElkShape & NodeResponse; @@ -392,6 +394,8 @@ export const generateElkGraph = ({ team: node.team, tooltip: node.tooltip, type: node.type, + uiColor: node.ui_color, + uiFgcolor: node.ui_fgcolor, width, }; }; diff --git a/airflow-core/src/airflow/ui/src/components/Graph/reactflowUtils.ts b/airflow-core/src/airflow/ui/src/components/Graph/reactflowUtils.ts index 67fe2c2d2406d..7a2667a7abc23 100644 --- a/airflow-core/src/airflow/ui/src/components/Graph/reactflowUtils.ts +++ b/airflow-core/src/airflow/ui/src/components/Graph/reactflowUtils.ts @@ -41,6 +41,8 @@ export type CustomNodeProps = { team?: string | null; tooltip?: string | null; type: string; + uiColor?: string | null; + uiFgcolor?: string | null; width?: number; }; diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_structure.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_structure.py index 0bde80a22dc37..5176bebe5c4ed 100644 --- a/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_structure.py +++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_structure.py @@ -33,6 +33,7 @@ from airflow.providers.standard.sensors.external_task import ExternalTaskSensor from airflow.sdk import Metadata, task from airflow.sdk.definitions.asset import Asset, AssetAlias, Dataset +from airflow.sdk.definitions.taskgroup import TaskGroup from tests_common.test_utils.asserts import assert_queries_count from tests_common.test_utils.db import clear_db_assets, clear_db_runs @@ -58,6 +59,8 @@ "team": None, "operator": "EmptyOperator", "asset_condition_type": None, + "ui_color": None, + "ui_fgcolor": None, }, { "children": None, @@ -70,6 +73,8 @@ "team": None, "operator": "EmptyOperator", "asset_condition_type": None, + "ui_color": None, + "ui_fgcolor": None, }, { "children": None, @@ -82,6 +87,8 @@ "team": None, "operator": "EmptyOperator", "asset_condition_type": None, + "ui_color": None, + "ui_fgcolor": None, }, ], } @@ -274,6 +281,8 @@ class TestStructureDataEndpoint: "nodes": [ { "asset_condition_type": None, + "ui_color": None, + "ui_fgcolor": None, "children": None, "id": "task_1", "is_mapped": None, @@ -286,6 +295,8 @@ class TestStructureDataEndpoint: }, { "asset_condition_type": None, + "ui_color": None, + "ui_fgcolor": None, "children": None, "id": "external_task_sensor", "is_mapped": None, @@ -298,6 +309,8 @@ class TestStructureDataEndpoint: }, { "asset_condition_type": None, + "ui_color": None, + "ui_fgcolor": None, "children": None, "id": "task_2", "is_mapped": None, @@ -332,6 +345,8 @@ class TestStructureDataEndpoint: "nodes": [ { "asset_condition_type": None, + "ui_color": None, + "ui_fgcolor": None, "children": None, "id": "task_1", "is_mapped": None, @@ -361,6 +376,8 @@ class TestStructureDataEndpoint: "nodes": [ { "asset_condition_type": None, + "ui_color": None, + "ui_fgcolor": None, "children": None, "id": "trigger_dag_run_operator", "is_mapped": None, @@ -373,6 +390,8 @@ class TestStructureDataEndpoint: }, { "asset_condition_type": None, + "ui_color": None, + "ui_fgcolor": None, "children": None, "id": "trigger:external_trigger:dag_with_multiple_versions:trigger_dag_run_operator", "is_mapped": None, @@ -480,6 +499,8 @@ def test_should_return_200_with_asset(self, test_client, asset1_id, asset2_id, a "team": None, "operator": "EmptyOperator", "asset_condition_type": None, + "ui_color": None, + "ui_fgcolor": None, }, { "children": None, @@ -492,6 +513,8 @@ def test_should_return_200_with_asset(self, test_client, asset1_id, asset2_id, a "team": None, "operator": "ExternalTaskSensor", "asset_condition_type": None, + "ui_color": None, + "ui_fgcolor": None, }, { "children": None, @@ -504,6 +527,8 @@ def test_should_return_200_with_asset(self, test_client, asset1_id, asset2_id, a "team": None, "operator": "EmptyOperator", "asset_condition_type": None, + "ui_color": None, + "ui_fgcolor": None, }, { "children": None, @@ -516,6 +541,8 @@ def test_should_return_200_with_asset(self, test_client, asset1_id, asset2_id, a "team": None, "operator": None, "asset_condition_type": None, + "ui_color": None, + "ui_fgcolor": None, }, { "children": None, @@ -528,6 +555,8 @@ def test_should_return_200_with_asset(self, test_client, asset1_id, asset2_id, a "team": None, "operator": None, "asset_condition_type": None, + "ui_color": None, + "ui_fgcolor": None, }, { "children": None, @@ -540,6 +569,8 @@ def test_should_return_200_with_asset(self, test_client, asset1_id, asset2_id, a "team": None, "operator": None, "asset_condition_type": None, + "ui_color": None, + "ui_fgcolor": None, }, { "children": None, @@ -552,6 +583,8 @@ def test_should_return_200_with_asset(self, test_client, asset1_id, asset2_id, a "team": None, "operator": None, "asset_condition_type": "and-gate", + "ui_color": None, + "ui_fgcolor": None, }, { "children": None, @@ -564,6 +597,8 @@ def test_should_return_200_with_asset(self, test_client, asset1_id, asset2_id, a "team": None, "operator": None, "asset_condition_type": None, + "ui_color": None, + "ui_fgcolor": None, }, { "children": None, @@ -576,6 +611,8 @@ def test_should_return_200_with_asset(self, test_client, asset1_id, asset2_id, a "team": None, "operator": None, "asset_condition_type": None, + "ui_color": None, + "ui_fgcolor": None, }, { "children": None, @@ -588,6 +625,8 @@ def test_should_return_200_with_asset(self, test_client, asset1_id, asset2_id, a "team": None, "operator": None, "asset_condition_type": None, + "ui_color": None, + "ui_fgcolor": None, }, ], } @@ -637,6 +676,8 @@ def test_should_return_200_with_resolved_asset_alias_attached_to_the_corrrect_pr "setup_teardown_type": None, "operator": "@task", "asset_condition_type": None, + "ui_color": None, + "ui_fgcolor": None, }, { "id": "task_2", @@ -649,6 +690,8 @@ def test_should_return_200_with_resolved_asset_alias_attached_to_the_corrrect_pr "setup_teardown_type": None, "operator": "EmptyOperator", "asset_condition_type": None, + "ui_color": None, + "ui_fgcolor": None, }, { "id": f"asset:{resolved_asset.id}", @@ -661,6 +704,8 @@ def test_should_return_200_with_resolved_asset_alias_attached_to_the_corrrect_pr "setup_teardown_type": None, "operator": None, "asset_condition_type": None, + "ui_color": None, + "ui_fgcolor": None, }, ], } @@ -834,6 +879,38 @@ def test_mapped_operator_in_task_group(self, dag_maker, test_client, session): assert mapped_in_group["is_mapped"] is True assert mapped_in_group["operator"] == "PythonOperator" + def test_ui_colors_only_exposed_for_chakra_tokens(self, dag_maker, test_client, session): + """Only Chakra palette tokens reach the graph; legacy hex/defaults are dropped to null.""" + + class ColoredOperator(EmptyOperator): + ui_color = "blue.500" + ui_fgcolor = "red.700" + + with dag_maker( + dag_id="test_ui_colors_dag", + serialized=True, + session=session, + start_date=pendulum.DateTime(2023, 2, 1, 0, 0, 0, tzinfo=pendulum.UTC), + ): + ColoredOperator(task_id="colored") + EmptyOperator(task_id="plain") # default ui_color "#fff" -> not a token + with TaskGroup(group_id="grp", ui_color="teal.400"): + EmptyOperator(task_id="inner") + + dag_maker.sync_dagbag_to_db() + response = test_client.get("/structure/structure_data", params={"dag_id": "test_ui_colors_dag"}) + assert response.status_code == 200 + nodes = {node["id"]: node for node in response.json()["nodes"]} + + assert nodes["colored"]["ui_color"] == "blue.500" + assert nodes["colored"]["ui_fgcolor"] == "red.700" + # Hex/default values are filtered out + assert nodes["plain"]["ui_color"] is None + assert nodes["plain"]["ui_fgcolor"] is None + # Task group keeps its token fill; the unset fgcolor default is dropped + assert nodes["grp"]["ui_color"] == "teal.400" + assert nodes["grp"]["ui_fgcolor"] is None + @pytest.mark.parametrize( ("params", "expected_task_ids", "description"), [ diff --git a/shared/dagnode/src/airflow_shared/dagnode/ui_color.py b/shared/dagnode/src/airflow_shared/dagnode/ui_color.py new file mode 100644 index 0000000000000..883dffdf16a43 --- /dev/null +++ b/shared/dagnode/src/airflow_shared/dagnode/ui_color.py @@ -0,0 +1,40 @@ +# 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. +"""Validation helper for the ``ui_color`` / ``ui_fgcolor`` graph node colors.""" + +from __future__ import annotations + +from typing import TypeGuard + + +def is_chakra_color_token(value: str | None) -> TypeGuard[str]: + """ + Return whether ``value`` is shaped like a Chakra UI color token (``family.shade``). + + A token is ``.`` such as ``blue.500`` -- a color family (built-in or one added + through custom UI theming) followed by an integer shade. Only the shape is checked, with plain + string operations rather than a regular expression so a Dag author cannot craft a value that + triggers pathological backtracking. Unknown families simply resolve to an undefined CSS variable + and are ignored by the UI. Raw hex (``#fff``), ``rgb()``/``hsl()`` and bare CSS names + (``CornflowerBlue``) are rejected. + """ + if not value or not value.isascii(): + return False + family, separator, shade = value.partition(".") + if not separator: + return False + return family[:1].isalpha() and family.isalnum() and shade.isdigit() diff --git a/shared/dagnode/tests/dagnode/test_ui_color.py b/shared/dagnode/tests/dagnode/test_ui_color.py new file mode 100644 index 0000000000000..3396fae7e4194 --- /dev/null +++ b/shared/dagnode/tests/dagnode/test_ui_color.py @@ -0,0 +1,49 @@ +# 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 + +import pytest + +from airflow_shared.dagnode.ui_color import is_chakra_color_token + + +@pytest.mark.parametrize( + ("value", "expected"), + [ + ("blue.500", True), + ("red.50", True), + ("corporate.500", True), + ("team1.300", True), + ("#fff", False), + ("#e8b7e4", False), + ("CornflowerBlue", False), + ("blue", False), + ("rgb(1, 2, 3)", False), + ("", False), + (None, False), + ("blue.", False), + (".500", False), + ("blue.500.foo", False), + ("blue.5e2", False), + # Non-ASCII digit/letter must not slip through the plain checks. + ("blue.²", False), + ("1blue.500", False), + ], +) +def test_is_chakra_color_token(value, expected): + assert is_chakra_color_token(value) is expected diff --git a/task-sdk/src/airflow/sdk/bases/operator.py b/task-sdk/src/airflow/sdk/bases/operator.py index 97da8869686cc..dd66830167cf1 100644 --- a/task-sdk/src/airflow/sdk/bases/operator.py +++ b/task-sdk/src/airflow/sdk/bases/operator.py @@ -38,6 +38,7 @@ import attrs from airflow.sdk import TriggerRule, timezone +from airflow.sdk._shared.dagnode.ui_color import is_chakra_color_token from airflow.sdk._shared.secrets_masker import redact from airflow.sdk.definitions._internal.abstractoperator import ( DEFAULT_EMAIL_ON_FAILURE, @@ -588,9 +589,32 @@ def __new__(cls, name, bases, namespace, **kwargs): if new_cls.__init__ is not first_superclass.__init__: new_cls.__init__ = cls._apply_defaults(new_cls.__init__) + _warn_on_invalid_ui_color(name, namespace, new_cls.__module__) + return new_cls +def _warn_on_invalid_ui_color(name: str, namespace: dict[str, Any], module: str) -> None: + """ + Warn when user code sets ``ui_color``/``ui_fgcolor`` to a non Chakra color token. + + Such values are ignored by the UI graph (only Chakra palette tokens like ``blue.500`` are + rendered). Built-in operators (``airflow.*``, including providers) are left alone so the many + that still carry legacy hex values do not flood Dag parsing with warnings. + """ + if module.startswith("airflow."): + return + for color_attr in ("ui_color", "ui_fgcolor"): + value = namespace.get(color_attr) + if isinstance(value, str) and not is_chakra_color_token(value): + warnings.warn( + f"{name}.{color_attr} value {value!r} is not a Chakra color token (e.g. 'blue.500') " + f"and will be ignored in the UI graph.", + UserWarning, + stacklevel=3, + ) + + # TODO: The following mapping is used to validate that the arguments passed to the BaseOperator are of the # correct type. This is a temporary solution until we find a more sophisticated method for argument # validation. One potential method is to use `get_type_hints` from the typing module. However, this is not diff --git a/task-sdk/src/airflow/sdk/definitions/taskgroup.py b/task-sdk/src/airflow/sdk/definitions/taskgroup.py index 14bb2fba31918..3414308832c36 100644 --- a/task-sdk/src/airflow/sdk/definitions/taskgroup.py +++ b/task-sdk/src/airflow/sdk/definitions/taskgroup.py @@ -21,6 +21,7 @@ import copy import re +import warnings import weakref from collections import deque from collections.abc import Generator, Iterator, Sequence @@ -29,6 +30,7 @@ import attrs from airflow.sdk import TriggerRule +from airflow.sdk._shared.dagnode.ui_color import is_chakra_color_token from airflow.sdk.definitions._internal.node import DAGNode, validate_group_key from airflow.sdk.exceptions import ( AirflowDagCycleException, @@ -76,6 +78,22 @@ def _validate_group_id(instance, attribute, value: str) -> None: validate_group_key(value) +def _warn_non_token_ui_color(instance, attribute, value: str) -> None: + """ + Warn when ``ui_color``/``ui_fgcolor`` is overridden with a non Chakra color token. + + Only Chakra palette tokens like ``blue.500`` are rendered by the UI graph; other values + (the legacy default, raw hex, CSS names) are ignored. + """ + if value != attribute.default and not is_chakra_color_token(value): + warnings.warn( + f"TaskGroup {attribute.name} value {value!r} is not a Chakra color token (e.g. 'blue.500') " + f"and will be ignored in the UI graph.", + UserWarning, + stacklevel=3, + ) + + def _convert_doc_md(doc_md: str | None) -> str | None: """Convert markdown file paths to file contents.""" if doc_md is None: @@ -149,8 +167,14 @@ class TaskGroup(DAGNode): on_setattr=attrs.setters.frozen, ) - ui_color: str = attrs.field(default="CornflowerBlue", validator=attrs.validators.instance_of(str)) - ui_fgcolor: str = attrs.field(default="#000", validator=attrs.validators.instance_of(str)) + ui_color: str = attrs.field( + default="CornflowerBlue", + validator=[attrs.validators.instance_of(str), _warn_non_token_ui_color], + ) + ui_fgcolor: str = attrs.field( + default="#000", + validator=[attrs.validators.instance_of(str), _warn_non_token_ui_color], + ) add_suffix_on_collision: bool = False diff --git a/task-sdk/tests/task_sdk/definitions/test_taskgroup.py b/task-sdk/tests/task_sdk/definitions/test_taskgroup.py index d11f8eb3632b3..f64a1238a1860 100644 --- a/task-sdk/tests/task_sdk/definitions/test_taskgroup.py +++ b/task-sdk/tests/task_sdk/definitions/test_taskgroup.py @@ -17,6 +17,8 @@ from __future__ import annotations +import warnings + import pendulum import pytest @@ -1169,3 +1171,25 @@ def spy(self, nodes, projected): assert called["value"] _assert_valid_topological_order(dag.task_group, order) + + +class TestTaskGroupUiColor: + @pytest.mark.parametrize( + ("kwargs", "expected_fields"), + [ + pytest.param({"ui_color": "red"}, ["ui_color"], id="non-token-color"), + pytest.param({"ui_fgcolor": "#000000"}, ["ui_fgcolor"], id="non-token-fgcolor"), + pytest.param({"ui_color": "teal.400", "ui_fgcolor": "gray.900"}, [], id="valid-tokens"), + pytest.param({}, [], id="defaults"), + ], + ) + def test_warns_on_non_token_colors(self, kwargs, expected_fields): + with DAG("d", schedule=None, start_date=DEFAULT_DATE): + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + TaskGroup("group", **kwargs) + + messages = [str(w.message) for w in caught if "Chakra color token" in str(w.message)] + assert len(messages) == len(expected_fields) + for field in expected_fields: + assert any(f"TaskGroup {field}" in message for message in messages) diff --git a/task-sdk/tests/task_sdk/definitions/test_ui_color.py b/task-sdk/tests/task_sdk/definitions/test_ui_color.py new file mode 100644 index 0000000000000..6f1a014588506 --- /dev/null +++ b/task-sdk/tests/task_sdk/definitions/test_ui_color.py @@ -0,0 +1,81 @@ +# 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 + +import warnings + +import pytest + +from airflow.sdk.bases.operator import BaseOperator, _warn_on_invalid_ui_color + + +@pytest.mark.parametrize( + ("module", "namespace", "expected_fields"), + [ + pytest.param("my_dags.operators", {"ui_color": "#fff"}, ["ui_color"], id="user-hex-color"), + pytest.param("my_dags.operators", {"ui_fgcolor": "black"}, ["ui_fgcolor"], id="user-named-fgcolor"), + pytest.param( + "my_dags.operators", + {"ui_color": "#fff", "ui_fgcolor": "#000"}, + ["ui_color", "ui_fgcolor"], + id="user-both", + ), + pytest.param("my_dags.operators", {"ui_color": "blue.500"}, [], id="user-valid-token"), + pytest.param("my_dags.operators", {}, [], id="user-no-override"), + pytest.param("airflow.providers.foo", {"ui_color": "#fff"}, [], id="provider-skipped"), + pytest.param("airflow.sdk.bases.sensor", {"ui_color": "#e6f1f2"}, [], id="core-skipped"), + ], +) +def test_warn_on_invalid_ui_color(module, namespace, expected_fields): + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + _warn_on_invalid_ui_color("MyOperator", namespace, module) + + messages = [str(w.message) for w in caught if issubclass(w.category, UserWarning)] + assert len(messages) == len(expected_fields) + for field in expected_fields: + assert any(f"MyOperator.{field}" in message for message in messages) + + +def test_operator_subclass_warns_on_invalid_ui_color(): + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + + class BadColorOperator(BaseOperator): + ui_color = "#e8b7e4" + + def execute(self, context): + pass + + messages = [str(w.message) for w in caught if issubclass(w.category, UserWarning)] + assert any("BadColorOperator.ui_color" in message for message in messages) + + +def test_operator_subclass_silent_on_valid_token(): + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + + class TokenColorOperator(BaseOperator): + ui_color = "blue.500" + ui_fgcolor = "gray.900" + + def execute(self, context): + pass + + messages = [str(w.message) for w in caught if "Chakra color token" in str(w.message)] + assert messages == []