[UI] Enhance editor with sqlglot (#513)

Author: mykalmax

.pre-commit-config.yaml

@@ -42,7 +42,7 @@ repos:
       - id: prettier
         name: prettier
         files: ^(web/client)
-        entry: prettier --write
+        entry: prettier --ignore-path web/client/.prettierignore
         exclude: ^(web/client/node_modules)
         require_serial: true
         language: node

web/client/.prettierignore

@@ -0,0 +1,2 @@
+**/*.py
+.prettierignore

web/client/src/api/index.ts

@@ -24,6 +24,8 @@ import {
   applyApiCommandsApplyPost,
   cancelPlanApiPlanCancelPost,
   type BodyApplyApiCommandsApplyPostCategories,
+  getModelsApiModelsGet,
+  type Models,
 } from './client'
 
 export function useApiDag(): UseQueryResult<DagApiCommandsDagGet200> {
@@ -45,6 +47,15 @@ export function useApiFileByPath(path?: string): UseQueryResult<File> {
   })
 }
 
+export function useApiModels(): UseQueryResult<Models> {
+  return useQuery({
+    queryKey: ['/api/models'],
+    queryFn: async ({ signal }) => await getModelsApiModelsGet({ signal }),
+    cacheTime: 0,
+    enabled: false,
+  })
+}
+
 export function useApiFiles(): UseQueryResult<Directory> {
   return useQuery({
     queryKey: ['/api/files'],
@@ -161,3 +172,7 @@ export function apiCancelGetEnvironments(client: QueryClient): void {
 export function apiCancelFiles(client: QueryClient): void {
   void client.cancelQueries({ queryKey: ['/api/files'] })
 }
+
+export function apiCancelModels(client: QueryClient): void {
+  void client.cancelQueries({ queryKey: ['/api/models'] })
+}

web/client/src/context/fileTree.ts

@@ -4,7 +4,7 @@ import { isArrayNotEmpty, isFalse } from '~/utils'
 import { ModelFile } from '../models'
 
 interface FileTreeStore {
-  files: ModelFile[]
+  files: Map<ID, ModelFile>
   activeFile: ModelFile
   openedFiles: Set<ModelFile>
   setOpenedFiles: (files: Set<ModelFile>) => void
@@ -20,24 +20,27 @@ const [getOpenedFilesIds, setOpenedFilesIds] = useLocalStorage<{ ids: ID[] }>(
 )
 
 export const useStoreFileTree = create<FileTreeStore>((set, get) => ({
-  files: [],
+  files: new Map(),
   activeFile: initialFile,
   openedFiles: new Set([initialFile]),
   setFiles(files: ModelFile[]) {
     set(() => {
       const openedFilesIds = getOpenedFilesIds()?.ids ?? []
       const openedFiles = new Set<ModelFile>([initialFile])
+      const output = new Map()
 
       if (isArrayNotEmpty(openedFilesIds)) {
         files.forEach(file => {
           if (openedFilesIds.includes(file.id)) {
             openedFiles.add(file)
           }
+
+          output.set(file.id, file)
         })
       }
 
       return {
-        files,
+        files: output,
         openedFiles,
       }
     })
@@ -70,4 +73,4 @@ export const useStoreFileTree = create<FileTreeStore>((set, get) => ({
       activeFile: file,
     }))
   },
-}))
+}))

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

@@ -0,0 +1,99 @@
+import { useCallback, useEffect, useMemo, useState } from 'react'
+import CodeMirror from '@uiw/react-codemirror'
+import { python } from '@codemirror/lang-python'
+import { StreamLanguage } from '@codemirror/language'
+
+import { yaml } from '@codemirror/legacy-modes/mode/yaml'
+import { type Extension } from '@codemirror/state'
+import { type Model } from '~/api/client'
+import { useStoreFileTree } from '~/context/fileTree'
+import { type ModelFile } from '~/models'
+import {
+  events,
+  SqlMeshModel,
+  HoverTooltip,
+  useSqlMeshExtension,
+} from './extensions'
+import { sqlglotWorker } from '~/library/components/editor/workers'
+import { dracula, tomorrow } from 'thememirror'
+import { useColorScheme, EnumColorScheme } from '~/context/theme'
+
+export default function CodeEditor({
+  file,
+  models,
+  dialect,
+  dialects,
+  onChange,
+}: {
+  file: ModelFile
+  models: Map<string, Model>
+  dialect?: string
+  dialects?: string[]
+  onChange: (value: string) => void
+}): JSX.Element {
+  const { mode } = useColorScheme()
+  const theme = mode === EnumColorScheme.Dark ? dracula : tomorrow
+
+  const files = useStoreFileTree(s => s.files)
+  const selectFile = useStoreFileTree(s => s.selectFile)
+
+  const [sqlDialectOptions, setSqlDialectOptions] = useState()
+
+  const [SqlMeshDialectExtension, SqlMeshDialectCleanUp] =
+    useSqlMeshExtension(dialects)
+
+  const extensions = useMemo(() => {
+    const showSqlSqlMeshDialect =
+      file.extension === '.sql' && models != null && sqlDialectOptions != null
+    return [
+      models != null && HoverTooltip(models),
+      models != null && events(models, files, selectFile),
+      models != null && SqlMeshModel(models),
+      showSqlSqlMeshDialect &&
+      SqlMeshDialectExtension(models, file, sqlDialectOptions),
+      file.extension === '.py' && python(),
+      file.extension === '.yaml' && StreamLanguage.define(yaml),
+      theme,
+    ].filter(Boolean) as Extension[]
+  }, [file, models, sqlDialectOptions])
+
+  const handleSqlGlotWorkerMessage = useCallback((e: MessageEvent): void => {
+    if (e.data.topic === 'dialect') {
+      setSqlDialectOptions(e.data.payload)
+    }
+  }, [])
+
+  useEffect(() => {
+    sqlglotWorker.addEventListener('message', handleSqlGlotWorkerMessage)
+
+    return () => {
+      sqlglotWorker.removeEventListener('message', handleSqlGlotWorkerMessage)
+      SqlMeshDialectCleanUp()
+    }
+  }, [])
+
+  useEffect(() => {
+    sqlglotWorker.postMessage({
+      topic: 'parse',
+      payload: file.content,
+    })
+  }, [file.content, sqlglotWorker])
+
+  useEffect(() => {
+    sqlglotWorker.postMessage({
+      topic: 'dialect',
+      payload: dialect,
+    })
+  }, [dialect, sqlglotWorker])
+
+  return (
+    <CodeMirror
+      value={file.content}
+      height="100%"
+      width="100%"
+      className="w-full h-full overflow-auto text-sm"
+      extensions={extensions}
+      onChange={onChange}
+    />
+  )
+}

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

@@ -26,7 +26,7 @@
 }
 
 .cm-editor .cm-line::selection,
-.cm-editor .cm-line > *::selection {
+.cm-editor .cm-line>*::selection {
   background-color: var(--color-brand-100) !important;
   color: var(--color-brand-700) !important;
 }
@@ -44,3 +44,22 @@
   background: var(--color-brand);
   border-radius: 1rem;
 }
+
+.sqlmesh-model {
+  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;
+}