Web console: fix time shifting (#15359)

* fix time shifting

web-console/src/views/explore-view/explore-view.tsx

@@ -32,7 +32,7 @@ import React, { useEffect, useMemo, useRef, useState } from 'react';
 
 import { ShowValueDialog } from '../../dialogs/show-value-dialog/show-value-dialog';
 import { useLocalStorageState, useQueryManager } from '../../hooks';
-import { deepGet, filterMap, LocalStorageKeys, oneOf, queryDruidSql } from '../../utils';
+import { deepGet, filterMap, findMap, LocalStorageKeys, oneOf, queryDruidSql } from '../../utils';
 
 import { ControlPane } from './control-pane/control-pane';
 import { DroppableContainer } from './droppable-container/droppable-container';
@@ -156,10 +156,19 @@ async function getMaxTimeForTable(tableName: string): Promise<Date | undefined>
   return maxTime;
 }
 
+function getFirstTableName(q: SqlQuery): string | undefined {
+  return (
+    findMap(q.getWithParts(), withPart => {
+      if (!(withPart.query instanceof SqlQuery)) return;
+      return getFirstTableName(withPart.query);
+    }) ?? q.getFirstTableName()
+  );
+}
+
 async function extendedQueryDruidSql<T = any>(sqlQueryPayload: Record<string, any>): Promise<T[]> {
   if (sqlQueryPayload.query.includes('MAX_DATA_TIME()')) {
     const parsed = SqlQuery.parse(sqlQueryPayload.query);
-    const tableName = parsed.getFirstTableName();
+    const tableName = getFirstTableName(parsed);
     if (tableName) {
       const maxTime = await getMaxTimeForTable(tableName);
       if (maxTime) {

web-console/src/views/explore-view/modules/table-react-module.tsx

@@ -21,7 +21,6 @@ import {
   C,
   F,
   SqlCase,
-  SqlColumn,
   SqlExpression,
   SqlFunction,
   SqlLiteral,
@@ -39,6 +38,7 @@ import { useQueryManager } from '../../../hooks';
 import { getInitQuery } from '../utils';
 
 import { GenericOutputTable } from './components';
+import { shiftTimeInWhere } from './utils/utils';
 
 import './table-react-module.scss';
 
@@ -131,20 +131,6 @@ function toShowColumnExpression(
   return ex.as(showColumn.name);
 }
 
-function shiftTimeInWhere(where: SqlExpression, period: string): SqlExpression {
-  return where.walk(q => {
-    if (
-      (q instanceof SqlColumn && q.getName() === '__time') ||
-      (q instanceof SqlFunction && q.getEffectiveFunctionName() === 'TIME_SHIFT') ||
-      (q instanceof SqlFunction && q.getEffectiveFunctionName() === 'MAX_DATA_TIME')
-    ) {
-      return SqlFunction.simple('TIME_SHIFT', [q, period, 1]);
-    } else {
-      return q;
-    }
-  }) as SqlExpression;
-}
-
 interface QueryAndHints {
   query: SqlQuery;
   groupHints: string[];

web-console/src/views/explore-view/modules/utils/utils.spec.ts

@@ -0,0 +1,45 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import { SqlExpression } from '@druid-toolkit/query';
+
+import { shiftTimeInWhere } from './utils';
+
+describe('shiftTimeInWhere', () => {
+  it('works with TIME_IN_INTERVAL', () => {
+    expect(
+      shiftTimeInWhere(
+        SqlExpression.parse(`TIME_IN_INTERVAL("__time", '2016-06-27/2016-06-28')`),
+        'P1D',
+      ).toString(),
+    ).toEqual(`TIME_IN_INTERVAL(TIME_SHIFT("__time", 'P1D', 1), '2016-06-27/2016-06-28')`);
+  });
+
+  it('works with relative time', () => {
+    expect(
+      shiftTimeInWhere(
+        SqlExpression.parse(
+          `(TIME_SHIFT(MAX_DATA_TIME(), 'PT1H', -1) <= "__time" AND "__time" < MAX_DATA_TIME())`,
+        ),
+        'PT1H',
+      ).toString(),
+    ).toEqual(
+      `(TIME_SHIFT(TIME_SHIFT(MAX_DATA_TIME(), 'PT1H', -1), 'PT1H', -1) <= "__time" AND "__time" < TIME_SHIFT(MAX_DATA_TIME(), 'PT1H', -1))`,
+    );
+  });
+});

web-console/src/views/explore-view/modules/utils/utils.ts

@@ -0,0 +1,40 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import type { SqlExpression } from '@druid-toolkit/query';
+import { F, SqlFunction } from '@druid-toolkit/query';
+
+export function shiftTimeInWhere(where: SqlExpression, period: string): SqlExpression {
+  return where.walk(ex => {
+    if (!(ex instanceof SqlFunction)) return ex;
+    const effectiveFunctionName = ex.getEffectiveFunctionName();
+
+    // Works with: TIME_IN_INTERVAL(__time, '<interval>')
+    if (effectiveFunctionName === 'TIME_IN_INTERVAL') {
+      return ex.changeArgs(ex.args!.change(0, F('TIME_SHIFT', ex.getArg(0), period, 1)));
+    }
+
+    // Works with: TIME_SHIFT(...) <= __time
+    //        and: __time < MAX_DATA_TIME()
+    if (effectiveFunctionName === 'TIME_SHIFT' || effectiveFunctionName === 'MAX_DATA_TIME') {
+      return F('TIME_SHIFT', ex, period, -1);
+    }
+
+    return ex;
+  }) as SqlExpression;
+}