trying to fix flaky ws failing tests

Author: KSDaemon

packages/cubejs-api-gateway/src/SubscriptionServer.ts

@@ -19,7 +19,7 @@ const calcMessageLength = (message: unknown) => Buffer.byteLength(
   typeof message === 'string' ? message : JSON.stringify(message)
 );
 
-export type WebSocketSendMessageFn = (connectionId: string, message: any) => void;
+export type WebSocketSendMessageFn = (connectionId: string, message: any) => Promise<void>;
 
 export class SubscriptionServer {
   public constructor(
@@ -31,7 +31,7 @@ export class SubscriptionServer {
   }
 
   public resultFn(connectionId: string, messageId: string, requestId: string | undefined) {
-    return (message, { status } = { status: 200 }) => {
+    return async (message, { status } = { status: 200 }) => {
       this.apiGateway.log({
         type: 'Outgoing network usage',
         service: 'api-ws',

packages/cubejs-api-gateway/src/gateway.ts

@@ -1832,17 +1832,15 @@ class ApiGateway {
       if (props.queryType === 'multi') {
         // We prepare the final json result on native side
         const resultMulti = new ResultMultiWrapper(results, { queryType, slowQuery });
-        return res(resultMulti);
+        await res(resultMulti);
       } else {
         // We prepare the full final json result on native side
-        return res(results[0]);
+        await res(results[0]);
       }
     } catch (e: any) {
       this.handleError({
         e, context, query, res, requestStarted
       });
-
-      return null;
     }
   }
 
@@ -1934,7 +1932,7 @@ class ApiGateway {
           }];
         }
 
-        return res(request.streaming ? results[0] : { results });
+        await res(request.streaming ? results[0] : { results });
       } else {
         results = await Promise.all(
           normalizedQueries.map(async (normalizedQuery, index) => {
@@ -1968,19 +1966,17 @@ class ApiGateway {
         );
 
         if (request.streaming) {
-          return res(results[0]);
+          await res(results[0]);
         } else {
           // We prepare the final json result on native side
           const resultArray = new ResultArrayWrapper(results);
-          return res(resultArray);
+          await res(resultArray);
         }
       }
     } catch (e: any) {
       this.handleError({
         e, context, query, res, requestStarted
       });
-
-      return null;
     }
   }