Fix: Encode model/column name for linage in case it has special chars (#1747)

mykalmax · web-flow · commit e5e191905c59 · 2023-11-22T14:47:00.000-08:00
diff --git a/web/client/src/library/components/editor/Editor.css b/web/client/src/library/components/editor/Editor.css
@@ -119,21 +119,23 @@
 }
 
 .sqlmesh-model__column.--is-original,
-.sqlmesh-model__column.--is-original > span {
+.sqlmesh-model__column.--is-original > span,
+.sqlmesh-model__column.--is-alias,
+.sqlmesh-model__column.--is-alias > span {
   display: inline-block;
   color: var(--color-accent-300);
   font-weight: bold;
 }
 
-.sqlmesh-model__column.--is-active-model.--is-derived,
-.sqlmesh-model__column.--is-active-model.--is-derived > span {
+.sqlmesh-model__column.--is-active-model.--is-alias,
+.sqlmesh-model__column.--is-active-model.--is-alias > span {
   display: inline-block;
   color: var(--color-brand-300);
   font-weight: bold;
   box-shadow: inset 0 -2px 0 0 transparent;
 }
-.sqlmesh-model__column.--is-active-model.--is-derived.--is-action-mode,
-.sqlmesh-model__column.--is-active-model.--is-derived.--is-action-mode > span {
+.sqlmesh-model__column.--is-active-model.--is-alias.--is-action-mode,
+.sqlmesh-model__column.--is-active-model.--is-alias.--is-action-mode > span {
   cursor: pointer;
   box-shadow: inset 0 -2px 0 0 var(--color-brand-300);
 }
diff --git a/web/client/src/library/components/editor/extensions/SQLMeshDialect.ts b/web/client/src/library/components/editor/extensions/SQLMeshDialect.ts
@@ -23,11 +23,11 @@ export const SQLMeshDialect: ExtensionSQLMeshDialect = function SQLMeshDialect(
   options = { types: '', keywords: '' },
   dialects,
 ): LanguageSupport {
-  const SQLKeywords = options.keywords
-  const SQLTypes = options.types
+  const SQLKeywords = options.keywords + ' coalesce sum count avg min max cast'
+  const SQLTypes = options.types + ' string'
   const SQLMeshModelDictionary = getSQLMeshModelKeywords(dialects)
   const SQLMeshKeywords =
-    'columns grain grains references metric tags audit model name kind owner cron start storage_format time_column partitioned_by pre post batch_size audits dialect'
+    'path threshold jinja_query_begin number_of_rows jinja_end not_null forall criteria length unique_values interval_unit unique_key columns grain grains references metric tags audit model name kind owner cron start storage_format time_column partitioned_by pre post batch_size audits dialect'
   const SQLMeshTypes =
     'expression seed full incremental_by_time_range incremental_by_unique_key view embedded'
   const lang = SQLDialect.define({
diff --git a/web/client/src/library/components/editor/extensions/help.ts b/web/client/src/library/components/editor/extensions/help.ts
@@ -1,25 +1,35 @@
-import { syntaxTree } from '@codemirror/language'
 import {
   type EditorView,
   type DecorationSet,
   Decoration,
 } from '@codemirror/view'
 import { type ModelSQLMeshModel } from '@models/sqlmesh-model'
-import { isFalse } from '@utils/index'
+import { isFalse, isNil, isNotNil } from '@utils/index'
 import { type Column } from '@api/client'
 import clsx from 'clsx'
 
+interface DecorationRange {
+  from: number
+  to: number
+}
+type MarkDecorationCallback = (
+  options: DecorationRange & {
+    name: string
+    key: string
+  },
+) => void
+
 export function findModel(
   event: MouseEvent,
   models: Map<string, ModelSQLMeshModel>,
 ): ModelSQLMeshModel | undefined {
-  if (event.target == null) return
+  if (isNil(event.target)) return
 
   const el = event.target as HTMLElement
   const modelName =
     el.getAttribute('model') ?? el.parentElement?.getAttribute('model')
 
-  if (modelName == null) return
+  if (isNil(modelName)) return
 
   return models.get(modelName)
 }
@@ -28,13 +38,13 @@ export function findColumn(
   event: MouseEvent,
   model: ModelSQLMeshModel,
 ): Column | undefined {
-  if (event.target == null) return
+  if (isNil(event.target)) return
 
   const el = event.target as HTMLElement
   const columnName =
     el.getAttribute('column') ?? el.parentElement?.getAttribute('column')
 
-  if (columnName == null) return
+  if (isNil(columnName)) return
 
   return model.columns.find(c => c.name === columnName)
 }
@@ -46,89 +56,155 @@ export function getDecorations(
   columns: Set<string>,
   isActionMode: boolean,
 ): DecorationSet {
-  const decorations: any = []
-  const modelColumns = model.columns.map(c => c.name)
+  const ranges: any = []
+  const columnNames = Array.from(columns)
+  const modelNames = Array.from(new Set(models.values())).map(m => m.name)
+  const modelColumns = model.columns.map(c => c.name.toLowerCase())
+  const validLeftCharColumn = new Set(['.', '(', '[', ' ', '\n'])
+  const validRightCharColumn = new Set([':', ')', ',', ']', ' ', '\n'])
 
   for (const range of view.visibleRanges) {
-    syntaxTree(view.state).iterate({
-      from: range.from,
-      to: range.to,
-      enter({ from, to }) {
-        // In case model name represented in qoutes
-        // like in python files, we need to remove qoutes
-        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,
+    getMarkDecorations(
+      modelNames,
+      view.state.doc.sliceString(range.from, range.to),
+      range,
+      ({ from, to, name }) => {
+        ranges.push(
+          createMarkDecorationModel({
+            model: name,
+            isActionMode,
+            isActiveModel: model.name === name,
+          }).range(from, to),
         )
+      },
+    )
+
+    getMarkDecorations(
+      columnNames,
+      view.state.doc.sliceString(range.from, range.to),
+      range,
+      ({ from, to, key }) => {
+        const column = view.state.doc.sliceString(from, to)
+        const word = view.state.doc.sliceString(from - 1, to + 1)
+        const leftChar = view.state.doc.sliceString(from - 1, from)
+        const rightChar = view.state.doc.sliceString(to, to + 1)
 
         if (
-          isFalse(isOriginal) &&
-          columns.has(maybeModelOrColumn) &&
-          !modelColumns.includes(maybeModelOrColumn)
-        ) {
-          isOriginal = true
-        }
-
-        let decoration
-
-        if (maybeModelOrColumn === model.name) {
-          decoration = Decoration.mark({
-            attributes: {
-              class: clsx(
-                'sqlmesh-model --is-active-model',
-                isActionMode && '--is-action-mode',
-              ),
-              model: maybeModelOrColumn,
-            },
-          }).range(from, to)
-        } else if (models.get(maybeModelOrColumn) != null) {
-          decoration = Decoration.mark({
-            attributes: {
-              class: clsx('sqlmesh-model', isActionMode && '--is-action-mode'),
-              model: maybeModelOrColumn,
-            },
-          }).range(from, to)
-        } else if (modelColumns.includes(maybeModelOrColumn)) {
-          decoration = Decoration.mark({
-            attributes: {
-              class: clsx(
-                'sqlmesh-model__column --is-active-model',
-                isOriginal ? '--is-original' : '--is-derived',
-                isActionMode && '--is-action-mode',
-              ),
-              column: maybeModelOrColumn,
-            },
-          }).range(from, to)
-        } else if (columns.has(maybeModelOrColumn)) {
-          decoration = Decoration.mark({
-            attributes: {
-              class: clsx(
-                'sqlmesh-model__column',
-                isOriginal ? '--is-original' : '--is-derived',
-                isActionMode && '--is-action-mode',
-              ),
-              column: maybeModelOrColumn,
-            },
-          }).range(from, to)
-        }
-
-        decoration != null && decorations.push(decoration)
+          (isNotNil(leftChar) && isFalse(validLeftCharColumn.has(leftChar))) ||
+          (isNotNil(rightChar) && isFalse(validRightCharColumn.has(rightChar)))
+        )
+          return
+
+        ranges.push(
+          createMarkDecorationColumn({
+            column,
+            isOriginalColumn: isOriginalColumn(word),
+            isActiveModel: modelColumns.includes(key),
+            isActionMode,
+          }).range(from, to),
+        )
       },
-    })
+    )
+  }
+
+  return Decoration.set(
+    ranges.sort((a: DecorationRange, b: DecorationRange) => a.from - b.from),
+  )
+}
+
+function getMarkDecorations(
+  list: string[],
+  doc: string,
+  range: DecorationRange,
+  callback: MarkDecorationCallback,
+): void {
+  const visisted = new Set()
+
+  for (const name of list) {
+    const key = name.toLowerCase()
+
+    if (visisted.has(key)) continue
+
+    visisted.add(key)
+
+    const regex = new RegExp(name, 'ig')
+    const regex_normalized = new RegExp(alternativeNameFormat(name), 'ig')
+    let found
+
+    while (isNotNil((found = regex_normalized.exec(doc) ?? regex.exec(doc)))) {
+      const from = range.from + found.index
+      const to = from + found[0].length
+      const options = {
+        from,
+        to,
+        name,
+        key,
+      }
+
+      isNotNil(callback) && callback(options)
+    }
   }
+}
+
+function createMarkDecorationModel({
+  model,
+  isActionMode,
+  isActiveModel,
+}: {
+  model: string
+  isActionMode: boolean
+  isActiveModel: boolean
+}): Decoration {
+  return Decoration.mark({
+    attributes: {
+      class: clsx(
+        'sqlmesh-model',
+        isActiveModel && '--is-active-model',
+        isActionMode && '--is-action-mode',
+      ),
+      model,
+    },
+  })
+}
+
+function createMarkDecorationColumn({
+  column,
+  isOriginalColumn,
+  isActionMode,
+  isActiveModel,
+}: {
+  column: string
+  isOriginalColumn: boolean
+  isActiveModel: boolean
+  isActionMode: boolean
+}): Decoration {
+  return Decoration.mark({
+    attributes: {
+      class: clsx(
+        'sqlmesh-model__column',
+        isOriginalColumn ? '--is-original' : '--is-alias',
+        isActiveModel && '--is-active-model',
+        isActionMode && '--is-action-mode',
+      ),
+      column,
+    },
+  })
+}
+
+function isOriginalColumn(column: string): boolean {
+  return (
+    column.startsWith('.') ||
+    column.endsWith(':') ||
+    column.startsWith('(') ||
+    column.endsWith(')')
+  )
+}
 
-  return Decoration.set(decorations)
+function alternativeNameFormat(modelName: string): string {
+  return modelName.includes('"')
+    ? modelName
+    : modelName
+        .split('.')
+        .map(name => `"${name}"`)
+        .join('.')
 }
diff --git a/web/client/src/library/components/graph/Graph.tsx b/web/client/src/library/components/graph/Graph.tsx
@@ -171,7 +171,7 @@ const ModelColumnDisplay = memo(function ModelColumnDisplay({
             className={clsx('flex items-center', disabled && 'opacity-50')}
           >
             {disabled && <NoSymbolIcon className="w-3 h-3 mr-2" />}
-            <b>{columnName}</b>
+            <b>{decodeURI(columnName)}</b>
           </span>
           <span className="inline-block text-neutral-400 dark:text-neutral-300 ml-2">
             {columnType}
@@ -326,7 +326,7 @@ const ModelNodeHeaderHandles = memo(function ModelNodeHeaderHandles({
             )}
             onClick={handleClick}
           >
-            {label}
+            {decodeURI(label)}
           </span>
           <span className="flex justify-between mx-2 px-2 rounded-full bg-neutral-10">
             {count}
diff --git a/web/client/src/library/components/graph/help.ts b/web/client/src/library/components/graph/help.ts
