AltimateAI · saravmajestic · Jun 13, 2024 · Jun 13, 2024 · Jun 13, 2024 · Jun 14, 2024
diff --git a/package.json b/package.json
@@ -1117,7 +1117,7 @@
   "scripts": {
     "prepare": "husky install",
     "vscode:prepublish": "npm run panel:webviews && npm run panel:lineage && webpack --mode production",
-    "webpack": "npm run panel:webviews && npm run panel:lineage && webpack --mode development",
+    "webpack": "npm run panel:webviews && webpack --mode development",
     "webpack-dev": "npm run panel:webviews && npm run panel:lineage && webpack --mode development --watch",
     "watch": "concurrently \"npm:watch --prefix ./new_lineage_panel\" \"npm:watch --prefix ./webview_panels\" \"webpack --mode development --watch\"",
     "test-compile": "tsc -p ./",

diff --git a/src/webview_provider/altimateWebviewProvider.ts b/src/webview_provider/altimateWebviewProvider.ts
@@ -473,6 +473,17 @@ export class AltimateWebviewProvider implements WebviewViewProvider {
         ),
       ),
     );
+    const LineageGif = webview.asWebviewUri(
+      Uri.file(
+        path.join(
+          extensionUri.fsPath,
+          "webview_panels",
+          "dist",
+          "assets",
+          "lineage.gif",
+        ),
+      ),
+    );
     const codiconsUri = webview.asWebviewUri(
       Uri.joinPath(
         extensionUri,
@@ -505,9 +516,11 @@ export class AltimateWebviewProvider implements WebviewViewProvider {
           <body>
             <div id="root"></div>
             <div id="sidebar"></div>
+            <div id="modal"></div>
             <script nonce="${nonce}" >
               window.viewPath = "${this.viewPath}";
               var spinnerUrl = "${SpinnerUrl}"
+              var lineageGif = "${LineageGif}"
             </script>
 
             <script nonce="${nonce}" type="module" src="${indexJs}"></script>

diff --git a/src/webview_provider/newLineagePanel.ts b/src/webview_provider/newLineagePanel.ts
@@ -1,3 +1,4 @@
+/* eslint-disable prettier/prettier */
 import { readFileSync, writeFileSync } from "fs";
 import * as path from "path";
 import {
@@ -9,12 +10,12 @@ import {
   TextEditor,
   Uri,
   Webview,
-  WebviewOptions,
   WebviewView,
   WebviewViewResolveContext,
   window,
   workspace,
   env,
+  WebviewOptions,
 } from "vscode";
 import { AltimateRequest, ModelNode } from "../altimate";
 import {
@@ -33,6 +34,10 @@ import { TelemetryService } from "../telemetry";
 import { PythonException } from "python-bridge";
 import { AbortError } from "node-fetch";
 import { DBTTerminal } from "../dbt_client/dbtTerminal";
+import { AltimateWebviewProvider } from "./altimateWebviewProvider";
+import { QueryManifestService } from "../services/queryManifestService";
+import { SharedStateService } from "../services/sharedStateService";
+import { UsersService } from "../services/usersService";
 
 type Table = {
   label: string;
@@ -70,19 +75,52 @@ class DerivedCancellationTokenSource extends CancellationTokenSource {
 }
 
 @provideSingleton(NewLineagePanel)
-export class NewLineagePanel implements LineagePanelView {
-  private _panel: WebviewView | undefined;
-  private eventMap: Map<string, ManifestCacheProjectAddedEvent> = new Map();
+export class NewLineagePanel
+  extends AltimateWebviewProvider
+  implements LineagePanelView
+{
+  protected viewPath = "/lineage";
+  protected panelDescription = "Lineage panel";
+  protected _panel: WebviewView | undefined;
+  protected eventMap: Map<string, ManifestCacheProjectAddedEvent> = new Map();
   // since lineage can be cancelled from 2 places: progress bar and panel actions
   private cancellationTokenSource: DerivedCancellationTokenSource | undefined;
   private cllProgressResolve: () => void = () => {};
 
   public constructor(
-    private dbtProjectContainer: DBTProjectContainer,
+    protected dbtProjectContainer: DBTProjectContainer,
     private altimate: AltimateRequest,
-    private telemetry: TelemetryService,
+    protected telemetry: TelemetryService,
     private terminal: DBTTerminal,
-  ) {}
+    private eventEmitterService: SharedStateService,
+    protected queryManifestService: QueryManifestService,
+    protected usersService: UsersService,
+  ) {
+    super(
+      dbtProjectContainer,
+      altimate,
+      telemetry,
+      eventEmitterService,
+      terminal,
+      queryManifestService,
+      usersService,
+    );
+
+    this._disposables.push(
+      workspace.onDidChangeConfiguration(
+        (e) => {
+          if (!e.affectsConfiguration("dbt.enableLineagePanelV2")) {
+            return;
+          }
+          if (this._panel) {
+            this.renderWebviewView(this._panel.webview);
+          }
+        },
+        this,
+        this._disposables,
+      ),
+    );
+  }
 
   public changedActiveTextEditor(event: TextEditor | undefined) {
     if (event === undefined) {
@@ -130,23 +168,9 @@ export class NewLineagePanel implements LineagePanelView {
     });
   }
 
-  resolveWebviewView(
-    panel: WebviewView,
-    context: WebviewViewResolveContext<unknown>,
-    _token: CancellationToken,
-  ): void | Thenable<void> {
-    this.terminal.debug(
-      "newLineagePanel:resolveWebviewView",
-      "onResolveWebviewView",
-    );
-    this._panel = panel;
-    this.setupWebviewOptions(context);
-    this.renderWebviewView(context);
-  }
-
-  async handleCommand(message: { command: string; args: any }): Promise<void> {
-    const { command, args } = message;
-    const { id, params } = args;
+  async handleCommand(message: { command: string; args: any, syncRequestId?: string }): Promise<void> {
+    const { command, args = {}, syncRequestId } = message;
+    const { id = syncRequestId, params } = args;
 
     if (command === "openProblemsTab") {
       commands.executeCommand("workbench.action.problems.focus");
@@ -156,7 +180,7 @@ export class NewLineagePanel implements LineagePanelView {
       const body = await this.getUpstreamTables(params);
       this._panel?.webview.postMessage({
         command: "response",
-        args: { id, body, status: true },
+        args: { id, syncRequestId, body, status: true },
       });
       return;
     }
@@ -165,7 +189,7 @@ export class NewLineagePanel implements LineagePanelView {
       const body = await this.getDownstreamTables(params);
       this._panel?.webview.postMessage({
         command: "response",
-        args: { id, body, status: true },
+        args: { id, syncRequestId, body, status: true },
       });
       return;
     }
@@ -174,7 +198,7 @@ export class NewLineagePanel implements LineagePanelView {
       const body = await this.getColumns(params);
       this._panel?.webview.postMessage({
         command: "response",
-        args: { id, body, status: true },
+        args: { id, syncRequestId, body, status: true },
       });
       return;
     }
@@ -183,7 +207,7 @@ export class NewLineagePanel implements LineagePanelView {
       const body = await this.getExposureDetails(params);
       this._panel?.webview.postMessage({
         command: "response",
-        args: { id, body, status: true },
+        args: { id, syncRequestId, body, status: true },
       });
       return;
     }
@@ -192,7 +216,7 @@ export class NewLineagePanel implements LineagePanelView {
       const body = await this.getConnectedColumns(params);
       this._panel?.webview.postMessage({
         command: "response",
-        args: { id, body, status: !!body },
+        args: { id, syncRequestId, body, status: !!body },
       });
       return;
     }
@@ -207,12 +231,12 @@ export class NewLineagePanel implements LineagePanelView {
         });
         this._panel?.webview.postMessage({
           command: "response",
-          args: { id, status: true },
+          args: { id, syncRequestId, status: true },
         });
       } catch (error) {
         this._panel?.webview.postMessage({
           command: "response",
-          args: { id, status: false },
+          args: { id, syncRequestId, status: false },
         });
         window.showErrorMessage(
           extendErrorWithSupportLinks(
@@ -253,6 +277,7 @@ export class NewLineagePanel implements LineagePanelView {
         command: "response",
         args: {
           id,
+          syncRequestId,
           status: true,
           body: {
             showSelectEdges: config.get("showSelectEdges", true),
@@ -273,6 +298,7 @@ export class NewLineagePanel implements LineagePanelView {
         command: "response",
         args: {
           id,
+          syncRequestId,
           status: true,
           body: { ok: true },
         },
@@ -285,6 +311,7 @@ export class NewLineagePanel implements LineagePanelView {
       "Unsupported command",
       message,
     );
+    super.handleCommand(message);
   }
 
   private async handleColumnLineage({ event }: { event: CllEvents }) {
@@ -927,14 +954,19 @@ export class NewLineagePanel implements LineagePanelView {
     return { node, aiEnabled };
   }
 
-  private setupWebviewOptions(context: WebviewViewResolveContext) {
-    this._panel!.description =
-      "Show table level and column level lineage SQL queries";
-    this._panel!.webview.options = <WebviewOptions>{ enableScripts: true };
-  }
+  protected renderWebviewView(webview: Webview) {
+    const enableLineagePanelV2 = workspace
+      .getConfiguration("dbt")
+      .get<boolean>("enableLineagePanelV2", false);
 
-  private renderWebviewView(context: WebviewViewResolveContext) {
-    const webview = this._panel!.webview!;
+    if (enableLineagePanelV2){
+      this._panel!.webview.html = super.getHtml(
+        webview,
+        this.dbtProjectContainer.extensionUri,
+      );
+      return;
+    }
+    this._panel!.webview.options = <WebviewOptions>{ enableScripts: true };
     this._panel!.webview.html = getHtml(
       webview,
       this.dbtProjectContainer.extensionUri,

diff --git a/webview_panels/src/AppConstants.tsx b/webview_panels/src/AppConstants.tsx
@@ -5,6 +5,7 @@ import DocumentationProvider from "@modules/documentationEditor/DocumentationPro
 import DataPilotPanel from "@modules/dataPilot";
 import DbtDocsView from "@modules/dbtDocs/DbtDocsView";
 import QueryPanelProvider from "@modules/queryPanel/QueryPanelProvider";
+import LineageView from "@modules/lineage/LineageView";
 
 // TODO: lazy loading breaks loading dynamic webviews when having css because of vite dynamic loading
 // research on how to fix that and then use lazy loading
@@ -25,4 +26,5 @@ export const AvailableRoutes = {
   },
   "/dbt-docs": { component: <DbtDocsView /> },
   "/query-panel": { component: <QueryPanelProvider /> },
+  "/lineage": { component: <LineageView /> },
 };
diff --git a/webview_panels/src/assets/icons/index.tsx b/webview_panels/src/assets/icons/index.tsx
@@ -18,7 +18,9 @@ export { default as EmptySquareIcon } from "./square.svg?react";
 export { default as CheckedSquareIcon } from "./checked-square.svg?react";
 export { default as TestsIcon } from "./tests.svg?react";
 export { default as FolderIcon } from "./folder.svg?react";
+export { default as PlayCircleIcon } from "./play-circle.svg?react";
 import LoadingSpinnerUrl from "./spinner.gif";
+import LineageGif from "./lineage.gif";
 import "./styles.css";
 
 interface Props {
@@ -124,6 +126,14 @@ export const LoadingSpinner = (): JSX.Element => (
   />
 );
 
+export const LineageDemo = (): JSX.Element => (
+  <img
+    // @ts-expect-error added in altimateWebViewProvider
+    src={(window.lineageGif as string) ?? LineageGif}
+    alt="Lineage demo"
+  />
+);
+
 export const LoadingIcon = (
   props: HTMLAttributes<HTMLElement>,
 ): JSX.Element => <Icon icon="loading" className="rotate" {...props} />;
diff --git a/webview_panels/src/assets/icons/lineage.gif b/webview_panels/src/assets/icons/lineage.gif
diff --git a/webview_panels/src/assets/icons/play-circle.svg b/webview_panels/src/assets/icons/play-circle.svg