[Vega] Restores signal values on refresh (#90774)

* Vega kibanaAddFilter() Resets Signal Values Back to Default

Closes: #88976

* fix ci

* introduce restoreSignalValuesOnRefresh option

* update docs

docs/user/dashboard/vega-reference.asciidoc

@@ -401,7 +401,9 @@ Vega-Lite compilation.
 [[vega-expression-functions]]
 ===== (Vega only) Expression functions which can update the time range and dashboard filters
 
-{kib} has extended the Vega expression language with these functions:
+{kib} has extended the Vega expression language with these functions. 
+These functions will trigger new data to be fetched, which by default will reset Vega signals. 
+To keep signal values set `restoreSignalValuesOnRefresh: true` in the Vega config.
 
 ```js
 /**
@@ -444,6 +446,8 @@ kibanaSetTimeFilter(start, end)
       hideWarnings: true
       // Vega renderer to use: `svg` or `canvas` (default)
       renderer: canvas
+      // Defaults to 'false', restores Vega signal values on refresh
+      restoreSignalValuesOnRefresh: false
     }
   }
 }

src/plugins/vis_type_vega/public/data_model/vega_parser.ts

@@ -53,6 +53,7 @@ const DEFAULT_PARSER: string = 'elasticsearch';
 export class VegaParser {
   spec: VegaSpec;
   hideWarnings: boolean;
+  restoreSignalValuesOnRefresh: boolean;
   error?: string;
   warnings: string[];
   _urlParsers: UrlParserConfig | undefined;
@@ -137,6 +138,8 @@ The URL is an identifier only. Kibana and your browser will never access this UR
 
     this._config = this._parseConfig();
     this.hideWarnings = !!this._config.hideWarnings;
+    this._parseBool('restoreSignalValuesOnRefresh', this._config, false);
+    this.restoreSignalValuesOnRefresh = this._config.restoreSignalValuesOnRefresh;
     this.useMap = this._config.type === 'map';
     this.renderer = this._config.renderer === 'svg' ? 'svg' : 'canvas';
     this.tooltips = this._parseTooltips();

src/plugins/vis_type_vega/public/lib/vega_state_restorer.test.ts

@@ -0,0 +1,93 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0 and the Server Side Public License, v 1; you may not use this file except
+ * in compliance with, at your election, the Elastic License 2.0 or the Server
+ * Side Public License, v 1.
+ */
+
+import { createVegaStateRestorer } from './vega_state_restorer';
+
+describe('extractIndexPatternsFromSpec', () => {
+  test('should create vega state restorer ', async () => {
+    expect(createVegaStateRestorer()).toMatchInlineSnapshot(`
+      Object {
+        "clear": [Function],
+        "restore": [Function],
+        "save": [Function],
+      }
+    `);
+  });
+
+  test('should save state', async () => {
+    const vegaStateRestorer = createVegaStateRestorer();
+
+    vegaStateRestorer.save({
+      signals: { foo: 'foo' },
+      data: { test: 'test' },
+    });
+
+    expect(vegaStateRestorer.restore()).toMatchInlineSnapshot(`
+      Object {
+        "signals": Object {
+          "foo": "foo",
+        },
+      }
+    `);
+  });
+
+  test('should restore of "data" if "restoreData" is true', () => {
+    const vegaStateRestorer = createVegaStateRestorer();
+
+    vegaStateRestorer.save({
+      signals: { foo: 'foo' },
+      data: { test: 'test' },
+    });
+
+    expect(vegaStateRestorer.restore(true)).toMatchInlineSnapshot(`
+      Object {
+        "data": Object {
+          "test": "test",
+        },
+        "signals": Object {
+          "foo": "foo",
+        },
+      }
+    `);
+  });
+
+  test('should clear saved state', () => {
+    const vegaStateRestorer = createVegaStateRestorer();
+
+    vegaStateRestorer.save({
+      signals: { foo: 'foo' },
+      data: { test: 'test' },
+    });
+    vegaStateRestorer.clear();
+
+    expect(vegaStateRestorer.restore(true)).toMatchInlineSnapshot(`null`);
+  });
+
+  test('should omit signals', () => {
+    const vegaStateRestorer = createVegaStateRestorer({ omitSignals: ['foo'] });
+
+    vegaStateRestorer.save({
+      signals: { foo: 'foo' },
+    });
+
+    expect(vegaStateRestorer.restore()).toMatchInlineSnapshot(`
+      Object {
+        "signals": Object {},
+      }
+    `);
+  });
+  test('should not save state if isActive is false', () => {
+    const vegaStateRestorer = createVegaStateRestorer({ isActive: () => false });
+
+    vegaStateRestorer.save({
+      signals: { foo: 'foo' },
+    });
+
+    expect(vegaStateRestorer.restore()).toMatchInlineSnapshot(`null`);
+  });
+});

src/plugins/vis_type_vega/public/lib/vega_state_restorer.ts

@@ -0,0 +1,69 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0 and the Server Side Public License, v 1; you may not use this file except
+ * in compliance with, at your election, the Elastic License 2.0 or the Server
+ * Side Public License, v 1.
+ */
+
+import { omit } from 'lodash';
+
+interface VegaStateRestorerOptions {
+  /**
+   *  List of excluded signals
+   *
+   *  By default, all Build-in signals (width,height,padding,autosize,background) were excluded
+   *  @see https://vega.github.io/vega/docs/signals/
+   */
+  omitSignals?: string[];
+  /**
+   * Gets a value that indicates whether the VegaStateRestorer is active.
+   */
+  isActive?: () => boolean;
+}
+
+type State = Partial<{
+  signals: Record<string, any>;
+  data: Record<string, any>;
+}>;
+
+export const createVegaStateRestorer = ({
+  omitSignals = ['width', 'height', 'padding', 'autosize', 'background'],
+  isActive = () => true,
+}: VegaStateRestorerOptions = {}) => {
+  let state: State | null;
+
+  return {
+    /**
+     * Save Vega state
+     * @public
+     * @param newState - new state value
+     */
+    save: (newState: State) => {
+      if (newState && isActive()) {
+        state = {
+          signals: omit(newState.signals, omitSignals || []),
+          data: newState.data,
+        };
+      }
+    },
+
+    /**
+     * Restore Vega state
+     * @public
+     * @param restoreData - by default, we only recover signals,
+     *        but if the data also needs to be recovered, this option should be set to true
+     */
+    restore: (restoreData = false) =>
+      isActive() && state ? omit(state, restoreData ? undefined : 'data') : null,
+
+    /**
+     *  Clear saved Vega state
+     *
+     *  @public
+     */
+    clear: () => {
+      state = null;
+    },
+  };
+};

src/plugins/vis_type_vega/public/vega_view/vega_base_view.d.ts

@@ -10,6 +10,7 @@ import { DataPublicPluginStart } from 'src/plugins/data/public';
 import { IInterpreterRenderHandlers } from 'src/plugins/expressions';
 import { IServiceSettings } from 'src/plugins/maps_legacy/public';
 import { VegaParser } from '../data_model/vega_parser';
+import { createVegaStateRestorer } from '../lib/vega_state_restorer';
 
 interface VegaViewParams {
   parentEl: HTMLDivElement;
@@ -18,6 +19,7 @@ interface VegaViewParams {
   serviceSettings: IServiceSettings;
   filterManager: DataPublicPluginStart['query']['filterManager'];
   timefilter: DataPublicPluginStart['query']['timefilter']['timefilter'];
+  vegaStateRestorer: ReturnType<typeof createVegaStateRestorer>;
 }
 
 export class VegaBaseView {
@@ -34,5 +36,6 @@ export class VegaBaseView {
   _$container: any;
   _parser: any;
   _vegaViewConfig: any;
-  _serviceSettings: any;
+  _serviceSettings: VegaViewParams['serviceSettings'];
+  _vegaStateRestorer: VegaViewParams['vegaStateRestorer'];
 }

src/plugins/vis_type_vega/public/vega_view/vega_base_view.js

@@ -62,6 +62,7 @@ export class VegaBaseView {
     this._destroyHandlers = [];
     this._initialized = false;
     this._enableExternalUrls = getEnableExternalUrls();
+    this._vegaStateRestorer = opts.vegaStateRestorer;
   }
 
   async init() {
@@ -103,6 +104,10 @@ export class VegaBaseView {
           this._$messages = null;
         }
         if (this._view) {
+          const state = this._view.getState();
+          if (state) {
+            this._vegaStateRestorer.save(state);
+          }
           this._view.finalize();
         }
         this._view = null;
@@ -262,7 +267,13 @@ export class VegaBaseView {
         this._addDestroyHandler(() => tthandler.hideTooltip());
       }
 
-      return view.runAsync(); // Allows callers to await rendering
+      const state = this._vegaStateRestorer.restore();
+
+      if (state) {
+        return view.setState(state);
+      } else {
+        return view.runAsync();
+      }
     }
   }
 

src/plugins/vis_type_vega/public/vega_view/vega_map_view/view.test.ts

@@ -100,13 +100,18 @@ describe('vega_map_view/view', () => {
 
     async function createVegaMapView() {
       await vegaParser.parseAsync();
-      return new VegaMapView({
+      return new VegaMapView(({
         vegaParser,
         filterManager: dataPluginStart.query.filterManager,
         timefilter: dataPluginStart.query.timefilter.timefilter,
         fireEvent: (event: any) => {},
         parentEl: document.createElement('div'),
-      } as VegaViewParams);
+        vegaStateRestorer: {
+          save: jest.fn(),
+          restore: jest.fn(),
+          clear: jest.fn(),
+        },
+      } as unknown) as VegaViewParams);
     }
 
     beforeEach(() => {

src/plugins/vis_type_vega/public/vega_view/vega_map_view/view.ts

@@ -68,11 +68,18 @@ export class VegaMapView extends VegaBaseView {
 
   private getMapParams(defaults: { maxZoom: number; minZoom: number }): Partial<MapboxOptions> {
     const { longitude, latitude, scrollWheelZoom } = this._parser.mapConfig;
-    const zoomSettings = validateZoomSettings(this._parser.mapConfig, defaults, this.onWarn);
+    const { zoom, maxZoom, minZoom } = validateZoomSettings(
+      this._parser.mapConfig,
+      defaults,
+      this.onWarn
+    );
+    const { signals } = this._vegaStateRestorer.restore() || {};
 
     return {
-      ...zoomSettings,
-      center: [longitude, latitude],
+      maxZoom,
+      minZoom,
+      zoom: signals?.zoom ?? zoom,
+      center: [signals?.longitude ?? longitude, signals?.latitude ?? latitude],
       scrollZoom: scrollWheelZoom,
     };
   }

src/plugins/vis_type_vega/public/vega_visualization.ts

@@ -12,6 +12,7 @@ import { VegaParser } from './data_model/vega_parser';
 import { VegaVisualizationDependencies } from './plugin';
 import { getNotifications, getData } from './services';
 import type { VegaView } from './vega_view/vega_view';
+import { createVegaStateRestorer } from './lib/vega_state_restorer';
 
 type VegaVisType = new (el: HTMLDivElement, fireEvent: IInterpreterRenderHandlers['event']) => {
   render(visData: VegaParser): Promise<void>;
@@ -24,6 +25,9 @@ export const createVegaVisualization = ({
   class VegaVisualization {
     private readonly dataPlugin = getData();
     private vegaView: InstanceType<typeof VegaView> | null = null;
+    private vegaStateRestorer = createVegaStateRestorer({
+      isActive: () => Boolean(this.vegaView?._parser?.restoreSignalValuesOnRefresh),
+    });
 
     constructor(
       private el: HTMLDivElement,
@@ -71,6 +75,7 @@ export const createVegaVisualization = ({
         const vegaViewParams = {
           parentEl: this.el,
           fireEvent: this.fireEvent,
+          vegaStateRestorer: this.vegaStateRestorer,
           vegaParser,
           serviceSettings,
           filterManager,
@@ -89,6 +94,7 @@ export const createVegaVisualization = ({
     }
 
     destroy() {
+      this.vegaStateRestorer.clear();
       this.vegaView?.destroy();
     }
   };