From 04eaedd13818f06d7ad8ae34d0c6c7d995500a87 Mon Sep 17 00:00:00 2001
From: Konstantin Burkalev <KSDaemon@gmail.com>
Date: Mon, 3 Mar 2025 15:59:00 +0200
Subject: [PATCH] fix(query-orchestrator): Fix dropping temp tables during
 pre-agg creation

---
 .../src/orchestrator/PreAggregationLoader.ts             | 9 ++++-----
 1 file changed, 4 insertions(+), 5 deletions(-)

diff --git a/packages/cubejs-query-orchestrator/src/orchestrator/PreAggregationLoader.ts b/packages/cubejs-query-orchestrator/src/orchestrator/PreAggregationLoader.ts
index f559998a568e9..716f5ba2b6090 100644
--- a/packages/cubejs-query-orchestrator/src/orchestrator/PreAggregationLoader.ts
+++ b/packages/cubejs-query-orchestrator/src/orchestrator/PreAggregationLoader.ts
@@ -618,7 +618,7 @@ export class PreAggregationLoader {
     dropSourceTempTable: boolean,
   ) {
     if (withTempTable && dropSourceTempTable) {
-      await this.withDropLock(false, async () => {
+      await this.withDropLock(`drop-temp-table:${this.preAggregation.dataSource}:${targetTableName}`, async () => {
         this.logger('Dropping source temp table', queryOptions);
 
         const actualTables = await client.getTablesQuery(this.preAggregation.preAggregationsSchema);
@@ -965,8 +965,7 @@ export class PreAggregationLoader {
     });
   }
 
-  private async withDropLock<T>(external: boolean, lockFn: () => MaybeCancelablePromise<T>): Promise<boolean> {
-    const lockKey = this.dropLockKey(external);
+  private async withDropLock<T>(lockKey: string, lockFn: () => MaybeCancelablePromise<T>): Promise<boolean> {
     return this.queryCache.withLock(lockKey, 60 * 5, lockFn);
   }
 
@@ -979,7 +978,7 @@ export class PreAggregationLoader {
   ) {
     await this.preAggregations.addTableUsed(justCreatedTable);
 
-    return this.withDropLock(external, async () => {
+    return this.withDropLock(this.dropOrphanedLockKey(external), async () => {
       this.logger('Dropping orphaned tables', { ...queryOptions, external });
       const actualTables = await client.getTablesQuery(
         this.preAggregation.preAggregationsSchema,
@@ -1039,7 +1038,7 @@ export class PreAggregationLoader {
     });
   }
 
-  private dropLockKey(external: boolean) {
+  private dropOrphanedLockKey(external: boolean) {
     return external
       ? 'drop-orphaned-tables-external'
       : `drop-orphaned-tables:${this.preAggregation.dataSource}`;