From 8f1fa775bad6baa077763712f2d93cd3ac6ef6f5 Mon Sep 17 00:00:00 2001
From: Konstantin Burkalev <KSDaemon@gmail.com>
Date: Wed, 26 Feb 2025 13:55:55 +0200
Subject: [PATCH] make loadBuildRange() timezone-aware

---
 packages/cubejs-backend-shared/src/time.ts    |  8 ++++-
 .../PreAggregationPartitionRangeLoader.ts     | 29 +++++++++++++------
 2 files changed, 27 insertions(+), 10 deletions(-)

diff --git a/packages/cubejs-backend-shared/src/time.ts b/packages/cubejs-backend-shared/src/time.ts
index f6277de9fae7b..6bb49a1f9f1bf 100644
--- a/packages/cubejs-backend-shared/src/time.ts
+++ b/packages/cubejs-backend-shared/src/time.ts
@@ -227,7 +227,13 @@ export const inDbTimeZone = (timezone: string, timestampFormat: string, timestam
   return moment.tz(timestamp, timezone).utc().format(timestampFormat);
 };
 
-export const utcToLocalTimeZone = (timezone: string, timestampFormat: string, timestamp: string): string => {
+/**
+ * Takes timestamp in UTC, treat it as local time in provided timezone and returns the corresponding timestamp in UTC
+ */
+export const utcToLocalTimeZoneInUtc = (timezone: string, timestampFormat: string, timestamp: string): string | null => {
+  if (!timestamp) {
+    return null;
+  }
   if (timestamp.length === 23) {
     const zone = moment.tz.zone(timezone);
     if (!zone) {
diff --git a/packages/cubejs-query-orchestrator/src/orchestrator/PreAggregationPartitionRangeLoader.ts b/packages/cubejs-query-orchestrator/src/orchestrator/PreAggregationPartitionRangeLoader.ts
index 07dd1902f52b8..8cbb25969087e 100644
--- a/packages/cubejs-query-orchestrator/src/orchestrator/PreAggregationPartitionRangeLoader.ts
+++ b/packages/cubejs-query-orchestrator/src/orchestrator/PreAggregationPartitionRangeLoader.ts
@@ -6,7 +6,7 @@ import {
   TO_PARTITION_RANGE,
   MAX_SOURCE_ROW_LIMIT,
   reformatInIsoLocal,
-  utcToLocalTimeZone,
+  utcToLocalTimeZoneInUtc,
   timeSeries,
   inDbTimeZone,
   extractDate
@@ -136,9 +136,9 @@ export class PreAggregationPartitionRangeLoader {
       return queryValues?.map(
         param => {
           if (param === BUILD_RANGE_START_LOCAL) {
-            return utcToLocalTimeZone(this.preAggregation.timezone, this.preAggregation.timestampFormat, buildRangeStart);
+            return utcToLocalTimeZoneInUtc(this.preAggregation.timezone, this.preAggregation.timestampFormat, buildRangeStart);
           } else if (param === BUILD_RANGE_END_LOCAL) {
-            return utcToLocalTimeZone(this.preAggregation.timezone, this.preAggregation.timestampFormat, buildRangeEnd);
+            return utcToLocalTimeZoneInUtc(this.preAggregation.timezone, this.preAggregation.timestampFormat, buildRangeEnd);
           } else {
             return param;
           }
@@ -396,12 +396,19 @@ export class PreAggregationPartitionRangeLoader {
     const { preAggregationStartEndQueries } = this.preAggregation;
     const [startDate, endDate] = await Promise.all(
       preAggregationStartEndQueries.map(
-        async rangeQuery => PreAggregationPartitionRangeLoader.extractDate(await this.loadRangeQuery(rangeQuery)),
+        async rangeQuery => utcToLocalTimeZoneInUtc(
+          this.preAggregation.timezone,
+          'YYYY-MM-DDTHH:mm:ss.SSS',
+          PreAggregationPartitionRangeLoader.extractDate(await this.loadRangeQuery(rangeQuery)),
+        )
       ),
     );
+
     if (!this.preAggregation.partitionGranularity) {
       return this.orNowIfEmpty([startDate, endDate]);
     }
+
+    // startDate & endDate are `localized` here
     const wholeSeriesRanges = PreAggregationPartitionRangeLoader.timeSeries(
       this.preAggregation.partitionGranularity,
       this.orNowIfEmpty([startDate, endDate]),
@@ -409,18 +416,22 @@ export class PreAggregationPartitionRangeLoader {
     );
     const [rangeStart, rangeEnd] = await Promise.all(
       preAggregationStartEndQueries.map(
-        async (rangeQuery, i) => PreAggregationPartitionRangeLoader.extractDate(
-          await this.loadRangeQuery(
-            rangeQuery, i === 0 ? wholeSeriesRanges[0] : wholeSeriesRanges[wholeSeriesRanges.length - 1],
+        async (rangeQuery, i) => utcToLocalTimeZoneInUtc(
+          this.preAggregation.timezone,
+          'YYYY-MM-DDTHH:mm:ss.SSS',
+          PreAggregationPartitionRangeLoader.extractDate(
+            await this.loadRangeQuery(
+              rangeQuery, i === 0 ? wholeSeriesRanges[0] : wholeSeriesRanges[wholeSeriesRanges.length - 1],
+            ),
           ),
-        ),
+        )
       ),
     );
     return this.orNowIfEmpty([rangeStart, rangeEnd]);
   }
 
   private now() {
-    return utcToLocalTimeZone(this.preAggregation.timezone, 'YYYY-MM-DDTHH:mm:ss.SSS', new Date().toJSON().substring(0, 23));
+    return utcToLocalTimeZoneInUtc(this.preAggregation.timezone, 'YYYY-MM-DDTHH:mm:ss.SSS', new Date().toJSON().substring(0, 23));
   }
 
   private orNowIfEmpty(dateRange: QueryDateRange): QueryDateRange {