Feat: Display node for CTE (#915)

Author: mykalmax

tests/web/test_main.py

@@ -484,13 +484,13 @@ def test_get_environments(project_context: Context) -> None:
 
 def test_get_lineage(web_sushi_context: Context) -> None:
     response = client.get("/api/lineage/sushi.waiters/ds")
-    assert response.status_code == 200
 
+    assert response.status_code == 200
     assert response.json() == {
         "sushi.waiters": {
             "ds": {
                 "source": """SELECT DISTINCT
-  <b>CAST(o.ds AS TEXT) AS ds</b>
+  CAST(o.ds AS TEXT) AS ds
 FROM (
   SELECT
     CAST(NULL AS INT) AS id,
@@ -504,12 +504,14 @@ def test_get_lineage(web_sushi_context: Context) -> None:
 ) AS o /* source: sushi.orders */
 WHERE
   o.ds <= '1970-01-01' AND o.ds >= '1970-01-01'""",
+                "expression": "CAST(o.ds AS TEXT) AS ds",
                 "models": {"sushi.orders": ["ds"]},
             }
         },
         "sushi.orders": {
             "ds": {
-                "source": "SELECT\n  <b>CAST(NULL AS TEXT) AS ds</b>\nFROM (VALUES\n  (1)) AS t(dummy)",
+                "source": "SELECT\n  CAST(NULL AS TEXT) AS ds\nFROM (VALUES\n  (1)) AS t(dummy)",
+                "expression": "CAST(NULL AS TEXT) AS ds",
                 "models": {},
             }
         },

web/client/openapi.json

@@ -985,6 +985,7 @@
         "type": "object",
         "properties": {
           "source": { "title": "Source", "type": "string" },
+          "expression": { "title": "Expression", "type": "string" },
           "models": {
             "title": "Models",
             "type": "object",

web/client/src/api/instance.ts

@@ -65,7 +65,6 @@ export async function fetchAPI<T = any, B extends object = any>(
     })
       .then(async response => {
         const headerContentType = response.headers.get('Content-Type')
-
         if (headerContentType == null)
           return { ok: false, message: 'Empty response' }
         if (response.status >= 400) {
@@ -75,10 +74,10 @@ export async function fetchAPI<T = any, B extends object = any>(
             json.status = json.status ?? response.status
 
             throw json
-          } catch (error) {
+          } catch (error: any) {
             throw {
               message: response.statusText,
-              status: response.status,
+              ...error,
             } as unknown as Error
           }
         }

web/client/src/index.css

@@ -7,6 +7,10 @@
     height: 0.2rem;
   }
 
+  .scrollbar--horizontal-md::-webkit-scrollbar {
+    height: 0.4rem;
+  }
+
   .scrollbar--vertical::-webkit-scrollbar {
     width: 0.2rem;
   }

web/client/src/library/components/documentation/Documentation.tsx

@@ -7,10 +7,9 @@ import { MinusCircleIcon, PlusCircleIcon } from '@heroicons/react/24/solid'
 import { EnumFileExtensions } from '@models/file'
 import { isFalse, isString, isTrue, toDateFormat } from '@utils/index'
 import clsx from 'clsx'
-import { useNavigate } from 'react-router-dom'
-import { EnumRoutes } from '~/routes'
-import { ModelSQLMeshModel } from '@models/sqlmesh-model'
+import { type ModelSQLMeshModel } from '@models/sqlmesh-model'
 import { ModelColumns } from '@components/graph/Graph'
+import { useLineageFlow } from '@components/graph/context'
 
 const Documentation = function Documentation({
   model,
@@ -27,12 +26,10 @@ const Documentation = function Documentation({
   withDescription?: boolean
   withColumns?: boolean
 }): JSX.Element {
-  const navigate = useNavigate()
+  const { handleClickModel } = useLineageFlow()
 
   const modelExtensions = useSQLMeshModelExtensions(model.path, model => {
-    navigate(
-      `${EnumRoutes.IdeDocsModels}/${ModelSQLMeshModel.encodeName(model.name)}`,
-    )
+    handleClickModel?.(model.name)
   })
 
   return (
@@ -86,6 +83,7 @@ const Documentation = function Documentation({
             columns={model.columns}
             disabled={model?.type === 'python'}
             withHandles={false}
+            withSource={false}
             limit={10}
           />
         </Section>

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

@@ -36,6 +36,14 @@
   width: 0.2rem;
 }
 
+.scrollbar--vertical-md .cm-scroller::-webkit-scrollbar {
+  width: 0.4rem;
+}
+
+.scrollbar--horizontal-md .cm-scroller::-webkit-scrollbar {
+  height: 0.4rem;
+}
+
 .cm-editor .cm-scroller::-webkit-scrollbar-track {
   background-color: transparent;
 }
@@ -51,6 +59,11 @@
   background: var(--color-theme) !important;
 }
 
+.sqlmesh-expression {
+  color: var(--color-primary-500);
+  background: var(--color-primary-10);
+}
+
 .sqlmesh-model,
 .sqlmesh-model > span {
   display: inline-block;

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

@@ -44,9 +44,11 @@ function CodeEditorDefault({
   type,
   content = '',
   children,
+  className,
 }: {
   type: FileExtensions
   content: string
+  className?: string
   children: (options: {
     extensions: Extension[]
     content: string
@@ -65,7 +67,7 @@ function CodeEditorDefault({
   }, [type, mode])
 
   return (
-    <div className="flex overflow-auto h-full">
+    <div className={clsx('flex w-full h-full', className)}>
       {children({ extensions, content })}
     </div>
   )
@@ -75,9 +77,11 @@ function CodeEditorSQLMesh({
   type,
   content = '',
   children,
+  className,
 }: {
   type: FileExtensions
   content?: string
+  className?: string
   children: (options: {
     extensions: Extension[]
     content: string
@@ -110,7 +114,6 @@ function CodeEditorSQLMesh({
 
   const extensions = useMemo(() => {
     return [
-      mode === EnumColorScheme.Dark ? dracula : tomorrow,
       mode === EnumColorScheme.Dark ? dracula : tomorrow,
       type === EnumFileExtensions.Python && python(),
       type === EnumFileExtensions.YAML && StreamLanguage.define(yaml),
@@ -140,7 +143,7 @@ function CodeEditorSQLMesh({
   }, [content])
 
   return (
-    <div className="flex overflow-auto h-full">
+    <div className={clsx('flex w-full h-full', className)}>
       {children({ extensions, content })}
     </div>
   )
@@ -214,9 +217,7 @@ export function useSQLMeshModelExtensions(
   const files = useStoreFileTree(s => s.files)
 
   const extensions = useMemo(() => {
-    if (path == null) return []
-
-    const model = models.get(path)
+    const model = path == null ? undefined : models.get(path)
     const columns =
       lineage == null
         ? new Set<string>()

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

@@ -1,5 +1,5 @@
 import { syntaxTree } from '@codemirror/language'
-import { type Extension } from '@codemirror/state'
+import { RangeSetBuilder, type Extension } from '@codemirror/state'
 import {
   ViewPlugin,
   type DecorationSet,
@@ -128,6 +128,53 @@ export function HoverTooltip(
   )
 }
 
+export function SqlMeshExpression(expression: string): Extension {
+  return ViewPlugin.fromClass(
+    class SqlMeshModelView {
+      decorations: DecorationSet = Decoration.set([])
+
+      constructor(readonly view: EditorView) {
+        this.decorations = markExpressionLine(expression, view)
+      }
+
+      update(viewUpdate: ViewUpdate): void {
+        this.decorations = markExpressionLine(expression, viewUpdate.view)
+      }
+    },
+    {
+      decorations: value => value.decorations,
+    },
+  )
+}
+
+function markExpressionLine(
+  expression: string,
+  view: EditorView,
+): DecorationSet {
+  const mark = Decoration.line({
+    attributes: {
+      id: expression,
+      class: 'sqlmesh-expression',
+    },
+  })
+
+  const builder = new RangeSetBuilder<Decoration>()
+
+  for (const { from, to } of view.visibleRanges) {
+    for (let pos = from; pos <= to; ) {
+      const line = view.state.doc.lineAt(pos)
+
+      if (line.text.includes(expression)) {
+        builder.add(line.from, line.from, mark)
+      }
+
+      pos = line.to + 1
+    }
+  }
+
+  return builder.finish()
+}
+
 function getDecorations(
   models: Map<string, ModelSQLMeshModel>,
   view: EditorView,