Merge pull request #704 from KxSystems/ee-conn

Implemented connect and run

Author: ecmel

src/classes/localConnection.ts

@@ -19,7 +19,6 @@ import { commands } from "vscode";
 import { ext } from "../extensionVariables";
 import { QueryResult, QueryResultType } from "../models/queryResult";
 import { ServerObject } from "../models/serverObject";
-import { delay } from "../utils/core";
 import { convertStringToArray, handleQueryResults } from "../utils/execution";
 import { MessageKind, notify } from "../utils/notifications";
 import { queryWrapper } from "../utils/queryUtils";
@@ -32,8 +31,6 @@ export class LocalConnection {
   public labels: string[];
   private options: nodeq.ConnectionParameters;
   private connection?: nodeq.Connection;
-  private isError: boolean = false;
-  private result?: string;
 
   constructor(
     connectionString: string,
@@ -80,35 +77,30 @@ export class LocalConnection {
     return this.connection;
   }
 
-  public async connect(
-    callback: nodeq.AsyncValueCallback<LocalConnection>,
-  ): Promise<void> {
-    if (this.connection && this.connected) return;
-
+  public async connect() {
     const options = await this.getCustomAuthOptions();
-
-    nodeq.connect(options, (err, conn) => {
-      if (err || !conn) {
-        ext.serverProvider.reload();
-        this.connection = undefined;
-        this.connected = false;
-        callback(err, this);
-        return;
-      }
-
-      conn.addListener("close", () => {
-        commands.executeCommand("kdb.connections.disconnect", this.connLabel);
-        notify(
-          `Connection closed: ${this.options.host}:${this.options.port}`,
-          MessageKind.DEBUG,
-          { logger },
-        );
+    return new Promise<nodeq.Connection>((resolve, reject) => {
+      nodeq.connect(options, (err, conn) => {
+        if (err || !conn) {
+          ext.serverProvider.reload();
+          this.connection = undefined;
+          this.connected = false;
+          reject(err);
+          return;
+        }
+        conn.addListener("close", () => {
+          commands.executeCommand("kdb.connections.disconnect", this.connLabel);
+          notify(
+            `Connection closed: ${this.options.host}:${this.options.port}`,
+            MessageKind.DEBUG,
+            { logger },
+          );
+        });
+        this.connection = conn;
+        this.connected = true;
+        this.update();
+        resolve(conn);
       });
-
-      this.connection = conn;
-      this.connected = true;
-      this.update();
-      callback(err, this);
     });
   }
 
@@ -125,131 +117,60 @@ export class LocalConnection {
     this.updateReservedKeywords();
   }
 
-  public async execute(query: string): Promise<string | Error> {
-    let result;
-    let error;
-    let retryCount = 0;
-    while (this.connection === undefined) {
-      if (retryCount > ext.maxRetryCount) {
-        return "timeout";
-      }
-      await delay(500);
-      retryCount++;
-    }
-
-    this.connection.k(query, function (err: Error, res: string) {
-      if (err) {
-        error = err;
-        result = "";
-        return;
-      }
-      result = res;
-    });
-
-    // wait for result (lack of await using callbacks)
-    while (result === undefined || result === null) {
-      await delay(500);
-    }
-
-    if (error) {
-      throw error;
-    }
-
-    return result;
-  }
-
   public async executeQuery(
     command: string,
     context?: string,
     stringify?: boolean,
     isPython?: boolean,
   ): Promise<any> {
-    let result;
-    await this.waitForConnection();
-
-    if (!this.connection) {
-      return "timeout";
-    }
-    const args: any[] = [];
-    const wrapper = queryWrapper(!!isPython);
-
-    if (isPython) {
-      args.push(stringify ? "text" : "serialized", command, "first", 10000);
-    } else {
-      args.push(context ?? ".", command, stringify ? "text" : "structuredText");
-    }
-
-    args.push((err: Error, res: QueryResult) => {
-      if (err) {
-        this.isError = true;
-        result = handleQueryResults(err.toString(), QueryResultType.Error);
-      }
-      if (res) {
-        if (res.errored) {
-          this.isError = true;
-          result = handleQueryResults(
-            res.error + (res.backtrace ? "\n" + res.backtrace : ""),
-            QueryResultType.Error,
-          );
+    return new Promise((resolve, reject) => {
+      if (this.connection) {
+        const args: any[] = [];
+        const wrapper = queryWrapper(!!isPython);
+        if (isPython) {
+          args.push(stringify ? "text" : "serialized", command, "first", 10000);
         } else {
-          result = res.result === null ? "" : res.result;
+          args.push(
+            context ?? ".",
+            command,
+            stringify ? "text" : "structuredText",
+          );
         }
-      }
+        this.connection.k(wrapper, ...args, (err: Error, res: QueryResult) => {
+          if (err) {
+            reject(handleQueryResults(err.toString(), QueryResultType.Error));
+          } else if (res.errored) {
+            resolve(
+              handleQueryResults(
+                res.error + (res.backtrace ? "\n" + res.backtrace : ""),
+                QueryResultType.Error,
+              ),
+            );
+          } else {
+            const result = res.result === null ? "" : res.result;
+            if (!stringify && !isPython) {
+              resolve(JSON.parse(result));
+            } else if (ext.isResultsTabVisible && stringify) {
+              resolve(convertStringToArray(result ? result : ""));
+            } else {
+              resolve(result);
+            }
+          }
+          this.updateGlobal();
+        });
+      } else reject(new Error("Not connected."));
     });
-
-    this.connection.k(wrapper, ...args);
-
-    while (result === undefined || result === null) {
-      await delay(50);
-    }
-
-    this.updateGlobal();
-
-    if (this.isError) {
-      this.isError = false;
-      return result;
-    }
-
-    if (!stringify && !isPython) {
-      return JSON.parse(result);
-    }
-
-    if (ext.isResultsTabVisible && stringify) {
-      return convertStringToArray(result ? result : "");
-    }
-
-    return result;
   }
 
   public async executeQueryRaw(command: string): Promise<any> {
-    let result;
-    let retryCount = 0;
-    let error;
-    while (this.connection === undefined) {
-      if (retryCount > ext.maxRetryCount) {
-        return "timeout";
-      }
-      await delay(500);
-      retryCount++;
-    }
-    this.connection.k(command, (err: Error, res: any) => {
-      if (err) {
-        error = err;
-        result = "";
-        return;
-      }
-      result = res;
+    return new Promise((resolve, reject) => {
+      if (this.connection) {
+        this.connection.k(command, (err: Error, res: any) => {
+          if (err) reject(err);
+          else resolve(res || "");
+        });
+      } else reject(new Error("Not connected."));
     });
-
-    while (result === undefined || result === null) {
-      await delay(500);
-    }
-
-    if (error) {
-      throw error;
-    }
-
-    return result;
   }
 
   public async loadServerObjects(): Promise<ServerObject[]> {
@@ -276,38 +197,6 @@ export class LocalConnection {
     }
   }
 
-  private async waitForConnection(): Promise<void> {
-    let retryCount = 0;
-    while (this.connection === undefined) {
-      if (retryCount > ext.maxRetryCount) {
-        throw new Error("timeout");
-      }
-      await delay(500);
-      retryCount++;
-    }
-  }
-
-  private handleQueryResult(res: QueryResult): void {
-    if (res.errored) {
-      this.isError = true;
-      this.result = handleQueryResults(
-        res.error + (res.backtrace ? "\n" + res.backtrace : ""),
-        QueryResultType.Error,
-      );
-    } else {
-      this.result = res.result;
-    }
-  }
-
-  private async waitForResult(): Promise<any> {
-    while (this.result === undefined || this.result === null) {
-      await delay(500);
-    }
-    const result = this.result;
-    this.result = undefined;
-    return result;
-  }
-
   private updateGlobal() {
     const globalQuery =
       '{[q] t:system"T";tm:@[{$[x>0;[system"T ",string x;1b];0b]};0;{0b}];r:$[tm;@[0;(q;::);{[tm; t; msgs] if[tm;system"T ",string t];\'msgs}[tm;t]];@[q;::;{\'x}]];if[tm;system"T ",string t];r}{do[1000;2+2];{@[{.z.ide.ns.r1:x;:.z.ide.ns.r1};x;{r:y;:r}[;x]]}({:x!{![sv[`;] each x cross `Tables`Functions`Variables; system each "afv" cross enlist[" "] cross enlist string x]} each x} [{raze x,.z.s\'[{x where{@[{1#get x};x;`]~1#.q}\'[x]}` sv\'x,\'key x]}`]),(enlist `.z)!flip (`.z.Tables`.z.Functions`.z.Variables)!(enlist 0#`;enlist `ac`bm`exit`pc`pd`pg`ph`pi`pm`po`pp`ps`pw`vs`ts`s`wc`wo`ws;enlist `a`b`e`f`h`i`k`K`l`o`q`u`w`W`x`X`n`N`p`P`z`Z`t`T`d`D`c`zd)}';

src/commands/workspaceCommand.ts

@@ -726,20 +726,18 @@ export async function findConnection(uri: Uri) {
       server = node.label;
       conn = connMngService.retrieveConnectedConnection(server);
       if (conn === undefined) {
-        offerConnectAction(server);
-        return;
+        const res = await offerConnectAction(server);
+        if (res) {
+          conn = connMngService.retrieveConnectedConnection(server);
+        }
       }
     } else {
       notify(`Connection ${server} not found.`, MessageKind.ERROR, {
         logger,
       });
-      return;
     }
-  } else if (ext.activeConnection) {
-    conn = ext.activeConnection;
   } else {
-    offerConnectAction();
-    return;
+    await offerConnectAction();
   }
   return conn;
   /* c8 ignore stop */

src/services/connectionManagerService.ts

@@ -146,28 +146,23 @@ export class ConnectionManagementService {
           authCredentials ? authCredentials.split(":") : undefined,
           connection.details.tls,
         );
-        await localConnection.connect((err, conn) => {
-          if (err) {
-            this.connectFailBehaviour(connLabel, err);
-            return;
-          }
-          if (conn) {
-            notify(
-              `Connection established successfully to: ${connLabel}`,
-              MessageKind.DEBUG,
-              {
-                logger,
-                telemetry:
-                  "Connection.Connected" +
-                  this.getTelemetryConnectionType(connLabel),
-              },
-            );
-
-            ext.connectedConnectionList.push(localConnection);
-
-            this.connectSuccessBehaviour(connection);
-          }
-        });
+        const conn = await localConnection.connect();
+        if (conn) {
+          notify(
+            `Connection established successfully to: ${connLabel}`,
+            MessageKind.DEBUG,
+            {
+              logger,
+              telemetry:
+                "Connection.Connected" +
+                this.getTelemetryConnectionType(connLabel),
+            },
+          );
+
+          ext.connectedConnectionList.push(localConnection);
+
+          this.connectSuccessBehaviour(connection);
+        }
       } else {
         ext.context.secrets.delete(connection.details.alias);
         const insightsConn: InsightsConnection = new InsightsConnection(
@@ -212,22 +207,8 @@ export class ConnectionManagementService {
       logger,
       telemetry: "Connection.Connected.Active",
     });
+    commands.executeCommand("setContext", "kdb.pythonEnabled", true);
     ext.activeConnection = connection;
-
-    if (node instanceof InsightsNode) {
-      commands.executeCommand("setContext", "kdb.pythonEnabled", true);
-    } else if (connection instanceof LocalConnection) {
-      // check if pykx namespace is defined
-      connection
-        .execute("`pykx in key`")
-        .then((res) => {
-          commands.executeCommand("setContext", "kdb.pythonEnabled", !!res);
-        })
-        .catch(() => {
-          commands.executeCommand("setContext", "kdb.pythonEnabled", false);
-        });
-    }
-
     ext.connectionNode = node;
     ext.serverProvider.reload();
   }