feat: support multiple time columns with time grain in Pivot Table v2 (…

…#21537)
apache · Sep 22, 2022 · e671d80 · e671d80
1 parent bbac67a
commit e671d80
Show file tree

Hide file tree

Showing 4 changed files with 94 additions and 9 deletions.
diff --git a/superset-frontend/packages/superset-ui-chart-controls/src/shared-controls/index.tsx b/superset-frontend/packages/superset-ui-chart-controls/src/shared-controls/index.tsx
@@ -45,6 +45,7 @@ import {
   ComparisionType,
   isAdhocColumn,
   isPhysicalColumn,
+  ensureIsArray,
 } from '@superset-ui/core';
 
 import {
@@ -57,7 +58,13 @@ import {
   DEFAULT_NUMBER_FORMAT,
 } from '../utils';
 import { TIME_FILTER_LABELS } from '../constants';
-import { SharedControlConfig, Dataset, ColumnMeta } from '../types';
+import {
+  SharedControlConfig,
+  Dataset,
+  ColumnMeta,
+  ControlState,
+  ControlPanelState,
+} from '../types';
 
 import {
   dndAdhocFilterControl,
@@ -340,6 +347,16 @@ const show_empty_columns: SharedControlConfig<'CheckboxControl'> = {
   description: t('Show empty columns'),
 };
 
+const datetime_columns_lookup: SharedControlConfig<'HiddenControl'> = {
+  type: 'HiddenControl',
+  initialValue: (control: ControlState, state: ControlPanelState) =>
+    Object.fromEntries(
+      ensureIsArray<Record<string, any>>(state?.datasource?.columns)
+        .filter(option => option.is_dttm)
+        .map(option => [option.column_name ?? option.name, option.is_dttm]),
+    ),
+};
+
 export default {
   metrics: dndAdhocMetricsControl,
   metric: dndAdhocMetricControl,
@@ -376,4 +393,5 @@ export default {
   truncate_metric,
   x_axis: dndXAxisControl,
   show_empty_columns,
+  datetime_columns_lookup,
 };
diff --git a/superset-frontend/plugins/plugin-chart-pivot-table/package.json b/superset-frontend/plugins/plugin-chart-pivot-table/package.json
@@ -32,7 +32,8 @@
     "@ant-design/icons": "^4.2.2",
     "react": "^16.13.1",
     "react-dom": "^16.13.1",
-    "prop-types": "*"
+    "prop-types": "*",
+    "lodash": "^4.17.11"
   },
   "devDependencies": {
     "@babel/types": "^7.13.12",

diff --git a/superset-frontend/plugins/plugin-chart-pivot-table/src/plugin/buildQuery.ts b/superset-frontend/plugins/plugin-chart-pivot-table/src/plugin/buildQuery.ts
@@ -16,9 +16,15 @@
  * specific language governing permissions and limitations
  * under the License.
  */
+import omit from 'lodash/omit';
+
 import {
+  AdhocColumn,
   buildQueryContext,
   ensureIsArray,
+  FeatureFlag,
+  isFeatureEnabled,
+  isPhysicalColumn,
   QueryFormColumn,
   QueryFormOrderBy,
 } from '@superset-ui/core';
@@ -27,10 +33,29 @@ import { PivotTableQueryFormData } from '../types';
 export default function buildQuery(formData: PivotTableQueryFormData) {
   const { groupbyColumns = [], groupbyRows = [] } = formData;
   // TODO: add deduping of AdhocColumns
-  const groupbySet = new Set([
-    ...ensureIsArray<QueryFormColumn>(groupbyColumns),
-    ...ensureIsArray<QueryFormColumn>(groupbyRows),
-  ]);
+  const columns = Array.from(
+    new Set([
+      ...ensureIsArray<QueryFormColumn>(groupbyColumns),
+      ...ensureIsArray<QueryFormColumn>(groupbyRows),
+    ]),
+  ).map(col => {
+    if (
+      isPhysicalColumn(col) &&
+      formData.time_grain_sqla &&
+      isFeatureEnabled(FeatureFlag.GENERIC_CHART_AXES) &&
+      formData?.datetime_columns_lookup?.[col]
+    ) {
+      return {
+        timeGrain: formData.time_grain_sqla,
+        columnType: 'BASE_AXIS',
+        sqlExpression: col,
+        label: col,
+        expressionType: 'SQL',
+      } as AdhocColumn;
+    }
+    return col;
+  });
+
   return buildQueryContext(formData, baseQueryObject => {
     const { series_limit_metric, metrics, order_desc } = baseQueryObject;
     let orderBy: QueryFormOrderBy[] | undefined;
@@ -41,9 +66,11 @@ export default function buildQuery(formData: PivotTableQueryFormData) {
     }
     return [
       {
-        ...baseQueryObject,
+        ...(isFeatureEnabled(FeatureFlag.GENERIC_CHART_AXES)
+          ? omit(baseQueryObject, ['extras.time_grain_sqla'])
+          : baseQueryObject),
         orderby: orderBy,
-        columns: [...groupbySet],
+        columns,
       },
     ];
   });

diff --git a/superset-frontend/plugins/plugin-chart-pivot-table/src/plugin/controlPanel.tsx b/superset-frontend/plugins/plugin-chart-pivot-table/src/plugin/controlPanel.tsx
@@ -19,6 +19,10 @@
 import React from 'react';
 import {
   ensureIsArray,
+  FeatureFlag,
+  isAdhocColumn,
+  isFeatureEnabled,
+  isPhysicalColumn,
   QueryFormMetric,
   smartDateFormatter,
   t,
@@ -38,7 +42,7 @@ import { MetricsLayoutEnum } from '../types';
 
 const config: ControlPanelConfig = {
   controlPanelSections: [
-    { ...sections.legacyTimeseriesTime, expanded: false },
+    { ...sections.genericTime, expanded: false },
     {
       label: t('Query'),
       expanded: true,
@@ -63,6 +67,41 @@ const config: ControlPanelConfig = {
             },
           },
         ],
+        [
+          isFeatureEnabled(FeatureFlag.GENERIC_CHART_AXES)
+            ? {
+                name: 'time_grain_sqla',
+                config: {
+                  ...sharedControls.time_grain_sqla,
+                  visibility: ({ controls }) => {
+                    const dttmLookup = Object.fromEntries(
+                      ensureIsArray(controls?.groupbyColumns?.options).map(
+                        option => [option.column_name, option.is_dttm],
+                      ),
+                    );
+
+                    return [
+                      ...ensureIsArray(controls?.groupbyColumns.value),
+                      ...ensureIsArray(controls?.groupbyRows.value),
+                    ]
+                      .map(selection => {
+                        if (isAdhocColumn(selection)) {
+                          return true;
+                        }
+                        if (isPhysicalColumn(selection)) {
+                          return !!dttmLookup[selection];
+                        }
+                        return false;
+                      })
+                      .some(Boolean);
+                  },
+                },
+              }
+            : null,
+          isFeatureEnabled(FeatureFlag.GENERIC_CHART_AXES)
+            ? 'datetime_columns_lookup'
+            : null,
+        ],
         [
           {
             name: 'metrics',