Adds changelog and UT

Signed-off-by: Bandini Bhopi <bandinib@amazon.com>
bandinib-amzn · Feb 16, 2024 · 3468a37 · 3468a37
1 parent 174cf73
commit 3468a37
Show file tree

Hide file tree

Showing 8 changed files with 141 additions and 6 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -39,6 +39,7 @@ Inspired from [Keep a Changelog](https://keepachangelog.com/en/1.0.0/)
 - [Multiple Datasource] Add datasource picker component and use it in devtools and tutorial page when multiple datasource is enabled ([#5756](https://github.com/opensearch-project/OpenSearch-Dashboards/pull/5756))
 - [Multiple Datasource] Add datasource picker to import saved object flyout when multiple data source is enabled ([#5781](https://github.com/opensearch-project/OpenSearch-Dashboards/pull/5781))
 - [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] Refactor client and legacy client to use authentication registry ([#5881](https://github.com/opensearch-project/OpenSearch-Dashboards/pull/5881))
 
 ### 🐛 Bug Fixes
 

diff --git a/src/plugins/data_source/server/auth_registry/authentication_methods_registry.mock.ts b/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 };
diff --git a/src/plugins/data_source/server/auth_registry/index.ts b/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';
diff --git a/src/plugins/data_source/server/client/configure_client.test.mocks.ts b/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,
+}));
diff --git a/src/plugins/data_source/server/client/configure_client.test.ts b/src/plugins/data_source/server/client/configure_client.test.ts
@@ -13,15 +13,23 @@ 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';
 // 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 {
+  IAuthenticationMethodRegistery,
+  authenticationMethodRegisteryMock,
+} from '../auth_registry';
 
 const DATA_SOURCE_ID = 'a54b76ec86771ee865a0f74a305dfff8';
 
@@ -38,12 +46,14 @@ describe('configureClient', () => {
   let dataSourceClientParams: DataSourceClientParams;
   let usernamePasswordAuthContent: UsernamePasswordTypedContent;
   let sigV4AuthContent: SigV4Content;
+  let authenticationMethodRegistery: jest.Mocked<IAuthenticationMethodRegistery>;
 
   beforeEach(() => {
     dsClient = opensearchClientMock.createInternalClient();
     logger = loggingSystemMock.createLogger();
     savedObjectsMock = savedObjectsClientMock.create();
     cryptographyMock = cryptographyServiceSetupMock.create();
+    authenticationMethodRegistery = authenticationMethodRegisteryMock.create();
 
     config = {
       enabled: true,
@@ -238,4 +248,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);
+  });
 });
diff --git a/src/plugins/data_source/server/legacy/configure_legacy_client.test.mocks.ts b/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,
+}));
diff --git a/src/plugins/data_source/server/legacy/configure_legacy_client.test.ts b/src/plugins/data_source/server/legacy/configure_legacy_client.test.ts
@@ -10,11 +10,19 @@ 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 {
+  IAuthenticationMethodRegistery,
+  authenticationMethodRegisteryMock,
+} from '../auth_registry';
 
 const DATA_SOURCE_ID = 'a54b76ec86771ee865a0f74a305dfff8';
 
@@ -28,6 +36,7 @@ describe('configureLegacyClient', () => {
   let configOptions: ConfigOptions;
   let dataSourceAttr: DataSourceAttributes;
   let sigV4AuthContent: SigV4Content;
+  let authenticationMethodRegistery: jest.Mocked<IAuthenticationMethodRegistery>;
 
   let mockOpenSearchClientInstance: {
     close: jest.Mock;
@@ -46,6 +55,7 @@ describe('configureLegacyClient', () => {
     logger = loggingSystemMock.createLogger();
     savedObjectsMock = savedObjectsClientMock.create();
     cryptographyMock = cryptographyServiceSetupMock.create();
+    authenticationMethodRegistery = authenticationMethodRegisteryMock.create();
     config = {
       enabled: true,
       clientPool: {
@@ -251,4 +261,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);
+  });
 });
diff --git a/src/plugins/data_source/server/legacy/configure_legacy_client.ts b/src/plugins/data_source/server/legacy/configure_legacy_client.ts
@@ -128,15 +128,18 @@ const getQueryClient = async (
       );
 
     case AuthType.UsernamePasswordType:
-      credential = await getCredential(dataSourceAttr, cryptography);
+      credential =
+        (credential as UsernamePasswordTypedContent) ??
+        (await getCredential(dataSourceAttr, cryptography));
 
       if (!rootClient) rootClient = new LegacyClient(clientOptions);
       addClientToPool(cacheKey, type, rootClient);
 
       return getBasicAuthClient(rootClient, { endpoint, clientParams, options }, credential);
 
     case AuthType.SigV4:
-      credential = await getAWSCredential(dataSourceAttr, cryptography);
+      credential =
+        (credential as SigV4Content) ?? (await getAWSCredential(dataSourceAttr, cryptography));
 
       const awsClient = rootClient ? rootClient : getAWSClient(credential, clientOptions);
       addClientToPool(cacheKey, type, awsClient);