feat: Remove job from pre-aggregations queue, Debug API

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

@@ -23,6 +23,7 @@ import {
   normalizeQuery,
   normalizeQueryPreAggregations,
   normalizeQueryPreAggregationPreview,
+  normalizeQueryCancelPreAggregations,
   QUERY_TYPE
 } from './query';
 import {
@@ -405,6 +406,14 @@ export class ApiGateway {
           res: this.resToResultFn(res)
         });
       }));
+
+      app.post('/cubejs-system/v1/pre-aggregations/cancel', jsonParser, systemMiddlewares, (async (req, res) => {
+        await this.cancelPreAggregationsFromQueue({
+          query: req.body.query,
+          context: req.context,
+          res: this.resToResultFn(res)
+        });
+      }));
     }
 
     app.get('/readyz', guestMiddlewares, cachedHandler(this.readiness));
@@ -602,6 +611,23 @@ export class ApiGateway {
     }
   }
 
+  public async cancelPreAggregationsFromQueue(
+    { query, context, res }: { query: any, context: RequestContext, res: ResponseResultFn }
+  ) {
+    const requestStarted = new Date();
+    try {
+      const { queryKeys, dataSource } = normalizeQueryCancelPreAggregations(this.parseQueryParam(query));
+      const orchestratorApi = this.getAdapterApi(context);
+      res({
+        result: await orchestratorApi.cancelPreAggregationQueriesFromQueue(queryKeys, dataSource)
+      });
+    } catch (e) {
+      this.handleError({
+        e, context, res, requestStarted
+      });
+    }
+  }
+
   protected async getNormalizedQueries(query, context: RequestContext): Promise<any> {
     query = this.parseQueryParam(query);
     let queryType = QUERY_TYPE.REGULAR_QUERY;

packages/cubejs-api-gateway/src/query.js

@@ -257,3 +257,17 @@ export const normalizeQueryPreAggregationPreview = (query) => {
 
   return query;
 };
+
+const queryCancelPreAggregationPreviewSchema = Joi.object().keys({
+  dataSource: Joi.string(),
+  queryKeys: Joi.array().items(Joi.string())
+});
+
+export const normalizeQueryCancelPreAggregations = query => {
+  const { error } = Joi.validate(query, queryCancelPreAggregationPreviewSchema);
+  if (error) {
+    throw new UserError(`Invalid query format: ${error.message || error.toString()}`);
+  }
+
+  return query;
+};

packages/cubejs-query-orchestrator/src/orchestrator/LocalQueueDriver.js

@@ -122,6 +122,21 @@ export class LocalQueueDriverConnection {
     return [query];
   }
 
+  async cancelQuery(queryKey) {
+    const query = await this.getQueryAndRemove(queryKey);
+
+    if (this.getQueueEventsBus) {
+      this.getQueueEventsBus().emit({
+        event: 'cancelQuery',
+        redisQueuePrefix: this.redisQueuePrefix,
+        queryKey: this.redisHash(queryKey),
+        payload: query
+      });
+    }
+
+    return true;
+  }
+
   async setResultAndRemoveQuery(queryKey, executionResult, processingId) {
     const key = this.redisHash(queryKey);
     if (this.processingLocks[key] !== processingId) {

packages/cubejs-query-orchestrator/src/orchestrator/PreAggregations.ts

@@ -1457,8 +1457,13 @@ export class PreAggregations {
     return data.filter(res => res);
   }
 
-  public async getQueueState(dataSource = undefined) {
+  public async getQueueState(dataSource: string) {
     const queries = await this.getQueue(dataSource).getQueries();
     return queries;
   }
+
+  public async cancelQueriesFromQueue(queryKeys: string[], dataSource: string) {
+    const queue = this.getQueue(dataSource);
+    return Promise.all(queryKeys.map(queryKey => queue.cancelQuery(queryKey)));
+  }
 }

packages/cubejs-query-orchestrator/src/orchestrator/QueryOrchestrator.ts

@@ -234,8 +234,12 @@ export class QueryOrchestrator {
     return this.preAggregations.expandPartitionsInPreAggregations(queryBody);
   }
 
-  public async getPreAggregationQueueStates() {
-    return this.preAggregations.getQueueState();
+  public async getPreAggregationQueueStates(dataSource = 'default') {
+    return this.preAggregations.getQueueState(dataSource);
+  }
+
+  public async cancelPreAggregationQueriesFromQueue(queryKeys: string[], dataSource = 'default') {
+    return this.preAggregations.cancelQueriesFromQueue(queryKeys, dataSource);
   }
 
   public async subscribeQueueEvents(id, callback) {

packages/cubejs-query-orchestrator/src/orchestrator/QueryQueue.js

@@ -205,6 +205,15 @@ export class QueryQueue {
     }
   }
 
+  async cancelQuery(queryKey) {
+    const redisClient = await this.queueDriver.createConnection();
+    try {
+      return redisClient.cancelQuery(queryKey);
+    } finally {
+      this.queueDriver.release(redisClient);
+    }
+  }
+
   async reconcileQueueImpl() {
     const redisClient = await this.queueDriver.createConnection();
     try {

packages/cubejs-query-orchestrator/src/orchestrator/RedisQueueDriver.js

@@ -53,6 +53,7 @@ export class RedisQueueDriverConnection {
       .zcard(this.toProcessRedisKey());
 
     if (this.getQueueEventsBus) {
+      console.log('addedToQueue', this.redisHash(queryKey));
       tx.publish(
         this.getQueueEventsBus().eventsChannel,
         JSON.stringify({
@@ -87,6 +88,24 @@ export class RedisQueueDriverConnection {
     return [JSON.parse(query), ...restResult];
   }
 
+  async cancelQuery(queryKey) {
+    const [query] = await this.getQueryAndRemove(queryKey);
+
+    if (this.getQueueEventsBus) {
+      await this.redisClient.publish(
+        this.getQueueEventsBus().eventsChannel,
+        JSON.stringify({
+          event: 'cancelQuery',
+          redisQueuePrefix: this.redisQueuePrefix,
+          queryKey: this.redisHash(queryKey),
+          payload: query
+        })
+      );
+    }
+
+    return true;
+  }
+
   async setResultAndRemoveQuery(queryKey, executionResult, processingId) {
     try {
       await this.redisClient.watchAsync(this.queryProcessingLockKey(queryKey));

packages/cubejs-server-core/src/core/OrchestratorApi.ts

@@ -189,6 +189,11 @@ export class OrchestratorApi {
     return result;
   }
 
+  public async cancelPreAggregationQueriesFromQueue(queryKeys: string[], dataSource: string) {
+    const result = await this.orchestrator.cancelPreAggregationQueriesFromQueue(queryKeys, dataSource);
+    return result;
+  }
+
   public async subscribeQueueEvents(id, callback) {
     return this.orchestrator.subscribeQueueEvents(id, callback);
   }