[MD] Integrate test connection to support SigV4 auth type

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

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(
@@ -78,9 +101,14 @@ export const configureTestClient = async (
       dataSourceId
     ) as Client;
 
-    if (type === AuthType.UsernamePasswordType && !credentials?.password && dataSourceId) {
-      dataSourceAttr = await getDataSource(dataSourceId, savedObjects);
-      requireDecryption = true;
+    if (dataSourceId) {
+      if (
+        (type === AuthType.UsernamePasswordType && !credentials?.password) ||
+        (type === AuthType.SigV4 && !credentials?.accessKey && !credentials?.secretKey)
+      ) {
+        dataSourceAttr = await getDataSource(dataSourceId, savedObjects);
+        requireDecryption = true;
+      }
     }
 
     return getQueryClient(
@@ -108,7 +136,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"
                 >

src/plugins/data_source_management/public/components/edit_data_source/components/edit_form/edit_data_source_form.tsx

@@ -346,25 +346,41 @@ export class EditDataSourceForm extends React.Component<EditDataSourceProps, Edi
 
   onClickTestConnection = async () => {
     this.setState({ isLoading: true });
-    const existingAuthType = this.props.existingDataSource.auth.type;
+    // const existingAuthType = this.props.existingDataSource.auth.type;
+    const isNewCredential = !!(this.state.auth.type !== this.props.existingDataSource.auth.type);
 
-    try {
-      const isNewCredential = !!(
-        existingAuthType === AuthType.NoAuth && this.state.auth.type !== existingAuthType
-      );
-      const formValues: DataSourceAttributes = {
-        title: this.state.title,
-        description: this.state.description,
-        endpoint: this.props.existingDataSource.endpoint,
-        auth: {
-          ...this.state.auth,
-          credentials: {
-            ...this.state.auth.credentials,
-            password: isNewCredential ? this.state.auth.credentials.password : '',
-          },
-        },
-      };
+    let credentials = this.state.auth.credentials;
 
+    switch (this.state.auth.type) {
+      case AuthType.UsernamePasswordType:
+        credentials = {
+          username: this.state.auth.credentials?.username,
+          password: isNewCredential ? this.state.auth.credentials?.password : '',
+        } as UsernamePasswordTypedContent;
+        break;
+      case AuthType.SigV4:
+        credentials = {
+          region: this.state.auth.credentials?.region,
+          accessKey: isNewCredential ? this.state.auth.credentials?.accessKey : '',
+          secretKey: isNewCredential ? this.state.auth.credentials?.secretKey : '',
+        } as SigV4Content;
+        break;
+      case AuthType.NoAuth:
+        credentials = undefined;
+        break;
+
+      default:
+        break;
+    }
+
+    const formValues: DataSourceAttributes = {
+      title: this.state.title,
+      description: this.state.description,
+      endpoint: this.state.endpoint,
+      auth: { ...this.state.auth, credentials },
+    };
+
+    try {
       await this.props.handleTestConnection(formValues);
 
       this.props.displayToastMessage({
@@ -532,9 +548,7 @@ export class EditDataSourceForm extends React.Component<EditDataSourceProps, Edi
     return (
       <Header
         showDeleteIcon={true}
-        isFormValid={
-          !!this.state.auth?.credentials?.username || this.state.auth.type === AuthType.NoAuth
-        }
+        isFormValid={this.isFormValid()}
         onClickDeleteIcon={this.onClickDeleteDataSource}
         dataSourceName={this.props.existingDataSource.title}
         onClickTestConnection={this.onClickTestConnection}

src/plugins/data_source_management/public/components/utils.test.ts

@@ -156,7 +156,7 @@ describe('DataSourceManagement: Utils.ts', () => {
           Array [
             "/internal/data-source-management/validate",
             Object {
-              "body": "{\\"dataSourceAttr\\":{\\"endpoint\\":\\"https://test.com\\",\\"auth\\":{\\"type\\":\\"no_auth\\",\\"credentials\\":null}}}",
+              "body": "{\\"dataSourceAttr\\":{\\"endpoint\\":\\"https://test.com\\",\\"auth\\":{\\"type\\":\\"no_auth\\"}}}",
             },
           ],
         ]
@@ -170,7 +170,7 @@ describe('DataSourceManagement: Utils.ts', () => {
           Array [
             "/internal/data-source-management/validate",
             Object {
-              "body": "{\\"id\\":\\"test1234\\",\\"dataSourceAttr\\":{\\"endpoint\\":\\"https://test.com\\",\\"auth\\":{\\"type\\":\\"no_auth\\",\\"credentials\\":null}}}",
+              "body": "{\\"id\\":\\"test1234\\",\\"dataSourceAttr\\":{\\"endpoint\\":\\"https://test.com\\",\\"auth\\":{\\"type\\":\\"no_auth\\"}}}",
             },
           ],
         ]

src/plugins/data_source_management/public/components/utils.ts

@@ -90,7 +90,7 @@ export async function testConnection(
       endpoint,
       auth: {
         type,
-        credentials: type === AuthType.NoAuth ? null : { ...credentials },
+        credentials,
       },
     },
   };