Fix: Tab not showing after refresh + store unsaved changes (#1164)

Author: mykalmax

web/client/src/context/editor.ts

@@ -1,6 +1,6 @@
 import { type LineageColumn } from '@api/client'
 import { type TableColumn, type TableRow } from '@components/table/help'
-import { uid } from '@utils/index'
+import { isFalse, isNil, isNotNil, uid } from '@utils/index'
 import { create } from 'zustand'
 import useLocalStorage from '~/hooks/useLocalStorage'
 import { type ErrorKey, type ErrorIDE } from '~/library/pages/ide/context'
@@ -17,9 +17,14 @@ export interface Lineage {
   columns?: Record<string, LineageColumn>
 }
 
+export interface StoredTab {
+  id?: ID
+  content?: string
+}
+
 interface EditorStore {
-  storedTabsId?: ID
-  storedTabsIds: ID[]
+  storedTabId?: ID
+  storedTabs: StoredTab[]
   tabs: Map<ModelFile, EditorTab>
   tab?: EditorTab
   engine: Worker
@@ -34,6 +39,7 @@ interface EditorStore {
   replaceTab: (from: EditorTab, to: EditorTab) => void
   updateStoredTabsIds: () => void
   addTab: (tab: EditorTab) => void
+  addTabs: (tabs: EditorTab[]) => void
   closeTab: (file: ModelFile) => void
   createTab: (file?: ModelFile) => EditorTab
   setDialects: (dialects: Dialect[]) => void
@@ -65,17 +71,23 @@ export interface EditorTab {
   el?: HTMLElement
 }
 
-const [getStoredTabs, setStoredTabs] = useLocalStorage<{ ids: ID[]; id: ID }>(
-  'tabs',
-)
+const [getStoredTabs, setStoredTabs] = useLocalStorage<{
+  tabs: StoredTab[]
+  id?: ID
+}>('tabs')
 
+const { tabs: storedTabs = [], id: storedTabId } = getStoredTabs() ?? {}
 const initialFile = createLocalFile()
 const initialTab: EditorTab = createTab(initialFile)
-const initialTabs = new Map([[initialFile, initialTab]])
+const initialTabs = new Map(
+  storedTabs.length > 0 && isNotNil(storedTabId)
+    ? []
+    : [[initialFile, initialTab]],
+)
 
 export const useStoreEditor = create<EditorStore>((set, get) => ({
-  storedTabsIds: getStoredTabs()?.ids ?? [],
-  storedTabsId: getStoredTabs()?.id,
+  storedTabs,
+  storedTabId,
   tab: initialTab,
   tabs: initialTabs,
   engine: sqlglotWorker,
@@ -101,19 +113,43 @@ export const useStoreEditor = create<EditorStore>((set, get) => ({
   },
   updateStoredTabsIds() {
     const s = get()
-    const id = s.tab?.file.id
-    const ids = Array.from(get().tabs.values())
-      .filter(tab => tab.file.isRemote)
-      .map(tab => tab.file.id)
+
+    if (isNil(s.tab)) {
+      setStoredTabs({
+        id: undefined,
+        tabs: [],
+      })
+
+      set(() => ({
+        storedTabId: undefined,
+        storedTabs: [],
+      }))
+
+      return
+    }
+
+    const tabs: StoredTab[] = []
+
+    for (const tab of s.tabs.values()) {
+      if (isFalse(tab.file.isChanged) && tab.file.isLocal) continue
+
+      tabs.push({
+        id: tab.file.id,
+        content: tab.file.isChanged ? tab.file.content : undefined,
+      })
+    }
+
+    const id =
+      s.tab.file.isChanged || s.tab.file.isRemote ? s.tab.file.id : undefined
 
     setStoredTabs({
       id,
-      ids,
+      tabs,
     })
 
     set(() => ({
-      storedTabsId: id,
-      storedTabsIds: ids,
+      storedTabId: id,
+      storedTabs: tabs,
     }))
   },
   refreshTab() {
@@ -129,7 +165,22 @@ export const useStoreEditor = create<EditorStore>((set, get) => ({
     }))
   },
   selectTab(tab) {
+    const s = get()
+
     set(() => ({ tab }))
+
+    s.updateStoredTabsIds()
+  },
+  addTabs(tabs) {
+    const s = get()
+
+    for (const tab of tabs) {
+      s.tabs.set(tab.file, tab)
+    }
+
+    set(() => ({
+      tabs: new Map(s.tabs),
+    }))
   },
   addTab(tab) {
     const s = get()
@@ -204,8 +255,9 @@ function createTab(file: ModelFile = createLocalFile()): EditorTab {
   }
 }
 
-function createLocalFile(): ModelFile {
+export function createLocalFile(id?: ID): ModelFile {
   return new ModelFile({
+    id,
     name: '',
     path: '',
     content:

web/client/src/context/project.ts

@@ -1,10 +1,9 @@
 import { create } from 'zustand'
 import { ModelDirectory, type ModelFile } from '../models'
-import { type Directory } from '~/api/client'
 
 interface ProjectStore {
   project: ModelDirectory
-  setProject: (project?: Directory) => void
+  setProject: (project?: ModelDirectory) => void
   files: Map<ID, ModelFile>
   setFiles: (files: ModelFile[]) => void
   selectedFile?: ModelFile
@@ -17,7 +16,7 @@ export const useStoreProject = create<ProjectStore>((set, get) => ({
   files: new Map(),
   setProject(project) {
     set(() => ({
-      project: new ModelDirectory(project),
+      project,
     }))
   },
   setFiles(files) {

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

@@ -3,7 +3,7 @@ import CodeMirror from '@uiw/react-codemirror'
 import { type KeyBinding, keymap } from '@codemirror/view'
 import { type Extension } from '@codemirror/state'
 import { useApiFileByPath } from '~/api'
-import { debounceAsync, isNil, isStringNotEmpty } from '~/utils'
+import { debounceAsync, isNil } from '~/utils'
 import { isCancelledError } from '@tanstack/react-query'
 import { useStoreContext } from '~/context/context'
 import { useStoreEditor } from '~/context/editor'
@@ -133,7 +133,7 @@ function CodeEditorRemoteFile({
     const tempFile = files.get(path)
 
     if (tempFile == null) return
-    if (isStringNotEmpty(tempFile.content)) {
+    if (tempFile.isSynced) {
       setFile(tempFile)
       return
     }

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

@@ -72,8 +72,7 @@ export const SQLMeshDialect: ExtensionSQLMeshDialect = function SQLMeshDialect(
         const keywordFrom = ctx.matchBefore(/from.+/i)
         const keywordKind = ctx.matchBefore(/kind.+/i)
         const keywordDialect = ctx.matchBefore(/dialect.+/i)
-        const matchModels =
-          text.match(/MODEL \((?:[^)(]+|\((?:[^)(]+|\([^)(]*\))*\))*\)/g) ?? []
+        const matchModels = text.match(/MODEL \(([\s\S]+?)\);/g) ?? []
         const isInsideModel = matchModels
           .filter(str => str.includes(match))
           .map<[number, number]>(str => [

web/client/src/library/pages/ide/IDE.tsx

@@ -17,7 +17,15 @@ import {
   type PlanTasks,
 } from '../../../context/plan'
 import { useChannelEvents } from '../../../api/channels'
-import { debounceAsync, isFalse, isObject, isObjectEmpty } from '~/utils'
+import {
+  debounceAsync,
+  isArrayEmpty,
+  isFalse,
+  isNil,
+  isNotNil,
+  isObject,
+  isObjectEmpty,
+} from '~/utils'
 import { useStoreContext } from '~/context/context'
 import { ArrowLongRightIcon } from '@heroicons/react/24/solid'
 import { EnumSize, EnumVariant } from '~/types/enum'
@@ -29,11 +37,12 @@ import { type Model } from '@api/client'
 import { Button } from '@components/button/Button'
 import { Divider } from '@components/divider/Divider'
 import Container from '@components/container/Container'
-import { useStoreEditor } from '@context/editor'
+import { useStoreEditor, createLocalFile } from '@context/editor'
 import { type ModelFile } from '@models/file'
 import ModalConfirmation, {
   type Confirmation,
 } from '@components/modal/ModalConfirmation'
+import { ModelDirectory } from '@models/directory'
 
 const ReportErrors = lazy(
   async () => await import('../../components/report/ReportErrors'),
@@ -73,11 +82,11 @@ export default function PageIDE(): JSX.Element {
   const setProject = useStoreProject(s => s.setProject)
   const setFiles = useStoreProject(s => s.setFiles)
 
-  const storedTabsIds = useStoreEditor(s => s.storedTabsIds)
-  const storedTabsId = useStoreEditor(s => s.storedTabsId)
+  const storedTabs = useStoreEditor(s => s.storedTabs)
+  const storedTabId = useStoreEditor(s => s.storedTabId)
   const selectTab = useStoreEditor(s => s.selectTab)
   const createTab = useStoreEditor(s => s.createTab)
-  const addTab = useStoreEditor(s => s.addTab)
+  const addTabs = useStoreEditor(s => s.addTabs)
 
   const subscribe = useChannelEvents()
 
@@ -125,7 +134,14 @@ export default function PageIDE(): JSX.Element {
     })
 
     void debouncedGetFiles().then(({ data }) => {
-      setProject(data)
+      if (isNil(data)) return
+
+      const project = new ModelDirectory(data)
+      const files = project.allFiles
+
+      setFiles(files)
+      restoreEditorTabsFromSaved(files)
+      setProject(project)
     })
 
     return () => {
@@ -152,13 +168,6 @@ export default function PageIDE(): JSX.Element {
     }
   }, [location])
 
-  useEffect(() => {
-    const files = project?.allFiles ?? []
-
-    setFiles(files)
-    restoreEditorTabsFromSaved(files)
-  }, [project])
-
   useEffect(() => {
     if (dataEnvironments == null || isObjectEmpty(dataEnvironments)) return
 
@@ -228,17 +237,23 @@ export default function PageIDE(): JSX.Element {
   }
 
   function restoreEditorTabsFromSaved(files: ModelFile[]): void {
-    files.forEach(file => {
-      if (storedTabsIds.includes(file.id)) {
-        const tab = createTab(file)
-
-        if (storedTabsId === file.id) {
-          selectTab(tab)
-        } else {
-          addTab(tab)
-        }
-      }
+    if (isArrayEmpty(storedTabs)) return
+
+    const tabs = storedTabs.map(({ id, content }) => {
+      const file = files.find(file => file.id === id) ?? createLocalFile(id)
+      const storedTab = createTab(file)
+
+      storedTab.file.content = content ?? storedTab.file.content ?? ''
+
+      return storedTab
     })
+    const tab = tabs.find(tab => tab.file.id === storedTabId)
+
+    addTabs(tabs)
+
+    if (isNotNil(tab)) {
+      selectTab(tab)
+    }
   }
 
   function closeModalConfirmation(confirmation?: Confirmation): void {

web/client/src/models/file.ts

@@ -1,7 +1,7 @@
 import { type File, FileType } from '../api/client'
 import { type ModelDirectory } from './directory'
 import { type InitialArtifact, ModelArtifact } from './artifact'
-import { isStringEmptyOrNil, toUniqueName } from '@utils/index'
+import { isFalse, isStringEmptyOrNil, toUniqueName } from '@utils/index'
 
 export const EnumFileExtensions = {
   SQL: '.sql',
@@ -57,8 +57,12 @@ export class ModelFile extends ModelArtifact<InitialFile> {
     this._type = newType ?? getFileType(this.path)
   }
 
+  get isSynced(): boolean {
+    return isFalse(isStringEmptyOrNil(this._content))
+  }
+
   get isEmpty(): boolean {
-    return this.content === ''
+    return isStringEmptyOrNil(this.content)
   }
 
   get isSupported(): boolean {
@@ -92,8 +96,16 @@ export class ModelFile extends ModelArtifact<InitialFile> {
   }
 
   updateContent(newContent: string = ''): void {
+    // When modifying a file locally, we only modify the content.
+    // Therefore, if we have content but the variable "_content" is empty,
+    // it is likely because we restored the file content from localStorage.
+    // After updating "_content", we still want to retain the content
+    // because it is unsaved changes.
+    if (this.isSynced) {
+      this.content = newContent
+    }
+
     this._content = newContent
-    this.content = newContent
   }
 
   update(newFile?: File): void {

web/client/src/models/initial.ts

@@ -1,4 +1,4 @@
-import { uid } from '@utils/index'
+import { isNotNil, uid } from '@utils/index'
 
 type Initial<T extends object> = T & { id?: ID }
 type InitialWithId<T extends object> = T & { id: ID }
@@ -9,19 +9,18 @@ export class ModelInitial<T extends object = any> {
   isModel = true
 
   constructor(initial: Initial<T> | InitialWithId<T>) {
-    this._initial =
-      'id' in initial
-        ? (initial as InitialWithId<T>)
-        : new Proxy<InitialWithId<T>>(
-            Object.assign(initial ?? {}, {
-              id: uid(),
-            }),
-            {
-              set() {
-                throw new Error('Cannot change initial file')
-              },
+    this._initial = isNotNil(initial?.id)
+      ? (initial as InitialWithId<T>)
+      : new Proxy<InitialWithId<T>>(
+          Object.assign(initial ?? {}, {
+            id: uid(),
+          }),
+          {
+            set() {
+              throw new Error('Cannot change initial file')
             },
-          )
+          },
+        )
   }
 
   get initial(): InitialWithId<T> {