[MD] Integrate test connection to support SigV4 auth type

Signed-off-by: Su <szhongna@amazon.com>

CHANGELOG.md

@@ -62,6 +62,7 @@ Inspired from [Keep a Changelog](https://keepachangelog.com/en/1.0.0/)
 - [Data] Update `createAggConfig` so that newly created configs can be added to beginning of `aggConfig` array ([#3160](https://github.com/opensearch-project/OpenSearch-Dashboards/pull/3160))
 - Add disablePrototypePoisoningProtection configuration to prevent JS client from erroring when cluster utilizes JS reserved words ([#2992](https://github.com/opensearch-project/OpenSearch-Dashboards/pull/2992))
 - [Multiple DataSource] Add support for SigV4 authentication ([#3058](https://github.com/opensearch-project/OpenSearch-Dashboards/issues/3058))
+- [Multiple DataSource] Refactor test connection to support SigV4 auth type ([#3456](https://github.com/opensearch-project/OpenSearch-Dashboards/issues/3456))
 
 ### 🐛 Bug Fixes
 

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

@@ -6,7 +6,12 @@
 import { SavedObjectsClientContract } from '../../../../core/server';
 import { loggingSystemMock, savedObjectsClientMock } from '../../../../core/server/mocks';
 import { DATA_SOURCE_SAVED_OBJECT_TYPE } from '../../common';
-import { DataSourceAttributes, AuthType } from '../../common/data_sources/types';
+import {
+  DataSourceAttributes,
+  AuthType,
+  UsernamePasswordTypedContent,
+  SigV4Content,
+} from '../../common/data_sources/types';
 import { DataSourcePluginConfigType } from '../../config';
 import { ClientMock, parseClientOptionsMock } from './configure_client.test.mocks';
 import { OpenSearchClientPoolSetup } from './client_pool';
@@ -31,6 +36,8 @@ describe('configureClient', () => {
   let dataSourceAttr: DataSourceAttributes;
   let dsClient: ReturnType<typeof opensearchClientMock.createInternalClient>;
   let dataSourceClientParams: DataSourceClientParams;
+  let usernamePasswordAuthContent: UsernamePasswordTypedContent;
+  let sigV4AuthContent: SigV4Content;
 
   beforeEach(() => {
     dsClient = opensearchClientMock.createInternalClient();
@@ -51,15 +58,24 @@ describe('configureClient', () => {
         rejectUnauthorized: true,
       },
     } as ClientOptions;
+
+    usernamePasswordAuthContent = {
+      username: 'username',
+      password: 'password',
+    };
+
+    sigV4AuthContent = {
+      region: 'us-east-1',
+      accessKey: 'accessKey',
+      secretKey: 'secretKey',
+    };
+
     dataSourceAttr = {
       title: 'title',
       endpoint: 'http://localhost',
       auth: {
         type: AuthType.UsernamePasswordType,
-        credentials: {
-          username: 'username',
-          password: 'password',
-        },
+        credentials: usernamePasswordAuthContent,
       },
     } as DataSourceAttributes;
 
@@ -126,6 +142,48 @@ describe('configureClient', () => {
     expect(client).toBe(dsClient.child.mock.results[0].value);
   });
 
+  test('configure client with auth.type == sigv4, will first call decodeAndDecrypt()', async () => {
+    savedObjectsMock.get.mockReset().mockResolvedValueOnce({
+      id: DATA_SOURCE_ID,
+      type: DATA_SOURCE_SAVED_OBJECT_TYPE,
+      attributes: {
+        ...dataSourceAttr,
+        auth: {
+          type: AuthType.SigV4,
+          credentials: sigV4AuthContent,
+        },
+      },
+      references: [],
+    });
+
+    const decodeAndDecryptSpy = jest.spyOn(cryptographyMock, 'decodeAndDecrypt').mockResolvedValue({
+      decryptedText: 'accessKey',
+      encryptionContext: { endpoint: 'http://localhost' },
+    });
+    await configureClient(dataSourceClientParams, clientPoolSetup, config, logger);
+
+    expect(ClientMock).toHaveBeenCalledTimes(1);
+    expect(savedObjectsMock.get).toHaveBeenCalledTimes(1);
+    expect(decodeAndDecryptSpy).toHaveBeenCalledTimes(2);
+  });
+
+  test('configure test client for non-exist datasource should not call saved object api, nor decode any credential', async () => {
+    const decodeAndDecryptSpy = jest.spyOn(cryptographyMock, 'decodeAndDecrypt').mockResolvedValue({
+      decryptedText: 'password',
+      encryptionContext: { endpoint: 'http://localhost' },
+    });
+    const testClientParams: DataSourceClientParams = {
+      ...dataSourceClientParams,
+      testClientDataSourceAttr: dataSourceAttr,
+      dataSourceId: undefined,
+    };
+    await configureClient(testClientParams, clientPoolSetup, config, logger);
+
+    expect(ClientMock).toHaveBeenCalledTimes(1);
+    expect(savedObjectsMock.get).not.toHaveBeenCalled();
+    expect(decodeAndDecryptSpy).not.toHaveBeenCalled();
+  });
+
   test('configure client with auth.type == username_password and password contaminated', async () => {
     const decodeAndDecryptSpy = jest
       .spyOn(cryptographyMock, 'decodeAndDecrypt')

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

@@ -29,13 +29,35 @@ import {
 } from './configure_client_utils';
 
 export const configureClient = async (
-  { dataSourceId, savedObjects, cryptography }: DataSourceClientParams,
+  { dataSourceId, savedObjects, cryptography, testClientDataSourceAttr }: DataSourceClientParams,
   openSearchClientPoolSetup: OpenSearchClientPoolSetup,
   config: DataSourcePluginConfigType,
   logger: Logger
 ): Promise<Client> => {
+  let dataSource;
+  let requireDecryption = true;
+
   try {
-    const dataSource = await getDataSource(dataSourceId!, savedObjects);
+    // configure test client
+    if (testClientDataSourceAttr) {
+      const {
+        auth: { type, credentials },
+      } = testClientDataSourceAttr;
+      // handle test connection case when changing non-credential field of existing data source
+      if (
+        dataSourceId &&
+        ((type === AuthType.UsernamePasswordType && !credentials?.password) ||
+          (type === AuthType.SigV4 && !credentials?.accessKey && !credentials?.secretKey))
+      ) {
+        dataSource = await getDataSource(dataSourceId, savedObjects);
+      } else {
+        dataSource = testClientDataSourceAttr;
+        requireDecryption = false;
+      }
+    } else {
+      dataSource = await getDataSource(dataSourceId!, savedObjects);
+    }
+
     const rootClient = getRootClient(
       dataSource,
       openSearchClientPoolSetup.getClientFromPool,
@@ -48,7 +70,8 @@ export const configureClient = async (
       config,
       cryptography,
       rootClient,
-      dataSourceId
+      dataSourceId,
+      requireDecryption
     );
   } catch (error: any) {
     logger.error(
@@ -59,46 +82,6 @@ export const configureClient = async (
   }
 };
 
-export const configureTestClient = async (
-  { savedObjects, cryptography, dataSourceId }: DataSourceClientParams,
-  dataSourceAttr: DataSourceAttributes,
-  openSearchClientPoolSetup: OpenSearchClientPoolSetup,
-  config: DataSourcePluginConfigType,
-  logger: Logger
-): Promise<Client> => {
-  try {
-    const {
-      auth: { type, credentials },
-    } = dataSourceAttr;
-    let requireDecryption = false;
-
-    const rootClient = getRootClient(
-      dataSourceAttr,
-      openSearchClientPoolSetup.getClientFromPool,
-      dataSourceId
-    ) as Client;
-
-    if (type === AuthType.UsernamePasswordType && !credentials?.password && dataSourceId) {
-      dataSourceAttr = await getDataSource(dataSourceId, savedObjects);
-      requireDecryption = true;
-    }
-
-    return getQueryClient(
-      dataSourceAttr,
-      openSearchClientPoolSetup.addClientToPool,
-      config,
-      cryptography,
-      rootClient,
-      dataSourceId,
-      requireDecryption
-    );
-  } catch (error: any) {
-    logger.error(`Failed to get test client. ${error}: ${error.stack}`);
-    // Re-throw as DataSourceError
-    throw createDataSourceError(error);
-  }
-};
-
 /**
  * Create a child client object with given auth info.
  *
@@ -108,7 +91,7 @@ export const configureTestClient = async (
  * @param config data source config
  * @param addClientToPool function to add client to client pool
  * @param dataSourceId id of data source saved Object
- * @param requireDecryption boolean
+ * @param requireDecryption false when creating test client before data source exists
  * @returns Promise of query client
  */
 const getQueryClient = async (

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

@@ -32,9 +32,17 @@ export const getRootClient = (
 ): Client | LegacyClient | undefined => {
   const {
     auth: { type },
+    lastUpdatedTime,
   } = dataSourceAttr;
+  let cachedClient;
   const cacheKey = generateCacheKey(dataSourceAttr, dataSourceId);
-  const cachedClient = getClientFromPool(cacheKey, type);
+
+  // return undefined when building SigV4 test client with new credentials
+  if (type === AuthType.SigV4) {
+    cachedClient = dataSourceId && lastUpdatedTime ? getClientFromPool(cacheKey, type) : undefined;
+  } else {
+    cachedClient = getClientFromPool(cacheKey, type);
+  }
 
   return cachedClient;
 };

src/plugins/data_source/server/data_source_service.ts

@@ -8,8 +8,7 @@ import { DataSourcePluginConfigType } from '../config';
 import { OpenSearchClientPool } from './client';
 import { configureLegacyClient } from './legacy';
 import { DataSourceClientParams } from './types';
-import { DataSourceAttributes } from '../common/data_sources';
-import { configureTestClient, configureClient } from './client/configure_client';
+import { configureClient } from './client/configure_client';
 export interface DataSourceServiceSetup {
   getDataSourceClient: (params: DataSourceClientParams) => Promise<OpenSearchClient>;
 
@@ -22,11 +21,6 @@ export interface DataSourceServiceSetup {
       options?: LegacyCallAPIOptions
     ) => Promise<unknown>;
   };
-
-  getTestingClient: (
-    params: DataSourceClientParams,
-    dataSource: DataSourceAttributes
-  ) => Promise<OpenSearchClient>;
 }
 export class DataSourceService {
   private readonly openSearchClientPool: OpenSearchClientPool;
@@ -49,19 +43,6 @@ export class DataSourceService {
       return configureClient(params, opensearchClientPoolSetup, config, this.logger);
     };
 
-    const getTestingClient = async (
-      params: DataSourceClientParams,
-      dataSource: DataSourceAttributes
-    ): Promise<OpenSearchClient> => {
-      return configureTestClient(
-        params,
-        dataSource,
-        opensearchClientPoolSetup,
-        config,
-        this.logger
-      );
-    };
-
     const getDataSourceLegacyClient = (params: DataSourceClientParams) => {
       return {
         callAPI: (
@@ -79,7 +60,7 @@ export class DataSourceService {
       };
     };
 
-    return { getDataSourceClient, getDataSourceLegacyClient, getTestingClient };
+    return { getDataSourceClient, getDataSourceLegacyClient };
   }
 
   start() {}

src/plugins/data_source/server/routes/test_connection.ts

@@ -5,7 +5,7 @@
 
 import { schema } from '@osd/config-schema';
 import { IRouter, OpenSearchClient } from 'opensearch-dashboards/server';
-import { DataSourceAttributes } from '../../common/data_sources';
+import { AuthType, DataSourceAttributes } from '../../common/data_sources';
 import { DataSourceConnectionValidator } from './data_source_connection_validator';
 import { DataSourceServiceSetup } from '../data_source_service';
 import { CryptographyServiceSetup } from '../cryptography_service';
@@ -26,14 +26,20 @@ export const registerTestConnectionRoute = (
             auth: schema.maybe(
               schema.object({
                 type: schema.oneOf([
-                  schema.literal('username_password'),
-                  schema.literal('no_auth'),
+                  schema.literal(AuthType.UsernamePasswordType),
+                  schema.literal(AuthType.NoAuth),
+                  schema.literal(AuthType.SigV4),
                 ]),
                 credentials: schema.oneOf([
                   schema.object({
                     username: schema.string(),
                     password: schema.string(),
                   }),
+                  schema.object({
+                    region: schema.string(),
+                    accessKey: schema.string(),
+                    secretKey: schema.string(),
+                  }),
                   schema.literal(null),
                 ]),
               })
@@ -46,13 +52,13 @@ export const registerTestConnectionRoute = (
       const { dataSourceAttr, id: dataSourceId } = request.body;
 
       try {
-        const dataSourceClient: OpenSearchClient = await dataSourceServiceSetup.getTestingClient(
+        const dataSourceClient: OpenSearchClient = await dataSourceServiceSetup.getDataSourceClient(
           {
-            dataSourceId,
             savedObjects: context.core.savedObjects.client,
             cryptography,
-          },
-          dataSourceAttr as DataSourceAttributes
+            dataSourceId,
+            testClientDataSourceAttr: dataSourceAttr as DataSourceAttributes,
+          }
         );
         const dsValidator = new DataSourceConnectionValidator(dataSourceClient);
 

src/plugins/data_source/server/types.ts

@@ -8,6 +8,7 @@ import {
   OpenSearchClient,
   SavedObjectsClientContract,
 } from 'src/core/server';
+import { DataSourceAttributes } from '../common/data_sources';
 
 import { CryptographyServiceSetup } from './cryptography_service';
 import { DataSourceError } from './lib/error';
@@ -19,11 +20,13 @@ export interface LegacyClientCallAPIParams {
 }
 
 export interface DataSourceClientParams {
-  // id is optional when creating test client
-  dataSourceId?: string;
-  // this saved objects client is used to fetch data source on behalf of users, caller should pass scoped saved objects client
+  // to fetch data source on behalf of users, caller should pass scoped saved objects client
   savedObjects: SavedObjectsClientContract;
   cryptography: CryptographyServiceSetup;
+  // optional when creating test client, required for normal client
+  dataSourceId?: string;
+  // required when creating test client
+  testClientDataSourceAttr?: DataSourceAttributes;
 }
 
 export interface DataSourcePluginRequestContext {

src/plugins/data_source_management/public/components/create_data_source_wizard/components/create_form/create_data_source_form.tsx

@@ -581,8 +581,7 @@ export class CreateDataSourceForm extends React.Component<
                 <EuiButton
                   type="submit"
                   fill={false}
-                  // TODO: remove the type check when working on https://github.com/opensearch-project/OpenSearch-Dashboards/issues/3418
-                  disabled={!this.isFormValid() || this.state.auth.type === AuthType.SigV4}
+                  disabled={!this.isFormValid()}
                   onClick={this.onClickTestConnection}
                   data-test-subj="createDataSourceTestConnectionButton"
                 >