feat(vscode): rerender on save model (#4637)

Author: benfdking

vscode/extension/src/commands/renderModel.ts

@@ -4,6 +4,51 @@ import { isErr } from '@bus/result'
 import { RenderModelEntry } from '../lsp/custom'
 import { RenderedModelProvider } from '../providers/renderedModelProvider'
 
+export async function reRenderModelForSourceFile(
+  sourceUri: string,
+  lspClient: LSPClient,
+  renderedModelProvider: RenderedModelProvider,
+): Promise<void> {
+  const renderedUri = renderedModelProvider.getRenderedUriForSource(sourceUri)
+  if (!renderedUri) {
+    return // No rendered model exists for this source file
+  }
+
+  // Call the render model API
+  const result = await lspClient.call_custom_method('sqlmesh/render_model', {
+    textDocumentUri: sourceUri,
+  })
+
+  if (isErr(result)) {
+    // Silently fail on auto-rerender errors to avoid spamming user
+    return
+  }
+
+  // Check if we got any models
+  if (!result.value.models || result.value.models.length === 0) {
+    return
+  }
+
+  // Get the originally rendered model information
+  const originalModelInfo =
+    renderedModelProvider.getModelInfoForRendered(renderedUri)
+
+  // Find the specific model that was originally rendered, or fall back to the first model
+  const selectedModel = originalModelInfo
+    ? result.value.models.find(
+        model =>
+          model.name === originalModelInfo.name &&
+          model.fqn === originalModelInfo.fqn,
+      ) || result.value.models[0]
+    : result.value.models[0]
+
+  // Update the existing rendered model content
+  renderedModelProvider.updateRenderedModel(
+    renderedUri,
+    selectedModel.rendered_query,
+  )
+}
+
 export function renderModel(
   lspClient?: LSPClient,
   renderedModelProvider?: RenderedModelProvider,
@@ -114,6 +159,8 @@ export function renderModel(
     const uri = renderedModelProvider.storeRenderedModel(
       selectedModel.name,
       selectedModel.rendered_query,
+      documentUri,
+      selectedModel,
     )
 
     // Open the virtual document

vscode/extension/src/extension.ts

@@ -12,7 +12,7 @@ import { AuthenticationProviderTobikoCloud } from './auth/auth'
 import { signOut } from './commands/signout'
 import { signIn } from './commands/signin'
 import { signInSpecifyFlow } from './commands/signinSpecifyFlow'
-import { renderModel } from './commands/renderModel'
+import { renderModel, reRenderModelForSourceFile } from './commands/renderModel'
 import { isErr } from '@bus/result'
 import {
   handleNotSginedInError,
@@ -23,6 +23,7 @@ import {
 import { selector, completionProvider } from './completion/completion'
 import { LineagePanel } from './webviews/lineagePanel'
 import { RenderedModelProvider } from './providers/renderedModelProvider'
+import { sleep } from './utilities/sleep'
 
 let lspClient: LSPClient | undefined
 
@@ -99,6 +100,25 @@ export async function activate(context: vscode.ExtensionContext) {
     ),
   )
 
+  // Add file save listener for auto-rerendering models
+  context.subscriptions.push(
+    vscode.workspace.onDidSaveTextDocument(async document => {
+      if (
+        lspClient &&
+        renderedModelProvider.hasRenderedModelForSource(
+          document.uri.toString(true),
+        )
+      ) {
+        await sleep(100)
+        await reRenderModelForSourceFile(
+          document.uri.toString(true),
+          lspClient,
+          renderedModelProvider,
+        )
+      }
+    }),
+  )
+
   const restart = async () => {
     if (lspClient) {
       traceVerbose('Restarting LSP client')

vscode/extension/src/providers/renderedModelProvider.ts

@@ -1,4 +1,11 @@
 import * as vscode from 'vscode'
+import { RenderModelEntry } from '../lsp/custom'
+
+interface RenderedModelInfo {
+  content: string
+  sourceUri?: string
+  modelInfo?: RenderModelEntry
+}
 
 /**
  * Content provider for read-only rendered SQL models
@@ -8,7 +15,10 @@ export class RenderedModelProvider
 {
   private static readonly scheme = 'sqlmesh-rendered'
 
-  private renderedModels = new Map<string, string>()
+  // Single map containing all rendered model information
+  private renderedModels = new Map<string, RenderedModelInfo>()
+  // Track which source file URIs are associated with rendered models
+  private sourceToRenderedUri = new Map<string, vscode.Uri>()
 
   // Event emitter for content changes
   private _onDidChange = new vscode.EventEmitter<vscode.Uri>()
@@ -19,13 +29,19 @@ export class RenderedModelProvider
    */
   provideTextDocumentContent(uri: vscode.Uri): string {
     const key = uri.toString()
-    return this.renderedModels.get(key) || ''
+    const modelInfo = this.renderedModels.get(key)
+    return modelInfo?.content || ''
   }
 
   /**
    * Store rendered model content and create a URI for it
    */
-  storeRenderedModel(modelName: string, content: string): vscode.Uri {
+  storeRenderedModel(
+    modelName: string,
+    content: string,
+    sourceUri?: string,
+    modelInfo?: RenderModelEntry,
+  ): vscode.Uri {
     const fileName = `${modelName} (rendered)`
     // Add a timestamp to make the URI unique for each render
     const timestamp = Date.now()
@@ -35,11 +51,78 @@ export class RenderedModelProvider
       path: fileName,
       fragment: timestamp.toString(),
     })
-    this.renderedModels.set(uri.toString(), content)
+
+    const uriString = uri.toString()
+
+    // Store all information in single map
+    this.renderedModels.set(uriString, {
+      content,
+      sourceUri,
+      modelInfo,
+    })
+
+    // Track the association between a source file and the rendered model
+    if (sourceUri) {
+      // Remove any existing mapping for this source file
+      const existingRenderedUri = this.sourceToRenderedUri.get(sourceUri)
+      if (existingRenderedUri) {
+        this.renderedModels.delete(existingRenderedUri.toString())
+      }
+
+      this.sourceToRenderedUri.set(sourceUri, uri)
+    }
+
     this._onDidChange.fire(uri)
     return uri
   }
 
+  /**
+   * Update an existing rendered model with new content
+   */
+  updateRenderedModel(uri: vscode.Uri, content: string): void {
+    const uriString = uri.toString()
+    const existingInfo = this.renderedModels.get(uriString)
+    if (existingInfo) {
+      this.renderedModels.set(uriString, {
+        ...existingInfo,
+        content,
+      })
+    }
+    this._onDidChange.fire(uri)
+  }
+
+  /**
+   * Get the rendered URI for a given source file URI
+   */
+  getRenderedUriForSource(sourceUri: string): vscode.Uri | undefined {
+    return this.sourceToRenderedUri.get(sourceUri)
+  }
+
+  /**
+   * Get the source URI for a given rendered model URI
+   */
+  getSourceUriForRendered(renderedUri: string): string | undefined {
+    const modelInfo = this.renderedModels.get(renderedUri)
+    return modelInfo?.sourceUri
+  }
+
+  /**
+   * Get the model information for a given rendered model URI
+   */
+  getModelInfoForRendered(
+    renderedUri: vscode.Uri,
+  ): RenderModelEntry | undefined {
+    const modelInfo = this.renderedModels.get(renderedUri.toString())
+    return modelInfo?.modelInfo
+  }
+
+  /**
+   * Check if a source file has an associated rendered model
+   */
+  hasRenderedModelForSource(sourceUri: string): boolean {
+    return this.sourceToRenderedUri.has(sourceUri)
+  }
+
   /**
    * Get the URI scheme for rendered models
    */
@@ -52,6 +135,7 @@ export class RenderedModelProvider
    */
   dispose() {
     this.renderedModels.clear()
+    this.sourceToRenderedUri.clear()
     this._onDidChange.dispose()
   }
 }