feat: add YAML documentation parsing for dbt models (#1686)

Author: mdesmet

package-lock.json

@@ -1345,6 +1345,7 @@
     "@types/glob": "^8.1.0",
     "@types/istanbul-lib-coverage": "^2.0.6",
     "@types/jest": "^29.5.0",
+    "@types/js-yaml": "^4.0.9",
     "@types/node": "^20.14.11",
     "@types/vscode": "^1.95.0",
     "@types/which": "^3.0.4",

package.json

@@ -9,7 +9,7 @@ import { createFullPathForNode } from ".";
 export class DocParser {
   constructor(private terminal: DBTTerminal) {}
 
-  createDocMetaMap(docs: any[], project: DBTProject): Promise<DocMetaMap> {
+  createDocMetaMap(docs: any, project: DBTProject): Promise<DocMetaMap> {
     return new Promise(async (resolve) => {
       this.terminal.debug(
         "DocParser",
@@ -22,12 +22,8 @@ export class DocParser {
         resolve(docMetaMap);
         return;
       }
-      if (typeof docs[Symbol.iterator] !== "function") {
-        resolve(docMetaMap);
-        return;
-      }
-      for (const doc of docs) {
-        const { package_name, name, original_file_path } = doc;
+      for (const doc of Object.values(docs)) {
+        const { package_name, name, original_file_path } = doc as any;
         const packageName = package_name;
         // TODO: these things can change so we should recreate them if project config changes
         const projectName = project.getProjectName();
@@ -47,6 +43,7 @@ export class DocParser {
           original_file_path,
         );
         if (!fullPath) {
+          resolve(docMetaMap);
           return;
         }
         try {

src/manifest/parsers/docParser.ts

@@ -1,4 +1,6 @@
 import path = require("path");
+import { promises as fs } from "fs";
+import * as yaml from "js-yaml";
 import {
   ProgressLocation,
   Uri,
@@ -13,7 +15,11 @@ import { RateLimitException } from "../exceptions";
 import { DBTProject } from "../manifest/dbtProject";
 import { DBTProjectContainer } from "../manifest/dbtProjectContainer";
 import { TelemetryService } from "../telemetry";
-import { extendErrorWithSupportLinks, provideSingleton } from "../utils";
+import {
+  extendErrorWithSupportLinks,
+  provideSingleton,
+  removeProtocol,
+} from "../utils";
 import {
   AIColumnDescription,
   DBTDocumentation,
@@ -33,7 +39,7 @@ interface DocumentationSchemaModel {
   name: string;
   description: string;
   tests: any;
-  columns: [];
+  columns: { name: string; description?: string; data_type?: string }[];
 }
 export interface DocumentationSchema {
   version: number;
@@ -77,6 +83,85 @@ export class DocGenService {
     private dbtTerminal: DBTTerminal,
   ) {}
 
+  private async getDocumentationFromYaml(
+    projectRoot: string,
+    filePath: string,
+    modelName: string,
+  ): Promise<DBTDocumentation | undefined> {
+    const eventResult = this.queryManifestService.getEventByCurrentProject();
+    if (!eventResult || !eventResult.event) {
+      return undefined;
+    }
+
+    const { event } = eventResult;
+    const currentNode = event.nodeMetaMap.lookupByBaseName(modelName);
+    if (!currentNode?.patch_path) {
+      return undefined;
+    }
+
+    try {
+      // Read and parse the YAML file
+      const yamlPath = path.join(
+        projectRoot,
+        removeProtocol(currentNode.patch_path),
+      );
+      const content = await fs.readFile(yamlPath, "utf8");
+      const parsedDoc = yaml.load(content) as DocumentationSchema;
+
+      // Find matching model definition
+      const modelDef = parsedDoc.models?.find((m) => m.name === modelName);
+      if (!modelDef) {
+        return undefined;
+      }
+
+      // Map to DBTDocumentation format
+      return {
+        aiEnabled: this.altimateRequest.enabled(),
+        name: modelName,
+        patchPath: currentNode.patch_path,
+        description: modelDef.description || "",
+        generated: false,
+        resource_type: currentNode.resource_type,
+        uniqueId: currentNode.uniqueId,
+        filePath,
+        columns: (modelDef.columns || []).map((column) => ({
+          name: column.name,
+          description: column.description || "",
+          generated: false,
+          source: Source.YAML,
+          type: column.data_type?.toLowerCase(),
+        })),
+      } as DBTDocumentation;
+    } catch (error) {
+      this.dbtTerminal.error(
+        "docGenService:getDocumentationYamlError",
+        `Error reading YAML documentation: ${error}`,
+        error,
+      );
+    }
+    // falling back on original implementation
+    const docColumns = currentNode.columns;
+    return {
+      aiEnabled: this.altimateRequest.enabled(),
+      name: modelName,
+      patchPath: currentNode.patch_path,
+      description: currentNode.description,
+      generated: false,
+      resource_type: currentNode.resource_type,
+      uniqueId: currentNode.uniqueId,
+      filePath,
+      columns: Object.values(docColumns).map((column) => {
+        return {
+          name: column.name,
+          description: column.description,
+          generated: false,
+          source: Source.YAML,
+          type: column.data_type?.toLowerCase(),
+        };
+      }),
+    } as DBTDocumentation;
+  }
+
   private async generateDocsForColumn(
     documentation: DBTDocumentation | undefined,
     compiledSql: string | undefined,
@@ -221,53 +306,33 @@ export class DocGenService {
     documentation: DBTDocumentation | undefined;
     message?: { message: string; type: string };
   }> {
-    const eventResult = this.queryManifestService.getEventByCurrentProject();
-    if (!eventResult) {
-      return {
-        documentation: undefined,
-        message: this.getMissingDocumentationMessage(filePath),
-      };
+    const messageObj = this.getMissingDocumentationMessage(filePath);
+    if (messageObj.type === "error") {
+      return { documentation: undefined, message: messageObj };
     }
-    const { event } = eventResult;
 
-    if (!event || !filePath) {
-      return {
-        documentation: undefined,
-        message: this.getMissingDocumentationMessage(filePath),
-      };
+    if (!filePath) {
+      return { documentation: undefined, message: messageObj };
     }
 
     const modelName = path.basename(filePath, ".sql");
-    const currentNode = event.nodeMetaMap.lookupByBaseName(modelName);
-    if (currentNode === undefined) {
+    const project = this.dbtProjectContainer.findDBTProject(Uri.file(filePath));
+
+    if (!project || !project.projectRoot) {
       return {
         documentation: undefined,
         message: this.getMissingDocumentationMessage(filePath),
       };
     }
 
-    const docColumns = currentNode.columns;
-    return {
-      documentation: {
-        aiEnabled: this.altimateRequest.enabled(),
-        name: modelName,
-        patchPath: currentNode.patch_path,
-        description: currentNode.description,
-        generated: false,
-        resource_type: currentNode.resource_type,
-        uniqueId: currentNode.uniqueId,
-        filePath,
-        columns: Object.values(docColumns).map((column) => {
-          return {
-            name: column.name,
-            description: column.description,
-            generated: false,
-            source: Source.YAML,
-            type: column.data_type?.toLowerCase(),
-          };
-        }),
-      } as DBTDocumentation,
-    };
+    // Get model directly from YAML
+    const documentation = await this.getDocumentationFromYaml(
+      project.projectRoot.fsPath,
+      filePath,
+      modelName,
+    );
+
+    return { documentation };
   }
 
   private chunk(a: string[], n: number) {

src/services/docGenService.ts

@@ -54,12 +54,7 @@ import { DbtTestService } from "../services/dbtTestService";
 import { gte } from "semver";
 import { TelemetryEvents } from "../telemetry/events";
 import { SendMessageProps } from "./altimateWebviewProvider";
-import {
-  CllEvents,
-  DbtLineageService,
-  Table,
-} from "../services/dbtLineageService";
-import { Model } from "@lib";
+import { DbtLineageService, Table } from "../services/dbtLineageService";
 
 export enum Source {
   YAML = "YAML",
@@ -181,10 +176,33 @@ export class DocsEditViewPanel implements WebviewViewProvider {
         collaborationEnabled: workspace
           .getConfiguration("dbt")
           .get<boolean>("enableCollaboration", false),
+        docBlocks: this.getDocBlocksForCurrentProject(),
       });
     }
   }
 
+  private getDocBlocksForCurrentProject(): Array<{
+    name: string;
+    path: string;
+  }> {
+    const project = this.getProject();
+    if (!project) {
+      return [];
+    }
+
+    const manifestEvent = this.eventMap.get(project.projectRoot.fsPath);
+    if (!manifestEvent?.docMetaMap) {
+      return [];
+    }
+
+    return Array.from(manifestEvent.docMetaMap.entries()).map(
+      ([name, metaData]) => ({
+        name,
+        path: metaData.path,
+      }),
+    );
+  }
+
   private async transmitColumns(columns: MetadataColumn[]) {
     if (this._panel) {
       await this._panel.webview.postMessage({

src/webview_provider/docsEditPanel.ts

@@ -30,12 +30,14 @@ import documentationSlice, {
   updateSelectedConversationGroup,
   updateSingleDocsPropRightPanel,
   updateUserInstructions,
+  setDocBlocks,
 } from "./state/documentationSlice";
 import {
   DBTDocumentation,
   DBTModelTest,
   DocsGenerateUserInstructions,
   MetadataColumn,
+  DocBlock,
 } from "./state/types";
 import { ContextProps } from "./types";
 import { getGenerationsInModel, isStateDirty } from "./utils";
@@ -57,6 +59,7 @@ type IncomingMessageEvent = MessageEvent<
     project?: string;
     columns?: DBTDocumentation["columns"];
     model?: string;
+    docBlocks?: DocBlock[];
     name?: string;
     description?: string;
     collaborationEnabled?: boolean;
@@ -132,6 +135,7 @@ const DocumentationProvider = (): JSX.Element => {
     dispatch(
       setMissingDocumentationMessage(event.data.missingDocumentationMessage),
     );
+    dispatch(setDocBlocks(event.data.docBlocks ?? []));
   };
 
   const onMessage = useCallback((event: IncomingMessageEvent) => {

webview_panels/src/modules/documentationEditor/DocumentationProvider.tsx

@@ -0,0 +1,39 @@
+.docBlockItem {
+  padding: 8px 12px;
+  cursor: pointer;
+  border-radius: 4px;
+  margin-bottom: 4px;
+  border: 1px solid transparent;
+  transition: background-color 0.15s ease, border-color 0.15s ease;
+
+  &:hover {
+    background-color: var(--background--01);
+    border-color: var(--background--02);
+  }
+}
+
+.popoverContent {
+  min-width: 320px;
+  padding: 8px;
+}
+
+.searchContainer {
+  margin-bottom: 12px;
+}
+
+.itemsList {
+  max-height: 240px;
+  overflow-y: auto;
+}
+
+.itemContent {
+  gap: 2px;
+}
+
+.itemName {
+  font-size: 14px;
+}
+
+.emptyState {
+  padding: 16px;
+}