checkpoint before ws datasource

KxSystems · May 1, 2024 · 65b3d0c · 65b3d0c
1 parent 032bbb0
commit 65b3d0c
Show file tree

Hide file tree

Showing 4 changed files with 64 additions and 70 deletions.
diff --git a/package.json b/package.json
@@ -190,9 +190,9 @@
           "description": "Enable linting for q files",
           "default": false
         },
-        "kdb.scratchpads": {
+        "kdb.connectionMap": {
           "type": "object",
-          "description": "Scratchpads",
+          "description": "Connection map for workspace files",
           "default": {},
           "scope": "resource"
         }
@@ -390,7 +390,7 @@
         "category": "KX",
         "command": "kdb.terminal.run",
         "title": "Run q file in a new q instance",
-        "icon": "$(run)"
+        "icon": "$(debug-alt)"
       },
       {
         "category": "KX",
@@ -401,7 +401,7 @@
         "category": "KX",
         "command": "kdb.execute.selectedQuery",
         "title": "Execute Current Selection",
-        "icon": "$(run)"
+        "icon": "$(run-below)"
       },
       {
         "category": "KX",
@@ -412,7 +412,8 @@
       {
         "category": "KX",
         "command": "kdb.execute.pythonScratchpadQuery",
-        "title": "Execute Current Selection in Insights Scratchpad"
+        "title": "Execute Current Selection in Insights Scratchpad",
+        "icon": "$(run-below)"
       },
       {
         "category": "KX",

diff --git a/src/commands/scratchpadCommand.ts b/src/commands/scratchpadCommand.ts
@@ -34,6 +34,11 @@ function setRunScratchpadItemText(text: string) {
   ext.runScratchpadItem.text = `$(run) ${text}`;
 }
 
+export function workspaceFoldersChanged() {
+  ext.dataSourceTreeProvider.reload();
+  ext.scratchpadTreeProvider.reload();
+}
+
 function getServers() {
   const conf = workspace.getConfiguration("kdb");
   const servers = conf.get<{ [key: string]: { serverAlias: string } }>(
@@ -51,39 +56,29 @@ function getServers() {
   ];
 }
 
-async function setServerForScratchpad(uri: Uri, server: string | undefined) {
+export function getServerForUri(uri: Uri) {
   const conf = workspace.getConfiguration("kdb", uri);
   const scratchpads = conf.get<{ [key: string]: string | undefined }>(
-    "scratchpads",
+    "connectionMap",
     {},
   );
-  scratchpads[uri.path] = server;
-  await conf.update("scratchpads", scratchpads);
+  return scratchpads[workspace.asRelativePath(uri)];
 }
 
-async function waitForConnection(name: string) {
-  return new Promise<void>((resolve, reject) => {
-    let count = 0;
-    const retry = () => {
-      count++;
-      setTimeout(() => {
-        if (connectionService.isConnected(name)) {
-          resolve();
-        } else if (count < 5) {
-          retry();
-        } else {
-          reject(`Can not connect to ${name}`);
-        }
-      }, 50);
-    };
-    retry();
-  });
+async function setServerForUri(uri: Uri, server: string | undefined) {
+  const conf = workspace.getConfiguration("kdb", uri);
+  const map = conf.get<{ [key: string]: string | undefined }>(
+    "connectionMap",
+    {},
+  );
+  map[workspace.asRelativePath(uri)] = server;
+  await conf.update("connectionMap", map);
 }
 
-async function getConnectionForServer(server: string) {
+export function getConnectionForUri(uri: Uri) {
+  const server = getServerForUri(uri);
   if (server) {
-    const servers = await ext.serverProvider.getChildren();
-    return servers.find((item) => {
+    return ext.connectionsList.find((item) => {
       if (item instanceof InsightsNode) {
         return item.details.alias === server;
       } else if (item instanceof KdbNode) {
@@ -94,19 +89,10 @@ async function getConnectionForServer(server: string) {
   }
 }
 
-export function getServerForUri(uri: Uri) {
-  const conf = workspace.getConfiguration("kdb", uri);
-  const scratchpads = conf.get<{ [key: string]: string | undefined }>(
-    "scratchpads",
-    {},
-  );
-  return scratchpads[uri.path];
-}
-
-export function getConnectionForUri(uri: Uri) {
-  const server = getServerForUri(uri);
+async function getConnectionForServer(server: string) {
   if (server) {
-    return ext.connectionsList.find((item) => {
+    const servers = await ext.serverProvider.getChildren();
+    return servers.find((item) => {
       if (item instanceof InsightsNode) {
         return item.details.alias === server;
       } else if (item instanceof KdbNode) {
@@ -117,9 +103,23 @@ export function getConnectionForUri(uri: Uri) {
   }
 }
 
-export function workspaceFoldersChanged() {
-  ext.dataSourceTreeProvider.reload();
-  ext.scratchpadTreeProvider.reload();
+async function waitForConnection(name: string) {
+  return new Promise<void>((resolve, reject) => {
+    let count = 0;
+    const retry = () => {
+      count++;
+      setTimeout(() => {
+        if (connectionService.isConnected(name)) {
+          resolve();
+        } else if (count < 5) {
+          retry();
+        } else {
+          reject(`Can not connect to ${name}`);
+        }
+      }, 50);
+    };
+    retry();
+  });
 }
 
 function setRealActiveTextEditor(editor?: TextEditor | undefined) {
@@ -138,8 +138,7 @@ export function activeEditorChanged(editor?: TextEditor | undefined) {
   const item = ext.runScratchpadItem;
   if (ext.activeTextEditor) {
     const uri = ext.activeTextEditor.document.uri;
-    const path = uri.path;
-    if (path.endsWith(".kdb.q") || path.endsWith(".kdb.py")) {
+    if (isScratchpad(uri)) {
       const server = getServerForUri(uri);
       setRunScratchpadItemText(server || "Run");
       item.show();
@@ -163,7 +162,7 @@ export async function pickConnection(uri: Uri) {
     if (picked === "(none)") {
       picked = undefined;
     }
-    await setServerForScratchpad(uri, picked);
+    await setServerForUri(uri, picked);
     setRunScratchpadItemText(picked || "Run");
   }
 
@@ -178,7 +177,7 @@ function isScratchpad(uri: Uri) {
   return uri.path.endsWith(".kdb.q") || uri.path.endsWith(".kdb.py");
 }
 
-export async function runScratchpad(type: ExecutionTypes) {
+export async function runActiveEditor(type?: ExecutionTypes) {
   if (ext.activeTextEditor) {
     const uri = ext.activeTextEditor.document.uri;
     if (isScratchpad(uri)) {
@@ -208,7 +207,12 @@ export async function runScratchpad(type: ExecutionTypes) {
         }
       }
     }
-    runQuery(type);
+
+    runQuery(
+      type || isPython(uri)
+        ? ExecutionTypes.PythonQueryFile
+        : ExecutionTypes.QueryFile,
+    );
   }
 }
 

diff --git a/src/extension.ts b/src/extension.ts
@@ -59,7 +59,6 @@ import {
   connect,
   disconnect,
   enableTLS,
-  executeQuery,
   removeConnection,
   rerunQuery,
   resetScratchPad,
@@ -104,7 +103,7 @@ import {
   activeEditorChanged,
   ConnectionLensProvider,
   pickConnection,
-  runScratchpad,
+  runActiveEditor,
   workspaceFoldersChanged,
 } from "./commands/scratchpadCommand";
 import { createDefaultDataSourceFile } from "./models/dataSource";
@@ -254,7 +253,6 @@ export async function activate(context: ExtensionContext) {
     ),
     commands.registerCommand("kdb.refreshServerObjects", () => {
       ext.serverProvider.reload();
-      ext.activeConnection?.update();
     }),
     commands.registerCommand(
       "kdb.queryHistory.rerun",
@@ -342,36 +340,24 @@ export async function activate(context: ExtensionContext) {
       }
     }),
     commands.registerCommand("kdb.runScratchpad", async () => {
-      if (ext.activeTextEditor) {
-        const path = ext.activeTextEditor.document.uri.path;
-        await runScratchpad(
-          path.endsWith(".kdb.py")
-            ? ExecutionTypes.PythonQueryFile
-            : ExecutionTypes.QuerySelection,
-        );
-        ext.activeConnection?.update();
-      }
+      await runActiveEditor();
     }),
     commands.registerCommand("kdb.execute.selectedQuery", async () => {
-      await runScratchpad(ExecutionTypes.QuerySelection);
-      ext.activeConnection?.update();
+      await runActiveEditor(ExecutionTypes.QuerySelection);
     }),
     commands.registerCommand("kdb.execute.fileQuery", async () => {
-      await runScratchpad(ExecutionTypes.QueryFile);
-      ext.activeConnection?.update();
+      await runActiveEditor(ExecutionTypes.QueryFile);
     }),
     commands.registerCommand("kdb.execute.pythonScratchpadQuery", async () => {
-      await runScratchpad(ExecutionTypes.PythonQuerySelection);
-      ext.activeConnection?.update();
+      await runActiveEditor(ExecutionTypes.PythonQuerySelection);
     }),
     commands.registerCommand("kdb.scratchpad.reset", async () => {
       await resetScratchPad();
     }),
     commands.registerCommand(
       "kdb.execute.pythonFileScratchpadQuery",
       async () => {
-        await runScratchpad(ExecutionTypes.PythonQueryFile);
-        ext.activeConnection?.update();
+        await runActiveEditor(ExecutionTypes.PythonQueryFile);
       },
     ),
     commands.registerCommand(

diff --git a/src/services/workspaceTreeProvider.ts b/src/services/workspaceTreeProvider.ts
@@ -24,6 +24,7 @@ import {
 import Path from "path";
 import { getServerIconState } from "../utils/core";
 import { getConnectionForUri } from "../commands/scratchpadCommand";
+import { ext } from "../extensionVariables";
 
 export class WorkspaceTreeProvider implements TreeDataProvider<FileTreeItem> {
   private _onDidChangeTreeData = new EventEmitter<void>();
@@ -32,7 +33,9 @@ export class WorkspaceTreeProvider implements TreeDataProvider<FileTreeItem> {
   constructor(
     private readonly glob: string,
     private readonly baseIcon: string,
-  ) {}
+  ) {
+    ext.serverProvider.onDidChangeTreeData(() => this.reload());
+  }
 
   reload() {
     this._onDidChangeTreeData.fire();