First pass at DHE connections (#79)

packages/require-jsapi/src/dhc.ts

@@ -4,6 +4,9 @@ import type { dh as DhType } from '@deephaven/jsapi-types';
 import { polyfillDh } from './polyfill';
 import { downloadFromURL, hasStatusCode } from './serverUtils';
 
+export const AUTH_HANDLER_TYPE_DHE =
+  'io.deephaven.enterprise.dnd.authentication.DheAuthenticationHandler';
+
 /**
  * Check if a given server is running by checking if the `dh-core.js` file is
  * accessible.

packages/require-jsapi/src/dhe.ts

@@ -1,4 +1,10 @@
 import { hasStatusCode } from './serverUtils';
+import * as fs from 'node:fs';
+import * as path from 'node:path';
+import type {
+  EnterpriseDhType as DheType,
+  EnterpriseClient,
+} from '@deephaven-enterprise/jsapi-types';
 
 /**
  * Check if a given server is running by checking if the `irisapi/irisapi.nocache.js`
@@ -15,3 +21,69 @@ export async function isDheServerRunning(serverUrl: URL): Promise<boolean> {
     return false;
   }
 }
+
+export async function createDheClient(
+  dhe: DheType,
+  serverUrl: URL
+): Promise<EnterpriseClient> {
+  const dheClient = new dhe.Client(serverUrl.toString());
+
+  return new Promise(resolve => {
+    const unsubscribe = dheClient.addEventListener(
+      dhe.Client.EVENT_CONNECT,
+      () => {
+        unsubscribe();
+        resolve(dheClient);
+      }
+    );
+  });
+}
+
+export async function initDheApi(serverUrl: URL): Promise<DheType> {
+  polyfillDh();
+  return getDhe(serverUrl, true);
+}
+
+declare global {
+  export const iris: DheType;
+}
+
+export async function getDhe(
+  serverUrl: URL,
+  download: boolean
+): Promise<DheType> {
+  const tmpDir = getTempDir(false, urlToDirectoryName(serverUrl));
+  const dheFilePath = path.join(tmpDir, 'irisapi.nocache.js');
+
+  if (download) {
+    const dhe = await downloadFromURL(
+      path.join(serverUrl.toString(), 'irisapi/irisapi.nocache.js')
+    );
+
+    fs.writeFileSync(dheFilePath, dhe);
+  }
+
+  require(dheFilePath);
+
+  return iris;
+}
+
+export async function getDheAuthToken(
+  client: EnterpriseClient
+): Promise<{ type: string; token: string }> {
+  const token = await client.createAuthToken('RemoteQueryProcessor');
+  return {
+    type: 'io.deephaven.proto.auth.Token',
+    token,
+  };
+}
+
+export function getWsUrl(serverUrl: URL): URL {
+  const url = new URL('/socket', serverUrl);
+  if (url.protocol === 'http:') {
+    url.protocol = 'ws:';
+  } else {
+    url.protocol = 'wss:';
+  }
+  return url;
+}

src/controllers/ExtensionController.ts

@@ -1,4 +1,6 @@
 import * as vscode from 'vscode';
+import type { dh as DhcType } from '@deephaven/jsapi-types';
+import type { EnterpriseDhType as DheType } from '@deephaven-enterprise/jsapi-types';
 import {
   CONNECT_TO_SERVER_CMD,
   CREATE_NEW_TEXT_DOC_CMD,
@@ -32,9 +34,16 @@ import {
   ServerConnectionPanelTreeProvider,
   runSelectedLinesHoverProvider,
 } from '../providers';
-import { DhcServiceFactory, PanelService, ServerManager } from '../services';
+import {
+  CacheByUrlService,
+  DhcServiceFactory,
+  PanelService,
+  ServerManager,
+  URLMap,
+} from '../services';
 import type {
   Disposable,
+  ICacheService,
   IConfigService,
   IDhService,
   IDhServiceFactory,
@@ -51,6 +60,7 @@ import { ServerConnectionTreeDragAndDropController } from './ServerConnectionTre
 import { ConnectionController } from './ConnectionController';
 import { PipServerController } from './PipServerController';
 import { PanelController } from './PanelController';
+import { initDheApi } from '../dh/dhe';
 
 const logger = new Logger('ExtensionController');
 
@@ -84,10 +94,12 @@ export class ExtensionController implements Disposable {
   readonly _config: IConfigService;
 
   private _connectionController: ConnectionController | null = null;
+  private _credentialsCache: URLMap<DhcType.LoginCredentials> | null = null;
   private _panelController: PanelController | null = null;
   private _panelService: IPanelService | null = null;
   private _pipServerController: PipServerController | null = null;
   private _dhcServiceFactory: IDhServiceFactory | null = null;
+  private _dheJsApiCache: ICacheService<URL, DheType> | null = null;
   private _serverManager: IServerManager | null = null;
 
   // Tree providers
@@ -243,19 +255,27 @@ export class ExtensionController implements Disposable {
     assertDefined(this._outputChannel, 'outputChannel');
     assertDefined(this._toaster, 'toaster');
 
+    this._credentialsCache = new URLMap<DhcType.LoginCredentials>();
+
     this._panelService = new PanelService();
     this._context.subscriptions.push(this._panelService);
 
     this._dhcServiceFactory = new DhcServiceFactory(
+      this._credentialsCache,
       this._panelService,
       this._pythonDiagnostics,
       this._outputChannel,
       this._toaster
     );
 
+    this._dheJsApiCache = new CacheByUrlService(initDheApi);
+    this._context.subscriptions.push(this._dheJsApiCache);
+
     this._serverManager = new ServerManager(
       this._config,
-      this._dhcServiceFactory
+      this._credentialsCache,
+      this._dhcServiceFactory,
+      this._dheJsApiCache
     );
     this._context.subscriptions.push(this._serverManager);
 

src/services/CacheByUrlService.ts

@@ -0,0 +1,31 @@
+import type { ICacheService } from '../types';
+import { URLMap } from './URLMap';
+
+/**
+ * Cache service that stores values by URL.
+ */
+export class CacheByUrlService<TValue> implements ICacheService<URL, TValue> {
+  constructor(loader: (url: URL) => Promise<TValue>) {
+    this._loader = loader;
+    this._promiseMap = new URLMap<Promise<TValue>>();
+  }
+
+  private readonly _loader: (url: URL) => Promise<TValue>;
+  private readonly _promiseMap = new URLMap<Promise<TValue>>();
+
+  dispose = async (): Promise<void> => {
+    this._promiseMap.clear();
+  };
+
+  get = async (url: URL): Promise<TValue> => {
+    if (!this._promiseMap.has(url)) {
+      this._promiseMap.set(url, this._loader(url));
+    }
+
+    return this._promiseMap.get(url)!;
+  };
+
+  invalidate = (url: URL): void => {
+    this._promiseMap.delete(url);
+  };
+}

src/services/DhService.ts

@@ -19,6 +19,7 @@ import {
 } from '../common';
 import type { ConnectionAndSession } from '../dh/dhc';
 import { NoConsoleTypesError, parseServerError } from '../dh/errorUtils';
+import type { URLMap } from './URLMap';
 
 const logger = new Logger('DhService');
 
@@ -27,11 +28,13 @@ export abstract class DhService<TDH = unknown, TClient = unknown>
 {
   constructor(
     serverUrl: URL,
+    credentialsCache: URLMap<DhcType.LoginCredentials>,
     panelService: IPanelService,
     diagnosticsCollection: vscode.DiagnosticCollection,
     outputChannel: vscode.OutputChannel,
     toaster: IToastService
   ) {
+    this.credentialsCache = credentialsCache;
     this.serverUrl = serverUrl;
     this.panelService = panelService;
     this.diagnosticsCollection = diagnosticsCollection;
@@ -45,6 +48,7 @@ export abstract class DhService<TDH = unknown, TClient = unknown>
   public readonly serverUrl: URL;
   protected readonly subscriptions: (() => void)[] = [];
 
+  protected readonly credentialsCache: URLMap<DhcType.LoginCredentials>;
   protected readonly outputChannel: vscode.OutputChannel;
   protected readonly toaster: IToastService;
   private readonly panelService: IPanelService;

src/services/DhcService.ts

@@ -13,14 +13,14 @@ import {
 const logger = new Logger('DhcService');
 
 export class DhcService extends DhService<typeof DhcType, DhcType.CoreClient> {
-  private _psk?: string;
+  getPsk(): string | null {
+    const credentials = this.credentialsCache.get(this.serverUrl);
 
-  getPsk(): string | undefined {
-    return this._psk;
-  }
+    if (credentials?.type !== AUTH_HANDLER_TYPE_PSK) {
+      return null;
+    }
 
-  setPsk(psk: string): void {
-    this._psk = psk;
+    return credentials.token ?? null;
   }
 
   protected async initApi(): Promise<typeof DhcType> {
@@ -45,29 +45,29 @@ export class DhcService extends DhService<typeof DhcType, DhcType.CoreClient> {
     dh: typeof DhcType,
     client: DhcType.CoreClient
   ): Promise<ConnectionAndSession<DhcType.IdeConnection, DhcType.IdeSession>> {
-    const authConfig = new Set(
-      (await client.getAuthConfigValues()).map(([, value]) => value)
-    );
+    if (!this.credentialsCache.has(this.serverUrl)) {
+      const authConfig = new Set(
+        (await client.getAuthConfigValues()).map(([, value]) => value)
+      );
 
-    if (authConfig.has(AUTH_HANDLER_TYPE_ANONYMOUS)) {
-      return initDhcSession(client, {
-        type: dh.CoreClient.LOGIN_TYPE_ANONYMOUS,
-      });
-    } else if (authConfig.has(AUTH_HANDLER_TYPE_PSK)) {
-      if (this._psk == null) {
-        this._psk = await vscode.window.showInputBox({
-          placeHolder: 'Pre-Shared Key',
-          prompt: 'Enter your Deephaven pre-shared key',
-          password: true,
+      if (authConfig.has(AUTH_HANDLER_TYPE_ANONYMOUS)) {
+        this.credentialsCache.set(this.serverUrl, {
+          type: dh.CoreClient.LOGIN_TYPE_ANONYMOUS,
+        });
+      } else if (authConfig.has(AUTH_HANDLER_TYPE_PSK)) {
+        this.credentialsCache.set(this.serverUrl, {
+          type: AUTH_HANDLER_TYPE_PSK,
+          token: await vscode.window.showInputBox({
+            placeHolder: 'Pre-Shared Key',
+            prompt: 'Enter your Deephaven pre-shared key',
+            password: true,
+          }),
         });
       }
+    }
 
-      const connectionAndSession = await initDhcSession(client, {
-        type: 'io.deephaven.authentication.psk.PskAuthenticationHandler',
-        token: this._psk,
-      });
-
-      return connectionAndSession;
+    if (this.credentialsCache.has(this.serverUrl)) {
+      return initDhcSession(client, this.credentialsCache.get(this.serverUrl)!);
     }
 
     throw new Error('No supported authentication methods found.');

src/services/DhcServiceFactory.ts

@@ -1,31 +1,31 @@
 import * as vscode from 'vscode';
+import type { dh as DhcType } from '@deephaven/jsapi-types';
 import { DhcService } from './DhcService';
 import type { IDhServiceFactory, IPanelService, IToastService } from '../types';
+import type { URLMap } from './URLMap';
 
 /**
  * Factory for creating DhcService instances.
  */
 export class DhcServiceFactory implements IDhServiceFactory {
   constructor(
+    private credentialsCache: URLMap<DhcType.LoginCredentials>,
     private panelService: IPanelService,
     private diagnosticsCollection: vscode.DiagnosticCollection,
     private outputChannel: vscode.OutputChannel,
     private toaster: IToastService
   ) {}
 
-  create = (serverUrl: URL, psk?: string): DhcService => {
+  create = (serverUrl: URL): DhcService => {
     const dhService = new DhcService(
       serverUrl,
+      this.credentialsCache,
       this.panelService,
       this.diagnosticsCollection,
       this.outputChannel,
       this.toaster
     );
 
-    if (psk != null) {
-      dhService.setPsk(psk);
-    }
-
     return dhService;
   };
 }