From 1654373bc2c637811eb887bb362b30f3e319ee32 Mon Sep 17 00:00:00 2001 From: Saravanan Date: Fri, 25 Oct 2024 10:43:15 +0530 Subject: [PATCH] fix: preserve comments and formatting in yml (#1468) * fix: preserve comments and formatting in yml * fix: schema save * chore: rename model * fix: doc edit save use cases * fix: formatting issue --- src/services/docGenService.ts | 18 ++ .../altimateWebviewProvider.ts | 21 +- src/webview_provider/docsEditPanel.ts | 250 ++++++++++++------ 3 files changed, 192 insertions(+), 97 deletions(-) diff --git a/src/services/docGenService.ts b/src/services/docGenService.ts index c13294c46..7f924299c 100644 --- a/src/services/docGenService.ts +++ b/src/services/docGenService.ts @@ -22,6 +22,24 @@ import { import { QueryManifestService } from "./queryManifestService"; import { TelemetryEvents } from "../telemetry/events"; +export interface DocumentationSchemaColumn { + name: string; + description: string; + data_type?: string; + quote?: boolean; + [key: string]: unknown; +} +interface DocumentationSchemaModel { + name: string; + description: string; + tests: any; + columns: []; +} +export interface DocumentationSchema { + version: number; + models: DocumentationSchemaModel[]; +} + interface GenerateDocsForColumnsProps { panel: WebviewView | WebviewPanel | undefined; message: any; diff --git a/src/webview_provider/altimateWebviewProvider.ts b/src/webview_provider/altimateWebviewProvider.ts index 5aa31e464..19013d653 100644 --- a/src/webview_provider/altimateWebviewProvider.ts +++ b/src/webview_provider/altimateWebviewProvider.ts @@ -243,17 +243,16 @@ export class AltimateWebviewProvider implements WebviewViewProvider { try { switch (command) { - case "getTeammatesStatus": - { - const isEnabled = workspace - .getConfiguration("dbt") - .get("enableTeammates", false); - this.sendResponseToWebview({ - command: "teammatesUpdated", - data: isEnabled, - }); - break; - } + case "getTeammatesStatus": { + const isEnabled = workspace + .getConfiguration("dbt") + .get("enableTeammates", false); + this.sendResponseToWebview({ + command: "teammatesUpdated", + data: isEnabled, + }); + break; + } case "configEnabled": this.handleSyncRequestFromWebview( syncRequestId, diff --git a/src/webview_provider/docsEditPanel.ts b/src/webview_provider/docsEditPanel.ts index 6fba0d419..15e797da5 100644 --- a/src/webview_provider/docsEditPanel.ts +++ b/src/webview_provider/docsEditPanel.ts @@ -32,10 +32,14 @@ import path = require("path"); import { PythonException } from "python-bridge"; import { TelemetryService } from "../telemetry"; import { AltimateRequest } from "../altimate"; -import { stringify, parse } from "yaml"; +import { stringify, parse, parseDocument, YAMLSeq, YAMLMap } from "yaml"; import { NewDocsGenPanel } from "./newDocsGenPanel"; import { DBTProject } from "../manifest/dbtProject"; -import { DocGenService } from "../services/docGenService"; +import { + DocGenService, + DocumentationSchema, + DocumentationSchemaColumn, +} from "../services/docGenService"; import { DBTTerminal } from "../dbt_client/dbtTerminal"; import { TestMetaData, @@ -465,6 +469,46 @@ export class DocsEditViewPanel implements WebviewViewProvider { return this.modifyColumnNames(columns, existingColumnNames); } + private setOrDeleteInParsedDocument( + doc: YAMLMap, + key: string, + value: any, + ) { + if (value) { + doc.set(key, value); + } else { + doc.delete(key); + } + } + + private findEntityInParsedDoc( + models: + | YAMLSeq + | YAMLSeq + | undefined, + predicate: (name: string) => boolean, + ) { + if (models && models.items) { + return ( + // @ts-ignore + (models.items.find( + ( + item: + | DocumentationSchema["models"]["0"] + | DocumentationSchemaColumn, + ) => { + if (item instanceof YAMLMap) { + const name = item.get("name"); + return name && predicate(name as string); + } + return false; + }, + ) as YAMLMap | undefined) || null + ); + } + return null; + } + private setupWebviewHooks(context: WebviewViewResolveContext) { // Clear this listener before subscribing again if (this.onMessageDisposable) { @@ -672,33 +716,125 @@ export class DocsEditViewPanel implements WebviewViewProvider { const docFile: string = readFileSync(patchPath).toString("utf8"); - const parsedDocFile = - parse(docFile, { - strict: false, - uniqueKeys: false, - maxAliasCount: -1, - }) || {}; - if (parsedDocFile.models === undefined) { - // this is a fresh file or one without models, so init the models - parsedDocFile.models = []; - } - if ( - parsedDocFile.models.find( - (model: any) => model.name === message.name, - ) === undefined - ) { + const parsedDocFile = parseDocument< + YAMLSeq + >(docFile, { + strict: false, + uniqueKeys: false, + }); + const existingModels = parsedDocFile.get("models") as + | YAMLSeq + | undefined; + + const model = this.findEntityInParsedDoc( + existingModels, + (name: string) => name === message.name, + ); + + if (!model) { // there is a models section but the model does not exist yet. - parsedDocFile.models.push({ + const newModelData = { name: message.name, - description: message.description || undefined, - columns: message.columns.map((column: any) => { + description: message.description?.trim() || undefined, + columns: message.columns.length + ? message.columns.map((column: any) => { + const name = getColumnNameByCase( + column.name, + projectByFilePath.getAdapterType(), + ); + return { + name, + description: + column.description?.trim() || undefined, + data_type: column.type?.toLowerCase(), + ...this.getTestDataByColumn( + message, + column.name, + project, + // passing column to get correct key: data_tests or tests + // https://github.com/AltimateAI/vscode-dbt-power-user/issues/1449 + { name: column.name }, + ), + ...(isQuotedIdentifier( + column.name, + projectByFilePath.getAdapterType(), + ) + ? { quote: true } + : undefined), + }; + }) + : undefined, + }; + // Models does not exist + if (existingModels?.items.length) { + parsedDocFile.addIn(["models"], newModelData); + } else { + // Models exist, but current one is new model + parsedDocFile.set("models", [newModelData]); + } + } else { + // The model already exists + this.setOrDeleteInParsedDocument( + model, + "description", + message.description?.trim(), + ); + const modelTests = this.getTestDataByModel( + message, + model.get("name") as string, + ); + this.setOrDeleteInParsedDocument( + model, + "tests", + modelTests, + ); + + message.columns.forEach((column: any) => { + const existingColumn = this.findEntityInParsedDoc( + model.get("columns") as + | YAMLSeq + | undefined, + (name: string) => isColumnNameEqual(name, column.name), + ); + + if (existingColumn) { + // ignore tests, data_tests from existing column, as it will be recreated in `getTestDataByColumn` + const { tests, data_tests, ...rest } = + existingColumn.toJSON(); + this.setOrDeleteInParsedDocument( + existingColumn, + "description", + column.description?.trim(), + ); + this.setOrDeleteInParsedDocument( + existingColumn, + "data_type", + (rest.data_type || column.type)?.toLowerCase(), + ); + const allTests = this.getTestDataByColumn( + message, + column.name, + project, + existingColumn.toJSON(), + ); + this.setOrDeleteInParsedDocument( + existingColumn, + "tests", + allTests?.tests, + ); + this.setOrDeleteInParsedDocument( + existingColumn, + "data_tests", + allTests?.data_tests, + ); + } else { const name = getColumnNameByCase( column.name, projectByFilePath.getAdapterType(), ); - return { + model.addIn(["columns"], { name, - description: column.description || undefined, + description: column.description?.trim() || undefined, data_type: column.type?.toLowerCase(), ...this.getTestDataByColumn( message, @@ -711,74 +847,16 @@ export class DocsEditViewPanel implements WebviewViewProvider { ) ? { quote: true } : undefined), - }; - }), + }); + } }); - } else { - // The model already exists - parsedDocFile.models = parsedDocFile.models.map( - (model: any) => { - if (model.name === message.name) { - model.description = message.description || undefined; - model.tests = this.getTestDataByModel( - message, - model.name, - ); - model.columns = message.columns.map((column: any) => { - const existingColumn = - model.columns && - model.columns.find((yamlColumn: any) => - isColumnNameEqual(yamlColumn.name, column.name), - ); - if (existingColumn !== undefined) { - // ignore tests, data_tests from existing column, as it will be recreated in `getTestDataByColumn` - const { tests, data_tests, ...rest } = - existingColumn; - return { - ...rest, - name: existingColumn.name, - data_type: ( - rest.data_type || column.type - )?.toLowerCase(), - description: column.description || undefined, - ...this.getTestDataByColumn( - message, - column.name, - project, - existingColumn, - ), - }; - } else { - const name = getColumnNameByCase( - column.name, - projectByFilePath.getAdapterType(), - ); - return { - name, - description: column.description || undefined, - data_type: column.type?.toLowerCase(), - ...this.getTestDataByColumn( - message, - column.name, - project, - ), - ...(isQuotedIdentifier( - column.name, - projectByFilePath.getAdapterType(), - ) - ? { quote: true } - : undefined), - }; - } - }); - } - return model; - }, - ); } // Force reload from manifest after manifest refresh this.loadedFromManifest = false; - writeFileSync(patchPath, stringify(parsedDocFile)); + writeFileSync( + patchPath, + stringify(parsedDocFile, { lineWidth: 0 }), + ); this.documentation = ( await this.docGenService.getDocumentationForCurrentActiveFile() ).documentation;