Use config values for refs in Preview (#1179)

* add commands for context menu

* clean stack trace for telemetry

* add the ability to get ref values from config

client/src/ExtensionClient.ts

@@ -20,6 +20,8 @@ import { EventEmitter } from 'node:events';
 import * as path from 'node:path';
 import { AnalyzeEntireProject } from './commands/AnalyzeEntireProject';
 import { CreateDbtProject } from './commands/CreateDbtProject/CreateDbtProject';
+import { UseConfigForRefsPreview } from './commands/UseConfigForRefsPreview';
+import { NotUseConfigForRefsPreview } from './commands/NotUseConfigForRefsPreview';
 
 export interface PackageJson {
   name: string;
@@ -123,6 +125,8 @@ export class ExtensionClient {
     this.commandManager.register(new OpenOrCreatePackagesYml());
     this.commandManager.register(new Restart(this.dbtLanguageClientManager));
     this.commandManager.register(new InstallDbtPackages(this.dbtLanguageClientManager, this.outputChannelProvider));
+    this.commandManager.register(new UseConfigForRefsPreview(this.previewContentProvider));
+    this.commandManager.register(new NotUseConfigForRefsPreview(this.previewContentProvider));
   }
 
   registerSqlPreviewContentProvider(context: ExtensionContext): void {

client/src/SqlPreviewContentProvider.ts

@@ -8,13 +8,16 @@ import {
   Range,
   TextDocumentContentProvider,
   Uri,
+  commands,
   languages,
   window,
 } from 'vscode';
 import { SQL_LANG_ID } from './Utils';
+import { RefReplacement } from 'dbt-language-server-common';
 
 interface PreviewInfo {
   previewText: string;
+  refReplacements: RefReplacement[];
   diagnostics: Diagnostic[];
   langId: string;
 }
@@ -23,6 +26,22 @@ export default class SqlPreviewContentProvider implements TextDocumentContentPro
   static readonly SCHEME = 'query-preview';
   static readonly URI = Uri.parse(`${SqlPreviewContentProvider.SCHEME}:Preview?dbt-language-server`);
 
+  private useConfigForRefs = false;
+
+  changeMode(useConfigForRefs: boolean): void {
+    this.setUseConfigForRefs(useConfigForRefs);
+    this.onDidChangeEmitter.fire(SqlPreviewContentProvider.URI);
+  }
+
+  setUseConfigForRefs(useConfigForRefs: boolean): void {
+    if (this.useConfigForRefs !== useConfigForRefs) {
+      commands
+        .executeCommand('setContext', 'WizardForDbtCore:useConfigForRefs', useConfigForRefs)
+        .then(undefined, e => console.log(e instanceof Error ? e.message : String(e)));
+      this.useConfigForRefs = useConfigForRefs;
+    }
+  }
+
   previewInfos = new Map<string, PreviewInfo>();
 
   activeDocUri: Uri = Uri.parse('');
@@ -33,11 +52,12 @@ export default class SqlPreviewContentProvider implements TextDocumentContentPro
     return uri.path === SqlPreviewContentProvider.URI.path;
   }
 
-  updateText(uri: string, previewText: string, langId: string): void {
+  updateText(uri: string, previewText: string, refReplacements: RefReplacement[], langId: string): void {
     const currentValue = this.previewInfos.get(uri);
 
     this.previewInfos.set(uri, {
       previewText,
+      refReplacements,
       diagnostics: currentValue?.diagnostics ?? [],
       langId,
     });
@@ -62,6 +82,7 @@ export default class SqlPreviewContentProvider implements TextDocumentContentPro
     const currentValue = this.previewInfos.get(uri);
     this.previewInfos.set(uri, {
       previewText: currentValue?.previewText ?? '',
+      refReplacements: currentValue?.refReplacements ?? [],
       diagnostics: diagnostics.map<Diagnostic>(d => {
         const diag = new Diagnostic(d.range, d.message, DiagnosticSeverity.Error);
         if (d.relatedInformation && d.relatedInformation.length > 0) {
@@ -93,6 +114,7 @@ export default class SqlPreviewContentProvider implements TextDocumentContentPro
   changeActiveDocument(uri: Uri): void {
     if (uri.toString() !== this.activeDocUri.toString()) {
       this.activeDocUri = uri;
+      this.setUseConfigForRefs(false);
       this.onDidChangeEmitter.fire(SqlPreviewContentProvider.URI);
     }
   }
@@ -108,6 +130,13 @@ export default class SqlPreviewContentProvider implements TextDocumentContentPro
   provideTextDocumentContent(): string {
     const previewInfo = this.previewInfos.get(this.activeDocUri.toString());
     this.updateLangId(previewInfo?.langId ?? SQL_LANG_ID).catch(e => console.log(e));
-    return previewInfo?.previewText ?? '';
+
+    let text = previewInfo?.previewText ?? '';
+    if (this.useConfigForRefs) {
+      for (const replacement of previewInfo?.refReplacements ?? []) {
+        text = text.replaceAll(replacement.from, replacement.to);
+      }
+    }
+    return text;
   }
 }

client/src/commands/NotUseConfigForRefsPreview.ts

@@ -0,0 +1,12 @@
+import SqlPreviewContentProvider from '../SqlPreviewContentProvider';
+import { Command } from './CommandManager';
+
+export class NotUseConfigForRefsPreview implements Command {
+  readonly id = 'WizardForDbtCore(TM).notUseConfigForRefsPreview';
+
+  constructor(private sqlPreviewContentProvider: SqlPreviewContentProvider) {}
+
+  execute(): void {
+    this.sqlPreviewContentProvider.changeMode(false);
+  }
+}

client/src/commands/UseConfigForRefsPreview.ts

@@ -0,0 +1,12 @@
+import SqlPreviewContentProvider from '../SqlPreviewContentProvider';
+import { Command } from './CommandManager';
+
+export class UseConfigForRefsPreview implements Command {
+  readonly id = 'WizardForDbtCore(TM).useConfigForRefsPreview';
+
+  constructor(private sqlPreviewContentProvider: SqlPreviewContentProvider) {}
+
+  execute(): void {
+    this.sqlPreviewContentProvider.changeMode(true);
+  }
+}

client/src/lsp_client/DbtLanguageClient.ts

@@ -1,4 +1,4 @@
-import { LspModeType, TelemetryEvent } from 'dbt-language-server-common';
+import { LspModeType, RefReplacement, TelemetryEvent } from 'dbt-language-server-common';
 import { EventEmitter } from 'node:events';
 import {
   ConfigurationChangeEvent,
@@ -130,10 +130,11 @@ export class DbtLanguageClient extends DbtWizardLanguageClient {
     super.initializeNotifications();
 
     this.disposables.push(
-      this.client.onNotification('custom/updateQueryPreview', ({ uri, previewText }) => {
+      this.client.onNotification('custom/updateQueryPreview', ({ uri, previewText, refReplacements }) => {
         this.previewContentProvider.updateText(
           uri as string,
           previewText as string,
+          refReplacements as RefReplacement[],
           this.canUseSnowflakeLangId() ? SNOWFLAKE_SQL_LANG_ID : SQL_LANG_ID,
         );
       }),

common/src/Interfaces.ts

@@ -17,3 +17,5 @@ export interface CustomInitParams {
   disableLogger?: boolean;
   profilesDir?: string;
 }
+
+export type RefReplacement = { from: string; to: string };

package.json

@@ -70,6 +70,16 @@
         "category": "WizardForDbtCore(TM)",
         "icon": "$(book)"
       },
+      {
+        "command": "WizardForDbtCore(TM).notUseConfigForRefsPreview",
+        "title": "✓ Use schema from config for refs",
+        "category": "WizardForDbtCore(TM)"
+      },
+      {
+        "command": "WizardForDbtCore(TM).useConfigForRefsPreview",
+        "title": "  Use schema from config for refs",
+        "category": "WizardForDbtCore(TM)"
+      },
       {
         "command": "WizardForDbtCore(TM).compile",
         "title": "dbt compile",
@@ -138,6 +148,16 @@
           "submenu": "dbt",
           "group": "navigation",
           "when": "editorLangId =~ /^sql$|^jinja-sql$|^snowflake-sql$|^sql-bigquery$/ && !editorReadonly"
+        },
+        {
+          "command": "WizardForDbtCore(TM).useConfigForRefsPreview",
+          "group": "1_common",
+          "when": "resource == query-preview:Preview?dbt-language-server && !WizardForDbtCore:useConfigForRefs"
+        },
+        {
+          "command": "WizardForDbtCore(TM).notUseConfigForRefsPreview",
+          "group": "1_common",
+          "when": "resource == query-preview:Preview?dbt-language-server && WizardForDbtCore:useConfigForRefs"
         }
       ]
     },

server/src/NotificationSender.ts

@@ -1,4 +1,4 @@
-import { StatusNotification, TelemetryEvent } from 'dbt-language-server-common';
+import { RefReplacement, StatusNotification, TelemetryEvent } from 'dbt-language-server-common';
 import { Diagnostic, PublishDiagnosticsParams, TelemetryEventNotification, _Connection } from 'vscode-languageserver';
 
 export class NotificationSender {
@@ -32,8 +32,8 @@ export class NotificationSender {
       .catch(e => console.log(`Failed to send notification: ${e instanceof Error ? e.message : String(e)}`));
   }
 
-  sendUpdateQueryPreview(uri: string, previewText: string): void {
-    this.sendNotification('custom/updateQueryPreview', { uri, previewText });
+  sendUpdateQueryPreview(uri: string, previewText: string, refReplacements: RefReplacement[]): void {
+    this.sendNotification('custom/updateQueryPreview', { uri, previewText, refReplacements });
   }
 
   sendStatus(statusNotification: StatusNotification): void {

server/src/dag/Dag.ts

@@ -14,4 +14,12 @@ export class Dag {
   getNodesByPackage(packageName: string): DagNode[] {
     return this.nodes.filter(n => n.getValue().packageName === packageName);
   }
+
+  getNodeByUri(uri: string): DagNode | undefined {
+    return this.nodes.find(n => uri.endsWith(n.getValue().originalFilePath));
+  }
+
+  getNodeByName(name: string): DagNode | undefined {
+    return this.nodes.find(n => n.getValue().name === name);
+  }
 }

server/src/document/DbtTextDocument.ts

@@ -39,6 +39,7 @@ import { DefinitionProvider } from '../definition/DefinitionProvider';
 import { getLineByPosition, getSignatureInfo } from '../utils/TextUtils';
 import { areRangesEqual, debounce, getIdentifierRangeAtPosition, getModelPathOrFullyQualifiedName } from '../utils/Utils';
 import { DbtDocumentKind } from './DbtDocumentKind';
+import { RefReplacement } from 'dbt-language-server-common';
 
 export interface QueryParseInformation {
   selects: {
@@ -69,6 +70,7 @@ export class DbtTextDocument {
 
   rawDocDiagnostics: Diagnostic[] = [];
   compiledDocDiagnostics: Diagnostic[] = [];
+  refReplacements: RefReplacement[] = [];
 
   constructor(
     doc: TextDocumentItem,
@@ -100,6 +102,8 @@ export class DbtTextDocument {
     this.modelCompiler.onCompilationError(this.onCompilationError.bind(this));
     this.modelCompiler.onCompilationFinished(async (compiledSql: string) => {
       projectChangeListener.updateManifest();
+
+      this.updateRefReplacements();
       await this.onCompilationFinished(compiledSql);
     });
     this.modelCompiler.onFinishAllCompilationJobs(this.onFinishAllCompilationTasks.bind(this));
@@ -108,6 +112,29 @@ export class DbtTextDocument {
     this.dbtCli.onDbtReady(this.onDbtReady.bind(this));
   }
 
+  updateRefReplacements(): void {
+    const refReplacements = [];
+    const currentNode = this.dbtRepository.dag.getNodeByUri(this.rawDocument.uri);
+    if (currentNode) {
+      const { refs } = currentNode.getValue();
+      for (const ref of refs) {
+        if ('name' in ref) {
+          const refNode = this.dbtRepository.dag.getNodeByName(ref.name);
+          if (refNode) {
+            const { config } = refNode.getValue();
+            if (config?.schema) {
+              refReplacements.push({
+                from: `\`${refNode.getValue().schema}\`.\`${refNode.getValue().name}\``,
+                to: `\`${config.schema}\`.\`${refNode.getValue().name}\``,
+              });
+            }
+          }
+        }
+      }
+    }
+    this.refReplacements = refReplacements;
+  }
+
   willSaveTextDocument(reason: TextDocumentSaveReason): void {
     // Document can be modified and not saved before language server initialized, in this case we need to compile it on first save command call (see unit test).
     if (
@@ -274,7 +301,7 @@ export class DbtTextDocument {
 
   async updateAndSendDiagnosticsAndPreview(dbtCompilationError?: string): Promise<void> {
     await this.updateDiagnostics(dbtCompilationError);
-    this.notificationSender.sendUpdateQueryPreview(this.rawDocument.uri, this.compiledDocument.getText());
+    this.notificationSender.sendUpdateQueryPreview(this.rawDocument.uri, this.compiledDocument.getText(), this.refReplacements);
   }
 
   async updateDiagnostics(dbtCompilationError?: string): Promise<void> {

server/src/lsp_server/LspServerBase.ts

@@ -3,6 +3,7 @@ import { InitializeError, InitializeParams, InitializeResult, ResponseError, _Co
 import { InstallUtils } from '../InstallUtils';
 import { NotificationSender } from '../NotificationSender';
 import { FeatureFinderBase } from '../feature_finder/FeatureFinderBase';
+import path from 'node:path';
 
 export abstract class LspServerBase<T extends FeatureFinderBase> {
   constructor(
@@ -14,17 +15,32 @@ export abstract class LspServerBase<T extends FeatureFinderBase> {
   abstract onInitialize(params: InitializeParams): InitializeResult<unknown> | ResponseError<InitializeError>;
 
   onUncaughtException(error: Error, _origin: 'uncaughtException' | 'unhandledRejection'): void {
-    console.log(error.stack);
+    const stack = LspServerBase.getCleanStackTrace(error.stack);
+    console.log(stack);
 
     this.notificationSender.sendTelemetry('error', {
       name: error.name,
       message: error.message,
-      stack: error.stack ?? '',
+      stack,
     });
 
     throw new Error('Uncaught exception. Server will be restarted.');
   }
 
+  private static getCleanStackTrace(stack: string | undefined): string {
+    if (!stack) {
+      return '';
+    }
+
+    let lines = stack.split('\n');
+    lines = lines.map(line => {
+      const match = line.match(/\((.*?dbt-language-server)/);
+      return (match ? line.replace(match[1], '') : line).replaceAll(path.sep, '__');
+    });
+
+    return lines.join('\n');
+  }
+
   initializeNotifications(): void {
     this.connection.onNotification('WizardForDbtCore(TM)/installDbtCore', (version: string) => this.installDbtCore(version));
     this.connection.onNotification('WizardForDbtCore(TM)/installDbtAdapter', (dbtAdapter: DbtAdapter) =>

server/src/manifest/ManifestJson.ts

@@ -20,6 +20,7 @@ export interface ManifestModel extends ManifestNode {
   config?: {
     sqlHeader?: string;
     materialized?: string;
+    schema?: string;
   };
 }
 

server/src/manifest/ManifestParser.ts

@@ -25,6 +25,7 @@ interface RawNode {
   config?: {
     sql_header?: string;
     materialized?: string;
+    schema?: string;
   };
 }
 
@@ -92,6 +93,7 @@ export class ManifestParser {
           config: {
             sqlHeader: n.config?.sql_header,
             materialized: n.config?.materialized,
+            schema: n.config?.schema,
           },
         }));
     }