test

Signed-off-by: Sajid Alam <sajid_alam@mckinsey.com>
kedro-org · Jan 31, 2025 · a916ec3 · a916ec3
1 parent 67b1dc0
commit a916ec3
Show file tree

Hide file tree

Showing 5 changed files with 88 additions and 3 deletions.
diff --git a/src/common/activationHelper.ts b/src/common/activationHelper.ts
@@ -5,6 +5,7 @@ import {
     executeServerCommand,
     executeServerDefinitionCommand,
     setKedroProjectPath,
+    filterPipelines
 } from './commands';
 
 import * as vscode from 'vscode';
@@ -89,6 +90,7 @@ export const registerCommandsAndEvents = (
     const CMD_DEFINITION_REQUEST = `${serverId}.sendDefinitionRequest`;
     const CMD_SHOW_OUTPUT_CHANNEL = `${serverId}.showOutputChannel`;
     const CMD_SET_PROJECT_PATH = `${serverId}.kedroProjectPath`;
+    const CMD_FILTER_PIPELINES = `${serverId}.filterPipelines`;
 
     (async () => {
         // Status Bar
@@ -162,6 +164,10 @@ export const registerCommandsAndEvents = (
             registerCommand(CMD_SET_PROJECT_PATH, () => {
                 setKedroProjectPath();
             }),
+            registerCommand(CMD_FILTER_PIPELINES, async () => {
+                const lsClient = getLSClient();
+                await filterPipelines(lsClient);
+            }),
         );
     })();
 };
diff --git a/src/common/commands.ts b/src/common/commands.ts
@@ -166,7 +166,7 @@ export async function executeGetProjectDataCommand(lsClient: LanguageClient | un
         return;
     }
     if (!lsClient.initializeResult) {
-        await vscode.window.showErrorMessage('The Language Server fail to initialise.');
+        await vscode.window.showErrorMessage('The Language Server failed to initialize.');
         return;
     }
 
@@ -175,3 +175,36 @@ export async function executeGetProjectDataCommand(lsClient: LanguageClient | un
     const result = await vscode.commands.executeCommand(commandName);
     return result;
 }
+
+export async function filterPipelines(lsClient?: LanguageClient) {
+    try {
+        // 1) Fetch pipeline data
+        const projectData: any = await executeGetProjectDataCommand(lsClient);
+        console.log('Project Data:', projectData); // Add this line to log the project data
+        const pipelines = projectData?.pipelines || [];
+        if (!pipelines.length) {
+            vscode.window.showInformationMessage('No pipelines found in this Kedro project.');
+            return;
+        }
+
+        // 2) Show a QuickPick
+        const pipelineItems: QuickPickItem[] = pipelines.map((p: any) => ({ label: p.name }));
+        const picked = await vscode.window.showQuickPick(pipelineItems, {
+            placeHolder: 'Select a pipeline to filter...',
+        });
+        if (!picked) {
+            // User hit ESC or no selection
+            return;
+        }
+
+        // 3) Send the pipeline name to the webview
+        vscode.commands.executeCommand('kedro.viz.sendMessage', {
+            command: 'filterPipeline',
+            pipelineName: picked.label,
+        });
+    } catch (err) {
+        vscode.window.showErrorMessage(
+            `Error filtering pipelines: ${err instanceof Error ? err.message : String(err)}`
+        );
+    }
+}
diff --git a/src/extension.ts b/src/extension.ts
@@ -16,6 +16,8 @@ import {
 } from './common/utilities';
 
 import { runServer, registerCommandsAndEvents } from './common/activationHelper';
+import { filterPipelines } from './common/commands';
+import KedroVizPanel from './webview/vizWebView';
 
 let lsClient: LanguageClient | undefined;
 let isCommandsAndEventsRegistered = false;
@@ -91,6 +93,18 @@ export async function activate(context: vscode.ExtensionContext): Promise<void>
         isCommandsAndEventsRegistered = true;
     }
 
+    context.subscriptions.push(
+        vscode.commands.registerCommand('kedro.viz.sendMessage', (message) => {
+            KedroVizPanel.sendMessage(message);
+        })
+    );
+
+    context.subscriptions.push(
+        vscode.commands.registerCommand('kedro.viz.sendMessage', (message) => {
+            KedroVizPanel.sendMessage(message);
+        })
+    );
+
     setImmediate(async () => {
         const interpreter = getInterpreterFromSetting(serverId);
         if (interpreter === undefined || interpreter.length === 0) {

diff --git a/src/webview/vizWebView.ts b/src/webview/vizWebView.ts
@@ -37,6 +37,12 @@ export default class KedroVizPanel {
         KedroVizPanel.currentPanel = new KedroVizPanel(panel, extensionUri);
     }
 
+    public static sendMessage(message: any) {
+        if (KedroVizPanel.currentPanel) {
+            KedroVizPanel.currentPanel._panel.webview.postMessage(message);
+        }
+    }
+
     public static revive(panel: vscode.WebviewPanel, extensionUri: vscode.Uri) {
         KedroVizPanel.currentPanel = new KedroVizPanel(panel, extensionUri);
     }

diff --git a/webview/src/App.jsx b/webview/src/App.jsx
@@ -30,16 +30,42 @@ function App() {
             setError(true);
           }
           break;
+        case "filterPipeline":
+          if (message.pipelineName) {
+            console.log("Filtering pipeline:", message.pipelineName);
+            // Implement filtering logic here
+            const filteredData = filterDataByPipeline(data, message.pipelineName);
+            console.log("Filtered Data:", filteredData);
+            setData(filteredData);
+          }
+          break;
         default:
           break;  
       }
     });
 
     return () => {
-      window.removeEventListener("message", () => {console.log("removed")});
+      window.removeEventListener("message", () => { console.log("removed") });
     };
 
-  }, []);
+  }, [data]);
+
+  const filterDataByPipeline = (data, pipelineName) => {
+    console.log("Original Data:", data);
+    // Implement the logic to filter the data based on the pipeline name
+    const filteredNodes = data.nodes.filter(node => node.pipelines.includes(pipelineName));
+    const filteredEdges = data.edges.filter(edge => {
+        const sourceNode = data.nodes.find(node => node.id === edge.source);
+        const targetNode = data.nodes.find(node => node.id === edge.target);
+        return sourceNode && targetNode && sourceNode.pipelines.includes(pipelineName) && targetNode.pipelines.includes(pipelineName);
+    });
+    console.log("Filtered Nodes:", filteredNodes);
+    console.log("Filtered Edges:", filteredEdges);
+    return {
+        nodes: filteredNodes,
+        edges: filteredEdges,
+    };
+  };
 
   const handleNodeClick = (node) => {
     if (node) {