fix: handle source node (#721)

Co-authored-by: Michiel De Smet <[email protected]>

Author: AdiGajbhiye

new_lineage_panel/src/TableDetails.tsx

@@ -130,7 +130,7 @@ const ColumnSection: FunctionComponent<{
               if (!selectedTable) {
                 return;
               }
-              getColumns(selectedTable.table, true).then((_data) => {
+              getColumns(selectedTable, true).then((_data) => {
                 setData(_data);
                 setFilteredColumn(_data.columns);
               });
@@ -233,7 +233,7 @@ const TableDetails = () => {
     if (!selectedTable) {
       return;
     }
-    getColumns(selectedTable.table, false).then((_data) => {
+    getColumns(selectedTable, false).then((_data) => {
       setData(_data);
       setFilteredColumn(_data.columns);
       setIsLoading(false);

new_lineage_panel/src/service.ts

@@ -45,8 +45,15 @@ export const downstreamTables = (tableKey: string) => {
   }>;
 };
 
-export const getColumns = (table: string, refresh: boolean) => {
-  return requestExecutor("getColumns", { table, refresh }) as Promise<
+export const getColumns = (
+  { key, nodeType }: { key: string; nodeType: string },
+  refresh: boolean,
+) => {
+  return requestExecutor("getColumns", {
+    tableKey: key,
+    refresh,
+    nodeType,
+  }) as Promise<
     Columns
   >;
 };

src/manifest/dbtProject.ts

@@ -570,7 +570,7 @@ export class DBTProject implements Disposable {
     return yamlString;
   }
 
-  async getColumnsInRelation(
+  async getColumnsOfModel(
     modelName: string,
   ): Promise<{ [key: string]: string }[]> {
     await this.blockUntilPythonBridgeIsInitalized();
@@ -585,18 +585,59 @@ export class DBTProject implements Disposable {
       return [];
     }
     // Get database and schema
-    const refNode = (await this.python?.lock(
+    const node = (await this.python?.lock(
       (python) => python!`to_dict(project.get_ref_node(${modelName}))`,
     )) as ResolveReferenceResult;
     // Get columns
-    if (!refNode) {
+    if (!node) {
       return [];
     }
+    return this.getColumsOfRelation(
+      node.database,
+      node.schema,
+      node.alias || modelName,
+    );
+  }
+
+  async getColumnsOfSource(
+    sourceName: string,
+    tableName: string,
+  ): Promise<{ [key: string]: string }[]> {
+    await this.blockUntilPythonBridgeIsInitalized();
+    if (!this.pythonBridgeInitialized) {
+      window.showErrorMessage(
+        "Could not execute query, because the Python bridge has not been initalized. If the issue persists, please open a Github issue.",
+      );
+      this.telemetry.sendTelemetryError(
+        "getColumnsInRelationPythonBridgeNotInitializedError",
+      );
+      // TODO: improve this, the errors should be captured at a higher level
+      return [];
+    }
+    // Get database and schema
+    const node = (await this.python?.lock(
+      (python) =>
+        python!`to_dict(project.get_source_node(${sourceName}, ${tableName}))`,
+    )) as ResolveReferenceResult;
+    // Get columns
+    if (!node) {
+      return [];
+    }
+    return this.getColumsOfRelation(
+      node.database,
+      node.schema,
+      node.alias || tableName,
+    );
+  }
+
+  async getColumsOfRelation(
+    database: string | undefined,
+    schema: string | undefined,
+    objectName: string,
+  ) {
     return await this.python?.lock(
       (python) =>
-        python!`to_dict(project.get_columns_in_relation(project.create_relation(${
-          refNode.database
-        }, ${refNode.schema}, ${refNode.alias || modelName})))`,
+        python!`to_dict(project.get_columns_in_relation(project.create_relation(${database}, ${schema}, ${objectName})))`,
     );
   }
 
@@ -664,7 +705,7 @@ export class DBTProject implements Disposable {
         this.telemetry.sendTelemetryEvent("generateSchemaYML", {
           adapter: this.adapterType,
         });
-        const columnsInRelation = await this.getColumnsInRelation(modelName);
+        const columnsInRelation = await this.getColumnsOfModel(modelName);
         // Generate yml file content
         const fileContents = this.createYMLContent(
           columnsInRelation,

src/manifest/dbtProjectContainer.ts

@@ -134,7 +134,7 @@ export class DBTProjectContainer implements Disposable {
     );
     if (answer === PromptAnswer.YES) {
       commands.executeCommand("dbtPowerUser.openSetupWalkthrough");
-    } 
+    }
     this.setToGlobalState("showSetupWalkthrough", false);
   }
 

src/webview_provider/docsEditPanel.ts

@@ -261,7 +261,7 @@ export class DocsEditViewPanel implements WebviewViewProvider {
                 const modelName = path.basename(currentFilePath.fsPath, ".sql");
                 try {
                   const columnsInRelation =
-                    await project.getColumnsInRelation(modelName);
+                    await project.getColumnsOfModel(modelName);
                   const columns = columnsInRelation.map((column) => {
                     return {
                       name: column.column,