refactor: replace Node class hierarchy with NodeData interface

mdesmet · mdesmet · commit bdd76f715e73 · 2025-06-27T18:55:30.000+05:30
diff --git a/src/dbt_integration/domain.ts b/src/dbt_integration/domain.ts
@@ -170,9 +170,16 @@ export interface ExposureMetaData {
   meta?: Record<string, unknown>;
 }
 
+export interface NodeData {
+  label: string;
+  key: string;
+  url?: string;
+  resourceType: string;
+}
+
 interface NodeGraphMetaData {
-  currentNode: Node;
-  nodes: Node[];
+  currentNode: NodeData;
+  nodes: NodeData[];
 }
 
 interface ModelGraphMetaData {
@@ -191,64 +198,6 @@ export interface GraphMetaMap {
   metrics: NodeGraphMap;
 }
 
-interface IconPath {
-  light: string;
-  dark: string;
-}
-
-export abstract class Node {
-  label: string;
-  key: string;
-  url: string | undefined;
-  iconPath: IconPath = {
-    light: path.join(
-      path.resolve(__dirname),
-      "../media/images/model_light.svg",
-    ),
-    dark: path.join(path.resolve(__dirname), "../media/images/model_dark.svg"),
-  };
-  displayInModelTree: boolean = true;
-
-  constructor(label: string, key: string, url?: string) {
-    this.label = label;
-    this.key = key;
-    this.url = url;
-  }
-}
-
-export class Model extends Node {}
-
-export class Seed extends Node {}
-export class Test extends Node {
-  // displayInModelTree = false;
-  iconPath = {
-    light: path.join(
-      path.resolve(__dirname),
-      "../media/images/source_light.svg",
-    ),
-    dark: path.join(path.resolve(__dirname), "../media/images/source_dark.svg"),
-  };
-}
-export class Analysis extends Node {
-  displayInModelTree = true;
-}
-export class Exposure extends Node {
-  displayInModelTree = true;
-}
-export class Metric extends Node {
-  displayInModelTree = false;
-}
-export class Snapshot extends Node {}
-export class Source extends Node {
-  iconPath = {
-    light: path.join(
-      path.resolve(__dirname),
-      "../media/images/source_light.svg",
-    ),
-    dark: path.join(path.resolve(__dirname), "../media/images/source_dark.svg"),
-  };
-}
-
 export enum RunModelType {
   RUN_PARENTS,
   RUN_CHILDREN,
diff --git a/src/dbt_integration/parsers/graphParser.ts b/src/dbt_integration/parsers/graphParser.ts
@@ -1,22 +1,15 @@
 import {
-  Analysis,
-  Exposure,
   GraphMetaMap,
-  Metric,
   MetricMetaMap,
-  Model,
-  Node,
+  NodeData,
   NodeGraphMap,
   NodeMetaMap,
-  Seed,
-  Snapshot,
-  Source,
   SourceMetaMap,
-  Test,
   TestMetaMap,
 } from "../domain";
 import { DBTTerminal } from "../terminal";
 import { DBTIntegrationAdapter } from "../dbtIntegrationAdapter";
+import { RESOURCE_TYPE_METRIC, RESOURCE_TYPE_TEST } from "../dbtIntegration";
 
 const notEmpty = <T>(value: T | null | undefined): value is T => {
   return value !== null && value !== undefined;
@@ -73,7 +66,7 @@ export class GraphParser {
               metricMetaMap,
             ),
           )
-          .filter((n) => !(n instanceof Test))
+          .filter((n) => n?.resourceType !== RESOURCE_TYPE_TEST)
           .filter(notEmpty);
         map.set(nodeName, { nodes: currentNodes });
         return map;
@@ -92,7 +85,7 @@ export class GraphParser {
               metricMetaMap,
             ),
           )
-          .filter((n) => n instanceof Test)
+          .filter((n) => n?.resourceType === RESOURCE_TYPE_TEST)
           .filter(notEmpty);
         map.set(nodeName, { nodes: currentNodes });
         return map;
@@ -111,7 +104,7 @@ export class GraphParser {
               metricMetaMap,
             ),
           )
-          .filter((n) => n instanceof Metric)
+          .filter((n) => n?.resourceType === RESOURCE_TYPE_METRIC)
           .filter(notEmpty);
         map.set(nodeName, { nodes: currentNodes });
         return map;
@@ -138,7 +131,7 @@ export class GraphParser {
     nodeMetaMap: NodeMetaMap,
     testMetaMap: TestMetaMap,
     metricMetaMap: MetricMetaMap,
-  ): (parentNodeName: string) => Node | undefined {
+  ): (parentNodeName: string) => NodeData | undefined {
     return (parentNodeName) => {
       // Support dots in model names
       const [nodeType, nodePackage, ...restNodeName] =
@@ -150,11 +143,12 @@ export class GraphParser {
           const url = sourceMetaMap
             .get(sourceName)
             ?.tables.find((table) => table.name === tableName)?.path!;
-          return new Source(
-            `${tableName} (${sourceName})`,
-            parentNodeName,
-            url,
-          );
+          return {
+            label: `${tableName} (${sourceName})`,
+            key: parentNodeName,
+            url: url,
+            resourceType: "source",
+          };
         }
         case "model": {
           // can this ever be not there?
@@ -163,7 +157,12 @@ export class GraphParser {
             return;
           }
           const url = model?.path!;
-          return new Model(model.alias, parentNodeName, url);
+          return {
+            label: model.alias,
+            key: parentNodeName,
+            url: url,
+            resourceType: "model",
+          };
         }
         case "seed": {
           // can this ever be not there?
@@ -172,28 +171,57 @@ export class GraphParser {
             return;
           }
           const url = model?.path!;
-          return new Seed(model.alias, parentNodeName, url);
+          return {
+            label: model.alias,
+            key: parentNodeName,
+            url: url,
+            resourceType: "seed",
+          };
         }
         case "test": {
           // nodeName => more interesting label possibilities?
           // console.log(`${nodeName} => (parent: ${parentNodeName})`);
           const url = testMetaMap.get(nodeName.split(".")[0])?.path;
-          return new Test(nodeName, parentNodeName, url ?? "");
+          return {
+            label: nodeName,
+            key: parentNodeName,
+            url: url,
+            resourceType: "test",
+          };
         }
         case "analysis": {
           const url = nodeMetaMap.lookupByBaseName(nodeName)?.path!;
-          return new Analysis(nodeName, parentNodeName, url);
+          return {
+            label: nodeName,
+            key: parentNodeName,
+            url: url,
+            resourceType: "analysis",
+          };
         }
         case "snapshot": {
           const url = nodeMetaMap.lookupByBaseName(nodeName)?.path!;
-          return new Snapshot(nodeName, parentNodeName, url);
+          return {
+            label: nodeName,
+            key: parentNodeName,
+            url: url,
+            resourceType: "snapshot",
+          };
         }
         case "exposure": {
           const url = nodeMetaMap.lookupByBaseName(nodeName)?.path!;
-          return new Exposure(nodeName, parentNodeName, url);
+          return {
+            label: nodeName,
+            key: parentNodeName,
+            url: url,
+            resourceType: "exposure",
+          };
         }
         case "semantic_model": {
-          return new Metric(nodeName, parentNodeName);
+          return {
+            label: nodeName,
+            key: parentNodeName,
+            resourceType: "semantic_model",
+          };
         }
         default:
           console.log(`Node Type '${nodeType}' not implemented!`);
diff --git a/src/treeview_provider/modelTreeviewProvider.ts b/src/treeview_provider/modelTreeviewProvider.ts
@@ -16,19 +16,12 @@ import {
   TreeItemCollapsibleState,
   Uri,
   window,
-  workspace,
 } from "vscode";
 import {
-  Analysis,
-  Exposure,
   GraphMetaMap,
-  Node,
+  NodeData,
   NodeMetaData,
   NodeMetaMap,
-  Seed,
-  Snapshot,
-  Source,
-  Test,
 } from "../dbt_integration/domain";
 import { DBTProjectContainer } from "../manifest/dbtProjectContainer";
 import {
@@ -42,6 +35,66 @@ import {
   getDepthColor,
 } from "../utils";
 
+interface IconPath {
+  light: string;
+  dark: string;
+}
+
+abstract class Node {
+  label: string;
+  key: string;
+  url: string | undefined;
+  iconPath: IconPath = {
+    light: path.join(
+      path.resolve(__dirname),
+      "../media/images/model_light.svg",
+    ),
+    dark: path.join(path.resolve(__dirname), "../media/images/model_dark.svg"),
+  };
+  displayInModelTree: boolean = true;
+
+  constructor(label: string, key: string, url?: string) {
+    this.label = label;
+    this.key = key;
+    this.url = url;
+  }
+}
+
+class Model extends Node {}
+
+class Seed extends Node {}
+class Test extends Node {
+  // displayInModelTree = false;
+  iconPath = {
+    light: path.join(
+      path.resolve(__dirname),
+      "../media/images/source_light.svg",
+    ),
+    dark: path.join(path.resolve(__dirname), "../media/images/source_dark.svg"),
+  };
+}
+class Analysis extends Node {
+  displayInModelTree = true;
+}
+class Exposure extends Node {
+  displayInModelTree = true;
+}
+class Metric extends Node {
+  displayInModelTree = false;
+}
+
+class Snapshot extends Node {}
+
+class Source extends Node {
+  iconPath = {
+    light: path.join(
+      path.resolve(__dirname),
+      "../media/images/source_light.svg",
+    ),
+    dark: path.join(path.resolve(__dirname), "../media/images/source_dark.svg"),
+  };
+}
+
 @provide(ModelTreeviewProvider)
 abstract class ModelTreeviewProvider
   implements TreeDataProvider<NodeTreeItem>, Disposable
@@ -129,6 +182,33 @@ abstract class ModelTreeviewProvider
     return Promise.resolve(this.getTreeItems(model.uniqueId, event));
   }
 
+  private nodeDataToNode(nodeData: NodeData): Node | undefined {
+    const resourceType = nodeData.resourceType;
+    switch (resourceType) {
+      case "snapshot":
+        return new Snapshot(nodeData.label, nodeData.key, nodeData.url);
+      case "exposure":
+        return new Exposure(nodeData.label, nodeData.key, nodeData.url);
+      case "analysis":
+        return new Analysis(nodeData.label, nodeData.key, nodeData.url);
+      case "test":
+        return new Test(nodeData.label, nodeData.key, nodeData.url);
+      case "source":
+        return new Source(nodeData.label, nodeData.key, nodeData.url);
+      case "seed":
+        return new Seed(nodeData.label, nodeData.key, nodeData.url);
+      case "semantic_model":
+        return new Metric(nodeData.label, nodeData.key, nodeData.url);
+      case "model":
+        return new Model(nodeData.label, nodeData.key, nodeData.url);
+      default:
+        console.log(
+          `Resource Type '${resourceType}' not implemented in ModelTreeviewProvider.nodeDataToNode`,
+        );
+        return undefined;
+    }
+  }
+
   private getNodeTreeItem(node: Node): NodeTreeItem {
     if (node instanceof Snapshot) {
       return new SnapshotTreeItem(node);
@@ -161,11 +241,15 @@ abstract class ModelTreeviewProvider
       return [];
     }
     return parentModels.nodes
-      .filter((node) => node.displayInModelTree)
+      .flatMap((nodeData) => {
+        const node = this.nodeDataToNode(nodeData);
+        return node && node.displayInModelTree ? [node] : [];
+      })
       .map((node) => {
         const childNodes = graphMetaMap[this.treeType]
           .get(node.key)
-          ?.nodes.filter((node) => node.displayInModelTree);
+          ?.nodes.map((nodeData) => this.nodeDataToNode(nodeData))
+          .filter((node) => node && node.displayInModelTree);
 
         const treeItem = this.getNodeTreeItem(node);
         treeItem.collapsibleState =
