fix: Pre-aggregation table is not found after it was successfully created for CUBEJS_DROP_PRE_AGG_WITHOUT_TOUCH strategy (#5858)

Author: paveltiunov

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

@@ -2,32 +2,30 @@ import crypto from 'crypto';
 import R from 'ramda';
 import {
   addSecondsToLocalTimestamp,
+  BUILD_RANGE_END_LOCAL,
+  BUILD_RANGE_START_LOCAL,
   extractDate,
   FROM_PARTITION_RANGE,
   getEnv,
   inDbTimeZone,
-  timeSeries,
   MAX_SOURCE_ROW_LIMIT,
+  timeSeries,
   TO_PARTITION_RANGE,
-  BUILD_RANGE_START_LOCAL,
-  BUILD_RANGE_END_LOCAL,
   utcToLocalTimeZone,
 } from '@cubejs-backend/shared';
 
-import { cancelCombinator, SaveCancelFn, DriverInterface, BaseDriver,
+import {
+  BaseDriver,
+  cancelCombinator,
   DownloadTableData,
+  DriverCapabilities,
+  DriverInterface,
   InlineTable,
+  SaveCancelFn,
   StreamOptions,
   UnloadOptions,
-  DriverCapabilities } from '@cubejs-backend/base-driver';
-import {
-  Query,
-  QueryCache,
-  QueryTuple,
-  QueryWithParams,
-  QueryBody,
-  PreAggTableToTempTable,
-} from './QueryCache';
+} from '@cubejs-backend/base-driver';
+import { PreAggTableToTempTable, Query, QueryBody, QueryCache, QueryTuple, QueryWithParams } from './QueryCache';
 import { ContinueWaitError } from './ContinueWaitError';
 import { DriverFactory, DriverFactoryByDataSource } from './DriverFactory';
 import { QueryQueue } from './QueryQueue';
@@ -866,6 +864,7 @@ export class PreAggregationLoader {
   }
 
   public refresh(newVersionEntry: VersionEntry, invalidationKeys: InvalidationKeys, client) {
+    this.updateLastTouch(this.targetTableName(newVersionEntry));
     let refreshStrategy = this.refreshStoreInSourceStrategy;
     if (this.preAggregation.external) {
       const readOnly =

packages/cubejs-query-orchestrator/test/unit/QueryOrchestrator.test.js

@@ -171,6 +171,7 @@ describe('QueryOrchestrator', () => {
   let externalMockDriver = null;
   let queryOrchestrator = null;
   let queryOrchestratorExternalRefresh = null;
+  let queryOrchestratorDropWithoutTouch = null;
   let testCount = 1;
 
   beforeEach(() => {
@@ -227,6 +228,14 @@ describe('QueryOrchestrator', () => {
           externalRefresh: true,
         },
       });
+    queryOrchestratorDropWithoutTouch =
+      new QueryOrchestrator(redisPrefix, driverFactory, logger, {
+        ...options,
+        preAggregationsOptions: {
+          ...options.preAggregationsOptions,
+          dropPreAggregationsWithoutTouch: true,
+        },
+      });
     mockDriver = mockDriverLocal;
     fooMockDriver = fooMockDriverLocal;
     barMockDriver = barMockDriverLocal;
@@ -238,6 +247,7 @@ describe('QueryOrchestrator', () => {
   afterEach(async () => {
     await queryOrchestrator.cleanup();
     await queryOrchestratorExternalRefresh.cleanup();
+    await queryOrchestratorDropWithoutTouch.cleanup();
   });
 
   test('basic', async () => {
@@ -1300,4 +1310,42 @@ describe('QueryOrchestrator', () => {
 
     expect(streamingSourceMockDriver.downloadTableStreamOffset).toBe('earliest');
   });
+
+  test('drop without touch does not affect tables in progress', async () => {
+    const firstQuery = queryOrchestratorDropWithoutTouch.fetchQuery({
+      query: 'SELECT * FROM stb_pre_aggregations.orders_delay_d20181102',
+      values: [],
+      cacheKeyQueries: {
+        renewalThreshold: 21600,
+        queries: []
+      },
+      preAggregations: [{
+        preAggregationsSchema: 'stb_pre_aggregations',
+        tableName: 'stb_pre_aggregations.orders_delay_d20181102',
+        loadSql: ['CREATE TABLE stb_pre_aggregations.orders_d20181102 AS SELECT * FROM public.orders_delay', []],
+        invalidateKeyQueries: [['SELECT 2', []]]
+      }],
+      requestId: 'drop without touch does not affect tables in progress'
+    });
+    const promises = [firstQuery];
+    for (let i = 0; i < 10; i++) {
+      promises.push(queryOrchestratorDropWithoutTouch.fetchQuery({
+        query: `SELECT * FROM stb_pre_aggregations.orders_d201811${i}`,
+        values: [],
+        cacheKeyQueries: {
+          renewalThreshold: 21600,
+          queries: []
+        },
+        preAggregations: [{
+          preAggregationsSchema: 'stb_pre_aggregations',
+          tableName: `stb_pre_aggregations.orders_d201811${i}`,
+          loadSql: [`CREATE TABLE stb_pre_aggregations.orders_d201811${i} AS SELECT * FROM public.orders`, []],
+          invalidateKeyQueries: [['SELECT 2', []]]
+        }],
+        requestId: 'drop without touch does not affect tables in progress'
+      }));
+    }
+    await Promise.all(promises);
+    expect(mockDriver.tables).toContainEqual(expect.stringMatching(/orders_delay/));
+  });
 });