From 84e062d5ca2fff6d40e4de28d82dff028720c9d0 Mon Sep 17 00:00:00 2001 From: Pavel Grivachev Date: Mon, 4 Sep 2023 13:07:55 +0200 Subject: [PATCH] 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 | 4 +++ client/src/SqlPreviewContentProvider.ts | 33 +++++++++++++++++-- .../commands/NotUseConfigForRefsPreview.ts | 12 +++++++ .../src/commands/UseConfigForRefsPreview.ts | 12 +++++++ client/src/lsp_client/DbtLanguageClient.ts | 5 +-- common/src/Interfaces.ts | 2 ++ package.json | 20 +++++++++++ server/src/NotificationSender.ts | 6 ++-- server/src/dag/Dag.ts | 8 +++++ server/src/document/DbtTextDocument.ts | 29 +++++++++++++++- server/src/lsp_server/LspServerBase.ts | 20 +++++++++-- server/src/manifest/ManifestJson.ts | 1 + server/src/manifest/ManifestParser.ts | 2 ++ 13 files changed, 144 insertions(+), 10 deletions(-) create mode 100644 client/src/commands/NotUseConfigForRefsPreview.ts create mode 100644 client/src/commands/UseConfigForRefsPreview.ts diff --git a/client/src/ExtensionClient.ts b/client/src/ExtensionClient.ts index dcd4e45f..a7c67c3d 100644 --- a/client/src/ExtensionClient.ts +++ b/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 { diff --git a/client/src/SqlPreviewContentProvider.ts b/client/src/SqlPreviewContentProvider.ts index 21e18065..7357b225 100644 --- a/client/src/SqlPreviewContentProvider.ts +++ b/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(); 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(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; } } diff --git a/client/src/commands/NotUseConfigForRefsPreview.ts b/client/src/commands/NotUseConfigForRefsPreview.ts new file mode 100644 index 00000000..945eff2e --- /dev/null +++ b/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); + } +} diff --git a/client/src/commands/UseConfigForRefsPreview.ts b/client/src/commands/UseConfigForRefsPreview.ts new file mode 100644 index 00000000..719b0155 --- /dev/null +++ b/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); + } +} diff --git a/client/src/lsp_client/DbtLanguageClient.ts b/client/src/lsp_client/DbtLanguageClient.ts index 5e44cddc..6216ab1d 100644 --- a/client/src/lsp_client/DbtLanguageClient.ts +++ b/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, ); }), diff --git a/common/src/Interfaces.ts b/common/src/Interfaces.ts index 9ff2c516..9e146144 100644 --- a/common/src/Interfaces.ts +++ b/common/src/Interfaces.ts @@ -17,3 +17,5 @@ export interface CustomInitParams { disableLogger?: boolean; profilesDir?: string; } + +export type RefReplacement = { from: string; to: string }; diff --git a/package.json b/package.json index 84a2edfb..31047540 100644 --- a/package.json +++ b/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" } ] }, diff --git a/server/src/NotificationSender.ts b/server/src/NotificationSender.ts index 6317098f..97e2a7c6 100644 --- a/server/src/NotificationSender.ts +++ b/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 { diff --git a/server/src/dag/Dag.ts b/server/src/dag/Dag.ts index bc383b2c..d6b463e8 100644 --- a/server/src/dag/Dag.ts +++ b/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); + } } diff --git a/server/src/document/DbtTextDocument.ts b/server/src/document/DbtTextDocument.ts index 00314319..73f5b7e9 100644 --- a/server/src/document/DbtTextDocument.ts +++ b/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 { 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 { diff --git a/server/src/lsp_server/LspServerBase.ts b/server/src/lsp_server/LspServerBase.ts index d18befac..1fdd5981 100644 --- a/server/src/lsp_server/LspServerBase.ts +++ b/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 { constructor( @@ -14,17 +15,32 @@ export abstract class LspServerBase { abstract onInitialize(params: InitializeParams): InitializeResult | ResponseError; 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) => diff --git a/server/src/manifest/ManifestJson.ts b/server/src/manifest/ManifestJson.ts index ae7b37dd..0464c91a 100644 --- a/server/src/manifest/ManifestJson.ts +++ b/server/src/manifest/ManifestJson.ts @@ -20,6 +20,7 @@ export interface ManifestModel extends ManifestNode { config?: { sqlHeader?: string; materialized?: string; + schema?: string; }; } diff --git a/server/src/manifest/ManifestParser.ts b/server/src/manifest/ManifestParser.ts index 42c390a4..6ea5d8bd 100644 --- a/server/src/manifest/ManifestParser.ts +++ b/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, }, })); }