diff --git a/src/treeview_provider/modelTreeviewProvider.ts b/src/treeview_provider/modelTreeviewProvider.ts index 663fc8abe..80d7be7f3 100755 --- a/src/treeview_provider/modelTreeviewProvider.ts +++ b/src/treeview_provider/modelTreeviewProvider.ts @@ -1,6 +1,7 @@ import { unmanaged } from "inversify"; import { provide } from "inversify-binding-decorators"; import * as path from "path"; + import { Command, Disposable, @@ -29,7 +30,7 @@ import { ManifestCacheChangedEvent, ManifestCacheProjectAddedEvent, } from "../manifest/event/manifestCacheChangedEvent"; -import { provideSingleton } from "../utils"; +import { getModelNameInActiveEditor, provideSingleton } from "../utils"; @provide(ModelTreeviewProvider) abstract class ModelTreeviewProvider @@ -54,6 +55,9 @@ abstract class ModelTreeviewProvider this.dbtProjectContainer.onManifestChanged((event) => this.onManifestCacheChanged(event), ), + window.onDidChangeTextEditorSelection(() => { + this._onDidChangeTreeData.fire(); + }), ); } @@ -104,15 +108,28 @@ abstract class ModelTreeviewProvider if (element) { return Promise.resolve(this.getTreeItems(element.key, event)); } - const fileName = path.basename( - window.activeTextEditor!.document.fileName, - ".sql", + + const model_by_file_content = event.nodeMetaMap.lookupByBaseName( + getModelNameInActiveEditor(), ); - const model = event.nodeMetaMap.lookupByBaseName(fileName); - if (!model) { - return Promise.resolve([]); + + if (model_by_file_content) { + return Promise.resolve( + this.getTreeItems(model_by_file_content.uniqueId, event), + ); } - return Promise.resolve(this.getTreeItems(model.uniqueId, event)); + + const fileName = path.parse( + window.activeTextEditor!.document.fileName, + ).name; + const model_by_file_name = event.nodeMetaMap.lookupByBaseName(fileName); + if (model_by_file_name) { + return Promise.resolve( + this.getTreeItems(model_by_file_name.uniqueId, event), + ); + } + + return Promise.resolve([]); } private getNodeTreeItem(node: Node): NodeTreeItem { @@ -180,6 +197,9 @@ class DocumentationTreeviewProvider implements TreeDataProvider { this.dbtProjectContainer.onManifestChanged((event) => this.onManifestCacheChanged(event), ), + window.onDidChangeTextEditorSelection(() => { + this._onDidChangeTreeData.fire(); + }), ); } @@ -222,14 +242,23 @@ class DocumentationTreeviewProvider implements TreeDataProvider { const { nodeMetaMap } = event; if (!element) { - const modelName = path.basename( + const fileName = path.parse( window.activeTextEditor!.document.fileName, - ".sql", + ).name; + // Try content-based lookup first + const currentNodeByFileContent = event.nodeMetaMap.lookupByBaseName( + getModelNameInActiveEditor(), ); - const currentNode = nodeMetaMap.lookupByBaseName(modelName); + // Fall back to filename-based lookup + const currentNode = + currentNodeByFileContent ?? + event.nodeMetaMap.lookupByBaseName(fileName); + if (currentNode === undefined) { return Promise.resolve([]); } + const modelName = currentNode.name; + const children = []; if (Object.keys(currentNode.columns).length !== 0) { diff --git a/src/utils.ts b/src/utils.ts index b6a0633e7..204005e76 100755 --- a/src/utils.ts +++ b/src/utils.ts @@ -8,9 +8,10 @@ import { TextDocument, Uri, workspace, + window, } from "vscode"; import { readFileSync } from "fs"; -import { parse } from "yaml"; +import { parse, parseDocument } from "yaml"; import { TestMetadataAcceptedValues, TestMetadataRelationships, @@ -359,3 +360,56 @@ export const getStringSizeInMb = (str: string): number => { const sizeInMB = sizeInBytes / (1024 * 1024); return sizeInMB; }; + +export function getModelNameInActiveEditor(): string { + if ( + window.activeTextEditor === undefined || + window.activeTextEditor.document.languageId !== "yaml" + ) { + return ""; + } + + try { + const parsedYaml = parseDocument( + window.activeTextEditor.document.getText(), + ); + if (parsedYaml.contents === null) { + return ""; + } + const cursorPosition = window.activeTextEditor.selection.active; + const offset = window.activeTextEditor.document.offsetAt(cursorPosition); + + const contents = parsedYaml.contents as { items?: Array }; + if (!contents.items) { + return ""; + } + + const modelsNode = contents.items.find( + (item) => item?.key?.value === "models", + ); + if (!modelsNode?.value?.items) { + return ""; + } + + // Find a model at the current position + for (const model of modelsNode.value.items) { + if (!model?.items) { + continue; + } + + const nameNode = model.items.find( + (item: any) => item?.key?.value === "name", + ); + if (!nameNode?.value) { + continue; + } + + if (model.range?.[0] < offset && offset < model.range?.[1]) { + return nameNode.value.toString(); + } + } + } catch (error) { + console.error("Error parsing YAML document:", error); + } + return ""; +}