Improve db statuses

Sema4AI · Dec 23, 2024 · 366aaeb · 366aaeb
1 parent acdd427
commit 366aaeb
Show file tree

Hide file tree

Showing 3 changed files with 70 additions and 11 deletions.
diff --git a/sema4ai/vscode-client/src/robo/dataSourceHandling.ts b/sema4ai/vscode-client/src/robo/dataSourceHandling.ts
@@ -1,6 +1,6 @@
 import { Uri, window, commands, ProgressLocation, Progress, CancellationToken } from "vscode";
 import { logError, OUTPUT_CHANNEL } from "../channel";
-import { RobotEntry } from "../viewsCommon";
+import { RobotEntry, treeViewIdToTreeDataProvider } from "../viewsCommon";
 import { DatasourceInfo, DataSourceState } from "../protocols";
 import { langServer } from "../extension";
 import {
@@ -13,6 +13,8 @@ import { sleep } from "../time";
 import { findActionPackagePath } from "../actionServer";
 import { SEMA4AI_LIST_ACTIONS_INTERNAL } from "../robocorpCommands";
 import { QuickPickItemWithAction, showSelectOneQuickPick } from "../ask";
+import { TREE_VIEW_SEMA4AI_TASK_PACKAGES_TREE } from "../robocorpViews";
+import { RobotsTreeDataProvider } from "../viewsRobots";
 
 function isExternalDatasource(datasource: DatasourceInfo): boolean {
     const externalEngines = ["custom", "files", "models"];
@@ -253,6 +255,11 @@ export const setupDataSource = async (entry?: RobotEntry) => {
     }
 
     await setupSingleDataSource(datasource, dataServerInfo, Uri.file(entry.robot.directory).toString());
+
+    const dataProvider = treeViewIdToTreeDataProvider.get(
+        TREE_VIEW_SEMA4AI_TASK_PACKAGES_TREE
+    ) as RobotsTreeDataProvider;
+    dataProvider.updateDatasourceStatuses();
 };
 
 export async function dropDataSource(entry?: RobotEntry) {
@@ -333,6 +340,11 @@ export async function dropDataSource(entry?: RobotEntry) {
         return;
     }
 
+    const dataProvider = treeViewIdToTreeDataProvider.get(
+        TREE_VIEW_SEMA4AI_TASK_PACKAGES_TREE
+    ) as RobotsTreeDataProvider;
+    dataProvider.updateDatasourceStatuses();
+
     window.showInformationMessage(result["message"]);
 }
 
@@ -420,6 +432,11 @@ export async function dropAllDataSources(entry?: RobotEntry) {
             window.showInformationMessage("Data Sources dropped succesfully.");
         }
     );
+
+    const dataProvider = treeViewIdToTreeDataProvider.get(
+        TREE_VIEW_SEMA4AI_TASK_PACKAGES_TREE
+    ) as RobotsTreeDataProvider;
+    dataProvider.updateDatasourceStatuses();
 }
 
 export async function setupAllDataSources(entry?: RobotEntry) {
@@ -512,4 +529,9 @@ export async function setupAllDataSources(entry?: RobotEntry) {
             }
         }
     );
+
+    const dataProvider = treeViewIdToTreeDataProvider.get(
+        TREE_VIEW_SEMA4AI_TASK_PACKAGES_TREE
+    ) as RobotsTreeDataProvider;
+    dataProvider.updateDatasourceStatuses();
 }
diff --git a/sema4ai/vscode-client/src/views.ts b/sema4ai/vscode-client/src/views.ts
@@ -232,11 +232,6 @@ export function registerViews(context: ExtensionContext) {
         "treeDataProvider": robotsTreeDataProvider,
     });
 
-    // Periodic refresh every 60 seconds
-    setInterval(() => {
-        robotsTreeDataProvider.updateDatasourceStatuses();
-    }, 5 * 1000);
-
     treeViewIdToTreeView.set(TREE_VIEW_SEMA4AI_TASK_PACKAGES_TREE, robotsTree);
     treeViewIdToTreeDataProvider.set(TREE_VIEW_SEMA4AI_TASK_PACKAGES_TREE, robotsTreeDataProvider);
 

diff --git a/sema4ai/vscode-client/src/viewsRobots.ts b/sema4ai/vscode-client/src/viewsRobots.ts
@@ -16,6 +16,37 @@ function empty<T>(array: T[]) {
     return array === undefined || array.length === 0;
 }
 
+function dataSourcesAreDifferent(
+    globalDataSourceState: Map<string, Map<string, DatasourceInfo>>,
+    currentDatasources: Map<string, Map<string, DatasourceInfo>>
+): boolean {
+    if (globalDataSourceState.size !== currentDatasources.size) {
+        return true;
+    }
+
+    for (const [actionPath, currentInnerMap] of currentDatasources) {
+        const globalInnerMap = globalDataSourceState.get(actionPath);
+
+        if (!globalInnerMap) {
+            return true;
+        }
+
+        if (globalInnerMap.size !== currentInnerMap.size) {
+            return true;
+        }
+
+        for (const [sourceName, currentInfo] of currentInnerMap) {
+            const globalInfo = globalInnerMap.get(sourceName);
+
+            if (!globalInfo || JSON.stringify(currentInfo) !== JSON.stringify(globalInfo)) {
+                return true;
+            }
+        }
+    }
+
+    return false;
+}
+
 function getRobotLabel(robotInfo: LocalPackageMetadataInfo): string {
     let label: string = undefined;
     if (robotInfo.yamlContents) {
@@ -69,7 +100,7 @@ export class RobotsTreeDataProvider implements vscode.TreeDataProvider<RobotEntr
             return;
         }
 
-        this.globalDataSourceState.clear();
+        const currentDatasources: Map<string, Map<string, DatasourceInfo>> = new Map();
 
         for (const actionPackage of localActionPackages) {
             const dataSourceStateResult = (await langServer.sendRequest("computeDataSourceState", {
@@ -80,15 +111,18 @@ export class RobotsTreeDataProvider implements vscode.TreeDataProvider<RobotEntr
             if (dataSourceStateResult.success) {
                 const dataSourcesState: DatasourceInfo[] = dataSourceStateResult.result.required_data_sources;
                 for (const item of dataSourcesState) {
-                    const innerMap = this.globalDataSourceState.get(actionPackage.filePath) ?? new Map();
+                    const innerMap = currentDatasources.get(actionPackage.filePath) ?? new Map();
                     innerMap.set(getDataSourceCaption(item), item);
 
-                    this.globalDataSourceState.set(actionPackage.filePath, innerMap);
+                    currentDatasources.set(actionPackage.filePath, innerMap);
                 }
             }
         }
 
-        this.fireRootChange();
+        if (currentDatasources.size > 0 && dataSourcesAreDifferent(this.globalDataSourceState, currentDatasources)) {
+            this.globalDataSourceState = currentDatasources;
+            this.fireRootChange();
+        }
     }
 
     /**
@@ -132,6 +166,7 @@ export class RobotsTreeDataProvider implements vscode.TreeDataProvider<RobotEntr
         if (element === undefined) {
             // i.e.: this is the root entry, so, we've
             // collected the actual robots here.
+            this.updateDatasourceStatuses();
 
             let notifySelection = false;
             if (empty(this.lastRoot) && empty(ret)) {
@@ -566,7 +601,14 @@ export class RobotsTreeDataProvider implements vscode.TreeDataProvider<RobotEntr
 
                     let name = `${dbName} ?`;
                     if (isConfigured !== undefined) {
-                        name = `${dbName} ${isConfigured ? "✓" : "✗"}`;
+                        if (
+                            datasource.model_state &&
+                            ["generating", "training", "creating"].includes(datasource.model_state)
+                        ) {
+                            name = `${dbName} ◌`;
+                        } else {
+                            name = `${dbName} ${isConfigured ? "✓" : "✗"}`;
+                        }
                     }
 
                     children.push({