fix(api-gateway): Allow messageId to be int in unsubscribe WS request (#10250)

Author: ovr

packages/cubejs-api-gateway/src/ws/local-subscription-store.ts

@@ -2,14 +2,16 @@ interface LocalSubscriptionStoreOptions {
   heartBeatInterval?: number;
 }
 
+export type SubscriptionId = string | number;
+
 export type LocalSubscriptionStoreSubscription = {
   message: any,
   state: any,
   timestamp: Date,
 };
 
 export type LocalSubscriptionStoreConnection = {
-  subscriptions: Map<string, LocalSubscriptionStoreSubscription>,
+  subscriptions: Map<SubscriptionId, LocalSubscriptionStoreSubscription>,
   authContext?: any,
 };
 
@@ -35,7 +37,7 @@ export class LocalSubscriptionStore {
     });
   }
 
-  public async unsubscribe(connectionId: string, subscriptionId: string) {
+  public async unsubscribe(connectionId: string, subscriptionId: SubscriptionId) {
     const connection = this.getConnectionOrCreate(connectionId);
     connection.subscriptions.delete(subscriptionId);
   }

packages/cubejs-api-gateway/src/ws/message-schema.ts

@@ -1,14 +1,14 @@
 import { z } from 'zod';
 
-const messageId = z.union([z.string().max(16), z.number()]);
+const messageId = z.union([z.string().max(16), z.int()]);
 const requestId = z.string().max(64).optional();
 
 export const authMessageSchema = z.object({
   authorization: z.string(),
 }).strict();
 
 export const unsubscribeMessageSchema = z.object({
-  unsubscribe: z.string().max(16),
+  unsubscribe: messageId,
 }).strict();
 
 const queryParams = z.object({

packages/cubejs-api-gateway/test/ws/subscription-server.test.ts

@@ -59,6 +59,15 @@ describe('SubscriptionServer', () => {
       expect(mockSubscriptionStore.unsubscribe).toHaveBeenCalledWith('conn-1', 'msg-1');
     });
 
+    it('should accept unsubscribe with numeric messageId', async () => {
+      const { mockApiGateway, mockSubscriptionStore, mockSendMessage, mockContextAcceptor } = createMocks();
+      const server = new SubscriptionServer(mockApiGateway, mockSendMessage, mockSubscriptionStore, mockContextAcceptor);
+
+      await server.processMessage('conn-1', JSON.stringify({ unsubscribe: 123 }));
+
+      expect(mockSubscriptionStore.unsubscribe).toHaveBeenCalledWith('conn-1', 123);
+    });
+
     it('should accept valid load message', async () => {
       const { mockApiGateway, mockSubscriptionStore, mockSendMessage, mockContextAcceptor, sentMessages } = createMocks();
       const server = new SubscriptionServer(mockApiGateway, mockSendMessage, mockSubscriptionStore, mockContextAcceptor);