From 0ec81a66ae2c1e345e4b76fc30524eade8a62c07 Mon Sep 17 00:00:00 2001 From: CTC97 Date: Wed, 10 Jun 2026 10:25:15 -0400 Subject: [PATCH 1/2] feat(ui): add tag-based node color coding for lineage DAG Signed-off-by: CTC97 --- sqlmesh/core/config/ui.py | 5 +++++ web/client/openapi.json | 6 ++++++ web/client/src/context/context.ts | 8 ++++++++ .../library/components/graph/ModelNode.tsx | 19 ++++++++++++++++++- .../src/library/components/graph/help.ts | 6 ++++++ web/client/src/library/pages/root/Root.tsx | 2 ++ web/server/api/endpoints/meta.py | 13 +++++++++++-- web/server/models.py | 1 + 8 files changed, 57 insertions(+), 3 deletions(-) diff --git a/sqlmesh/core/config/ui.py b/sqlmesh/core/config/ui.py index e24dc43375..478c724be6 100644 --- a/sqlmesh/core/config/ui.py +++ b/sqlmesh/core/config/ui.py @@ -1,5 +1,7 @@ from __future__ import annotations +import typing as t + from sqlmesh.core.config.base import BaseConfig @@ -8,6 +10,9 @@ class UIConfig(BaseConfig): Args: format_on_save: Whether to format the SQL code on save or not. + node_colors: A mapping of model tags to hex color strings used + to color-code nodes in the lineage DAG visualization. """ format_on_save: bool = True + node_colors: t.Dict[str, str] = {} diff --git a/web/client/openapi.json b/web/client/openapi.json index bf1cef0809..56e25dda52 100644 --- a/web/client/openapi.json +++ b/web/client/openapi.json @@ -1371,6 +1371,12 @@ "type": "boolean", "title": "Has Running Task", "default": false + }, + "node_colors": { + "additionalProperties": { "type": "string" }, + "type": "object", + "title": "Node Colors", + "default": {} } }, "additionalProperties": false, diff --git a/web/client/src/context/context.ts b/web/client/src/context/context.ts index eedf9976ff..ca3a53926e 100644 --- a/web/client/src/context/context.ts +++ b/web/client/src/context/context.ts @@ -12,6 +12,7 @@ import { isNil, isStringEmptyOrNil, isTrue } from '~/utils' interface ContextStore { version?: string + nodeColors: Record showConfirmation: boolean confirmations: Confirmation[] environment: ModelEnvironment @@ -23,6 +24,7 @@ interface ContextStore { setLastSelectedModel: (model?: ModelSQLMeshModel) => void setSplitPaneSizes: (splitPaneSizes: number[]) => void setModules: (modules: ModelModuleController) => void + setNodeColors: (nodeColors: Record) => void setVersion: (version?: string) => void setShowConfirmation: (showConfirmation: boolean) => void addConfirmation: (confirmation: Confirmation) => void @@ -55,6 +57,7 @@ const environment = export const useStoreContext = create((set, get) => ({ version: undefined, + nodeColors: {}, modules: new ModelModuleController(), splitPaneSizes: [20, 80], showConfirmation: false, @@ -83,6 +86,11 @@ export const useStoreContext = create((set, get) => ({ splitPaneSizes, })) }, + setNodeColors(nodeColors) { + set(() => ({ + nodeColors, + })) + }, setVersion(version) { set(() => ({ version, diff --git a/web/client/src/library/components/graph/ModelNode.tsx b/web/client/src/library/components/graph/ModelNode.tsx index cfa17f7639..0fe2f27d51 100644 --- a/web/client/src/library/components/graph/ModelNode.tsx +++ b/web/client/src/library/components/graph/ModelNode.tsx @@ -2,6 +2,7 @@ import { isNil, isArrayNotEmpty, isNotNil, toID, isFalse } from '@utils/index' import clsx from 'clsx' import { useMemo, useCallback, useState, useRef } from 'react' import { ModelType } from '@api/client' +import { useStoreContext } from '@context/context' import { useLineageFlow } from './context' import { type GraphNodeData } from './help' import { Position, type NodeProps, NodeResizeControl } from 'reactflow' @@ -44,6 +45,7 @@ export default function ModelNode({ highlightedNodes, activeNodes, } = useLineageFlow() + const nodeColors = useStoreContext(s => s.nodeColors) const columns: Column[] = useMemo(() => { const model = models.get(id) @@ -113,6 +115,16 @@ export default function ModelNode({ [setSelectedNodes, highlightedNodeModels], ) + const tagColor = useMemo(() => { + const tags = nodeData.tags + if (isNil(tags) || Object.keys(nodeColors).length === 0) return undefined + for (const tag of tags) { + const color = nodeColors[tag] + if (color) return color + } + return undefined + }, [nodeData.tags, nodeColors]) + const splat = highlightedNodes['*'] const hasSelectedColumns = columns.some(({ name }) => connections.get(toID(id, name)), @@ -183,7 +195,12 @@ export default function ModelNode({ ? 'ring-8 ring-neutral-50' : isSelected && 'ring-8 ring-secondary-50 dark:ring-primary-50', )} - style={{ width: '100%' }} + style={{ + width: '100%', + ...(tagColor != null + ? { borderColor: tagColor, backgroundColor: tagColor, color: tagColor } + : {}), + }} > , modelName: string) => { const model = models.get(modelName) + const tagsStr = model?.details?.tags + const tags = tagsStr + ? tagsStr.split(',').map(t => t.trim()).filter(Boolean) + : undefined const node = createGraphNode(modelName, { label: model?.displayName ?? modelName, withColumns, + tags, type: isNotNil(model) ? (model.type as LineageNodeModelType) : // If model name present in lineage but not in global models diff --git a/web/client/src/library/pages/root/Root.tsx b/web/client/src/library/pages/root/Root.tsx index 6f609d73a3..b39f3750f7 100644 --- a/web/client/src/library/pages/root/Root.tsx +++ b/web/client/src/library/pages/root/Root.tsx @@ -106,6 +106,7 @@ export default function Root({ const closeTab = useStoreEditor(s => s.closeTab) const inTabs = useStoreEditor(s => s.inTabs) const setVersion = useStoreContext(s => s.setVersion) + const setNodeColors = useStoreContext(s => s.setNodeColors) const { refetch: getMeta, cancel: cancelRequestMeta } = useApiMeta() const { refetch: getModels, cancel: cancelRequestModels } = useApiModels() @@ -313,6 +314,7 @@ export default function Root({ useEffect(() => { void getMeta().then(({ data }) => { setVersion(data?.version) + setNodeColors(data?.node_colors ?? {}) if (isTrue(data?.has_running_task)) { setPlanAction( diff --git a/web/server/api/endpoints/meta.py b/web/server/api/endpoints/meta.py index 430bf0d7a1..f159f1992a 100644 --- a/web/server/api/endpoints/meta.py +++ b/web/server/api/endpoints/meta.py @@ -5,7 +5,7 @@ from sqlmesh.cli.main import _sqlmesh_version from web.server import models from web.server.console import api_console -from web.server.settings import Settings, get_settings +from web.server.settings import Settings, get_context, get_settings router = APIRouter() @@ -32,4 +32,13 @@ def get_api_meta( if not has_running_task and api_console.is_cancelling_plan(): api_console.finish_plan_cancellation() - return models.Meta(version=_sqlmesh_version(), has_running_task=has_running_task) + node_colors: dict[str, str] = {} + context = get_context(settings) + if context: + node_colors = context.config.ui.node_colors + + return models.Meta( + version=_sqlmesh_version(), + has_running_task=has_running_task, + node_colors=node_colors, + ) diff --git a/web/server/models.py b/web/server/models.py index c81fd8eb5a..935822cf46 100644 --- a/web/server/models.py +++ b/web/server/models.py @@ -133,6 +133,7 @@ class Directory(PydanticModel): class Meta(PydanticModel): version: str has_running_task: bool = False + node_colors: t.Dict[str, str] = {} class Reference(PydanticModel): From ae2b8596980aa3989a4db0da08bdf09ef90e4954 Mon Sep 17 00:00:00 2001 From: CTC97 Date: Thu, 11 Jun 2026 09:43:54 -0400 Subject: [PATCH 2/2] feat(ui): add tag-based node color coding for lineage DAG Signed-off-by: CTC97 --- tests/web/test_main.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/tests/web/test_main.py b/tests/web/test_main.py index cf2220ad6d..b20947c49d 100644 --- a/tests/web/test_main.py +++ b/tests/web/test_main.py @@ -415,7 +415,11 @@ def test_meta(client: TestClient) -> None: response = client.get("/api/meta") assert response.status_code == 200 - assert response.json() == {"version": _sqlmesh_version(), "has_running_task": False} + assert response.json() == { + "version": _sqlmesh_version(), + "has_running_task": False, + "node_colors": {}, + } def test_modules(client: TestClient) -> None: