Feat: Highlight know columns in editor (#761)

Author: mykalmax

web/client/src/index.css

@@ -49,6 +49,10 @@
     --color-neutral: hsla(0, 0%, 50%, 1);
 
     /* Brand */
+    --color-brand-10: hsla(24, 100%, 60%, 0.1);
+    --color-brand-20: hsla(24, 100%, 60%, 0.2);
+    --color-brand-30: hsla(24, 100%, 60%, 0.3);
+    --color-brand-40: hsla(24, 100%, 60%, 0.4);
     --color-brand-100: hsla(37, 100%, 92%, 1);
     --color-brand-200: hsla(34, 100%, 84%, 1);
     --color-brand-300: hsla(31, 100%, 76%, 1);

web/client/src/library/components/editor/Editor.css

@@ -45,21 +45,44 @@
   border-radius: 1rem;
 }
 
-.sqlmesh-model {
+.cm-tooltip {
+  border: none !important;
+  outline: none !important;
+  background: var(--color-theme) !important;
+}
+
+.sqlmesh-model,
+.sqlmesh-model > span {
+  display: inline-block;
   color: var(--color-primary-500);
   background: var(--color-primary-10);
   box-shadow: inset 0 -2px 0 0 var(--color-primary-300);
-  display: inline-block;
   cursor: pointer;
   font-weight: bold;
 }
 
-.cm-tooltip {
-  border: none !important;
-  outline: none !important;
-  background: var(--color-theme) !important;
-}
-
 .sqlmesh-model:hover {
   opacity: 0.8;
 }
+
+.sqlmesh-model.--is-active-model,
+.sqlmesh-model.--is-active-model:hover,
+.sqlmesh-model.--is-active-model > span,
+.sqlmesh-model.--is-active-model:hover > span {
+  box-shadow: none;
+  cursor: default;
+  opacity: 1;
+}
+
+.sqlmesh-model__column.--is-original,
+.sqlmesh-model__column.--is-original > span {
+  display: inline-block;
+  color: var(--color-accent-300);
+}
+
+.sqlmesh-model__column.--is-active-model.--is-derived,
+.sqlmesh-model__column.--is-active-model.--is-derived > span {
+  display: inline-block;
+  color: var(--color-brand-300);
+  background: var(--color-brand-10);
+}

web/client/src/library/components/editor/EditorCode.tsx

@@ -38,10 +38,10 @@ export default function CodeEditor({ tab }: { tab: EditorTab }): JSX.Element {
   const [SqlMeshDialect, SqlMeshDialectCleanUp] = useSqlMeshExtension()
 
   const models = useStoreContext(s => s.models)
-
   const files = useStoreFileTree(s => s.files)
   const selectFile = useStoreFileTree(s => s.selectFile)
 
+  const previewLineage = useStoreEditor(s => s.previewLineage)
   const tabs = useStoreEditor(s => s.tabs)
   const dialects = useStoreEditor(s => s.dialects)
   const engine = useStoreEditor(s => s.engine)
@@ -73,11 +73,17 @@ export default function CodeEditor({ tab }: { tab: EditorTab }): JSX.Element {
   )
 
   const extensions = useMemo(() => {
+    const model = models.get(tab.file.path)
+    const columns = Object.keys(previewLineage ?? {})
+      .map(model => models.get(model)?.columns.map(c => c.name))
+      .flat()
+      .filter(Boolean) as string[]
+
     return [
       mode === EnumColorScheme.Dark ? dracula : tomorrow,
       HoverTooltip(models),
       events(models, files, selectFile),
-      SqlMeshModel(models),
+      model != null && SqlMeshModel(models, model, columns),
       tab.file.extension === '.py' && python(),
       tab.file.extension === '.yaml' && StreamLanguage.define(yaml),
       tab.file.extension === '.sql' &&

web/client/src/library/components/editor/extensions/SqlMeshDialect.ts

@@ -128,14 +128,13 @@ export function useSqlMeshExtension(): [
       ])
     }
 
-  const SqlMeshDialectCleanUp: ExtensionCleanUp =
-    function SqlMeshDialectCleanUp(): void {
-      const handler = cache.get('message')
+  function SqlMeshDialectCleanUp(): void {
+    const handler = cache.get('message')
 
-      if (handler == null) return
+    if (handler == null) return
 
-      sqlglotWorker.removeEventListener('message', handler)
-    }
+    sqlglotWorker.removeEventListener('message', handler)
+  }
 
   return [SqlMeshDialectExtension, SqlMeshDialectCleanUp]
 }

web/client/src/library/components/editor/extensions/index.ts

@@ -11,24 +11,30 @@ import {
 } from '@codemirror/view'
 import { type Model } from '~/api/client'
 import { type ModelFile } from '~/models'
-import { isNil } from '~/utils'
 
 import { useSqlMeshExtension } from './SqlMeshDialect'
 
 export { useSqlMeshExtension }
 
-export function SqlMeshModel(models: Map<string, Model>): Extension {
+export function SqlMeshModel(
+  models: Map<string, Model>,
+  model: Model,
+  columns: string[],
+): Extension {
   return ViewPlugin.fromClass(
     class SqlMeshModelView {
       decorations: DecorationSet = Decoration.set([])
       constructor(readonly view: EditorView) {
-        this.decorations = getDecorations(models, view)
+        this.decorations = getDecorations(models, view, model, columns)
       }
 
       update(viewUpdate: ViewUpdate): void {
-        if (viewUpdate.docChanged) {
-          this.decorations = getDecorations(models, viewUpdate.view)
-        }
+        this.decorations = getDecorations(
+          models,
+          viewUpdate.view,
+          model,
+          columns,
+        )
       }
     },
     {
@@ -75,11 +81,11 @@ export function HoverTooltip(models: Map<string, Model>): Extension {
         create() {
           const dom = document.createElement('div')
           const template = `
-          <div class="flex items-center">
-            <span>Model Name:</span>
-            <span class="px-2 py-1 inline-block ml-1 bg-alternative-100 text-alternative-500 rounded">${model.name}</span>
-          </div>
-        `
+            <div class="flex items-center">
+              <span>Model Name:</span>
+              <span class="px-2 py-1 inline-block ml-1 bg-alternative-100 text-alternative-500 rounded">${model.name}</span>
+            </div>
+          `
 
           dom.className =
             'text-xs font-bold px-3 py-3 bg-white border-2 border-secondary-100 rounded outline-none shadow-lg mb-2'
@@ -97,8 +103,11 @@ export function HoverTooltip(models: Map<string, Model>): Extension {
 function getDecorations(
   models: Map<string, Model>,
   view: EditorView,
+  model: Model,
+  columns: string[],
 ): DecorationSet {
   const decorations: any = []
+  const modelColumns = model.columns.map(c => c.name)
 
   for (const range of view.visibleRanges) {
     syntaxTree(view.state).iterate({
@@ -107,21 +116,68 @@ function getDecorations(
       enter({ from, to }) {
         // In case model name represented in qoutes
         // like in python files, we need to remove qoutes
-        const model = view.state.doc
-          .sliceString(from, to)
+        let maybeModelOrColumn = view.state.doc
+          .sliceString(from - 1, to + 1)
           .replaceAll('"', '')
           .replaceAll("'", '')
+        let isOriginal = false
+
+        if (
+          maybeModelOrColumn.startsWith('.') ||
+          maybeModelOrColumn.endsWith(':')
+        ) {
+          isOriginal = true
+        }
+
+        maybeModelOrColumn = maybeModelOrColumn.slice(
+          1,
+          maybeModelOrColumn.length - 1,
+        )
 
-        if (isNil(models.get(model))) return true
+        if (
+          columns.includes(maybeModelOrColumn) &&
+          !modelColumns.includes(maybeModelOrColumn)
+        ) {
+          isOriginal = true
+        }
 
-        const decoration = Decoration.mark({
-          attributes: {
-            class: 'sqlmesh-model',
-            model,
-          },
-        }).range(from, to)
+        let decoration
+
+        if (maybeModelOrColumn === model.name) {
+          decoration = Decoration.mark({
+            attributes: {
+              class: 'sqlmesh-model --is-active-model',
+              model: maybeModelOrColumn,
+            },
+          }).range(from, to)
+        } else if (models.get(maybeModelOrColumn) != null) {
+          decoration = Decoration.mark({
+            attributes: {
+              class: 'sqlmesh-model',
+              model: maybeModelOrColumn,
+            },
+          }).range(from, to)
+        } else if (modelColumns.includes(maybeModelOrColumn)) {
+          decoration = Decoration.mark({
+            attributes: {
+              class: `sqlmesh-model__column --is-active-model ${
+                isOriginal ? '--is-original' : ' --is-derived'
+              }`,
+              column: maybeModelOrColumn,
+            },
+          }).range(from, to)
+        } else if (columns.includes(maybeModelOrColumn)) {
+          decoration = Decoration.mark({
+            attributes: {
+              class: `sqlmesh-model__column ${
+                isOriginal ? '--is-original' : ' --is-derived'
+              }`,
+              column: maybeModelOrColumn,
+            },
+          }).range(from, to)
+        }
 
-        decorations.push(decoration)
+        decoration != null && decorations.push(decoration)
       },
     })
   }

web/client/tailwind.config.js

@@ -43,6 +43,7 @@ module.exports = {
         DEFAULT: 'var(--color-divider)',
       },
       brand: {
+        10: 'var(--color-brand-10)',
         100: 'var(--color-brand-100)',
         200: 'var(--color-brand-200)',
         300: 'var(--color-brand-300)',