Move decideClient and decideLegacy client to data source plugin util

Signed-off-by: Zhongnan Su <[email protected]>

src/plugins/data/server/index.ts

@@ -147,7 +147,6 @@ export {
   FieldDescriptor as IndexPatternFieldDescriptor,
   shouldReadFieldFromDocValues, // used only in logstash_fields fixture
   FieldDescriptor,
-  decideClient,
 } from './index_patterns';
 
 export {

src/plugins/data/server/index_patterns/index.ts

@@ -31,4 +31,3 @@
 export * from './utils';
 export { IndexPatternsFetcher, FieldDescriptor, shouldReadFieldFromDocValues } from './fetcher';
 export { IndexPatternsService, IndexPatternsServiceStart } from './index_patterns_service';
-export { decideClient } from './routes';

src/plugins/data/server/index_patterns/routes.ts

@@ -29,12 +29,9 @@
  */
 
 import { schema } from '@osd/config-schema';
-import {
-  HttpServiceSetup,
-  LegacyAPICaller,
-  RequestHandlerContext,
-} from 'opensearch-dashboards/server';
+import { HttpServiceSetup, RequestHandlerContext } from 'opensearch-dashboards/server';
 import { IndexPatternsFetcher } from './fetcher';
+import { decideLegacyClient } from '../../../data_source/server';
 
 export function registerRoutes(http: HttpServiceSetup) {
   const parseMetaFields = (metaFields: string | string[]) => {
@@ -62,7 +59,7 @@ export function registerRoutes(http: HttpServiceSetup) {
       },
     },
     async (context, request, response) => {
-      const callAsCurrentUser = await decideClient(context, request);
+      const callAsCurrentUser = await decideLegacyClient(context, request);
       const indexPatterns = new IndexPatternsFetcher(callAsCurrentUser);
       const { pattern, meta_fields: metaFields } = request.query;
 
@@ -122,7 +119,7 @@ export function registerRoutes(http: HttpServiceSetup) {
       },
     },
     async (context: RequestHandlerContext, request: any, response: any) => {
-      const callAsCurrentUser = await decideClient(context, request);
+      const callAsCurrentUser = await decideLegacyClient(context, request);
 
       const indexPatterns = new IndexPatternsFetcher(callAsCurrentUser);
       const { pattern, interval, look_back: lookBack, meta_fields: metaFields } = request.query;
@@ -154,13 +151,3 @@ export function registerRoutes(http: HttpServiceSetup) {
     }
   );
 }
-
-export const decideClient = async (
-  context: RequestHandlerContext,
-  request: any
-): Promise<LegacyAPICaller> => {
-  const dataSourceId = request.query.data_source;
-  return dataSourceId
-    ? (context.dataSource.opensearch.legacy.getClient(dataSourceId).callAPI as LegacyAPICaller)
-    : context.core.opensearch.legacy.client.callAsCurrentUser;
-};

src/plugins/data/server/search/opensearch_search/opensearch_search_strategy.ts

@@ -43,7 +43,7 @@ import {
   getShardTimeout,
   shimAbortSignal,
 } from '..';
-import { decideClient } from './decide_client';
+import { decideClient } from '../../../../data_source/server';
 
 export const opensearchSearchStrategyProvider = (
   config$: Observable<SharedGlobalConfig>,

src/plugins/data_source/server/index.ts

@@ -24,5 +24,6 @@ export {
   DataSourcePluginSetup,
   DataSourcePluginStart,
   DataSourcePluginRequestContext,
-  DataSourceError,
 } from './types';
+
+export { decideClient, decideLegacyClient } from './util/decide_client';

src/plugins/data_source/server/util/decide_client.test.ts

@@ -0,0 +1,74 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { coreMock } from '../../../../core/server/mocks';
+import { RequestHandlerContext } from '../../../../core/server';
+import { IOpenSearchSearchRequest } from '../../../data/server';
+import { decideClient, decideLegacyClient } from './decide_client';
+
+const context: RequestHandlerContext = {
+  core: {
+    ...coreMock.createRequestHandlerContext(),
+  },
+  dataSource: {
+    opensearch: {
+      getClient: jest.fn(),
+      legacy: {
+        getClient: jest.fn(),
+      },
+    },
+  },
+};
+
+describe('decideClient', () => {
+  const request: IOpenSearchSearchRequest = {
+    dataSourceId: 'dataSourceId',
+  };
+
+  it('should return defaultOpenSearchClient when dataSourceId is not provided', async () => {
+    const result = await decideClient(context, { ...request, dataSourceId: undefined });
+    expect(result).toBe(context.core.opensearch.client.asCurrentUser);
+  });
+
+  it('should return defaultOpenSearchClientWithLongNumeralsSupport when withLongNumeralsSupport is true', async () => {
+    const result = await decideClient(context, { ...request, dataSourceId: undefined }, true);
+    expect(result).toBe(context.core.opensearch.client.asCurrentUserWithLongNumeralsSupport);
+  });
+
+  it('should return client from dataSource when dataSourceId is provided', async () => {
+    const dataSourceClient = jest.fn();
+    (context.dataSource.opensearch.getClient as jest.Mock).mockResolvedValueOnce(dataSourceClient);
+
+    const result = await decideClient(context, request);
+    expect(result).toBe(dataSourceClient);
+    expect(context.dataSource.opensearch.getClient).toHaveBeenCalledWith(request.dataSourceId);
+  });
+});
+
+describe('decideLegacyClient', () => {
+  const request = {
+    query: {
+      data_source: 'dataSourceId',
+    },
+  };
+
+  it('should return callAsCurrentUser when dataSourceId is not provided', async () => {
+    const result = await decideLegacyClient(context, { ...request, query: {} });
+    expect(result).toBe(context.core.opensearch.legacy.client.callAsCurrentUser);
+  });
+
+  it('should return legacy client from dataSource when dataSourceId is provided', async () => {
+    const dataSourceClient = jest.fn();
+    (context.dataSource.opensearch.legacy.getClient as jest.Mock).mockReturnValueOnce({
+      callAPI: dataSourceClient,
+    });
+
+    const result = await decideLegacyClient(context, request);
+    expect(result).toBe(dataSourceClient);
+    expect(context.dataSource.opensearch.legacy.getClient).toHaveBeenCalledWith(
+      request.query.data_source
+    );
+  });
+});

src/plugins/data_source/server/util/decide_client.ts

@@ -0,0 +1,37 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ *
+ * Any modifications Copyright OpenSearch Contributors. See
+ * GitHub history for details.
+ */
+
+import { LegacyAPICaller, OpenSearchClient, RequestHandlerContext } from '../../../../core/server';
+import { IOpenSearchSearchRequest } from '../../../data/server';
+
+export const decideClient = async (
+  context: RequestHandlerContext,
+  request: IOpenSearchSearchRequest,
+  withLongNumeralsSupport: boolean = false
+): Promise<OpenSearchClient> => {
+  const defaultOpenSearchClient = withLongNumeralsSupport
+    ? context.core.opensearch.client.asCurrentUserWithLongNumeralsSupport
+    : context.core.opensearch.client.asCurrentUser;
+
+  return request.dataSourceId && context.dataSource
+    ? await context.dataSource.opensearch.getClient(request.dataSourceId)
+    : defaultOpenSearchClient;
+};
+
+export const decideLegacyClient = async (
+  context: RequestHandlerContext,
+  request: any
+): Promise<LegacyAPICaller> => {
+  const dataSourceId = request.query.data_source;
+  return dataSourceId
+    ? (context.dataSource.opensearch.legacy.getClient(dataSourceId).callAPI as LegacyAPICaller)
+    : context.core.opensearch.legacy.client.callAsCurrentUser;
+};

src/plugins/vis_type_timeseries/server/lib/get_fields.ts

@@ -37,9 +37,9 @@ import {
   indexPatterns,
   IndexPatternFieldDescriptor,
   IndexPatternsFetcher,
-  decideClient,
 } from '../../../data/server';
 import { ReqFacade } from './search_strategies/strategies/abstract_search_strategy';
+import { decideLegacyClient } from '../../../data_source/server';
 
 export async function getFields(
   requestContext: RequestHandlerContext,
@@ -51,7 +51,7 @@ export async function getFields(
   // removes the need to refactor many layers of dependencies on "req", and instead just augments the top
   // level object passed from here. The layers should be refactored fully at some point, but for now
   // this works and we are still using the New Platform services for these vis data portions.
-  const client = await decideClient(requestContext, request);
+  const client = await decideLegacyClient(requestContext, request);
 
   const reqFacade: ReqFacade = {
     requestContext,