[Vega] Add back setMapView function (#128914)

src/plugins/vis_types/vega/public/vega_view/vega_map_view/view.test.ts

@@ -31,30 +31,43 @@ import { initVegaLayer, initTmsRasterLayer } from './layers';
 
 import { mapboxgl } from '@kbn/mapbox-gl';
 
-jest.mock('@kbn/mapbox-gl', () => ({
-  mapboxgl: {
-    setRTLTextPlugin: jest.fn(),
-    Map: jest.fn().mockImplementation(() => ({
-      getLayer: () => '',
-      removeLayer: jest.fn(),
-      once: (eventName: string, handler: Function) => handler(),
-      remove: () => jest.fn(),
-      getCanvas: () => ({ clientWidth: 512, clientHeight: 512 }),
-      getCenter: () => ({ lat: 20, lng: 20 }),
-      getZoom: () => 3,
-      addControl: jest.fn(),
-      addLayer: jest.fn(),
-      dragRotate: {
-        disable: jest.fn(),
+jest.mock('@kbn/mapbox-gl', () => {
+  const zoomTo = jest.fn();
+  const setCenter = jest.fn();
+  const fitBounds = jest.fn();
+  return {
+    mapboxgl: {
+      mocks: {
+        zoomTo,
+        setCenter,
+        fitBounds,
       },
-      touchZoomRotate: {
-        disableRotation: jest.fn(),
-      },
-    })),
-    MapboxOptions: jest.fn(),
-    NavigationControl: jest.fn(),
-  },
-}));
+      setRTLTextPlugin: jest.fn(),
+      Map: jest.fn().mockImplementation(() => ({
+        getLayer: () => '',
+        removeLayer: jest.fn(),
+        once: (eventName: string, handler: Function) => handler(),
+        remove: () => jest.fn(),
+        getCanvas: () => ({ clientWidth: 512, clientHeight: 512 }),
+        getCenter: () => ({ lat: 20, lng: 20 }),
+        getZoom: () => 3,
+        zoomTo,
+        setCenter,
+        fitBounds,
+        addControl: jest.fn(),
+        addLayer: jest.fn(),
+        dragRotate: {
+          disable: jest.fn(),
+        },
+        touchZoomRotate: {
+          disableRotation: jest.fn(),
+        },
+      })),
+      MapboxOptions: jest.fn(),
+      NavigationControl: jest.fn(),
+    },
+  };
+});
 
 jest.mock('./layers', () => ({
   initVegaLayer: jest.fn(),
@@ -206,5 +219,57 @@ describe('vega_map_view/view', () => {
 
       expect(mapboxgl.NavigationControl).toHaveBeenCalled();
     });
+
+    describe('setMapView', () => {
+      let vegaMapView: VegaMapView;
+      beforeEach(async () => {
+        vegaMapView = await createVegaMapView();
+        await vegaMapView.init();
+        mapboxgl.mocks.setCenter.mockReset();
+        mapboxgl.mocks.zoomTo.mockReset();
+        mapboxgl.mocks.fitBounds.mockReset();
+      });
+
+      test('should set just lat lng', async () => {
+        vegaMapView.setMapViewHandler(1, 2);
+        expect(mapboxgl.mocks.setCenter).toHaveBeenCalledWith({ lat: 1, lng: 2 });
+      });
+
+      test('should set just lng lat via array', async () => {
+        vegaMapView.setMapViewHandler([1, 2]);
+        expect(mapboxgl.mocks.setCenter).toHaveBeenCalledWith({ lat: 2, lng: 1 });
+      });
+
+      test('should set lat lng and zoom', async () => {
+        vegaMapView.setMapViewHandler(1, 2, 6);
+        expect(mapboxgl.mocks.setCenter).toHaveBeenCalledWith({ lat: 1, lng: 2 });
+        expect(mapboxgl.mocks.zoomTo).toHaveBeenCalledWith(6);
+      });
+
+      test('should set bounds', async () => {
+        vegaMapView.setMapViewHandler([
+          [1, 2],
+          [6, 7],
+        ]);
+        expect(mapboxgl.mocks.fitBounds).toHaveBeenCalledWith([
+          { lat: 2, lng: 1 },
+          { lat: 7, lng: 6 },
+        ]);
+      });
+
+      test('should throw on invalid input', async () => {
+        expect(() => {
+          vegaMapView.setMapViewHandler(undefined);
+        }).toThrow();
+
+        expect(() => {
+          vegaMapView.setMapViewHandler(['a', 'b'] as unknown as [number, number]);
+        }).toThrow();
+
+        expect(() => {
+          vegaMapView.setMapViewHandler([1, 2, 3, 4, 5] as unknown as [number, number]);
+        }).toThrow();
+      });
+    });
   });
 });

src/plugins/vis_types/vega/public/vega_view/vega_map_view/view.ts

@@ -9,7 +9,7 @@
 import { i18n } from '@kbn/i18n';
 import type { Map, Style, MapboxOptions } from '@kbn/mapbox-gl';
 
-import { View, parse } from 'vega';
+import { View, parse, expressionFunction } from 'vega';
 
 import { mapboxgl } from '@kbn/mapbox-gl';
 
@@ -45,6 +45,30 @@ async function updateVegaView(mapBoxInstance: Map, vegaView: View) {
   }
 }
 
+type SetMapViewArgs =
+  | [number, number, number]
+  | [number, number]
+  | [[number, number], number]
+  | [[number, number]]
+  | [[[number, number], [number, number]]];
+
+expressionFunction(
+  'setMapView',
+  function handlerFwd(
+    this: {
+      context: { dataflow: { _kibanaView: VegaMapView; runAfter: (fn: () => void) => void } };
+    },
+    ...args: SetMapViewArgs
+  ) {
+    const view = this.context.dataflow;
+    if (!('setMapViewHandler' in view._kibanaView)) {
+      // not a map view, don't do anything
+      return;
+    }
+    view.runAfter(() => view._kibanaView.setMapViewHandler(...args));
+  }
+);
+
 export class VegaMapView extends VegaBaseView {
   private mapBoxInstance?: Map;
 
@@ -200,4 +224,86 @@ export class VegaMapView extends VegaBaseView {
   protected async onViewContainerResize() {
     this.mapBoxInstance?.resize();
   }
+
+  public setMapViewHandler(...args: SetMapViewArgs) {
+    if (!this.mapBoxInstance) {
+      return;
+    }
+    function throwError() {
+      throw new Error(
+        i18n.translate('visTypeVega.visualization.setMapViewErrorMessage', {
+          defaultMessage:
+            'Unexpected setMapView() parameters. It could be called with a bounding box setMapView([[longitude1,latitude1],[longitude2,latitude2]]), or it could be the center point setMapView([longitude, latitude], optional_zoom), or it can be used as setMapView(latitude, longitude, optional_zoom)',
+        })
+      );
+    }
+
+    function checkArray(
+      val: number | [number, number] | [[number, number], [number, number]]
+    ): [number, number] {
+      if (
+        !Array.isArray(val) ||
+        val.length !== 2 ||
+        typeof val[0] !== 'number' ||
+        typeof val[1] !== 'number'
+      ) {
+        throwError();
+      }
+      return val as [number, number];
+    }
+
+    let lng: number | undefined;
+    let lat: number | undefined;
+    let zoom: number | undefined;
+    switch (args.length) {
+      default:
+        throwError();
+        break;
+      case 1: {
+        const arg = args[0];
+        if (
+          Array.isArray(arg) &&
+          arg.length === 2 &&
+          Array.isArray(arg[0]) &&
+          Array.isArray(arg[1])
+        ) {
+          // called with a bounding box, need to reverse order
+          const [lng1, lat1] = checkArray(arg[0]);
+          const [lng2, lat2] = checkArray(arg[1]);
+          this.mapBoxInstance.fitBounds([
+            { lat: lat1, lng: lng1 },
+            { lat: lat2, lng: lng2 },
+          ]);
+        } else {
+          // called with a center point and no zoom
+          [lng, lat] = checkArray(arg);
+        }
+        break;
+      }
+      case 2:
+        if (Array.isArray(args[0])) {
+          [lng, lat] = checkArray(args[0]);
+          zoom = args[1];
+        } else {
+          [lat, lng] = args;
+        }
+        break;
+      case 3:
+        [lat, lng, zoom] = args;
+        break;
+    }
+
+    if (lat !== undefined && lng !== undefined) {
+      if (typeof lat !== 'number' || typeof lng !== 'number') {
+        throwError();
+      }
+      if (zoom !== undefined && typeof zoom !== 'number') {
+        throwError();
+      }
+      this.mapBoxInstance.setCenter({ lat, lng });
+      if (zoom !== undefined) {
+        this.mapBoxInstance.zoomTo(zoom);
+      }
+    }
+  }
 }