feat(cubestore-driver): Queue - use exclusive flag with new Cube Store (cube-js#10488)

Author: ovr

packages/cubejs-backend-shared/src/index.ts

@@ -15,6 +15,7 @@ export * from './helpers';
 export * from './machine-id';
 export * from './type-helpers';
 export * from './shared-types';
+export * from './semver';
 export * from './http-utils';
 export * from './cli';
 export * from './proxy';

packages/cubejs-backend-shared/src/semver.ts

@@ -0,0 +1,51 @@
+interface VersionPart {
+  num: number;
+  pre: string;
+}
+
+function parseVersionParts(v: string): VersionPart[] {
+  return v.split('.').map((segment) => {
+    const idx = segment.indexOf('-');
+    if (idx === -1) {
+      return { num: parseInt(segment, 10) || 0, pre: '' };
+    }
+
+    return {
+      num: parseInt(segment.substring(0, idx), 10) || 0,
+      pre: segment.substring(idx + 1),
+    };
+  });
+}
+
+export function isVersionGte(version: string | null, minVersion: string): boolean {
+  if (!version) {
+    return false;
+  }
+
+  const parts = parseVersionParts(version);
+  const minParts = parseVersionParts(minVersion);
+
+  for (let i = 0; i < Math.max(parts.length, minParts.length); i++) {
+    const a = parts[i] || { num: 0, pre: '' };
+    const b = minParts[i] || { num: 0, pre: '' };
+
+    if (a.num > b.num) {
+      return true;
+    }
+
+    if (a.num < b.num) {
+      return false;
+    }
+
+    // The same numeric part — pre-release is less than no pre-release
+    if (a.pre && !b.pre) {
+      return false;
+    }
+
+    if (!a.pre && b.pre) {
+      return true;
+    }
+  }
+
+  return true;
+}

packages/cubejs-backend-shared/test/semver.test.ts

@@ -0,0 +1,36 @@
+import { isVersionGte } from '../src';
+
+describe('isVersionGte', () => {
+  test.each([
+    // null version
+    [null, '1.0.0', false],
+    // equal versions
+    ['1.6.22', '1.6.22', true],
+    ['0.0.0', '0.0.0', true],
+    // greater (major, minor, patch)
+    ['2.0.0', '1.6.22', true],
+    ['1.7.0', '1.6.22', true],
+    ['1.6.23', '1.6.22', true],
+    // lesser (major, minor, patch)
+    ['0.9.99', '1.0.0', false],
+    ['1.5.99', '1.6.0', false],
+    ['1.6.21', '1.6.22', false],
+    // different segment lengths
+    ['1.6', '1.6.0', true],
+    ['1.6.0', '1.6', true],
+    ['1.5', '1.6.0', false],
+    // pre-release is less than the same clean version
+    ['1.6.22-alpha', '1.6.22', false],
+    ['1.6.22-beta', '1.6.22', false],
+    ['1.6.22-rc.1', '1.6.22', false],
+    // pre-release of a higher version still passes
+    ['1.6.23-alpha', '1.6.22', true],
+    ['2.0.0-beta', '1.6.22', true],
+    // clean version is greater than pre-release
+    ['1.6.22', '1.6.22-alpha', true],
+    // both pre-release, same numeric — equal
+    ['1.6.22-alpha', '1.6.22-beta', true],
+  ])('isVersionGte(%j, %j) === %j', (version, minVersion, expected) => {
+    expect(isVersionGte(version, minVersion)).toBe(expected);
+  });
+});

packages/cubejs-cubestore-driver/src/CubeStoreDriver.ts

@@ -14,13 +14,19 @@ import {
   QueryOptions,
   ExternalDriverCompatibilities, TableStructure, TableColumnQueryResult,
 } from '@cubejs-backend/base-driver';
-import { AsyncDebounce, getEnv } from '@cubejs-backend/shared';
+import { AsyncDebounce, getEnv, isVersionGte } from '@cubejs-backend/shared';
 import { format as formatSql, escape } from 'sqlstring';
 import fetch from 'node-fetch';
 
 import { ConnectionConfig } from './types';
 import { WebSocketConnection } from './WebSocketConnection';
 
+type CubeStoreCapability = 'queueExclusive';
+
+const CubeStoreCapabilityMinVersion: Record<CubeStoreCapability, string> = {
+  queueExclusive: '1.6.22',
+};
+
 const GenericTypeToCubeStore: Record<string, string> = {
   string: 'varchar(255)',
   text: 'varchar(255)',
@@ -76,6 +82,12 @@ export class CubeStoreDriver extends BaseDriver implements DriverInterface {
     this.connection = new WebSocketConnection(`${this.baseUrl}/ws`);
   }
 
+  public async hasCapability(capability: CubeStoreCapability): Promise<boolean> {
+    const minVersion = CubeStoreCapabilityMinVersion[capability];
+
+    return isVersionGte(await this.connection.getCubeStoreVersion(), minVersion);
+  }
+
   public async testConnection() {
     await this.query('SELECT 1', []);
   }

packages/cubejs-cubestore-driver/src/CubeStoreQueueDriver.ts

@@ -82,7 +82,8 @@ class CubestoreQueueDriverConnection implements QueueDriverConnectionInterface {
     values.push(this.prefixKey(this.redisHash(queryKey)));
     values.push(JSON.stringify(data));
 
-    const rows = await this.driver.query(`QUEUE ADD PRIORITY ?${options.orphanedTimeout ? ' ORPHANED ?' : ''} ? ?`, values);
+    const exclusive = queryKey.persistent && await this.driver.hasCapability('queueExclusive');
+    const rows = await this.driver.query(`QUEUE ADD${exclusive ? ' EXCLUSIVE' : ''} PRIORITY ?${options.orphanedTimeout ? ' ORPHANED ?' : ''} ? ?`, values);
     if (rows && rows.length) {
       return [
         rows[0].added === 'true' ? 1 : 0,

packages/cubejs-cubestore-driver/src/WebSocketConnection.ts

@@ -42,6 +42,8 @@ export class WebSocketConnection {
 
   private readonly connectionId: string;
 
+  private cubeStoreVersion: string | null = null;
+
   public constructor(url: string) {
     this.url = url;
     this.messageCounter = 1;
@@ -57,6 +59,10 @@ export class WebSocketConnection {
       headers['x-process-id'] = getProcessUid();
 
       const webSocket = new WebSocket(this.url, { headers }) as CubeStoreWebSocket;
+      webSocket.on('upgrade', (response: any) => {
+        this.cubeStoreVersion = response.headers['x-cubestore-version'] || null;
+      });
+
       webSocket.readyPromise = new Promise<CubeStoreWebSocket>((resolve, reject) => {
         webSocket.lastHeartBeat = new Date();
         const pingInterval = setInterval(() => {
@@ -284,6 +290,14 @@ export class WebSocketConnection {
     return this.sendMessage(messageId, builder.asUint8Array());
   }
 
+  public async getCubeStoreVersion(): Promise<string> {
+    if (this.webSocket) {
+      await this.webSocket.readyPromise;
+    }
+
+    return this.cubeStoreVersion ?? '0.0.0';
+  }
+
   public close() {
     if (this.webSocket) {
       this.webSocket.close();

rust/cubestore/cubestore/src/http/mod.rs

@@ -147,7 +147,7 @@ impl HttpServer {
             .and_then(move |tx: mpsc::Sender<(mpsc::Sender<Arc<HttpMessage>>, SqlQueryContext, HttpMessage)>, sql_query_context: SqlQueryContext, ws: Ws| async move {
                 let tx_to_move = tx.clone();
                 let sql_query_context = sql_query_context.clone();
-                Result::<_, Rejection>::Ok(ws.max_frame_size(max_frame_size).max_message_size(max_message_size).on_upgrade(async move |mut web_socket| {
+                let reply = ws.max_frame_size(max_frame_size).max_message_size(max_message_size).on_upgrade(async move |mut web_socket| {
                     let process_id = sql_query_context.process_id.as_deref().unwrap_or("None");
                     trace!("WebSocket connection established (process_id: {})", process_id);
                     let (response_tx, mut response_rx) = mpsc::channel::<Arc<HttpMessage>>(10000);
@@ -207,7 +207,8 @@ impl HttpServer {
                             }
                         };
                     };
-                }))
+                });
+                Result::<_, Rejection>::Ok(warp::reply::with_header(reply, "X-CubeStore-Version", env!("CARGO_PKG_VERSION")))
             });
 
         let auth_filter_to_move = auth_filter.clone();