feat(vscode): rerender on save model

Author: benfdking

vscode/extension/src/commands/renderModel.ts

@@ -4,6 +4,39 @@ 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
+  }
+
+  // For auto-rerender, we'll use the first model found
+  // In the future, we could store which specific model was originally rendered
+  const selectedModel = 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 +147,7 @@ export function renderModel(
     const uri = renderedModelProvider.storeRenderedModel(
       selectedModel.name,
       selectedModel.rendered_query,
+      documentUri,
     )
 
     // 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

@@ -9,6 +9,10 @@ export class RenderedModelProvider
   private static readonly scheme = 'sqlmesh-rendered'
 
   private renderedModels = new Map<string, string>()
+  // Track which source file URIs are associated with rendered models
+  private sourceToRenderedUri = new Map<string, vscode.Uri>()
+  // Track which rendered URIs are associated with source files
+  private renderedToSourceUri = new Map<string, string>()
 
   // Event emitter for content changes
   private _onDidChange = new vscode.EventEmitter<vscode.Uri>()
@@ -25,7 +29,7 @@ export class RenderedModelProvider
   /**
    * Store rendered model content and create a URI for it
    */
-  storeRenderedModel(modelName: string, content: string): vscode.Uri {
+  storeRenderedModel(modelName: string, content: string, sourceUri?: string): vscode.Uri {
     const fileName = `${modelName} (rendered)`
     // Add a timestamp to make the URI unique for each render
     const timestamp = Date.now()
@@ -35,11 +39,56 @@ export class RenderedModelProvider
       path: fileName,
       fragment: timestamp.toString(),
     })
-    this.renderedModels.set(uri.toString(), content)
+    
+    const uriString = uri.toString()
+    this.renderedModels.set(uriString, content)
+    
+    // Track the association between source file and rendered model
+    if (sourceUri) {
+      // Remove any existing mapping for this source file
+      const existingRenderedUri = this.sourceToRenderedUri.get(sourceUri)
+      if (existingRenderedUri) {
+        this.renderedToSourceUri.delete(existingRenderedUri.toString())
+      }
+      
+      this.sourceToRenderedUri.set(sourceUri, uri)
+      this.renderedToSourceUri.set(uriString, sourceUri)
+    }
+    
     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()
+    this.renderedModels.set(uriString, 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 {
+    return this.renderedToSourceUri.get(renderedUri)
+  }
+
+  /**
+   * 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 +101,8 @@ export class RenderedModelProvider
    */
   dispose() {
     this.renderedModels.clear()
+    this.sourceToRenderedUri.clear()
+    this.renderedToSourceUri.clear()
     this._onDidChange.dispose()
   }
 }