Correctly sets timerange

Signed-off-by: Ashwin P Chandran <ashwinpc@amazon.com>

src/plugins/vis_builder/public/application/components/data_tab/use/use_dropbox.tsx

@@ -8,17 +8,17 @@ import { cloneDeep } from 'lodash';
 import { BucketAggType, IndexPatternField, propFilter } from '../../../../../../data/common';
 import { Schema } from '../../../../../../vis_default_editor/public';
 import { COUNT_FIELD, FieldDragDataType } from '../../../utils/drag_drop/types';
-import { useTypedDispatch, useTypedSelector } from '../../../utils/state_management';
+import { useTypedDispatch } from '../../../utils/state_management';
 import { DropboxDisplay, DropboxProps } from '../dropbox';
 import { useDrop } from '../../../utils/drag_drop';
 import {
   editDraftAgg,
   reorderAgg,
   updateAggConfigParams,
 } from '../../../utils/state_management/visualization_slice';
-import { useIndexPatterns } from '../../../utils/use/use_index_pattern';
 import { useOpenSearchDashboards } from '../../../../../../opensearch_dashboards_react/public';
 import { WizardServices } from '../../../../types';
+import { useAggs } from '../../../utils/use';
 
 const filterByName = propFilter('name');
 const filterByType = propFilter('type');
@@ -30,52 +30,33 @@ export interface UseDropboxProps extends Pick<DropboxProps, 'id' | 'label'> {
 export const useDropbox = (props: UseDropboxProps): DropboxProps => {
   const { id: dropboxId, label, schema } = props;
   const [validAggTypes, setValidAggTypes] = useState<string[]>([]);
+  const { aggConfigs, indexPattern, aggs, timeRange } = useAggs();
   const dispatch = useTypedDispatch();
-  const indexPattern = useIndexPatterns().selected;
   const {
     services: {
       data: {
         search: { aggs: aggService },
-        query: {
-          timefilter: { timefilter },
-        },
       },
     },
   } = useOpenSearchDashboards<WizardServices>();
-  const aggConfigParams = useTypedSelector(
-    (state) => state.visualization.activeVisualization?.aggConfigParams
-  );
-
-  const aggConfigs = useMemo(() => {
-    return indexPattern && aggService.createAggConfigs(indexPattern, cloneDeep(aggConfigParams));
-  }, [aggConfigParams, aggService, indexPattern]);
-
-  const aggs = useMemo(() => aggConfigs?.aggs ?? [], [aggConfigs?.aggs]);
 
   const dropboxAggs = useMemo(() => aggs.filter((agg) => agg.schema === schema.name), [
     aggs,
     schema.name,
   ]);
 
-  // This useEffect is required to update the timeRange value and initiate rerender to keep labels up to date (Issue #2531).
-  useEffect(() => {
-    const timeRange = timefilter.getTime();
-    dropboxAggs.forEach((agg) => {
-      if (timeRange && agg.type?.name === 'date_histogram') {
-        agg.aggConfigs.setTimeRange(timeRange);
-      }
-    });
-  }, [dropboxAggs, timefilter]);
-
   const displayFields: DropboxDisplay[] = useMemo(
     () =>
       dropboxAggs?.map(
-        (agg): DropboxDisplay => ({
-          id: agg.id,
-          label: agg.makeLabel(),
-        })
-      ) || [],
-    [dropboxAggs]
+        (agg): DropboxDisplay => {
+          agg.aggConfigs.setTimeRange(timeRange);
+          return {
+            id: agg.id,
+            label: agg.makeLabel(),
+          };
+        }
+      ) ?? [],
+    [dropboxAggs, timeRange]
   );
 
   // Event handlers for each dropbox action type

src/plugins/vis_builder/public/application/components/workspace.tsx

@@ -41,7 +41,7 @@ export const Workspace: FC = ({ children }) => {
   useEffect(() => {
     async function loadExpression() {
       const schemas = ui.containerConfig.data.schemas;
-      const [valid, errorMsg] = validateSchemaState(schemas, rootState);
+      const [valid, errorMsg] = validateSchemaState(schemas, rootState.visualization);
 
       if (!valid) {
         if (errorMsg) {
@@ -50,12 +50,13 @@ export const Workspace: FC = ({ children }) => {
         setExpression(undefined);
         return;
       }
-      const exp = await toExpression(rootState);
+
+      const exp = await toExpression(rootState, searchContext);
       setExpression(exp);
     }
 
     loadExpression();
-  }, [rootState, toExpression, toasts, ui.containerConfig.data.schemas]);
+  }, [rootState, toExpression, toasts, ui.containerConfig.data.schemas, searchContext]);
 
   useLayoutEffect(() => {
     const subscription = data.query.state$.subscribe(({ state }) => {

src/plugins/vis_builder/public/application/utils/state_management/store.ts

@@ -72,6 +72,7 @@ export const getPreloadedStore = async (services: WizardServices) => {
 
 // Infer the `RootState` and `AppDispatch` types from the store itself
 export type RootState = ReturnType<typeof rootReducer>;
+export type RenderState = Omit<RootState, 'metadata'>;
 type Store = ReturnType<typeof configurePreloadedStore>;
 export type AppDispatch = Store['dispatch'];
 

src/plugins/vis_builder/public/application/utils/use/index.ts

@@ -3,6 +3,7 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
+export { useAggs } from './use_aggs';
 export { useVisualizationType } from './use_visualization_type';
 export { useIndexPatterns } from './use_index_pattern';
 export { useSavedWizardVis } from './use_saved_vis_builder_vis';

src/plugins/vis_builder/public/application/utils/use/use_aggs.ts

@@ -0,0 +1,53 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { cloneDeep } from 'lodash';
+import { useLayoutEffect, useMemo, useState } from 'react';
+import { useOpenSearchDashboards } from '../../../../../opensearch_dashboards_react/public';
+import { WizardServices } from '../../../types';
+import { useTypedSelector, useTypedDispatch } from '../state_management';
+import { useIndexPatterns } from './use_index_pattern';
+
+export const useAggs = () => {
+  const {
+    services: {
+      data: {
+        search: { aggs: aggService },
+        query: {
+          timefilter: { timefilter },
+        },
+      },
+    },
+  } = useOpenSearchDashboards<WizardServices>();
+  const indexPattern = useIndexPatterns().selected;
+  const [timeRange, setTimeRange] = useState(timefilter.getTime());
+  const aggConfigParams = useTypedSelector(
+    (state) => state.visualization.activeVisualization?.aggConfigParams
+  );
+  const dispatch = useTypedDispatch();
+
+  const aggConfigs = useMemo(() => {
+    const configs =
+      indexPattern && aggService.createAggConfigs(indexPattern, cloneDeep(aggConfigParams));
+    return configs;
+  }, [aggConfigParams, aggService, indexPattern]);
+
+  useLayoutEffect(() => {
+    const subscription = timefilter.getTimeUpdate$().subscribe(() => {
+      setTimeRange(timefilter.getTime());
+    });
+
+    return () => {
+      subscription.unsubscribe();
+    };
+  }, [dispatch, timefilter]);
+
+  return {
+    indexPattern,
+    aggConfigs,
+    aggs: aggConfigs?.aggs ?? [],
+    timeRange,
+  };
+};

src/plugins/vis_builder/public/application/utils/validate_schema_state.ts

@@ -5,10 +5,19 @@
 
 import { countBy } from 'lodash';
 import { Schemas } from '../../../../vis_default_editor/public';
-import { RootState } from './state_management';
+import { VisualizationState } from './state_management';
 
-export const validateSchemaState = (schemas: Schemas, state: RootState): [boolean, string?] => {
-  const activeViz = state.visualization.activeVisualization;
+/**
+ * Validate if the visualization state fits the  vis type schema criteria
+ * @param schemas Visualization type config Schema objects
+ * @param state visualization state
+ * @returns [Validity, 'Message']
+ */
+export const validateSchemaState = (
+  schemas: Schemas,
+  state: VisualizationState
+): [boolean, string?] => {
+  const activeViz = state.activeVisualization;
   const vizName = activeViz?.name;
   const aggs = activeViz?.aggConfigParams;
 

src/plugins/vis_builder/public/embeddable/vis_builder_embeddable.tsx

@@ -30,6 +30,7 @@ import {
 import { validateSchemaState } from '../application/utils/validate_schema_state';
 import { getExpressionLoader, getTypeService } from '../plugin_services';
 import { PersistedState } from '../../../visualizations/public';
+import { RenderState, VisualizationState } from '../application/utils/state_management';
 
 // Apparently this needs to match the saved object type for the clone and replace panel actions to work
 export const WIZARD_EMBEDDABLE = WIZARD_SAVED_OBJECT;
@@ -121,24 +122,24 @@ export class WizardEmbeddable extends Embeddable<SavedObjectEmbeddableInput, Wiz
     const { visualization, style } = this.serializedState;
 
     const vizStateWithoutIndex = JSON.parse(visualization);
-    const visualizationState = {
+    const visualizationState: VisualizationState = {
       searchField: vizStateWithoutIndex.searchField,
       activeVisualization: vizStateWithoutIndex.activeVisualization,
       indexPattern: this.savedWizard?.searchSourceFields?.index,
     };
-    const rootState = {
+    const renderState: RenderState = {
       visualization: visualizationState,
       style: JSON.parse(style),
     };
-    const visualizationName = rootState.visualization?.activeVisualization?.name ?? '';
+    const visualizationName = renderState.visualization?.activeVisualization?.name ?? '';
     const visualizationType = getTypeService().get(visualizationName);
     if (!visualizationType) {
       this.onContainerError(new Error(`Invalid visualization type ${visualizationName}`));
       return;
     }
     const { toExpression, ui } = visualizationType;
     const schemas = ui.containerConfig.data.schemas;
-    const [valid, errorMsg] = validateSchemaState(schemas, rootState);
+    const [valid, errorMsg] = validateSchemaState(schemas, visualizationState);
 
     if (!valid) {
       if (errorMsg) {
@@ -147,7 +148,11 @@ export class WizardEmbeddable extends Embeddable<SavedObjectEmbeddableInput, Wiz
       }
     } else {
       // TODO: handle error in Expression creation
-      const exp = await toExpression(rootState);
+      const exp = await toExpression(renderState, {
+        filters: this.filters,
+        query: this.query,
+        timeRange: this.timeRange,
+      });
       return exp;
     }
   };
@@ -268,12 +273,15 @@ export class WizardEmbeddable extends Embeddable<SavedObjectEmbeddableInput, Wiz
     // Check if rootState has changed
     if (!isEqual(this.getSerializedState(), this.serializedState)) {
       this.serializedState = this.getSerializedState();
-      this.expression = (await this.getExpression()) ?? '';
       dirty = true;
     }
 
-    if (this.handler && dirty) {
-      this.updateHandler();
+    if (dirty) {
+      this.expression = (await this.getExpression()) ?? '';
+
+      if (this.handler) {
+        this.updateHandler();
+      }
     }
   }
 

src/plugins/vis_builder/public/services/type_service/types.ts

@@ -4,8 +4,9 @@
  */
 import { ReactElement } from 'react';
 import { IconType } from '@elastic/eui';
-import { RootState } from '../../application/utils/state_management';
+import { RenderState } from '../../application/utils/state_management';
 import { Schemas } from '../../../../vis_default_editor/public';
+import { IExpressionLoaderParams } from '../../../../expressions/public';
 
 export interface DataTabConfig {
   schemas: Schemas;
@@ -28,5 +29,8 @@ export interface VisualizationTypeOptions<T = any> {
       style: StyleTabConfig<T>;
     };
   };
-  readonly toExpression: (state: RootState) => Promise<string | undefined>;
+  readonly toExpression: (
+    state: RenderState,
+    searchContext: IExpressionLoaderParams['searchContext']
+  ) => Promise<string | undefined>;
 }

src/plugins/vis_builder/public/services/type_service/visualization_type.tsx

@@ -3,7 +3,8 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 import { IconType } from '@elastic/eui';
-import { RootState } from '../../application/utils/state_management';
+import { IExpressionLoaderParams } from '../../../../expressions/public';
+import { RenderState } from '../../application/utils/state_management';
 import { VisualizationTypeOptions } from './types';
 
 type IVisualizationType = VisualizationTypeOptions;
@@ -15,7 +16,10 @@ export class VisualizationType implements IVisualizationType {
   public readonly icon: IconType;
   public readonly stage: 'experimental' | 'production';
   public readonly ui: IVisualizationType['ui'];
-  public readonly toExpression: (state: RootState) => Promise<string | undefined>;
+  public readonly toExpression: (
+    state: RenderState,
+    searchContext: IExpressionLoaderParams['searchContext']
+  ) => Promise<string | undefined>;
 
   constructor(options: VisualizationTypeOptions) {
     this.name = options.name;

src/plugins/vis_builder/public/visualizations/vislib/common/create_vis.ts

@@ -3,12 +3,10 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { AggConfigs, IndexPattern, search } from '../../../../../data/public';
+import { AggConfigs, IndexPattern } from '../../../../../data/public';
 import { Vis } from '../../../../../visualizations/public';
 import { getSearchService, getTimeFilter } from '../../../plugin_services';
 
-const { isDateHistogramBucketAggConfig } = search.aggs;
-
 export const createVis = async (
   type: string,
   aggConfigs: AggConfigs,
@@ -21,12 +19,8 @@ export const createVis = async (
 
   const responseAggs = vis.data.aggs.getResponseAggs().filter((agg) => agg.enabled);
   responseAggs.forEach((agg) => {
-    if (isDateHistogramBucketAggConfig(agg)) {
-      const timeFilter = getTimeFilter();
-
-      agg.params.timeRange = timeFilter.getTime();
-    }
+    const timeFilter = getTimeFilter();
+    agg.params.timeRange = timeFilter.getTime();
   });
-
   return vis;
 };