[Manual Backport 2.x] Refactor getClient and getLegacyClient to use authentication method registry (#5881)

CHANGELOG.md

@@ -60,6 +60,7 @@ Inspired from [Keep a Changelog](https://keepachangelog.com/en/1.0.0/)
 - [Multiple Datasource] Add interfaces to register add-on authentication method from plug-in module ([#5851](https://github.com/opensearch-project/OpenSearch-Dashboards/pull/5851)) 
 - [Multiple Datasource] Able to Hide "Local Cluster" option from datasource DropDown ([#5827](https://github.com/opensearch-project/OpenSearch-Dashboards/pull/5827))
 - [Multiple Datasource] Add api registry and allow it to be added into client config in data source plugin ([#5895](https://github.com/opensearch-project/OpenSearch-Dashboards/pull/5895))
+- [Multiple Datasource] Refactor client and legacy client to use authentication registry ([#5881](https://github.com/opensearch-project/OpenSearch-Dashboards/pull/5881))
 
 ### 🐛 Bug Fixes
 

src/plugins/data_source/common/data_sources/types.ts

@@ -14,6 +14,7 @@ export interface DataSourceAttributes extends SavedObjectAttributes {
     credentials: UsernamePasswordTypedContent | SigV4Content | undefined | AuthTypeContent;
   };
   lastUpdatedTime?: string;
+  name: AuthType | string;
 }
 
 export interface AuthTypeContent {
@@ -30,6 +31,7 @@ export interface SigV4Content extends SavedObjectAttributes {
   secretKey: string;
   region: string;
   service?: SigV4ServiceName;
+  sessionToken?: string;
 }
 
 export interface UsernamePasswordTypedContent extends SavedObjectAttributes {

src/plugins/data_source/server/auth_registry/authentication_methods_registry.mock.ts

@@ -0,0 +1,14 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { IAuthenticationMethodRegistery } from './authentication_methods_registry';
+
+const create = () =>
+  (({
+    getAllAuthenticationMethods: jest.fn(),
+    getAuthenticationMethod: jest.fn(),
+  } as unknown) as jest.Mocked<IAuthenticationMethodRegistery>);
+
+export const authenticationMethodRegisteryMock = { create };

src/plugins/data_source/server/auth_registry/index.ts

@@ -7,3 +7,5 @@ export {
   IAuthenticationMethodRegistery,
   AuthenticationMethodRegistery,
 } from './authentication_methods_registry';
+
+export { authenticationMethodRegisteryMock } from './authentication_methods_registry.mock';

src/plugins/data_source/server/client/configure_client.test.mocks.ts

@@ -16,3 +16,8 @@ export const parseClientOptionsMock = jest.fn();
 jest.doMock('./client_config', () => ({
   parseClientOptions: parseClientOptionsMock,
 }));
+
+export const authRegistryCredentialProviderMock = jest.fn();
+jest.doMock('../util/credential_provider', () => ({
+  authRegistryCredentialProvider: authRegistryCredentialProviderMock,
+}));

src/plugins/data_source/server/client/configure_client.test.ts

@@ -13,16 +13,24 @@ import {
   SigV4Content,
 } from '../../common/data_sources/types';
 import { DataSourcePluginConfigType } from '../../config';
-import { ClientMock, parseClientOptionsMock } from './configure_client.test.mocks';
+import {
+  ClientMock,
+  parseClientOptionsMock,
+  authRegistryCredentialProviderMock,
+} from './configure_client.test.mocks';
 import { OpenSearchClientPoolSetup } from './client_pool';
 import { configureClient } from './configure_client';
 import { ClientOptions } from '@opensearch-project/opensearch-next';
 // eslint-disable-next-line @osd/eslint/no-restricted-paths
 import { opensearchClientMock } from '../../../../core/server/opensearch/client/mocks';
 import { cryptographyServiceSetupMock } from '../cryptography_service.mocks';
 import { CryptographyServiceSetup } from '../cryptography_service';
-import { DataSourceClientParams } from '../types';
+import { DataSourceClientParams, AuthenticationMethod } from '../types';
 import { CustomApiSchemaRegistry } from '../schema_registry';
+import {
+  IAuthenticationMethodRegistery,
+  authenticationMethodRegisteryMock,
+} from '../auth_registry';
 
 const DATA_SOURCE_ID = 'a54b76ec86771ee865a0f74a305dfff8';
 
@@ -40,13 +48,15 @@ describe('configureClient', () => {
   let usernamePasswordAuthContent: UsernamePasswordTypedContent;
   let sigV4AuthContent: SigV4Content;
   let customApiSchemaRegistry: CustomApiSchemaRegistry;
+  let authenticationMethodRegistery: jest.Mocked<IAuthenticationMethodRegistery>;
 
   beforeEach(() => {
     dsClient = opensearchClientMock.createInternalClient();
     logger = loggingSystemMock.createLogger();
     savedObjectsMock = savedObjectsClientMock.create();
     cryptographyMock = cryptographyServiceSetupMock.create();
     customApiSchemaRegistry = new CustomApiSchemaRegistry();
+    authenticationMethodRegistery = authenticationMethodRegisteryMock.create();
 
     config = {
       enabled: true,
@@ -242,4 +252,46 @@ describe('configureClient', () => {
     expect(savedObjectsMock.get).toHaveBeenCalledTimes(1);
     expect(decodeAndDecryptSpy).toHaveBeenCalledTimes(1);
   });
+
+  test('configureClient should retunrn client from authentication registery if method present in registry', async () => {
+    const name = 'typeA';
+    const customAuthContent = {
+      region: 'us-east-1',
+      roleARN: 'test-role',
+    };
+    savedObjectsMock.get.mockReset().mockResolvedValueOnce({
+      id: DATA_SOURCE_ID,
+      type: DATA_SOURCE_SAVED_OBJECT_TYPE,
+      attributes: {
+        ...dataSourceAttr,
+        auth: {
+          type: AuthType.SigV4,
+          credentials: customAuthContent,
+        },
+      },
+      references: [],
+    });
+    const authMethod: AuthenticationMethod = {
+      name,
+      authType: AuthType.SigV4,
+      credentialProvider: jest.fn(),
+    };
+    authenticationMethodRegistery.getAuthenticationMethod.mockImplementation(() => authMethod);
+
+    authRegistryCredentialProviderMock.mockReturnValue({
+      credential: sigV4AuthContent,
+      type: AuthType.SigV4,
+    });
+
+    await configureClient(
+      { ...dataSourceClientParams, authRegistry: authenticationMethodRegistery },
+      clientPoolSetup,
+      config,
+      logger
+    );
+    expect(authRegistryCredentialProviderMock).toHaveBeenCalled();
+    expect(authenticationMethodRegistery.getAuthenticationMethod).toHaveBeenCalledTimes(1);
+    expect(ClientMock).toHaveBeenCalledTimes(1);
+    expect(savedObjectsMock.get).toHaveBeenCalledTimes(1);
+  });
 });

src/plugins/data_source/server/client/configure_client.ts

@@ -7,7 +7,7 @@ import { Client, ClientOptions } from '@opensearch-project/opensearch-next';
 import { Client as LegacyClient } from 'elasticsearch';
 import { Credentials } from 'aws-sdk';
 import { AwsSigv4Signer } from '@opensearch-project/opensearch-next/aws';
-import { Logger } from '../../../../../src/core/server';
+import { Logger, OpenSearchDashboardsRequest } from '../../../../../src/core/server';
 import {
   AuthType,
   DataSourceAttributes,
@@ -27,6 +27,8 @@ import {
   getDataSource,
   generateCacheKey,
 } from './configure_client_utils';
+import { IAuthenticationMethodRegistery } from '../auth_registry';
+import { authRegistryCredentialProvider } from '../util/credential_provider';
 
 export const configureClient = async (
   {
@@ -35,6 +37,8 @@ export const configureClient = async (
     cryptography,
     testClientDataSourceAttr,
     customApiSchemaRegistryPromise,
+    request,
+    authRegistry,
   }: DataSourceClientParams,
   openSearchClientPoolSetup: OpenSearchClientPoolSetup,
   config: DataSourcePluginConfigType,
@@ -80,6 +84,8 @@ export const configureClient = async (
       cryptography,
       rootClient,
       dataSourceId,
+      request,
+      authRegistry,
       requireDecryption
     );
   } catch (error: any) {
@@ -101,6 +107,8 @@ export const configureClient = async (
  * @param config data source config
  * @param addClientToPool function to add client to client pool
  * @param dataSourceId id of data source saved Object
+ * @param request OpenSearch Dashboards incoming request to read client parameters from header.
+ * @param authRegistry registry to retrieve the credentials provider for the authentication method in order to return the client
  * @param requireDecryption false when creating test client before data source exists
  * @returns Promise of query client
  */
@@ -112,15 +120,31 @@ const getQueryClient = async (
   cryptography?: CryptographyServiceSetup,
   rootClient?: Client,
   dataSourceId?: string,
+  request?: OpenSearchDashboardsRequest,
+  authRegistry?: IAuthenticationMethodRegistery,
   requireDecryption: boolean = true
 ): Promise<Client> => {
-  const {
+  let credential;
+  let {
     auth: { type },
-    endpoint,
+    name,
   } = dataSourceAttr;
+  const { endpoint } = dataSourceAttr;
+  name = name ?? type;
   const clientOptions = parseClientOptions(config, endpoint, registeredSchema);
   const cacheKey = generateCacheKey(dataSourceAttr, dataSourceId);
 
+  const authenticationMethod = authRegistry?.getAuthenticationMethod(name);
+  if (authenticationMethod !== undefined) {
+    const credentialProvider = await authRegistryCredentialProvider(authenticationMethod, {
+      dataSourceAttr,
+      request,
+      cryptography,
+    });
+    credential = credentialProvider.credential;
+    type = credentialProvider.type;
+  }
+
   switch (type) {
     case AuthType.NoAuth:
       if (!rootClient) rootClient = new Client(clientOptions);
@@ -129,21 +153,25 @@ const getQueryClient = async (
       return rootClient.child();
 
     case AuthType.UsernamePasswordType:
-      const credential = requireDecryption
-        ? await getCredential(dataSourceAttr, cryptography!)
-        : (dataSourceAttr.auth.credentials as UsernamePasswordTypedContent);
+      credential =
+        (credential as UsernamePasswordTypedContent) ??
+        (requireDecryption
+          ? await getCredential(dataSourceAttr, cryptography!)
+          : (dataSourceAttr.auth.credentials as UsernamePasswordTypedContent));
 
       if (!rootClient) rootClient = new Client(clientOptions);
       addClientToPool(cacheKey, type, rootClient);
 
       return getBasicAuthClient(rootClient, credential);
 
     case AuthType.SigV4:
-      const awsCredential = requireDecryption
-        ? await getAWSCredential(dataSourceAttr, cryptography!)
-        : (dataSourceAttr.auth.credentials as SigV4Content);
+      credential =
+        (credential as SigV4Content) ??
+        (requireDecryption
+          ? await getAWSCredential(dataSourceAttr, cryptography!)
+          : (dataSourceAttr.auth.credentials as SigV4Content));
 
-      const awsClient = rootClient ? rootClient : getAWSClient(awsCredential, clientOptions);
+      const awsClient = rootClient ? rootClient : getAWSClient(credential, clientOptions);
       addClientToPool(cacheKey, type, awsClient);
 
       return awsClient;

src/plugins/data_source/server/legacy/configure_legacy_client.test.mocks.ts

@@ -16,3 +16,8 @@ export const parseClientOptionsMock = jest.fn();
 jest.doMock('./client_config', () => ({
   parseClientOptions: parseClientOptionsMock,
 }));
+
+export const authRegistryCredentialProviderMock = jest.fn();
+jest.doMock('../util/credential_provider', () => ({
+  authRegistryCredentialProvider: authRegistryCredentialProviderMock,
+}));

src/plugins/data_source/server/legacy/configure_legacy_client.test.ts

@@ -10,12 +10,20 @@ import { AuthType, DataSourceAttributes, SigV4Content } from '../../common/data_
 import { DataSourcePluginConfigType } from '../../config';
 import { cryptographyServiceSetupMock } from '../cryptography_service.mocks';
 import { CryptographyServiceSetup } from '../cryptography_service';
-import { DataSourceClientParams, LegacyClientCallAPIParams } from '../types';
+import { DataSourceClientParams, LegacyClientCallAPIParams, AuthenticationMethod } from '../types';
 import { OpenSearchClientPoolSetup } from '../client';
 import { ConfigOptions } from 'elasticsearch';
-import { ClientMock, parseClientOptionsMock } from './configure_legacy_client.test.mocks';
+import {
+  ClientMock,
+  parseClientOptionsMock,
+  authRegistryCredentialProviderMock,
+} from './configure_legacy_client.test.mocks';
 import { configureLegacyClient } from './configure_legacy_client';
 import { CustomApiSchemaRegistry } from '../schema_registry';
+import {
+  IAuthenticationMethodRegistery,
+  authenticationMethodRegisteryMock,
+} from '../auth_registry';
 
 const DATA_SOURCE_ID = 'a54b76ec86771ee865a0f74a305dfff8';
 
@@ -29,6 +37,7 @@ describe('configureLegacyClient', () => {
   let configOptions: ConfigOptions;
   let dataSourceAttr: DataSourceAttributes;
   let sigV4AuthContent: SigV4Content;
+  let authenticationMethodRegistery: jest.Mocked<IAuthenticationMethodRegistery>;
 
   let mockOpenSearchClientInstance: {
     close: jest.Mock;
@@ -48,6 +57,7 @@ describe('configureLegacyClient', () => {
     logger = loggingSystemMock.createLogger();
     savedObjectsMock = savedObjectsClientMock.create();
     cryptographyMock = cryptographyServiceSetupMock.create();
+    authenticationMethodRegistery = authenticationMethodRegisteryMock.create();
     config = {
       enabled: true,
       clientPool: {
@@ -254,4 +264,47 @@ describe('configureLegacyClient', () => {
     expect(mockOpenSearchClientInstance.ping).toHaveBeenCalledTimes(1);
     expect(mockOpenSearchClientInstance.ping).toHaveBeenLastCalledWith(mockParams);
   });
+
+  test('configureLegacyClient should retunrn client from authentication registery if method present in registry', async () => {
+    const name = 'typeA';
+    const customAuthContent = {
+      region: 'us-east-1',
+      roleARN: 'test-role',
+    };
+    savedObjectsMock.get.mockReset().mockResolvedValueOnce({
+      id: DATA_SOURCE_ID,
+      type: DATA_SOURCE_SAVED_OBJECT_TYPE,
+      attributes: {
+        ...dataSourceAttr,
+        auth: {
+          type: AuthType.SigV4,
+          credentials: customAuthContent,
+        },
+      },
+      references: [],
+    });
+    const authMethod: AuthenticationMethod = {
+      name,
+      authType: AuthType.SigV4,
+      credentialProvider: jest.fn(),
+    };
+    authenticationMethodRegistery.getAuthenticationMethod.mockImplementation(() => authMethod);
+
+    authRegistryCredentialProviderMock.mockReturnValue({
+      credential: sigV4AuthContent,
+      type: AuthType.SigV4,
+    });
+
+    await configureLegacyClient(
+      { ...dataSourceClientParams, authRegistry: authenticationMethodRegistery },
+      callApiParams,
+      clientPoolSetup,
+      config,
+      logger
+    );
+    expect(authRegistryCredentialProviderMock).toHaveBeenCalled();
+    expect(authenticationMethodRegistery.getAuthenticationMethod).toHaveBeenCalledTimes(1);
+    expect(ClientMock).toHaveBeenCalledTimes(1);
+    expect(savedObjectsMock.get).toHaveBeenCalledTimes(1);
+  });
 });