feat: add depth decorator (#1683)

Author: mdesmet

package.json

@@ -275,6 +275,11 @@
             "default": [],
             "description": "Add additional params to `sqlfmt`"
           },
+          "dbt.disableDepthsCalculation": {
+            "type": "boolean",
+            "description": "Disable DAG depth calculation and decoration",
+            "default": false
+          },
           "dbt.deferConfigPerProject": {
             "type": "object",
             "description": "Run subset of models without building their parent models",

src/hover_provider/depthDecorationProvider.ts

@@ -0,0 +1,193 @@
+import {
+  DecorationOptions,
+  Disposable,
+  Hover,
+  HoverProvider,
+  MarkdownString,
+  Position,
+  Range,
+  TextDocument,
+  TextEditor,
+  TextEditorDecorationType,
+  window,
+  workspace,
+} from "vscode";
+import { DBTProjectContainer } from "../manifest/dbtProjectContainer";
+import { provideSingleton } from "../utils";
+import { getDepthColor } from "../utils";
+
+@provideSingleton(DepthDecorationProvider)
+export class DepthDecorationProvider implements HoverProvider, Disposable {
+  private disposables: Disposable[] = [];
+  private readonly REF_PATTERN =
+    /\{\{\s*ref\s*\(\s*['"]([^'"]+)['"]\s*\)\s*\}\}/;
+  private readonly decorationType: TextEditorDecorationType;
+  private projectToDepthMap: Map<string, Map<string, number>> = new Map();
+
+  constructor(private dbtProjectContainer: DBTProjectContainer) {
+    this.decorationType = window.createTextEditorDecorationType({
+      after: {
+        margin: "0 0 0 5px",
+        textDecoration: "none",
+      },
+    });
+
+    this.dbtProjectContainer.onManifestChanged((event) => {
+      event.added?.forEach((added) => {
+        this.projectToDepthMap.set(
+          added.project.projectRoot.fsPath,
+          added.modelDepthMap,
+        );
+      });
+      event.removed?.forEach((removed) => {
+        this.projectToDepthMap.delete(removed.projectRoot.fsPath);
+      });
+
+      window.visibleTextEditors.forEach((editor) => {
+        this.updateDecorations(editor);
+      });
+    });
+
+    this.disposables.push(
+      window.onDidChangeActiveTextEditor((editor) => {
+        if (editor) {
+          this.updateDecorations(editor);
+        }
+      }),
+      workspace.onDidChangeTextDocument((event) => {
+        if (
+          window.activeTextEditor &&
+          event.document === window.activeTextEditor.document
+        ) {
+          this.updateDecorations(window.activeTextEditor);
+        }
+      }),
+      workspace.onDidChangeConfiguration((event) => {
+        if (event.affectsConfiguration("dbt")) {
+          window.visibleTextEditors.forEach((editor) => {
+            this.updateDecorations(editor);
+          });
+        }
+      }),
+    );
+  }
+
+  dispose() {
+    this.decorationType.dispose();
+    while (this.disposables.length) {
+      const x = this.disposables.pop();
+      if (x) {
+        x.dispose();
+      }
+    }
+  }
+
+  private updateDecorations(editor: TextEditor): void {
+    const project = this.dbtProjectContainer.findDBTProject(
+      editor.document.uri,
+    );
+    if (!project) {
+      return;
+    }
+
+    const projectRoot = project.projectRoot.fsPath;
+    const depthMapForProject = this.projectToDepthMap.get(projectRoot);
+    if (!depthMapForProject) {
+      return;
+    }
+
+    const text = editor.document.getText();
+    const decorations: DecorationOptions[] = [];
+    const pattern = new RegExp(this.REF_PATTERN, "g");
+    let match;
+
+    while ((match = pattern.exec(text)) !== null) {
+      const startPos = editor.document.positionAt(match.index);
+      const endPos = editor.document.positionAt(match.index + match[0].length);
+      const refRange = new Range(startPos, endPos);
+
+      const modelName = match[1];
+      const depth = depthMapForProject.get(modelName);
+
+      if (depth !== undefined) {
+        decorations.push(this.createDepthDecoration(refRange, depth));
+      }
+    }
+
+    editor.setDecorations(this.decorationType, decorations);
+  }
+
+  private createDepthDecoration(
+    refRange: Range,
+    depth: number,
+  ): DecorationOptions {
+    const color = getDepthColor(depth);
+    const depthText = `(${depth})`;
+
+    return {
+      range: refRange,
+      renderOptions: {
+        after: {
+          contentText: depthText,
+          backgroundColor: color,
+          color: "white",
+          fontWeight: "bold",
+          margin: "0 0 0 5px",
+        },
+      },
+    };
+  }
+
+  public provideHover(
+    document: TextDocument,
+    position: Position,
+  ): Hover | undefined {
+    const project = this.dbtProjectContainer.findDBTProject(document.uri);
+    if (!project) {
+      return;
+    }
+
+    const projectRoot = project.projectRoot.fsPath;
+    const depthMapForProject = this.projectToDepthMap.get(projectRoot);
+    if (!depthMapForProject) {
+      return;
+    }
+
+    const text = document.getText();
+    let matches: RegExpMatchArray[] = [];
+    try {
+      const pattern = new RegExp(this.REF_PATTERN, "g");
+      matches = Array.from(text.matchAll(pattern));
+    } catch (error) {
+      console.error("Error matching ref pattern in provideHover:", error);
+      return undefined;
+    }
+
+    for (const match of matches) {
+      if (match.index === undefined) {
+        continue;
+      }
+      const startPos = document.positionAt(match.index);
+      const endPos = document.positionAt(match.index + match[0].length);
+      const refRange = new Range(startPos, endPos);
+
+      if (refRange.contains(position)) {
+        const modelName = match[1];
+        const depth = depthMapForProject.get(modelName);
+
+        if (depth !== undefined) {
+          const color = getDepthColor(depth);
+          const markdown = new MarkdownString(
+            `**DAG Depth:** <span style="color:${color}">${depth}</span>\n\n` +
+              `The longest path of models between a source and this model is ${depth} nodes long.`,
+          );
+          markdown.isTrusted = true;
+          markdown.supportHtml = true;
+          return new Hover(markdown, refRange);
+        }
+      }
+    }
+
+    return undefined;
+  }
+}

src/hover_provider/index.ts

@@ -4,6 +4,7 @@ import { provideSingleton } from "../utils";
 import { ModelHoverProvider } from "./modelHoverProvider";
 import { SourceHoverProvider } from "./sourceHoverProvider";
 import { MacroHoverProvider } from "./macroHoverProvider";
+import { DepthDecorationProvider } from "./depthDecorationProvider";
 
 @provideSingleton(HoverProviders)
 export class HoverProviders implements Disposable {
@@ -13,6 +14,7 @@ export class HoverProviders implements Disposable {
     private modelHoverProvider: ModelHoverProvider,
     private sourceHoverProvider: SourceHoverProvider,
     private macroHoverProvider: MacroHoverProvider,
+    private depthDecorationProvider: DepthDecorationProvider,
   ) {
     this.disposables.push(
       languages.registerHoverProvider(
@@ -32,6 +34,12 @@ export class HoverProviders implements Disposable {
         this.macroHoverProvider,
       ),
     );
+    this.disposables.push(
+      languages.registerHoverProvider(
+        DBTPowerUserExtension.DBT_SQL_SELECTOR,
+        this.depthDecorationProvider,
+      ),
+    );
   }
 
   dispose() {

src/manifest/event/manifestCacheChangedEvent.ts

@@ -21,6 +21,7 @@ export interface ManifestCacheProjectAddedEvent {
   testMetaMap: TestMetaMap;
   docMetaMap: DocMetaMap;
   exposureMetaMap: ExposureMetaMap;
+  modelDepthMap: Map<string, number>;
 }
 
 export interface ManifestCacheProjectRemovedEvent {

src/manifest/parsers/index.ts

@@ -1,4 +1,4 @@
-import { existsSync, readFileSync } from "fs";
+import { readFileSync } from "fs";
 import { provide } from "inversify-binding-decorators";
 import * as path from "path";
 import { Uri } from "vscode";
@@ -15,6 +15,8 @@ import { TelemetryService } from "../../telemetry";
 import { ExposureParser } from "./exposureParser";
 import { MetricParser } from "./metricParser";
 import { ChildrenParentParser } from "./childrenParentParser";
+import { ModelDepthParser } from "./modelDepthParser";
+import { createFullPathForNode } from "./utils";
 
 @provide(ManifestParser)
 export class ManifestParser {
@@ -33,6 +35,7 @@ export class ManifestParser {
     private docParser: DocParser,
     private terminal: DBTTerminal,
     private telemetry: TelemetryService,
+    private modelDepthParser: ModelDepthParser,
   ) {}
 
   public async parseManifest(project: DBTProject) {
@@ -71,6 +74,7 @@ export class ManifestParser {
             },
             docMetaMap: new Map(),
             exposureMetaMap: new Map(),
+            modelDepthMap: new Map(),
           },
         ],
       };
@@ -130,6 +134,13 @@ export class ManifestParser {
       exposuresMetaMapPromise,
     ]);
 
+    // Calculate model depths
+    const modelDepthMap = this.modelDepthParser.createModelDepthsMap(
+      nodes,
+      parentMetaMap,
+      childMetaMap,
+    );
+
     const graphMetaMap = this.graphParser.createGraphMetaMap(
       project,
       parentMetaMap,
@@ -181,6 +192,7 @@ export class ManifestParser {
           testMetaMap: testMetaMap,
           docMetaMap: docMetaMap,
           exposureMetaMap: exposureMetaMap,
+          modelDepthMap,
         },
       ],
     };
@@ -216,29 +228,4 @@ export class ManifestParser {
   }
 }
 
-export const createFullPathForNode: (
-  projectName: string,
-  rootPath: string,
-  packageName: string,
-  packagePath: string,
-  relativeFilePath: string,
-) => string | undefined = (
-  projectName,
-  rootPath,
-  packageName,
-  packagePath,
-  relativeFilePath,
-) => {
-  if (packageName !== projectName) {
-    const rootPathWithPackage = path.join(
-      packagePath,
-      packageName,
-      relativeFilePath,
-    );
-    if (existsSync(rootPathWithPackage)) {
-      return rootPathWithPackage;
-    }
-    return undefined;
-  }
-  return path.join(rootPath, relativeFilePath);
-};
+export { createFullPathForNode } from "./utils";