feat: enhance session recovery and retry configuration handling in WebSocket client

Author: qevolg

nodejs/src/client/wsClient.ts

@@ -91,11 +91,20 @@ export class WsClient {
         return this._dsn.endpoint === WS_SQL_ENDPOINT;
     }
 
+    private normalizeConnectedDatabase(database?: string | null): string | null {
+        if (database && database.length > 0) {
+            return database;
+        }
+        return this.isSqlPath() ? "information_schema" : null;
+    }
+
     private async recoverSqlSessionContext(): Promise<void> {
         if (!this._wsConnector) {
             return;
         }
-        const connMsg = this.buildConnMessage(this._connectedDatabase);
+        const connMsg = this.buildConnMessage(
+            this.normalizeConnectedDatabase(this._connectedDatabase)
+        );
         await this.sendMsgDirect(JSON.stringify(connMsg), false);
 
         if (this._connectionOptions.size <= 0) {
@@ -138,14 +147,14 @@ export class WsClient {
         );
         this.bindReconnectRecoveryHook();
         if (this._wsConnector.readyState() === w3cwebsocket.OPEN) {
-            this._connectedDatabase = database ?? null;
+            this._connectedDatabase = this.normalizeConnectedDatabase(database ?? null);
             return;
         }
         try {
             await this._wsConnector.ready();
             let result: any = await this._wsConnector.sendMsg(JSON.stringify(connMsg));
             if (result.msg.code == 0) {
-                this._connectedDatabase = database ?? null;
+                this._connectedDatabase = this.normalizeConnectedDatabase(database ?? null);
                 return;
             }
             await this.close();

nodejs/src/client/wsConnector.ts

@@ -155,7 +155,7 @@ export class WebSocketConnector {
     private readonly _poolKey: string;
     private readonly _dsn: Dsn;
     private _currentAddress: Address;
-    private readonly _retryConfig: RetryConfig;
+    private _retryConfig: RetryConfig;
     private _inflightStore: InflightRequestStore;
     private readonly _suppressedSockets: WeakSet<w3cwebsocket> = new WeakSet();
     private _reconnectLock: Promise<void> | null = null;
@@ -188,6 +188,10 @@ export class WebSocketConnector {
         this.createConnection();
     }
 
+    public refreshRetryConfig(dsn: Dsn): void {
+        this._retryConfig = RetryConfig.fromDsn(dsn);
+    }
+
     private buildUrl(addr: Address): string {
         const path = this._dsn.path();
         const url = new URL(`${this._dsn.scheme}://${addr.host}:${addr.port}/${path}`);

nodejs/src/client/wsConnectorPool.ts

@@ -33,10 +33,15 @@ export class WebSocketConnectionPool {
         return poolKey.replace(/#auth=[^#]*/, "#auth=[REDACTED]");
     }
 
-    private buildAuth(dsn: Dsn): string {
+    private buildAuthScope(dsn: Dsn): string {
         const token = dsn.params.get("token") || "";
         const bearerToken = dsn.params.get("bearer_token") || "";
-        const raw = `${dsn.username}:${dsn.password}:${token}:${bearerToken}`;
+        const raw = JSON.stringify({
+            username: dsn.username,
+            password: dsn.password,
+            token,
+            bearerToken,
+        });
         return createHash("sha256").update(raw).digest("hex");
     }
 
@@ -45,7 +50,7 @@ export class WebSocketConnectionPool {
             .sort((a, b) => `${a.host}:${a.port}`.localeCompare(`${b.host}:${b.port}`))
             .map((addr) => `${addr.host}:${addr.port}`)
             .join(",");
-        const auth = this.buildAuth(dsn);
+        const auth = this.buildAuthScope(dsn);
         const path = dsn.path();
         return `${dsn.scheme}://${addrs}/${path}#auth=${auth}`;
     }
@@ -67,6 +72,7 @@ export class WebSocketConnectionPool {
                         continue;
                     }
                     if (candidate.readyState() === w3cwebsocket.OPEN) {
+                        candidate.refreshRetryConfig(dsn);
                         connector = candidate;
                         break;
                     }

nodejs/src/common/dsn.ts

@@ -102,13 +102,27 @@ export function parse(url: string): Dsn {
     const scheme = schemeMatch[1].toLowerCase();
     let remainder = url.slice(schemeMatch[0].length);
 
-    // Extract username:password@ if present
+    // Isolate authority from path/query before parsing userinfo.
+    const slashIndex = remainder.indexOf("/");
+    const queryMarkIndex = remainder.indexOf("?");
+    let authorityEndIndex = remainder.length;
+    if (slashIndex !== -1) {
+        authorityEndIndex = Math.min(authorityEndIndex, slashIndex);
+    }
+    if (queryMarkIndex !== -1) {
+        authorityEndIndex = Math.min(authorityEndIndex, queryMarkIndex);
+    }
+    const authority = remainder.slice(0, authorityEndIndex);
+    const suffix = remainder.slice(authorityEndIndex);
+
+    // Extract username:password@ from authority only.
     let username = "";
     let password = "";
-    const atIndex = remainder.indexOf("@");
+    const atIndex = authority.lastIndexOf("@");
+    let hostPort = authority;
     if (atIndex !== -1) {
-        const userInfo = remainder.slice(0, atIndex);
-        remainder = remainder.slice(atIndex + 1);
+        const userInfo = authority.slice(0, atIndex);
+        hostPort = authority.slice(atIndex + 1);
         const colonIndex = userInfo.indexOf(":");
         if (colonIndex !== -1) {
             username = userInfo.slice(0, colonIndex);
@@ -117,6 +131,7 @@ export function parse(url: string): Dsn {
             username = userInfo;
         }
     }
+    remainder = `${hostPort}${suffix}`;
 
     // Extract query params (after ?)
     let params = new Map<string, string>();

nodejs/test/client/wsClient.recovery.test.ts

@@ -115,4 +115,26 @@ describe("WsClient recovery hook", () => {
         await expect(hook()).rejects.toBeInstanceOf(WebSocketQueryError);
         expect(customRecoveryHook).not.toHaveBeenCalled();
     });
+
+    test("restores default information_schema during sql recovery when no db provided", async () => {
+        const dsn = parse("ws://root:taosdata@localhost:6041");
+        const connector = createMockConnector();
+        jest
+            .spyOn(WebSocketConnectionPool.instance(), "getConnection")
+            .mockResolvedValue(connector as any);
+
+        const client = new WsClient(dsn, 5000);
+        await client.connect();
+
+        const hookCalls = connector.setSessionRecoveryHook.mock.calls;
+        const hook = hookCalls[hookCalls.length - 1]?.[0];
+        expect(hook).toBeTruthy();
+        await hook();
+
+        expect(connector.sendMsgDirect).toHaveBeenCalledTimes(1);
+        const firstCall = connector.sendMsgDirect.mock.calls[0] as [string];
+        const connMsg = JSON.parse(firstCall[0]);
+        expect(connMsg.action).toBe("conn");
+        expect(connMsg.args.db).toBe("information_schema");
+    });
 });

nodejs/test/client/wsConnectorPool.key.test.ts

@@ -1,8 +1,10 @@
 import { WebSocketConnectionPool } from "@src/client/wsConnectorPool";
+import { RetryConfig } from "@src/client/wsConnector";
 import { WSConfig } from "@src/common/config";
 import { parse, WS_TMQ_ENDPOINT } from "@src/common/dsn";
 import { WsSql } from "@src/sql/wsSql";
 import { testPassword, testUsername } from "@test-helpers/utils";
+import { w3cwebsocket } from "websocket";
 
 function resetPoolSingleton() {
     const PoolClass = WebSocketConnectionPool as any;
@@ -52,6 +54,66 @@ describe("WebSocketConnectionPool key generation", () => {
         expect(keyA).not.toBe(keyB);
     });
 
+    test("does not collide auth scope when credentials contain colon", () => {
+        const pool = WebSocketConnectionPool.instance();
+
+        const dsnA = parse("ws://host1:6041/mydb");
+        dsnA.username = "a:b";
+        dsnA.password = "c";
+
+        const dsnB = parse("ws://host1:6041/mydb");
+        dsnB.username = "a";
+        dsnB.password = "b:c";
+
+        const keyA = (pool as any).getPoolKey(dsnA);
+        const keyB = (pool as any).getPoolKey(dsnB);
+        expect(keyA).not.toBe(keyB);
+    });
+
+    test("does not split pool key by reconnect policy", () => {
+        const pool = WebSocketConnectionPool.instance();
+        const lowRetryDsn = parse(
+            "ws://root:taosdata@host1:6041/mydb?retries=1&retry_backoff_ms=10&retry_backoff_max_ms=20"
+        );
+        const highRetryDsn = parse(
+            "ws://root:taosdata@host1:6041/mydb?retries=60&retry_backoff_ms=100&retry_backoff_max_ms=500"
+        );
+
+        const lowRetryKey = (pool as any).getPoolKey(lowRetryDsn);
+        const highRetryKey = (pool as any).getPoolKey(highRetryDsn);
+
+        expect(lowRetryKey).toBe(highRetryKey);
+    });
+
+    test("updates connector retry policy when reusing pooled connector", async () => {
+        const pool = WebSocketConnectionPool.instance();
+        const lowRetryDsn = parse(
+            "ws://root:taosdata@host1:6041/mydb?retries=1&retry_backoff_ms=10&retry_backoff_max_ms=20"
+        );
+        const highRetryDsn = parse(
+            "ws://root:taosdata@host1:6041/mydb?retries=60&retry_backoff_ms=100&retry_backoff_max_ms=500"
+        );
+        const poolKey = (pool as any).getPoolKey(lowRetryDsn);
+
+        let retries = 1;
+        const connector = {
+            readyState: jest.fn(() => w3cwebsocket.OPEN),
+            close: jest.fn(),
+            refreshRetryConfig: jest.fn((dsn) => {
+                retries = RetryConfig.fromDsn(dsn).retries;
+            }),
+            getReconnectRetries: jest.fn(() => retries),
+            getPoolKey: jest.fn(() => poolKey),
+        };
+
+        (pool as any).pool.set(poolKey, [connector]);
+        const reused = await pool.getConnection(highRetryDsn, 3000);
+
+        expect(reused).toBe(connector);
+        expect(connector.refreshRetryConfig).toHaveBeenCalledWith(highRetryDsn);
+        expect((reused as any).getReconnectRetries()).toBe(60);
+    });
+
     test("includes endpoint-derived websocket path in the pool key scope", () => {
         const pool = WebSocketConnectionPool.instance();
         const sqlDsn = parse("ws://root:taosdata@host1:6041/mydb");

nodejs/test/common/dsn.test.ts

@@ -238,6 +238,16 @@ describe("dsn", () => {
             expect(result.params.get("token")).toBe("abc123");
         });
 
+        test("does not parse @ in query parameter as userinfo delimiter", () => {
+            const result = parse("ws://localhost:6041?bearer_token=a@b");
+            expect(result.scheme).toBe("ws");
+            expect(result.username).toBe("");
+            expect(result.password).toBe("");
+            expect(result.addresses).toEqual([{ host: "localhost", port: 6041 }]);
+            expect(result.database).toBe("");
+            expect(result.params.get("bearer_token")).toBe("a@b");
+        });
+
         test("database with query parameters", () => {
             const result = parse("ws://root:taosdata@localhost:6041/mydb?timezone=UTC");
             expect(result.scheme).toBe("ws");