fix: DH-21024: Grizzly worker kind fixes (#266)

DH-21024: Grizzly worker kind fixes
This PR includes a couple of fixes + configuration improvements:
- Worker kinds are now filtered to only use Core+ workers. This fixes
bug where legacy worker sorted first would get picked
- Fixed a bug where `operateAs` was redeclared in an if block instead of
being set in the outer. This was causing operateAs to be ignored. Added
respective unit tests
- Added additional `experimentalWorkerConfig` settings:
   - `additionalMemory`
   - `classPaths`
   - `engine`
   - `envVars`
- I did not add support for `enableGcLogs` or `python environment` as
they require additional logic I didn't want to have to replicate from
the web

### Testing Config
Configure a grizzly server such that legacy workers are sorted to
beginning of worker kind list. See description in the Jira or you can
use https://bmingles-f3.int.illumon.com:8123/iriside/ is already which
is already setup for this.

#### VS Code `deephaven.enterpriseServers` settings for different tests
- Default config
```json
{
    "url": "https://bmingles-f3.int.illumon.com:8123/",
    "experimentalWorkerConfig": {
      "heapSize": 0.5
    }
},
```
Connecting to server should successfully create a connection
- Named config
```json
{
    "url": "https://bmingles-f3.int.illumon.com:8123/",
    "experimentalWorkerConfig": {
      "engine": "Core+",
      "heapSize": 0.5
    }
},
```
Connecting to server should successfully create a connection
- Additional config
```json
{
    "url": "https://bmingles-f3.int.illumon.com:8123/",
    "experimentalWorkerConfig": {
      "additionalMemory": 1,
      "classPaths": "",
      "envVars": "FOO=BAR",
      "heapSize": 0.5
    }
},
```
   - Connecting to server should successfully create a connection.
   - additionalMemory can be seen in query monitor column
   - envVars
   ```py
   import os
   print(os.environ.get('FOO'))
   ```
   - classPaths - not sure how to test this one

### Testing operateAs
- Disconnect from grizzly server if connect.
- Right-click on the server in VS Code, and `Connect to Server as
Another User`
- Login as iris, and give another user to operateAs
- Should see the `operateAs` user in the `owner` column of query monitor
for the interactive console query that gets created

Author: bmingles

__mocks__/vscode.ts

@@ -178,6 +178,7 @@ export const window = {
   registerFileDecorationProvider: vi
     .fn()
     .mockName('registerFileDecorationProvider'),
+  showInputBox: vi.fn().mockName('showInputBox'),
 };
 
 export const workspace = {

package.json

@@ -110,9 +110,23 @@
                     "type": "object",
                     "description": "(experimental) Worker configuration used when creating new connections to the server",
                     "properties": {
+                      "additionalMemory": {
+                        "type": "number",
+                        "description": "(GB) The amount of additional memory allocated to this worker process beyond the JVM heap memory. This allows Deephaven to account for memory allocated by Python native code, Java direct memory, and other types of non-heap memory."
+                      },
+                      "classPaths": {
+                        "type": "string"
+                      },
                       "dbServerName": {
                         "type": "string"
                       },
+                      "engine": {
+                        "type": "string",
+                        "description": "Core+ engine configuration name"
+                      },
+                      "envVars": {
+                        "type": "string"
+                      },
                       "heapSize": {
                         "type": "number"
                       },

releases/vscode-deephaven-latest.vsix

@@ -192,3 +192,11 @@ export const DH_PYTHON_REMOTE_SOURCE_PLUGIN_NAME =
 export const DH_PROTECTED_VARIABLE_NAMES: Set<string> = new Set([
   DH_PYTHON_REMOTE_SOURCE_PLUGIN_VARIABLE,
 ]);
+
+export const PROTOCOL = {
+  /* eslint-disable @typescript-eslint/naming-convention */
+  COMMUNITY: 'Community',
+  ENTERPRISE_COMM: 'EnterpriseComm',
+  ENTERPRISE_WEBSOCKET: 'EnterpriseWebsocket',
+  /* eslint-enable @typescript-eslint/naming-convention */
+} as const;

src/common/constants.ts

@@ -326,7 +326,7 @@ export class UserLoginController extends ControllerBase {
           title,
           userLoginPreferences,
           privateKeyUserNames,
-          showOperatesAs: operateAsAnotherUser,
+          showOperateAs: operateAsAnotherUser,
         });
 
         // Cancelled by user

src/controllers/UserLoginController.ts

@@ -5,6 +5,7 @@ import type {
   EnterpriseClient,
   QueryInfo,
   TypeSpecificFields,
+  WorkerKind,
 } from '@deephaven-enterprise/jsapi-types';
 import { DraftQuery, QueryScheduler } from '@deephaven-enterprise/query-utils';
 import type {
@@ -36,6 +37,7 @@ import {
   DHE_FEATURES_URL_PATH,
   INTERACTIVE_CONSOLE_QUERY_TYPE,
   INTERACTIVE_CONSOLE_TEMPORARY_QUEUE_NAME,
+  PROTOCOL,
   UnsupportedFeatureQueryError,
 } from '../common';
 import { withResolvers } from '../util';
@@ -203,7 +205,7 @@ export async function createInteractiveConsoleQuery(
   consoleType?: ConsoleType
 ): Promise<QuerySerial> {
   const userInfo = await dheClient.getUserInfo();
-  const owner = userInfo.username;
+  const owner = userInfo.operateAs;
   const type = INTERACTIVE_CONSOLE_QUERY_TYPE;
   const queueName = INTERACTIVE_CONSOLE_TEMPORARY_QUEUE_NAME;
   const autoDeleteTimeoutMs = DEFAULT_TEMPORARY_QUERY_AUTO_TIMEOUT_MS;
@@ -215,6 +217,27 @@ export async function createInteractiveConsoleQuery(
     dheClient.getServerConfigValues(),
   ]);
 
+  const coreWorkerKinds = serverConfigValues.workerKinds.filter(
+    isCommunityWorkerKind
+  );
+
+  const engine = (
+    workerConfig.engine == null
+      ? coreWorkerKinds[0]
+      : coreWorkerKinds.find(({ name, title }) =>
+          // Match on title first since it's the most likely to be configured by
+          // a user, but fallback to name just in case someone is using the actual
+          // worker kind name.
+          [title, name].includes(workerConfig.engine!)
+        )
+  )?.name;
+
+  if (engine == null) {
+    throw new Error(
+      `No Core+ engine configurations were found${workerConfig.engine == null ? '' : ` matching '${workerConfig.engine}'`}.`
+    );
+  }
+
   const name = createQueryName(tagId);
   const dbServerName =
     workerConfig?.dbServerName ?? dbServers[0]?.name ?? 'Query 1';
@@ -230,13 +253,13 @@ export async function createInteractiveConsoleQuery(
 
   const jvmProfile =
     workerConfig?.jvmProfile ?? serverConfigValues.jvmProfileDefault;
+
   const scriptLanguage =
     workerConfig?.scriptLanguage ??
     serverConfigValues.scriptSessionProviders?.find(
       p => p.toLocaleLowerCase() === consoleType
     ) ??
     'Python';
-  const workerKind = serverConfigValues.workerKinds?.[0]?.name;
 
   const autoDelete = autoDeleteTimeoutMs > 0;
 
@@ -256,21 +279,26 @@ export async function createInteractiveConsoleQuery(
     queueName,
   });
 
+  const { additionalMemory, classPaths, envVars } = workerConfig;
+
   const draftQuery = new DraftQuery({
-    isClientSide: true,
-    draftOwner: owner,
-    name,
-    type,
-    owner,
+    additionalMemory: additionalMemory ?? 0,
     dbServerName,
+    draftOwner: owner,
+    envVars: envVars ?? '',
+    extraClasspaths: classPaths ?? '',
     heapSize,
-    scheduling,
+    isClientSide: true,
     jvmArgs,
     jvmProfile,
+    name,
+    owner,
+    scheduling,
     scriptLanguage,
     timeout,
+    type,
     typeSpecificFields,
-    workerKind,
+    workerKind: engine,
   });
 
   if (draftQuery.serial == null) {
@@ -461,6 +489,14 @@ export async function getWorkerInfoFromQuery(
   };
 }
 
+/**
+ * Determine if a given worker kind is a Community worker.
+ * @returns True if the worker kind is a Community worker, false otherwise.
+ */
+export function isCommunityWorkerKind(workerKind?: WorkerKind): boolean {
+  return workerKind?.protocols?.includes(PROTOCOL.COMMUNITY) ?? false;
+}
+
 /**
  * Polyfill some things needed by the DHE jsapi. The need for this should go
  * away once DH-17942 is completed and the jsapi no longer relies on `window`

src/dh/dhe.ts

@@ -130,7 +130,11 @@ export type ServerConnectionConfig =
   | URL;
 
 export interface WorkerConfig {
+  additionalMemory?: number;
+  classPaths?: string;
   dbServerName?: string;
+  engine?: string;
+  envVars?: string;
   heapSize?: number;
   jvmArgs?: string;
   jvmProfile?: string;

src/types/commonTypes.d.ts

@@ -1,12 +1,14 @@
 import * as vscode from 'vscode';
-import { describe, it, expect, vi } from 'vitest';
+import { beforeEach, describe, it, expect, vi } from 'vitest';
+import type { Username } from '@deephaven-enterprise/auth-nodejs';
 import {
   createConnectText,
   ConnectionOption,
   createConnectionOption,
   updateConnectionStatusBarItem,
   createConnectionQuickPickOptions,
   createSeparatorPickItem,
+  promptForCredentials,
 } from './uiUtils';
 import type {
   ConnectionState,
@@ -28,6 +30,10 @@ const groovyServerConfig: CoreConnectionConfig = {
   url: new URL('http://localhost:10001'),
 };
 
+beforeEach(() => {
+  vi.clearAllMocks();
+});
+
 describe('createConnectionOption', () => {
   it.each([
     ['DHC', pythonServerConfig],
@@ -138,3 +144,99 @@ describe('createSeparatorPickItem', () => {
     });
   });
 });
+
+describe('promptForCredentials', () => {
+  const title = 'mock.title';
+  const username = 'mock.username' as Username;
+  const token = 'mock.token';
+  const operateAs = 'mock.operateAs';
+
+  it.each([
+    [
+      'password',
+      { title },
+      [username, token],
+      { type: 'password', token, username },
+    ],
+    ['password / cancelled username', { title }, [undefined], undefined],
+    ['password / cancelled token', { title }, [username, undefined], undefined],
+    [
+      'password operateAs',
+      { title, showOperateAs: true },
+      [username, token, operateAs],
+      { type: 'password', operateAs, token, username },
+    ],
+    [
+      'password operateAs / cancelled username',
+      { title, showOperateAs: true },
+      [undefined],
+      undefined,
+    ],
+    [
+      'password operateAs / cancelled token',
+      { title, showOperateAs: true },
+      [username, undefined],
+      undefined,
+    ],
+    [
+      'password operateAs / cancelled operateAs',
+      { title, showOperateAs: true },
+      [username, token, undefined],
+      undefined,
+    ],
+    [
+      'privateKey',
+      { title, privateKeyUserNames: [username] },
+      [username],
+      { type: 'keyPair', username },
+    ],
+    [
+      'privateKey / cancelled username',
+      { title, privateKeyUserNames: [username] },
+      [undefined],
+      undefined,
+    ],
+    [
+      'privateKey operateAs',
+      {
+        title,
+        privateKeyUserNames: [username],
+        showOperateAs: true,
+      },
+      [username, operateAs],
+      { type: 'keyPair', username, operateAs },
+    ],
+    [
+      'privateKey operateAs / cancelled username',
+      {
+        title,
+        privateKeyUserNames: [username],
+        showOperateAs: true,
+      },
+      [undefined],
+      undefined,
+    ],
+    [
+      'privateKey operateAs / cancelled operateAs',
+      {
+        title,
+        privateKeyUserNames: [username],
+        showOperateAs: true,
+      },
+      [username, undefined],
+      undefined,
+    ],
+  ])(
+    'should prompt for username, password, and operateAs: %s',
+    async (_label, arg, promptResponses, expected) => {
+      for (const promptResponse of promptResponses) {
+        vi.mocked(vscode.window.showInputBox).mockResolvedValueOnce(
+          promptResponse
+        );
+      }
+
+      const actual = await promptForCredentials(arg);
+      expect(actual).toEqual(expected);
+    }
+  );
+});

src/util/uiUtils.spec.ts

@@ -168,31 +168,31 @@ export async function promptForAuthFlow(
  * @param privateKeyUserNames Optional list of private key user names. If provided,
  * the authentication method will be prompted to determine if user wants to use
  * one of these private keys or username/password.
- * @param showOperatesAs Whether to show the operate as prompt.
+ * @param showOperateAs Whether to show the operate as prompt.
  */
 export async function promptForCredentials(args: {
   title: string;
   userLoginPreferences?: UserLoginPreferences;
   privateKeyUserNames?: undefined | [];
-  showOperatesAs?: boolean;
+  showOperateAs?: boolean;
 }): Promise<PasswordCredentials | undefined>;
 export async function promptForCredentials(args: {
   title: string;
   userLoginPreferences?: UserLoginPreferences;
   privateKeyUserNames?: Username[];
-  showOperatesAs?: boolean;
+  showOperateAs?: boolean;
 }): Promise<LoginPromptCredentials | undefined>;
 export async function promptForCredentials(args: {
   title: string;
   userLoginPreferences?: UserLoginPreferences;
   privateKeyUserNames?: Username[];
-  showOperatesAs?: boolean;
+  showOperateAs?: boolean;
 }): Promise<LoginPromptCredentials | undefined> {
   const {
     title,
     userLoginPreferences,
     privateKeyUserNames = [],
-    showOperatesAs,
+    showOperateAs,
   } = args;
 
   const username = await promptForUsername(
@@ -220,10 +220,10 @@ export async function promptForCredentials(args: {
   }
 
   // Operate As
-  if (showOperatesAs) {
+  if (showOperateAs) {
     const defaultValue = username as unknown as OperateAsUsername | undefined;
 
-    const operateAs = await promptForOperateAs(
+    operateAs = await promptForOperateAs(
       title,
       userLoginPreferences?.operateAsUser[username] ?? defaultValue
     );