[TS] Allow brotli to be specified for compression and reorganize some websocket stuff

Author: coolreader18

crates/bindings-typescript/src/sdk/db_connection_builder.ts

@@ -8,6 +8,7 @@ import type {
 } from '../';
 import { ensureMinimumVersionOrThrow } from './version';
 import { WebsocketDecompressAdapter } from './websocket_decompress_adapter';
+import type { WebSocketFactory } from './ws';
 
 /**
  * The database client connection to a SpacetimeDB server.
@@ -23,10 +24,10 @@ export class DbConnectionBuilder<DbConnection extends DbConnectionImpl<any>> {
   #identity?: Identity;
   #token?: string;
   #emitter: EventEmitter<ConnectionEvent> = new EventEmitter();
-  #compression: 'gzip' | 'none' = 'gzip';
+  #compression: 'gzip' | 'brotli' | 'none' = 'gzip';
   #lightMode: boolean = false;
   #confirmedReads?: boolean;
-  #createWSFn: typeof WebsocketDecompressAdapter.createWebSocketFn;
+  #createWSFn: WebSocketFactory;
 
   /**
    * Creates a new `DbConnectionBuilder` database client and set the initial parameters.
@@ -42,7 +43,7 @@ export class DbConnectionBuilder<DbConnection extends DbConnectionImpl<any>> {
       config: DbConnectionConfig<RemoteModuleOf<DbConnection>>
     ) => DbConnection
   ) {
-    this.#createWSFn = WebsocketDecompressAdapter.createWebSocketFn;
+    this.#createWSFn = WebsocketDecompressAdapter.openWebSocket;
   }
 
   /**
@@ -82,9 +83,7 @@ export class DbConnectionBuilder<DbConnection extends DbConnectionImpl<any>> {
     return this;
   }
 
-  withWSFn(
-    createWSFn: typeof WebsocketDecompressAdapter.createWebSocketFn
-  ): this {
+  withWSFn(createWSFn: WebSocketFactory): this {
     this.#createWSFn = createWSFn;
     return this;
   }
@@ -94,7 +93,17 @@ export class DbConnectionBuilder<DbConnection extends DbConnectionImpl<any>> {
    *
    * @param compression The compression algorithm to use for the connection.
    */
-  withCompression(compression: 'gzip' | 'none'): this {
+  withCompression(compression: 'gzip' | 'brotli' | 'none'): this {
+    if (compression === 'brotli') {
+      try {
+        new DecompressionStream('brotli' as CompressionFormat);
+      } catch (e) {
+        throw new TypeError(
+          `Brotli compression is not supported by the runtime. Please choose a different compression method.`,
+          { cause: e }
+        );
+      }
+    }
     this.#compression = compression;
     return this;
   }

crates/bindings-typescript/src/sdk/db_connection_impl.ts

@@ -37,10 +37,6 @@ import {
   type PendingCallback,
   type TableUpdate as CacheTableUpdate,
 } from './table_cache.ts';
-import {
-  WebsocketDecompressAdapter,
-  type WebsocketAdapter,
-} from './websocket_decompress_adapter.ts';
 import {
   SubscriptionBuilderImpl,
   SubscriptionHandleImpl,
@@ -60,6 +56,7 @@ import type { ProceduresView } from './procedures.ts';
 import type { Values } from '../lib/type_util.ts';
 import type { TransactionUpdate } from './client_api/types.ts';
 import { InternalError, SenderError } from '../lib/errors.ts';
+import type { WebSocketAdapter, WebSocketFactory } from './ws.ts';
 
 export {
   DbConnectionBuilder,
@@ -89,8 +86,8 @@ export type DbConnectionConfig<RemoteModule extends UntypedRemoteModule> = {
   identity?: Identity;
   token?: string;
   emitter: EventEmitter<ConnectionEvent>;
-  createWSFn: typeof WebsocketDecompressAdapter.createWebSocketFn;
-  compression: 'gzip' | 'none';
+  createWSFn: WebSocketFactory;
+  compression: 'gzip' | 'brotli' | 'none';
   lightMode: boolean;
   confirmedReads?: boolean;
   remoteModule: RemoteModule;
@@ -173,8 +170,8 @@ export class DbConnectionImpl<RemoteModule extends UntypedRemoteModule>
   // private fields.
   // We use them in testing.
   private clientCache: ClientCache<RemoteModule>;
-  private ws?: WebsocketAdapter;
-  private wsPromise: Promise<WebsocketAdapter | undefined>;
+  private ws?: WebSocketAdapter;
+  private wsPromise: Promise<WebSocketAdapter | undefined>;
 
   constructor({
     uri,
@@ -541,7 +538,7 @@ export class DbConnectionImpl<RemoteModule extends UntypedRemoteModule>
     return this.#mergeTableUpdates(updates);
   }
 
-  #flushOutboundQueue(wsResolved: WebsocketAdapter): void {
+  #flushOutboundQueue(wsResolved: WebSocketAdapter): void {
     const pending = this.#outboundQueue.splice(0);
     for (const message of pending) {
       wsResolved.send(message);

crates/bindings-typescript/src/sdk/decompress.ts

@@ -1,12 +1,11 @@
 export async function decompress(
-  buffer: Uint8Array,
-  // Leaving it here to expand to brotli when it lands in the browsers and NodeJS
-  type: 'gzip',
+  buffer: Uint8Array<ArrayBuffer>,
+  type: CompressionFormat,
   chunkSize: number = 128 * 1024 // 128KB
 ): Promise<Uint8Array> {
   // Create a single ReadableStream to handle chunks
   let offset = 0;
-  const readableStream = new ReadableStream({
+  const readableStream = new ReadableStream<BufferSource>({
     pull(controller) {
       if (offset < buffer.length) {
         // Slice a chunk of the buffer and enqueue it
@@ -29,24 +28,6 @@ export async function decompress(
   const decompressedStream = readableStream.pipeThrough(decompressionStream);
 
   // Collect the decompressed chunks efficiently
-  const reader = decompressedStream.getReader();
-  const chunks: Uint8Array[] = [];
-  let totalLength = 0;
-  let result: any;
-
-  while (!(result = await reader.read()).done) {
-    chunks.push(result.value);
-    totalLength += result.value.length;
-  }
-
-  // Allocate a single Uint8Array for the decompressed data
-  const decompressedArray = new Uint8Array(totalLength);
-  let chunkOffset = 0;
-
-  for (const chunk of chunks) {
-    decompressedArray.set(chunk, chunkOffset);
-    chunkOffset += chunk.length;
-  }
-
-  return decompressedArray;
+  const chunks = await Array.fromAsync(decompressedStream);
+  return new Blob(chunks).bytes();
 }

crates/bindings-typescript/src/sdk/websocket_decompress_adapter.ts

@@ -1,17 +1,7 @@
 import { decompress } from './decompress';
-import { resolveWS } from './ws';
+import { openWebSocket, type WebSocketAdapter, type WebSocketArgs } from './ws';
 
-export interface WebsocketAdapter {
-  send(msg: Uint8Array): void;
-  close(): void;
-
-  set onclose(handler: (ev: CloseEvent) => void);
-  set onopen(handler: () => void);
-  set onmessage(handler: (msg: { data: Uint8Array }) => void);
-  set onerror(handler: (msg: ErrorEvent) => void);
-}
-
-export class WebsocketDecompressAdapter implements WebsocketAdapter {
+export class WebsocketDecompressAdapter implements WebSocketAdapter {
   set onclose(handler: (ev: CloseEvent) => void) {
     this.#ws.onclose = handler;
   }
@@ -30,16 +20,17 @@ export class WebsocketDecompressAdapter implements WebsocketAdapter {
 
   #ws: WebSocket;
 
-  async #decompress(buffer: Uint8Array): Promise<Uint8Array> {
+  async #decompress(buffer: Uint8Array<ArrayBuffer>): Promise<Uint8Array> {
     const tag = buffer[0];
     const data = buffer.subarray(1);
     switch (tag) {
       case 0:
         return data;
       case 1:
-        throw new Error(
-          'Brotli Compression not supported. Please use gzip or none compression in withCompression method on DbConnection.'
-        );
+        // Some runtimes support brotli, but it's not yet defined in `lib.dom.d.ts`.
+        // We assert runtime support in `DbConnectionBuilder.withCompression`, so
+        // this cast is safe.
+        return await decompress(data, 'brotli' as CompressionFormat);
       case 2:
         return await decompress(data, 'gzip');
       default:
@@ -58,68 +49,12 @@ export class WebsocketDecompressAdapter implements WebsocketAdapter {
   }
 
   constructor(ws: WebSocket) {
-    ws.binaryType = 'arraybuffer';
-
     this.#ws = ws;
   }
 
-  static async createWebSocketFn({
-    url,
-    nameOrAddress,
-    wsProtocol,
-    authToken,
-    compression,
-    lightMode,
-    confirmedReads,
-  }: {
-    url: URL;
-    wsProtocol: string;
-    nameOrAddress: string;
-    authToken?: string;
-    compression: 'gzip' | 'none';
-    lightMode: boolean;
-    confirmedReads?: boolean;
-  }): Promise<WebsocketDecompressAdapter> {
-    const headers = new Headers();
-
-    const WS = await resolveWS();
-
-    // We swap our original token to a shorter-lived token
-    // to avoid sending the original via query params.
-    let temporaryAuthToken: string | undefined = undefined;
-    if (authToken) {
-      headers.set('Authorization', `Bearer ${authToken}`);
-      const tokenUrl = new URL('v1/identity/websocket-token', url);
-      tokenUrl.protocol = url.protocol === 'wss:' ? 'https:' : 'http:';
-
-      const response = await fetch(tokenUrl, { method: 'POST', headers });
-      if (response.ok) {
-        const { token } = await response.json();
-        temporaryAuthToken = token;
-      } else {
-        return Promise.reject(
-          new Error(`Failed to verify token: ${response.statusText}`)
-        );
-      }
-    }
-
-    const databaseUrl = new URL(`v1/database/${nameOrAddress}/subscribe`, url);
-    if (temporaryAuthToken) {
-      databaseUrl.searchParams.set('token', temporaryAuthToken);
-    }
-    databaseUrl.searchParams.set(
-      'compression',
-      compression === 'gzip' ? 'Gzip' : 'None'
-    );
-    if (lightMode) {
-      databaseUrl.searchParams.set('light', 'true');
-    }
-    if (confirmedReads !== undefined) {
-      databaseUrl.searchParams.set('confirmed', confirmedReads.toString());
-    }
-
-    const ws = new WS(databaseUrl.toString(), wsProtocol);
-
-    return new WebsocketDecompressAdapter(ws);
+  static async openWebSocket(
+    args: WebSocketArgs
+  ): Promise<WebsocketDecompressAdapter> {
+    return new this(await openWebSocket(args));
   }
 }

crates/bindings-typescript/src/sdk/websocket_test_adapter.ts

@@ -1,8 +1,8 @@
 import { BinaryReader, BinaryWriter } from '../';
 import { ClientMessage, ServerMessage } from './client_api/types';
-import type { WebsocketAdapter } from './websocket_decompress_adapter';
+import type { WebSocketAdapter, WebSocketFactory } from './ws';
 
-class WebsocketTestAdapter implements WebsocketAdapter {
+class WebsocketTestAdapter implements WebSocketAdapter {
   onclose: any;
   // eslint-disable-next-line @typescript-eslint/no-unsafe-function-type
   onopen!: () => void;
@@ -47,18 +47,7 @@ class WebsocketTestAdapter implements WebsocketAdapter {
     this.onmessage({ data: rawBytes });
   }
 
-  async createWebSocketFn(_args: {
-    url: URL;
-    wsProtocol: string;
-    nameOrAddress: string;
-    authToken?: string;
-    compression: 'gzip' | 'none';
-    lightMode: boolean;
-    confirmedReads?: boolean;
-  }): Promise<WebsocketTestAdapter> {
-    return this;
-  }
+  openWebSocket: WebSocketFactory = async () => this;
 }
 
-export type { WebsocketTestAdapter };
 export default WebsocketTestAdapter;

crates/bindings-typescript/src/sdk/ws.ts

@@ -1,9 +1,9 @@
 import { stdbLogger } from './logger';
 
-export async function resolveWS(): Promise<typeof WebSocket> {
+async function resolveWS(): Promise<typeof WebSocket> {
   // Browser or Node >= 22 (or any env that exposes global WebSocket)
-  if (typeof (globalThis as any).WebSocket !== 'undefined') {
-    return (globalThis as any).WebSocket as typeof WebSocket;
+  if (typeof WebSocket !== 'undefined') {
+    return WebSocket;
   }
 
   // Node without a global WebSocket: lazily load undici's polyfill.
@@ -25,3 +25,80 @@ export async function resolveWS(): Promise<typeof WebSocket> {
     throw err;
   }
 }
+
+export interface WebSocketAdapter {
+  send(msg: Uint8Array): void;
+  close(): void;
+
+  set onclose(handler: (ev: CloseEvent) => void);
+  set onopen(handler: () => void);
+  set onmessage(handler: (msg: { data: Uint8Array }) => void);
+  set onerror(handler: (msg: ErrorEvent) => void);
+}
+
+export interface WebSocketArgs {
+  url: URL;
+  wsProtocol: string;
+  nameOrAddress: string;
+  authToken?: string;
+  compression: 'gzip' | 'brotli' | 'none';
+  lightMode: boolean;
+  confirmedReads?: boolean;
+}
+export type WebSocketFactory = (
+  args: WebSocketArgs
+) => Promise<WebSocketAdapter>;
+
+/**
+ * Open a WebSocket to the database specified by the given `WebSocketArgs`.
+ * @returns a WebSocket with `binaryType` set to `arraybuffer`.
+ */
+export async function openWebSocket({
+  url,
+  nameOrAddress,
+  wsProtocol,
+  authToken,
+  compression,
+  lightMode,
+  confirmedReads,
+}: WebSocketArgs): Promise<WebSocket> {
+  const headers = new Headers();
+
+  const WS = await resolveWS();
+
+  // We swap our original token to a shorter-lived token
+  // to avoid sending the original via query params.
+  let temporaryAuthToken: string | undefined;
+  if (authToken) {
+    headers.set('Authorization', `Bearer ${authToken}`);
+    const tokenUrl = new URL('v1/identity/websocket-token', url);
+    tokenUrl.protocol = url.protocol === 'wss:' ? 'https:' : 'http:';
+
+    const response = await fetch(tokenUrl, { method: 'POST', headers });
+    if (response.ok) {
+      const { token } = await response.json();
+      temporaryAuthToken = token;
+    } else {
+      throw new Error(`Failed to verify token: ${response.statusText}`);
+    }
+  }
+
+  const databaseUrl = new URL(`v1/database/${nameOrAddress}/subscribe`, url);
+  if (temporaryAuthToken) {
+    databaseUrl.searchParams.set('token', temporaryAuthToken);
+  }
+  databaseUrl.searchParams.set(
+    'compression',
+    { gzip: 'Gzip', brotli: 'Brotli', none: 'None' }[compression] ?? 'None'
+  );
+  if (lightMode) {
+    databaseUrl.searchParams.set('light', 'true');
+  }
+  if (confirmedReads !== undefined) {
+    databaseUrl.searchParams.set('confirmed', confirmedReads.toString());
+  }
+
+  const ws = new WS(databaseUrl.toString(), wsProtocol);
+  ws.binaryType = 'arraybuffer';
+  return ws;
+}