feat(ui): add tag-based node color coding for lineage DAG

Signed-off-by: CTC97 <ctcurnin@gmail.com>

Author: CTC97

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] = {}

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,

web/client/src/context/context.ts

@@ -12,6 +12,7 @@ import { isNil, isStringEmptyOrNil, isTrue } from '~/utils'
 
 interface ContextStore {
   version?: string
+  nodeColors: Record<string, string>
   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<string, string>) => void
   setVersion: (version?: string) => void
   setShowConfirmation: (showConfirmation: boolean) => void
   addConfirmation: (confirmation: Confirmation) => void
@@ -55,6 +57,7 @@ const environment =
 
 export const useStoreContext = create<ContextStore>((set, get) => ({
   version: undefined,
+  nodeColors: {},
   modules: new ModelModuleController(),
   splitPaneSizes: [20, 80],
   showConfirmation: false,
@@ -83,6 +86,11 @@ export const useStoreContext = create<ContextStore>((set, get) => ({
       splitPaneSizes,
     }))
   },
+  setNodeColors(nodeColors) {
+    set(() => ({
+      nodeColors,
+    }))
+  },
   setVersion(version) {
     set(() => ({
       version,

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 }
+          : {}),
+      }}
     >
       <NodeResizeControl
         minWidth={150}

web/client/src/library/components/graph/help.ts

@@ -23,6 +23,7 @@ export interface GraphNodeData {
   label: string
   type: LineageNodeModelType
   withColumns: boolean
+  tags?: string[]
   [key: string]: any
 }
 
@@ -172,9 +173,14 @@ function getNodeMap({
 
   return modelNames.reduce((acc: Record<string, Node>, 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

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(

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,
+    )

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):