From d273c57081dd3e4652d8d40e613b8f25c72bd25f Mon Sep 17 00:00:00 2001 From: Bandini Bhopi Date: Thu, 14 Mar 2024 09:40:02 +0000 Subject: [PATCH] Added UT for client pool Signed-off-by: Bandini Bhopi --- .../authentication_methods_registry.test.ts | 8 + .../authentication_methods_registry.ts | 10 + .../data_source/server/client/client_pool.ts | 6 +- .../server/client/configure_client.test.ts | 278 +++++++++++++- .../server/client/configure_client.ts | 38 +- .../server/client/configure_client_utils.ts | 45 ++- .../legacy/configure_legacy_client.test.ts | 357 +++++++++++++++++- .../server/legacy/configure_legacy_client.ts | 50 ++- 8 files changed, 758 insertions(+), 34 deletions(-) diff --git a/src/plugins/data_source/server/auth_registry/authentication_methods_registry.test.ts b/src/plugins/data_source/server/auth_registry/authentication_methods_registry.test.ts index c7692acee782..ce8e13f7cdff 100644 --- a/src/plugins/data_source/server/auth_registry/authentication_methods_registry.test.ts +++ b/src/plugins/data_source/server/auth_registry/authentication_methods_registry.test.ts @@ -6,6 +6,12 @@ import { AuthenticationMethodRegistery } from './authentication_methods_registry'; import { AuthenticationMethod } from '../../server/types'; import { AuthType } from '../../common/data_sources'; +import { OpenSearchClientPoolSetup } from '../client'; + +const clientPoolSetup: OpenSearchClientPoolSetup = { + getClientFromPool: jest.fn(), + addClientToPool: jest.fn(), +}; const createAuthenticationMethod = ( authMethod: Partial @@ -13,6 +19,8 @@ const createAuthenticationMethod = ( name: 'unknown', authType: AuthType.NoAuth, credentialProvider: jest.fn(), + clientPoolSetup, + legacyClientPoolSetup: clientPoolSetup, ...authMethod, }); diff --git a/src/plugins/data_source/server/auth_registry/authentication_methods_registry.ts b/src/plugins/data_source/server/auth_registry/authentication_methods_registry.ts index e2f39498e007..9fe2eb1e37e3 100644 --- a/src/plugins/data_source/server/auth_registry/authentication_methods_registry.ts +++ b/src/plugins/data_source/server/auth_registry/authentication_methods_registry.ts @@ -5,6 +5,7 @@ import { deepFreeze } from '@osd/std'; import { AuthenticationMethod } from '../../server/types'; +import { AuthType } from '../../common/data_sources'; export type IAuthenticationMethodRegistery = Omit< AuthenticationMethodRegistery, @@ -18,6 +19,15 @@ export class AuthenticationMethodRegistery { * Authentication Method can only be registered once. subsequent calls with the same method name will throw an error. */ public registerAuthenticationMethod(method: AuthenticationMethod) { + if ( + method.name === AuthType.NoAuth || + method.name === AuthType.UsernamePasswordType || + method.name === AuthType.SigV4 + ) { + throw new Error( + `Must not be no_auth or username_password or sigv4 for registered auth types` + ); + } if (this.authMethods.has(method.name)) { throw new Error(`Authentication method '${method.name}' is already registered`); } diff --git a/src/plugins/data_source/server/client/client_pool.ts b/src/plugins/data_source/server/client/client_pool.ts index 02ad665718df..0231833bdda9 100644 --- a/src/plugins/data_source/server/client/client_pool.ts +++ b/src/plugins/data_source/server/client/client_pool.ts @@ -82,7 +82,11 @@ export class OpenSearchClientPool { }); this.logger.info(`Created data source aws client pool of size ${size}`); - const getClientFromPool = (key: string, authType: AuthType) => { + const getClientFromPool = ( + key: string, + authType: AuthType, + request?: OpenSearchDashboardsRequest + ) => { const selectedCache = authType === AuthType.SigV4 ? this.awsClientCache : this.clientCache; return selectedCache!.get(key); diff --git a/src/plugins/data_source/server/client/configure_client.test.ts b/src/plugins/data_source/server/client/configure_client.test.ts index 5596a3258676..2030793c2d61 100644 --- a/src/plugins/data_source/server/client/configure_client.test.ts +++ b/src/plugins/data_source/server/client/configure_client.test.ts @@ -19,7 +19,7 @@ import { parseClientOptionsMock, authRegistryCredentialProviderMock, } from './configure_client.test.mocks'; -import { OpenSearchClientPoolSetup } from './client_pool'; +import { OpenSearchClientPool, OpenSearchClientPoolSetup } from './client_pool'; import { configureClient } from './configure_client'; import { ClientOptions } from '@opensearch-project/opensearch'; // eslint-disable-next-line @osd/eslint/no-restricted-paths @@ -361,4 +361,280 @@ describe('configureClient', () => { }, }); }); + + test('configure client with auth method from registry, service == aoss, should successfully call new Client()', async () => { + savedObjectsMock.get.mockReset().mockResolvedValueOnce({ + id: DATA_SOURCE_ID, + type: DATA_SOURCE_SAVED_OBJECT_TYPE, + attributes: { + ...dataSourceAttr, + auth: { + type: AuthType.SigV4, + credentials: { ...customAuthContent, service: 'aoss' }, + }, + }, + references: [], + }); + + authRegistryCredentialProviderMock.mockReturnValue({ + credential: sigV4AuthContent, + type: AuthType.SigV4, + }); + + const client = await configureClient( + { ...dataSourceClientParams, authRegistry: authenticationMethodRegistery }, + clientPoolSetup, + config, + logger + ); + + expect(ClientMock).toHaveBeenCalledTimes(1); + expect(client).toBe(dsClient.child.mock.results[0].value); + expect(dsClient.child).toBeCalledWith({ + auth: { + credentials: { + accessKeyId: sigV4AuthContent.accessKey, + secretAccessKey: sigV4AuthContent.secretKey, + sessionToken: '', + }, + region: sigV4AuthContent.region, + service: 'aoss', + }, + }); + }); + + describe('Client Pool', () => { + let opensearchClientPoolSetup: OpenSearchClientPoolSetup; + let openSearchClientPool: OpenSearchClientPool; + beforeEach(() => { + openSearchClientPool = new OpenSearchClientPool(logger); + opensearchClientPoolSetup = openSearchClientPool.setup(config); + }); + + describe('NoAuth', () => { + beforeEach(() => { + savedObjectsMock.get.mockReset().mockResolvedValue({ + id: DATA_SOURCE_ID, + type: DATA_SOURCE_SAVED_OBJECT_TYPE, + attributes: { + ...dataSourceAttr, + auth: { + type: AuthType.NoAuth, + }, + }, + references: [], + }); + }); + + test('For same endpoint only one client object should be created', async () => { + await configureClient(dataSourceClientParams, opensearchClientPoolSetup, config, logger); + await configureClient(dataSourceClientParams, opensearchClientPoolSetup, config, logger); + + expect(ClientMock).toHaveBeenCalledTimes(1); + }); + + test('For different endpoints multiple client objects should be created', async () => { + await configureClient(dataSourceClientParams, opensearchClientPoolSetup, config, logger); + + const mockDataSourceAttr = { ...dataSourceAttr, endpoint: 'http://test.com' }; + + savedObjectsMock.get.mockReset().mockResolvedValueOnce({ + id: DATA_SOURCE_ID, + type: DATA_SOURCE_SAVED_OBJECT_TYPE, + attributes: { + ...mockDataSourceAttr, + auth: { + type: AuthType.NoAuth, + }, + }, + references: [], + }); + + await configureClient(dataSourceClientParams, opensearchClientPoolSetup, config, logger); + + expect(ClientMock).toHaveBeenCalledTimes(2); + }); + }); + + describe('UserNamePassword', () => { + beforeEach(() => { + savedObjectsMock.get.mockReset().mockResolvedValue({ + id: DATA_SOURCE_ID, + type: DATA_SOURCE_SAVED_OBJECT_TYPE, + attributes: dataSourceAttr, + references: [], + }); + jest.spyOn(cryptographyMock, 'decodeAndDecrypt').mockResolvedValue({ + decryptedText: 'password', + encryptionContext: { endpoint: 'http://localhost' }, + }); + }); + + test('For same endpoint only one client object should be created', async () => { + await configureClient(dataSourceClientParams, opensearchClientPoolSetup, config, logger); + await configureClient(dataSourceClientParams, opensearchClientPoolSetup, config, logger); + + expect(ClientMock).toHaveBeenCalledTimes(1); + }); + + test('For different endpoints multiple client objects should be created', async () => { + await configureClient(dataSourceClientParams, opensearchClientPoolSetup, config, logger); + + const mockDataSourceAttr = { ...dataSourceAttr, endpoint: 'http://test.com' }; + savedObjectsMock.get.mockReset().mockResolvedValue({ + id: DATA_SOURCE_ID, + type: DATA_SOURCE_SAVED_OBJECT_TYPE, + attributes: mockDataSourceAttr, + references: [], + }); + jest.spyOn(cryptographyMock, 'decodeAndDecrypt').mockResolvedValue({ + decryptedText: 'password', + encryptionContext: { endpoint: 'http://test.com' }, + }); + + await configureClient(dataSourceClientParams, opensearchClientPoolSetup, config, logger); + + expect(ClientMock).toHaveBeenCalledTimes(2); + }); + }); + + describe('AWSSigV4', () => { + beforeEach(() => { + savedObjectsMock.get.mockReset().mockResolvedValue({ + id: DATA_SOURCE_ID, + type: DATA_SOURCE_SAVED_OBJECT_TYPE, + attributes: { + ...dataSourceAttr, + auth: { + type: AuthType.SigV4, + credentials: sigV4AuthContent, + }, + }, + references: [], + }); + + jest.spyOn(cryptographyMock, 'decodeAndDecrypt').mockResolvedValue({ + decryptedText: 'accessKey', + encryptionContext: { endpoint: 'http://localhost' }, + }); + }); + test('For same endpoint only one client object should be created', async () => { + await configureClient(dataSourceClientParams, opensearchClientPoolSetup, config, logger); + await configureClient(dataSourceClientParams, opensearchClientPoolSetup, config, logger); + + expect(ClientMock).toHaveBeenCalledTimes(1); + }); + + test('For different endpoints multiple client objects should be created', async () => { + await configureClient(dataSourceClientParams, opensearchClientPoolSetup, config, logger); + + const mockDataSourceAttr = { ...dataSourceAttr, endpoint: 'http://test.com' }; + savedObjectsMock.get.mockReset().mockResolvedValue({ + id: DATA_SOURCE_ID, + type: DATA_SOURCE_SAVED_OBJECT_TYPE, + attributes: { + ...mockDataSourceAttr, + auth: { + type: AuthType.SigV4, + credentials: sigV4AuthContent, + }, + }, + references: [], + }); + + jest.spyOn(cryptographyMock, 'decodeAndDecrypt').mockResolvedValue({ + decryptedText: 'accessKey', + encryptionContext: { endpoint: 'http://test.com' }, + }); + await configureClient(dataSourceClientParams, opensearchClientPoolSetup, config, logger); + + expect(ClientMock).toHaveBeenCalledTimes(2); + }); + }); + + describe('Auth Method from Registry', () => { + beforeEach(() => { + const authMethodWithClientPool: AuthenticationMethod = { + name: 'clientPoolTest', + authType: AuthType.SigV4, + credentialProvider: jest.fn(), + clientPoolSetup: opensearchClientPoolSetup, + legacyClientPoolSetup: clientPoolSetup, + }; + authenticationMethodRegistery.getAuthenticationMethod + .mockReset() + .mockImplementation(() => authMethodWithClientPool); + const mockDataSourceAttr = { ...dataSourceAttr, name: 'custom_auth' }; + savedObjectsMock.get.mockReset().mockResolvedValue({ + id: DATA_SOURCE_ID, + type: DATA_SOURCE_SAVED_OBJECT_TYPE, + attributes: { + ...mockDataSourceAttr, + auth: { + type: AuthType.SigV4, + credentials: customAuthContent, + }, + }, + references: [], + }); + authRegistryCredentialProviderMock.mockReturnValue({ + credential: sigV4AuthContent, + type: AuthType.SigV4, + }); + }); + test('Auth Method from Registry: If endpoint is same for multiple requests client pool size should be 1', async () => { + await configureClient( + { ...dataSourceClientParams, authRegistry: authenticationMethodRegistery }, + clientPoolSetup, + config, + logger + ); + + await configureClient( + { ...dataSourceClientParams, authRegistry: authenticationMethodRegistery }, + clientPoolSetup, + config, + logger + ); + + expect(ClientMock).toHaveBeenCalledTimes(1); + }); + + test('Auth Method from Registry: If endpoint is different for two requests client pool size should be 2', async () => { + await configureClient( + { ...dataSourceClientParams, authRegistry: authenticationMethodRegistery }, + clientPoolSetup, + config, + logger + ); + + const mockDataSourceAttr = { + ...dataSourceAttr, + endpoint: 'http://test.com', + name: 'custom_auth', + }; + savedObjectsMock.get.mockReset().mockResolvedValue({ + id: DATA_SOURCE_ID, + type: DATA_SOURCE_SAVED_OBJECT_TYPE, + attributes: { + ...mockDataSourceAttr, + auth: { + type: AuthType.SigV4, + credentials: customAuthContent, + }, + }, + references: [], + }); + + await configureClient( + { ...dataSourceClientParams, authRegistry: authenticationMethodRegistery }, + clientPoolSetup, + config, + logger + ); + + expect(ClientMock).toHaveBeenCalledTimes(2); + }); + }); + }); }); diff --git a/src/plugins/data_source/server/client/configure_client.ts b/src/plugins/data_source/server/client/configure_client.ts index 3c9c0825dca0..a0592c8799aa 100644 --- a/src/plugins/data_source/server/client/configure_client.ts +++ b/src/plugins/data_source/server/client/configure_client.ts @@ -16,11 +16,16 @@ import { import { DataSourcePluginConfigType } from '../../config'; import { CryptographyServiceSetup } from '../cryptography_service'; import { createDataSourceError } from '../lib/error'; -import { DataSourceClientParams } from '../types'; +import { AuthenticationMethod, DataSourceClientParams } from '../types'; import { parseClientOptions } from './client_config'; import { OpenSearchClientPoolSetup } from './client_pool'; -import { getAWSCredential, getCredential, getDataSource } from './configure_client_utils'; -import { IAuthenticationMethodRegistery } from '../auth_registry'; +import { + getRootClient, + getAWSCredential, + getCredential, + getDataSource, + getAuthenticationMethod, +} from './configure_client_utils'; import { authRegistryCredentialProvider } from '../util/credential_provider'; export const configureClient = async ( @@ -61,17 +66,25 @@ export const configureClient = async ( dataSource = await getDataSource(dataSourceId!, savedObjects); } + let clientPool = openSearchClientPoolSetup; + const authenticationMethod = getAuthenticationMethod(dataSource, authRegistry); + if (authenticationMethod !== undefined) { + clientPool = authenticationMethod.clientPoolSetup; + } + const rootClient = getRootClient(dataSource, clientPool.getClientFromPool, request) as Client; + const registeredSchema = (await customApiSchemaRegistryPromise).getAll(); return await getQueryClient( dataSource, - openSearchClientPoolSetup, + clientPool, config, registeredSchema, cryptography, + rootClient, dataSourceId, request, - authRegistry, + authenticationMethod, requireDecryption ); } catch (error: any) { @@ -100,26 +113,23 @@ export const configureClient = async ( */ const getQueryClient = async ( dataSourceAttr: DataSourceAttributes, - openSearchClientPool: OpenSearchClientPoolSetup, + clientPool: OpenSearchClientPoolSetup, config: DataSourcePluginConfigType, registeredSchema: any[], cryptography?: CryptographyServiceSetup, + rootClient?: Client, dataSourceId?: string, request?: OpenSearchDashboardsRequest, - authRegistry?: IAuthenticationMethodRegistery, + authenticationMethod?: AuthenticationMethod, requireDecryption: boolean = true ): Promise => { let credential; - let clientPool = openSearchClientPool; let { auth: { type }, - name, } = dataSourceAttr; const { endpoint } = dataSourceAttr; - name = name ?? type; const clientOptions = parseClientOptions(config, endpoint, registeredSchema); - const authenticationMethod = authRegistry?.getAuthenticationMethod(name); if (authenticationMethod !== undefined) { const credentialProvider = await authRegistryCredentialProvider(authenticationMethod, { dataSourceAttr, @@ -128,11 +138,13 @@ const getQueryClient = async ( }); credential = credentialProvider.credential; type = credentialProvider.type; - clientPool = authenticationMethod.clientPoolSetup; + + if (credential.service === undefined) { + credential = { ...credential, service: dataSourceAttr.auth.credentials?.service }; + } } const cacheKey = endpoint; - let rootClient = clientPool.getClientFromPool(cacheKey, type, request) as Client; switch (type) { case AuthType.NoAuth: diff --git a/src/plugins/data_source/server/client/configure_client_utils.ts b/src/plugins/data_source/server/client/configure_client_utils.ts index 412dae5080a1..6fcdc1e2565e 100644 --- a/src/plugins/data_source/server/client/configure_client_utils.ts +++ b/src/plugins/data_source/server/client/configure_client_utils.ts @@ -3,15 +3,50 @@ * SPDX-License-Identifier: Apache-2.0 */ -import { SavedObjectsClientContract } from '../../../../../src/core/server'; +import { Client } from '@opensearch-project/opensearch'; +import { Client as LegacyClient } from 'elasticsearch'; +import { + OpenSearchDashboardsRequest, + SavedObjectsClientContract, +} from '../../../../../src/core/server'; import { DATA_SOURCE_SAVED_OBJECT_TYPE } from '../../common'; import { DataSourceAttributes, UsernamePasswordTypedContent, SigV4Content, + AuthType, } from '../../common/data_sources'; import { CryptographyServiceSetup } from '../cryptography_service'; import { createDataSourceError } from '../lib/error'; +import { IAuthenticationMethodRegistery } from '../auth_registry'; +import { AuthenticationMethod } from '../types'; + +/** + * Get the root client of datasource from + * client cache. If there's a cache miss, return undefined. + * + * @param dataSourceAttr data source saved objects attributes + * @param dataSourceId id of data source saved Object + * @param addClientToPool function to get client from client pool + * @returns cached OpenSearch client, or undefined if cache miss + */ +export const getRootClient = ( + dataSourceAttr: DataSourceAttributes, + getClientFromPool: ( + endpoint: string, + authType: AuthType, + request?: OpenSearchDashboardsRequest + ) => Client | LegacyClient | undefined, + request?: OpenSearchDashboardsRequest +): Client | LegacyClient | undefined => { + const { + auth: { type }, + endpoint, + } = dataSourceAttr; + const cacheKey = endpoint; + + return getClientFromPool(cacheKey, type, request); +}; export const getDataSource = async ( dataSourceId: string, @@ -93,3 +128,11 @@ export const getAWSCredential = async ( return credential; }; + +export const getAuthenticationMethod = ( + dataSourceAttr: DataSourceAttributes, + authRegistry?: IAuthenticationMethodRegistery +): AuthenticationMethod => { + const name = dataSourceAttr.name ?? dataSourceAttr.auth.type; + return authRegistry?.getAuthenticationMethod(name) as AuthenticationMethod; +}; 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 index be54ac0bb4d6..c35434417456 100644 --- a/src/plugins/data_source/server/legacy/configure_legacy_client.test.ts +++ b/src/plugins/data_source/server/legacy/configure_legacy_client.test.ts @@ -16,7 +16,7 @@ import { DataSourcePluginConfigType } from '../../config'; import { cryptographyServiceSetupMock } from '../cryptography_service.mocks'; import { CryptographyServiceSetup } from '../cryptography_service'; import { DataSourceClientParams, LegacyClientCallAPIParams, AuthenticationMethod } from '../types'; -import { OpenSearchClientPoolSetup } from '../client'; +import { OpenSearchClientPool, OpenSearchClientPoolSetup } from '../client'; import { ConfigOptions } from 'elasticsearch'; import { ClientMock, parseClientOptionsMock } from './configure_legacy_client.test.mocks'; import { authRegistryCredentialProviderMock } from '../client/configure_client.test.mocks'; @@ -383,4 +383,359 @@ describe('configureLegacyClient', () => { }, }); }); + test('configureLegacyClient with auth method from registry, service == aoss, should successfully call new Client()', async () => { + savedObjectsMock.get.mockReset().mockResolvedValueOnce({ + id: DATA_SOURCE_ID, + type: DATA_SOURCE_SAVED_OBJECT_TYPE, + attributes: { + ...dataSourceAttr, + auth: { + type: AuthType.SigV4, + credentials: { ...customAuthContent, service: 'aoss' }, + }, + }, + references: [], + }); + + 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); + expect(mockOpenSearchClientInstance.ping).toHaveBeenCalledTimes(1); + expect(mockOpenSearchClientInstance.ping).toHaveBeenLastCalledWith({ + headers: { + auth: { + credentials: { + accessKeyId: sigV4AuthContent.accessKey, + secretAccessKey: sigV4AuthContent.secretKey, + sessionToken: '', + }, + region: sigV4AuthContent.region, + service: 'aoss', + }, + }, + }); + }); + + describe('Client Pool', () => { + let opensearchClientPoolSetup: OpenSearchClientPoolSetup; + let openSearchClientPool: OpenSearchClientPool; + beforeEach(() => { + openSearchClientPool = new OpenSearchClientPool(logger); + opensearchClientPoolSetup = openSearchClientPool.setup(config); + }); + + describe('NoAuth', () => { + beforeEach(() => { + savedObjectsMock.get.mockReset().mockResolvedValue({ + id: DATA_SOURCE_ID, + type: DATA_SOURCE_SAVED_OBJECT_TYPE, + attributes: { + ...dataSourceAttr, + auth: { + type: AuthType.NoAuth, + }, + }, + references: [], + }); + }); + + test('For same endpoint only one client object should be created', async () => { + await configureLegacyClient( + dataSourceClientParams, + callApiParams, + opensearchClientPoolSetup, + config, + logger + ); + await configureLegacyClient( + dataSourceClientParams, + callApiParams, + opensearchClientPoolSetup, + config, + logger + ); + + expect(ClientMock).toHaveBeenCalledTimes(1); + }); + + test('For different endpoints multiple client objects should be created', async () => { + await configureLegacyClient( + dataSourceClientParams, + callApiParams, + opensearchClientPoolSetup, + config, + logger + ); + + const mockDataSourceAttr = { ...dataSourceAttr, endpoint: 'http://test.com' }; + + savedObjectsMock.get.mockReset().mockResolvedValueOnce({ + id: DATA_SOURCE_ID, + type: DATA_SOURCE_SAVED_OBJECT_TYPE, + attributes: { + ...mockDataSourceAttr, + auth: { + type: AuthType.NoAuth, + }, + }, + references: [], + }); + + await configureLegacyClient( + dataSourceClientParams, + callApiParams, + opensearchClientPoolSetup, + config, + logger + ); + + expect(ClientMock).toHaveBeenCalledTimes(2); + }); + }); + + describe('UserNamePassword', () => { + beforeEach(() => { + savedObjectsMock.get.mockReset().mockResolvedValue({ + id: DATA_SOURCE_ID, + type: DATA_SOURCE_SAVED_OBJECT_TYPE, + attributes: dataSourceAttr, + references: [], + }); + jest.spyOn(cryptographyMock, 'decodeAndDecrypt').mockResolvedValue({ + decryptedText: 'password', + encryptionContext: { endpoint: 'http://localhost' }, + }); + }); + + test('For same endpoint only one client object should be created', async () => { + await configureLegacyClient( + dataSourceClientParams, + callApiParams, + opensearchClientPoolSetup, + config, + logger + ); + await configureLegacyClient( + dataSourceClientParams, + callApiParams, + opensearchClientPoolSetup, + config, + logger + ); + + expect(ClientMock).toHaveBeenCalledTimes(1); + }); + + test('For different endpoints multiple client objects should be created', async () => { + await configureLegacyClient( + dataSourceClientParams, + callApiParams, + opensearchClientPoolSetup, + config, + logger + ); + + const mockDataSourceAttr = { ...dataSourceAttr, endpoint: 'http://test.com' }; + savedObjectsMock.get.mockReset().mockResolvedValue({ + id: DATA_SOURCE_ID, + type: DATA_SOURCE_SAVED_OBJECT_TYPE, + attributes: mockDataSourceAttr, + references: [], + }); + jest.spyOn(cryptographyMock, 'decodeAndDecrypt').mockResolvedValue({ + decryptedText: 'password', + encryptionContext: { endpoint: 'http://test.com' }, + }); + + await configureLegacyClient( + dataSourceClientParams, + callApiParams, + opensearchClientPoolSetup, + config, + logger + ); + + expect(ClientMock).toHaveBeenCalledTimes(2); + }); + }); + + describe('AWSSigV4', () => { + beforeEach(() => { + savedObjectsMock.get.mockReset().mockResolvedValue({ + id: DATA_SOURCE_ID, + type: DATA_SOURCE_SAVED_OBJECT_TYPE, + attributes: { + ...dataSourceAttr, + auth: { + type: AuthType.SigV4, + credentials: sigV4AuthContent, + }, + }, + references: [], + }); + + jest.spyOn(cryptographyMock, 'decodeAndDecrypt').mockResolvedValue({ + decryptedText: 'accessKey', + encryptionContext: { endpoint: 'http://localhost' }, + }); + }); + test('For same endpoint only one client object should be created', async () => { + await configureLegacyClient( + dataSourceClientParams, + callApiParams, + opensearchClientPoolSetup, + config, + logger + ); + await configureLegacyClient( + dataSourceClientParams, + callApiParams, + opensearchClientPoolSetup, + config, + logger + ); + + expect(ClientMock).toHaveBeenCalledTimes(1); + }); + + test('For different endpoints multiple client objects should be created', async () => { + await configureLegacyClient( + dataSourceClientParams, + callApiParams, + opensearchClientPoolSetup, + config, + logger + ); + const mockDataSourceAttr = { ...dataSourceAttr, endpoint: 'http://test.com' }; + savedObjectsMock.get.mockReset().mockResolvedValue({ + id: DATA_SOURCE_ID, + type: DATA_SOURCE_SAVED_OBJECT_TYPE, + attributes: { + ...mockDataSourceAttr, + auth: { + type: AuthType.SigV4, + credentials: sigV4AuthContent, + }, + }, + references: [], + }); + + jest.spyOn(cryptographyMock, 'decodeAndDecrypt').mockResolvedValue({ + decryptedText: 'accessKey', + encryptionContext: { endpoint: 'http://test.com' }, + }); + await configureLegacyClient( + dataSourceClientParams, + callApiParams, + opensearchClientPoolSetup, + config, + logger + ); + + expect(ClientMock).toHaveBeenCalledTimes(2); + }); + }); + + describe('Auth Method from Registry', () => { + beforeEach(() => { + const authMethodWithClientPool: AuthenticationMethod = { + name: 'clientPoolTest', + authType: AuthType.SigV4, + credentialProvider: jest.fn(), + clientPoolSetup, + legacyClientPoolSetup: opensearchClientPoolSetup, + }; + authenticationMethodRegistery.getAuthenticationMethod + .mockReset() + .mockImplementation(() => authMethodWithClientPool); + const mockDataSourceAttr = { ...dataSourceAttr, name: 'custom_auth' }; + savedObjectsMock.get.mockReset().mockResolvedValue({ + id: DATA_SOURCE_ID, + type: DATA_SOURCE_SAVED_OBJECT_TYPE, + attributes: { + ...mockDataSourceAttr, + auth: { + type: AuthType.SigV4, + credentials: customAuthContent, + }, + }, + references: [], + }); + authRegistryCredentialProviderMock.mockReturnValue({ + credential: sigV4AuthContent, + type: AuthType.SigV4, + }); + }); + test('Auth Method from Registry: If endpoint is same for multiple requests client pool size should be 1', async () => { + await configureLegacyClient( + { ...dataSourceClientParams, authRegistry: authenticationMethodRegistery }, + callApiParams, + clientPoolSetup, + config, + logger + ); + + await configureLegacyClient( + { ...dataSourceClientParams, authRegistry: authenticationMethodRegistery }, + callApiParams, + clientPoolSetup, + config, + logger + ); + + expect(ClientMock).toHaveBeenCalledTimes(1); + }); + + test('Auth Method from Registry: If endpoint is different for two requests client pool size should be 2', async () => { + await configureLegacyClient( + { ...dataSourceClientParams, authRegistry: authenticationMethodRegistery }, + callApiParams, + clientPoolSetup, + config, + logger + ); + + const mockDataSourceAttr = { + ...dataSourceAttr, + endpoint: 'http://test.com', + name: 'custom_auth', + }; + savedObjectsMock.get.mockReset().mockResolvedValue({ + id: DATA_SOURCE_ID, + type: DATA_SOURCE_SAVED_OBJECT_TYPE, + attributes: { + ...mockDataSourceAttr, + auth: { + type: AuthType.SigV4, + credentials: customAuthContent, + }, + }, + references: [], + }); + + await configureLegacyClient( + { ...dataSourceClientParams, authRegistry: authenticationMethodRegistery }, + callApiParams, + clientPoolSetup, + config, + logger + ); + + expect(ClientMock).toHaveBeenCalledTimes(2); + }); + }); + }); }); diff --git a/src/plugins/data_source/server/legacy/configure_legacy_client.ts b/src/plugins/data_source/server/legacy/configure_legacy_client.ts index 48ef8f6b57ce..8cd46be10e21 100644 --- a/src/plugins/data_source/server/legacy/configure_legacy_client.ts +++ b/src/plugins/data_source/server/legacy/configure_legacy_client.ts @@ -24,12 +24,17 @@ import { } from '../../common/data_sources'; import { DataSourcePluginConfigType } from '../../config'; import { CryptographyServiceSetup } from '../cryptography_service'; -import { DataSourceClientParams, LegacyClientCallAPIParams } from '../types'; +import { AuthenticationMethod, DataSourceClientParams, LegacyClientCallAPIParams } from '../types'; import { OpenSearchClientPoolSetup } from '../client'; import { parseClientOptions } from './client_config'; import { createDataSourceError } from '../lib/error'; -import { getAWSCredential, getCredential, getDataSource } from '../client/configure_client_utils'; -import { IAuthenticationMethodRegistery } from '../auth_registry'; +import { + getRootClient, + getAWSCredential, + getCredential, + getDataSource, + getAuthenticationMethod, +} from '../client/configure_client_utils'; import { authRegistryCredentialProvider } from '../util/credential_provider'; export const configureLegacyClient = async ( @@ -49,18 +54,30 @@ export const configureLegacyClient = async ( try { const dataSourceAttr = await getDataSource(dataSourceId!, savedObjects); + let clientPool = openSearchClientPoolSetup; + const authenticationMethod = getAuthenticationMethod(dataSourceAttr, authRegistry); + if (authenticationMethod !== undefined) { + clientPool = authenticationMethod.legacyClientPoolSetup; + } + const rootClient = getRootClient( + dataSourceAttr, + clientPool.getClientFromPool, + request + ) as LegacyClient; + const registeredSchema = (await customApiSchemaRegistryPromise).getAll(); return await getQueryClient( dataSourceAttr, cryptography, callApiParams, - openSearchClientPoolSetup, + clientPool, config, registeredSchema, + rootClient, dataSourceId, request, - authRegistry + authenticationMethod ); } catch (error: any) { logger.debug( @@ -87,24 +104,21 @@ const getQueryClient = async ( dataSourceAttr: DataSourceAttributes, cryptography: CryptographyServiceSetup, { endpoint, clientParams, options }: LegacyClientCallAPIParams, - openSearchClientPool: OpenSearchClientPoolSetup, + clientPool: OpenSearchClientPoolSetup, config: DataSourcePluginConfigType, registeredSchema: any[], + rootClient?: LegacyClient, dataSourceId?: string, request?: OpenSearchDashboardsRequest, - authRegistry?: IAuthenticationMethodRegistery + authenticationMethod?: AuthenticationMethod ) => { let credential; - let clientPool = openSearchClientPool; let { auth: { type }, - name, } = dataSourceAttr; const { endpoint: nodeUrl } = dataSourceAttr; - name = name ?? type; const clientOptions = parseClientOptions(config, nodeUrl, registeredSchema); - const authenticationMethod = authRegistry?.getAuthenticationMethod(name); if (authenticationMethod !== undefined) { const credentialProvider = await authRegistryCredentialProvider(authenticationMethod, { dataSourceAttr, @@ -113,16 +127,18 @@ const getQueryClient = async ( }); credential = credentialProvider.credential; type = credentialProvider.type; - clientPool = authenticationMethod.legacyClientPoolSetup; + + if (credential.service === undefined) { + credential = { ...credential, service: dataSourceAttr.auth.credentials?.service }; + } } - const cacheKey = endpoint; - let rootClient = clientPool.getClientFromPool(cacheKey, type, request) as LegacyClient; + const cacheKey = nodeUrl; switch (type) { case AuthType.NoAuth: if (!rootClient) rootClient = new LegacyClient(clientOptions); - clientPool.addClientToPool(cacheKey, type, rootClient); + clientPool.addClientToPool(cacheKey, type, rootClient, request); return await (callAPI.bind(null, rootClient) as LegacyAPICaller)( endpoint, @@ -136,7 +152,7 @@ const getQueryClient = async ( (await getCredential(dataSourceAttr, cryptography)); if (!rootClient) rootClient = new LegacyClient(clientOptions); - clientPool.addClientToPool(cacheKey, type, rootClient); + clientPool.addClientToPool(cacheKey, type, rootClient, request); return getBasicAuthClient(rootClient, { endpoint, clientParams, options }, credential); @@ -147,7 +163,7 @@ const getQueryClient = async ( if (!rootClient) { rootClient = getAWSClient(credential, clientOptions); } - clientPool.addClientToPool(cacheKey, type, rootClient); + clientPool.addClientToPool(cacheKey, type, rootClient, request); return await getAWSChildClient(rootClient, { endpoint, clientParams, options }, credential);