Support multi data source in Region map (#614) (#615)

(cherry picked from commit 61d9703)

Co-authored-by: Junqiu Lei <junqiu@amazon.com>

CHANGELOG.md

@@ -15,6 +15,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 ## [Unreleased 2.x](https://github.com/opensearch-project/dashboards-maps/compare/2.13...2.x)
 ### Features
 Support multi data source display in Maps app([#611](https://github.com/opensearch-project/dashboards-maps/pull/611))
+Support multi data source in Region map ([#614](https://github.com/opensearch-project/dashboards-maps/pull/614))
 ### Enhancements
 ### Bug Fixes
 * Fix zoom level type error in custom layer ([#605](https://github.com/opensearch-project/dashboards-maps/pull/605))

public/components/vector_upload_options.test.tsx

@@ -27,7 +27,17 @@ jest.mock('../../../../src/plugins/opensearch_dashboards_react/public', () => ({
 }));
 
 describe('vector_upload_options', () => {
-  const props = {};
+  const props = {
+    vis: {
+      data: {
+        indexPattern: {
+          dataSourceRef: {
+            id: 'mock-data-source-id',
+          },
+        },
+      },
+    },
+  };
 
   const getIndexResponseWhenIndexIsNotPresent = {
     ok: false,

public/components/vector_upload_options.tsx

@@ -78,6 +78,8 @@ const VectorUploadOptions = (props: RegionMapOptionsProps) => {
     return document.getElementsByName(elementName)[0];
   };
 
+  const dataSourceRefId = props.vis.data.indexPattern?.dataSourceRef?.id || '';
+
   const validateIndexName = (typedIndexName: string, isIndexNameWithSuffix: boolean) => {
     const error = [];
     const errorIndexNameDiv = fetchElementByName('errorIndexName');
@@ -175,7 +177,7 @@ const VectorUploadOptions = (props: RegionMapOptionsProps) => {
 
   const checkIfIndexExists = async (indexName: string) => {
     try {
-      const result = await getIndex(indexName, http);
+      const result = await getIndex(indexName, http, dataSourceRefId);
       return result.ok;
     } catch (e) {
       return false;
@@ -280,7 +282,7 @@ const VectorUploadOptions = (props: RegionMapOptionsProps) => {
       type: GEO_SHAPE_TYPE,
       data: [JSON.parse(fileData || null)],
     };
-    const result = await postGeojson(JSON.stringify(bodyData), http);
+    const result = await postGeojson(JSON.stringify(bodyData), http, dataSourceRefId);
     // error handling logic that displays correct toasts for the end users
     if (result?.ok) {
       parsePostGeojsonResult(result, indexName);

public/services.ts

@@ -7,25 +7,34 @@ import { CoreStart } from '../../../src/core/public';
 import { createGetterSetter } from '../../../src/plugins/opensearch_dashboards_utils/common';
 import { TimefilterContract } from '../../../src/plugins/data/public';
 
-export const postGeojson = async (requestBody: any, http: CoreStart['http']) => {
+export const postGeojson = async (
+  requestBody: any,
+  http: CoreStart['http'],
+  dataSourceRefId: string
+) => {
   try {
-    const response = await http.post('../api/custom_import_map/_upload', {
+    const query = dataSourceRefId ? { dataSourceId: dataSourceRefId } : undefined;
+
+    return await http.post('../api/custom_import_map/_upload', {
       body: requestBody,
+      ...(query && { query }),
     });
-    return response;
   } catch (e) {
     return e;
   }
 };
 
-export const getIndex = async (indexName: string, http: CoreStart['http']) => {
+export const getIndex = async (
+  indexName: string,
+  http: CoreStart['http'],
+  dataSourceRefId: string
+) => {
   try {
-    const response = await http.post('../api/custom_import_map/_indices', {
-      body: JSON.stringify({
-        index: indexName,
-      }),
+    const query = dataSourceRefId ? { dataSourceId: dataSourceRefId } : undefined;
+    return await http.post('../api/custom_import_map/_indices', {
+      body: JSON.stringify({ index: indexName }),
+      ...(query && { query }),
     });
-    return response;
   } catch (e) {
     return e;
   }

server/plugin.ts

@@ -26,10 +26,10 @@ import { geospatial, opensearch, statsRoute } from '../server/routes';
 import { mapSavedObjectsType } from './saved_objects';
 import { capabilitiesProvider } from './saved_objects/capabilities_provider';
 import { ConfigSchema } from '../common/config';
+import GeospatialPlugin from './clusters/geospatial_plugin';
 
 export class CustomImportMapPlugin
-  implements Plugin<CustomImportMapPluginSetup, CustomImportMapPluginStart>
-{
+  implements Plugin<CustomImportMapPluginSetup, CustomImportMapPluginStart> {
   private readonly logger: Logger;
   private readonly globalConfig$;
   private readonly config$;
@@ -58,7 +58,11 @@ export class CustomImportMapPlugin
     const opensearchService = new OpensearchService(geospatialClient);
 
     const router = core.http.createRouter();
-    const { home } = plugins;
+    const { home, dataSource } = plugins;
+
+    if (dataSource) {
+      dataSource.registerCustomApiSchema(GeospatialPlugin);
+    }
 
     // Register server side APIs
     geospatial(geospatialService, router);

server/routes/geospatial.ts

@@ -13,6 +13,7 @@ export default function (services, router) {
       path: '/api/custom_import_map/_upload',
       validate: {
         body: schema.any(),
+        query: schema.maybe(schema.object({}, { unknowns: 'allow' })),
       },
       options: {
         body: {

server/routes/opensearch.ts

@@ -14,6 +14,7 @@ export default function (services, router) {
         body: schema.object({
           index: schema.string(),
         }),
+        query: schema.maybe(schema.object({}, { unknowns: 'allow' })),
       },
     },
     services.getIndex

server/services/geospatial_service.js

@@ -9,11 +9,23 @@ export default class GeospatialService {
   }
 
   uploadGeojson = async (context, req, res) => {
+    const dataSourceRefId = req.query.dataSourceId;
+    let uploadResponse;
     try {
-      const { callAsCurrentUser } = await this.driver.asScoped(req);
-      const uploadResponse = await callAsCurrentUser('geospatial.geospatialQuery', {
-        body: req.body,
-      });
+      if (dataSourceRefId) {
+        const remoteDataSourceClient = context.dataSource.opensearch.legacy.getClient(
+          dataSourceRefId
+        ).callAPI;
+        uploadResponse = await remoteDataSourceClient('geospatial.geospatialQuery', {
+          body: req.body,
+        });
+      } else {
+        const { callAsCurrentUser } = await this.driver.asScoped(req);
+        uploadResponse = await callAsCurrentUser('geospatial.geospatialQuery', {
+          body: req.body,
+        });
+      }
+
       return res.ok({
         body: {
           ok: true,

server/services/opensearch_service.js

@@ -9,14 +9,23 @@ export default class OpensearchService {
   }
 
   getIndex = async (context, req, res) => {
+    const dataSourceRefId = req.query.dataSourceId;
     try {
-      const { index } = req.body;
-      const { callAsCurrentUser } = this.driver.asScoped(req);
-      const indices = await callAsCurrentUser('cat.indices', {
-        index,
-        format: 'json',
-        h: 'health,index,status',
-      });
+      if (dataSourceRefId) {
+        const remoteDataSourceClient = context.dataSource.opensearch.legacy.getClient(
+          dataSourceRefId
+        ).callAPI;
+        const { index } = req.body;
+        const indices = await remoteDataSourceClient('cat.indices', {
+          index,
+          format: 'json',
+          h: 'health,index,status',
+        });
+      } else {
+        const { callAsCurrentUser } = this.driver.asScoped(req);
+        const { index } = req.body;
+        const indices = await callAsCurrentUser;
+      }
       return res.ok({
         body: {
           ok: true,

server/types.ts

@@ -3,6 +3,7 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 import { HomeServerPluginSetup } from '../../../src/plugins/home/server';
+import { DataSourcePluginSetup } from '../../../src/plugins/data_source/server';
 
 // eslint-disable-next-line @typescript-eslint/no-empty-interface
 export interface CustomImportMapPluginSetup {}
@@ -11,4 +12,5 @@ export interface CustomImportMapPluginStart {}
 
 export interface AppPluginSetupDependencies {
   home?: HomeServerPluginSetup;
+  dataSource: DataSourcePluginSetup;
 }