feat: bulk docs prop (#1568)

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

src/altimate.ts

@@ -290,6 +290,11 @@ interface FeedbackResponse {
   ok: boolean;
 }
 
+interface BulkDocsPropRequest {
+  num_columns: number;
+  session_id: string;
+}
+
 interface AltimateConfig {
   key: string;
   instance: string;
@@ -954,6 +959,13 @@ export class AltimateRequest {
     });
   }
 
+  async bulkDocsPropCredit(req: BulkDocsPropRequest) {
+    return this.fetch<FeedbackResponse>("dbt/v4/bulk-docs-prop-credits", {
+      method: "POST",
+      body: JSON.stringify(req),
+    });
+  }
+
   async getPreConfiguredNotebooks() {
     return this.fetch<PreconfiguredNotebookItem[]>(
       "notebook/preconfigured/list",

src/services/dbtLineageService.ts

@@ -51,66 +51,15 @@ export type Table = {
   packageName?: string;
 };
 
-class DerivedCancellationTokenSource extends CancellationTokenSource {
-  constructor(linkedToken: CancellationToken) {
-    super();
-    linkedToken.onCancellationRequested(() => {
-      super.cancel();
-    });
-  }
-}
-
 @provideSingleton(DbtLineageService)
 export class DbtLineageService {
-  private cllProgressResolve: () => void = () => {};
-
   public constructor(
     private altimateRequest: AltimateRequest,
     protected telemetry: TelemetryService,
     private dbtTerminal: DBTTerminal,
     private queryManifestService: QueryManifestService,
   ) {}
 
-  async handleColumnLineage(
-    { event }: { event: CllEvents },
-    onCancel: () => void,
-  ) {
-    if (event === CllEvents.START) {
-      window.withProgress(
-        {
-          title: "Retrieving column level lineage",
-          location: ProgressLocation.Notification,
-          cancellable: true,
-        },
-        async (_, token) => {
-          await new Promise<void>((resolve) => {
-            this.cancellationTokenSource = new DerivedCancellationTokenSource(
-              token,
-            );
-            this.cllProgressResolve = resolve;
-            token.onCancellationRequested(() => {
-              onCancel();
-            });
-          });
-        },
-      );
-      return;
-    }
-    this.cancellationTokenSource?.token.onCancellationRequested((e) => {
-      console.log(e);
-    });
-    if (event === CllEvents.END) {
-      this.cllProgressResolve();
-      this.cancellationTokenSource?.dispose();
-      return;
-    }
-    if (event === CllEvents.CANCEL) {
-      this.cllProgressResolve();
-      this.cancellationTokenSource?.cancel();
-      return;
-    }
-  }
-
   getUpstreamTables({ table }: { table: string }) {
     return { tables: this.getConnectedTables("children", table) };
   }
@@ -260,23 +209,25 @@ export class DbtLineageService {
     return g.get(key)?.nodes.length || 0;
   }
 
-  private cancellationTokenSource: CancellationTokenSource | undefined;
-  async getConnectedColumns({
-    targets,
-    upstreamExpansion,
-    currAnd1HopTables,
-    selectedColumn,
-    showIndirectEdges,
-    eventType,
-  }: {
-    targets: [string, string][];
-    upstreamExpansion: boolean;
-    currAnd1HopTables: string[];
-    // select_column is used for pricing not business logic
-    selectedColumn: { name: string; table: string };
-    showIndirectEdges: boolean;
-    eventType: string;
-  }) {
+  async getConnectedColumns(
+    {
+      targets,
+      upstreamExpansion,
+      currAnd1HopTables,
+      selectedColumn,
+      showIndirectEdges,
+      eventType,
+    }: {
+      targets: [string, string][];
+      upstreamExpansion: boolean;
+      currAnd1HopTables: string[];
+      // select_column is used for pricing not business logic
+      selectedColumn: { name: string; table: string };
+      showIndirectEdges: boolean;
+      eventType: string;
+    },
+    cancellationTokenSource: CancellationTokenSource,
+  ) {
     const _event = this.queryManifestService.getEventByCurrentProject();
     if (!_event) {
       return;
@@ -317,15 +268,15 @@ export class DbtLineageService {
       await project.getNodesWithDBColumns(
         event,
         modelsToFetch,
-        this.cancellationTokenSource!.token,
+        cancellationTokenSource.token,
       );
 
     const selected_column = {
       model_node: mappedNode[selectedColumn.table],
       column: selectedColumn.name,
     };
 
-    if (this.cancellationTokenSource?.token.isCancellationRequested) {
+    if (cancellationTokenSource.token.isCancellationRequested) {
       return { column_lineage: [] };
     }
 
@@ -411,7 +362,7 @@ export class DbtLineageService {
 
     const modelDialect = project.getAdapterType();
     try {
-      if (this.cancellationTokenSource?.token.isCancellationRequested) {
+      if (cancellationTokenSource.token.isCancellationRequested) {
         return { column_lineage: [] };
       }
       const sessionId = `${env.sessionId}-${selectedColumn.table}-${selectedColumn.name}`;

src/telemetry/events.ts

@@ -22,6 +22,7 @@ export enum TelemetryEvents {
   "DocumentationEditor/FeedbackClick" = "DocumentationEditor/FeedbackClick",
   "DocumentationEditor/BulkGenerateTests" = "DocumentationEditor/BulkGenerateTests",
   "DocumentationEditor/SaveClick" = "DocumentationEditor/SaveClick",
+  "DocumentationEditor/SaveBulk" = "DocumentationEditor/SaveBulk",
   "DocumentationEditor/SaveError" = "DocumentationEditor/SaveError",
   "DocumentationEditor/SaveNewFilePathSelect" = "DocumentationEditor/SaveNewFilePathSelect",
   "DocumentationEditor/BulkGenerateAllClick" = "DocumentationEditor/BulkGenerateAllClick",

src/webview_provider/docsEditPanel.ts

@@ -1,6 +1,7 @@
 import { existsSync, readFileSync, writeFileSync } from "fs";
 import {
   CancellationToken,
+  CancellationTokenSource,
   ColorThemeKind,
   Disposable,
   ProgressLocation,
@@ -13,6 +14,7 @@ import {
   WebviewViewResolveContext,
   window,
   workspace,
+  env,
 } from "vscode";
 import { DBTProjectContainer } from "../manifest/dbtProjectContainer";
 import {
@@ -110,6 +112,7 @@ export class DocsEditViewPanel implements WebviewViewProvider {
   private eventMap: Map<string, ManifestCacheProjectAddedEvent> = new Map();
   private _disposables: Disposable[] = [];
   private onMessageDisposable: Disposable | undefined;
+  private cancellationTokenSource: CancellationTokenSource | undefined;
 
   public constructor(
     private dbtProjectContainer: DBTProjectContainer,
@@ -682,15 +685,6 @@ export class DocsEditViewPanel implements WebviewViewProvider {
               panel: this._panel,
             });
             break;
-          case "columnLineageBase": {
-            this.dbtLineageService.handleColumnLineage(params, () => {
-              this._panel?.webview.postMessage({
-                command: "columnLineage",
-                args: { event: CllEvents.CANCEL },
-              });
-            });
-            break;
-          }
           case "getDownstreamColumns": {
             const targets = params.targets as [string, string][];
             const testsResult = await Promise.all(
@@ -720,7 +714,7 @@ export class DocsEditViewPanel implements WebviewViewProvider {
               this.handleSyncRequestFromWebview(
                 syncRequestId,
                 () => ({ column_lineage: [], tables: [], tests }),
-                command,
+                "response",
               );
               return;
             }
@@ -729,23 +723,29 @@ export class DocsEditViewPanel implements WebviewViewProvider {
               name: params.column as string,
             };
             const currAnd1HopTables = [...tables, ...targets.map((t) => t[0])];
-            const columns = await this.dbtLineageService.getConnectedColumns({
-              targets,
-              currAnd1HopTables,
-              selectedColumn,
-              upstreamExpansion: true,
-              showIndirectEdges: false,
-              eventType: "documentation_propagation",
-            });
+            this.cancellationTokenSource = new CancellationTokenSource();
+            const columns = await this.dbtLineageService.getConnectedColumns(
+              {
+                targets,
+                currAnd1HopTables,
+                selectedColumn,
+                upstreamExpansion: true,
+                showIndirectEdges: false,
+                eventType: "documentation_propagation",
+              },
+              this.cancellationTokenSource!,
+            );
             this.handleSyncRequestFromWebview(
               syncRequestId,
-              () => {
-                return { ...columns, tables: _tables, tests };
-              },
-              command,
+              () => ({ ...columns, tables: _tables, tests }),
+              "response",
             );
             break;
           }
+          case "cancelColumnLineage": {
+            this.cancellationTokenSource?.cancel();
+            break;
+          }
           case "saveDocumentation":
             this.telemetry.sendTelemetryEvent(
               TelemetryEvents["DocumentationEditor/SaveClick"],
@@ -761,20 +761,24 @@ export class DocsEditViewPanel implements WebviewViewProvider {
             break;
           case "saveDocumentationBulk": {
             this.telemetry.sendTelemetryEvent(
-              TelemetryEvents["DocumentationEditor/SaveClick"],
-            );
-            window.withProgress(
-              {
-                title: "Saving documentation",
-                location: ProgressLocation.Notification,
-                cancellable: false,
-              },
-              async () => {
-                for (const item of message.models) {
-                  await this.saveDocumentation(item, syncRequestId);
-                }
-              },
+              TelemetryEvents["DocumentationEditor/SaveBulk"],
             );
+            const successfulSaves: string[] = [];
+            for (const item of message.models) {
+              const model = await this.saveDocumentation(item, syncRequestId);
+              if (model) {
+                successfulSaves.push(model);
+              }
+            }
+            if (successfulSaves.length > 0) {
+              window.showInformationMessage(
+                `Successfully propagated to: ${Array.from(new Set(successfulSaves)).join(", ")}`,
+              );
+              this.altimateRequest.bulkDocsPropCredit({
+                num_columns: message.numColumns,
+                session_id: env.sessionId,
+              });
+            }
             break;
           }
         }
@@ -983,6 +987,7 @@ export class DocsEditViewPanel implements WebviewViewProvider {
           },
         });
       }
+      return this.documentation?.name;
     } catch (error) {
       this.transmitError();
       this.telemetry.sendTelemetryError(