Upgrade sdk v2 to sdk v3 - DataZoneCustomClient

Author: PotatoWKY

package-lock.json

@@ -582,6 +582,7 @@
         "@amzn/codewhisperer-streaming": "file:../../src.gen/@amzn/codewhisperer-streaming",
         "@amzn/sagemaker-client": "file:../../src.gen/@amzn/sagemaker-client/1.0.0.tgz",
         "@amzn/glue-catalog-client": "file:../../src.gen/@amzn/glue-catalog-client/0.0.1.tgz",
+        "@amzn/datazone-custom-client": "file:../../src.gen/@amzn/datazone-custom-client/3.0.0.tgz",
         "@aws-sdk/client-accessanalyzer": "^3.888.0",
         "@aws-sdk/client-api-gateway": "<3.731.0",
         "@aws-sdk/client-apprunner": "<3.731.0",

packages/core/package.json

@@ -245,10 +245,6 @@ void (async () => {
             serviceJsonPath: 'src/sagemakerunifiedstudio/shared/client/sqlworkbench.json',
             serviceName: 'SQLWorkbench',
         },
-        {
-            serviceJsonPath: 'src/sagemakerunifiedstudio/shared/client/datazonecustomclient.json',
-            serviceName: 'DataZoneCustomClient',
-        },
     ]
     await generateServiceClients(serviceClientDefinitions)
 })()

packages/core/scripts/build/generateServiceClient.ts

@@ -4,12 +4,19 @@
  */
 
 import { getLogger } from '../../../shared/logger/logger'
-import apiConfig = require('./datazonecustomclient.json')
-import globals from '../../../shared/extensionGlobals'
-import { Service } from 'aws-sdk'
-import { ServiceConfigurationOptions } from 'aws-sdk/lib/service'
-import * as DataZoneCustomClient from './datazonecustomclient'
-import { adaptConnectionCredentialsProvider } from './credentialsAdapter'
+import {
+    DataZone,
+    ListDomainsCommand,
+    GetDomainCommand,
+    SearchGroupProfilesCommand,
+    SearchUserProfilesCommand,
+    DomainSummary,
+    GetDomainOutput,
+    GroupProfileSummary,
+    UserProfileSummary,
+    GroupSearchType,
+    UserSearchType,
+} from '@amzn/datazone-custom-client'
 import { CredentialsProvider } from '../../../auth/providers/credentials'
 import { ToolkitError } from '../../../shared/errors'
 import { SmusUtils, isIamDomain } from '../smusUtils'
@@ -30,7 +37,7 @@ export const DataZoneErrorCode = {
  * Helper client for interacting with AWS DataZone Custom API
  */
 export class DataZoneCustomClientHelper {
-    private datazoneCustomClient: DataZoneCustomClient | undefined
+    private datazoneCustomClient: DataZone | undefined
     private static instances = new Map<string, DataZoneCustomClientHelper>()
     private readonly logger = getLogger('smus')
 
@@ -91,33 +98,38 @@ export class DataZoneCustomClientHelper {
     /**
      * Gets the DataZone client, initializing it if necessary
      */
-    private async getDataZoneCustomClient(): Promise<DataZoneCustomClient> {
+    private async getDataZoneCustomClient(): Promise<DataZone> {
         if (!this.datazoneCustomClient) {
             try {
                 this.logger.info('DataZoneCustomClientHelper: Creating authenticated DataZone client')
 
-                // Use user setting for endpoint if provided, otherwise use default
+                // Create credential provider function for auto-refresh
+                const awsCredentialProvider = async () => {
+                    const credentials = await this.credentialProvider.getCredentials()
+                    return {
+                        accessKeyId: credentials.accessKeyId,
+                        secretAccessKey: credentials.secretAccessKey,
+                        sessionToken: credentials.sessionToken,
+                        expiration: credentials.expiration,
+                    }
+                }
+
+                const clientConfig: any = {
+                    region: this.region,
+                    credentials: awsCredentialProvider,
+                }
+
+                // Use user setting for endpoint if provided
                 const devSettings = DevSettings.instance
                 const customEndpoint = devSettings.get('endpoints', {})['datazone']
-                const endpoint = customEndpoint || `https://datazone.${this.region}.api.aws`
-
                 if (customEndpoint) {
+                    clientConfig.endpoint = customEndpoint
                     this.logger.debug(
-                        `DataZoneCustomClientHelper: Using custom DataZone endpoint from settings: ${endpoint}`
+                        `DataZoneCustomClientHelper: Using custom DataZone endpoint from settings: ${customEndpoint}`
                     )
                 }
 
-                this.datazoneCustomClient = (await globals.sdkClientBuilder.createAwsService(
-                    Service,
-                    {
-                        apiConfig: apiConfig,
-                        endpoint: endpoint,
-                        region: this.region,
-                        credentialProvider: adaptConnectionCredentialsProvider(this.credentialProvider),
-                    } as ServiceConfigurationOptions,
-                    undefined,
-                    false
-                )) as DataZoneCustomClient
+                this.datazoneCustomClient = new DataZone(clientConfig)
 
                 this.logger.info('DataZoneCustomClientHelper: Successfully created authenticated DataZone client')
             } catch (err) {
@@ -137,20 +149,19 @@ export class DataZoneCustomClientHelper {
         maxResults?: number
         status?: string
         nextToken?: string
-    }): Promise<{ domains: DataZoneCustomClient.Types.DomainSummary[]; nextToken?: string }> {
+    }): Promise<{ domains: DomainSummary[]; nextToken?: string }> {
         try {
             this.logger.info(`DataZoneCustomClientHelper: Listing domains in region ${this.region}`)
 
             const datazoneCustomClient = await this.getDataZoneCustomClient()
 
             // Call DataZone API to list domains with pagination
-            const response = await datazoneCustomClient
-                .listDomains({
-                    maxResults: options?.maxResults,
-                    status: options?.status,
-                    nextToken: options?.nextToken,
-                })
-                .promise()
+            const command = new ListDomainsCommand({
+                maxResults: options?.maxResults,
+                status: options?.status as any,
+                nextToken: options?.nextToken,
+            })
+            const response = await datazoneCustomClient.send(command)
 
             const domains = response.items || []
 
@@ -172,9 +183,9 @@ export class DataZoneCustomClientHelper {
      * @param options Options for listing domains (excluding nextToken which is handled internally)
      * @returns Promise resolving to an array of all DataZone domains
      */
-    public async fetchAllDomains(options?: { status?: string }): Promise<DataZoneCustomClient.Types.DomainSummary[]> {
+    public async fetchAllDomains(options?: { status?: string }): Promise<DomainSummary[]> {
         try {
-            let allDomains: DataZoneCustomClient.Types.DomainSummary[] = []
+            let allDomains: DomainSummary[] = []
             let nextToken: string | undefined
             do {
                 const maxResultsPerPage = 25
@@ -199,7 +210,7 @@ export class DataZoneCustomClientHelper {
      * Gets the domain with IAM authentication mode using pagination with early termination
      * @returns Promise resolving to the DataZone domain or undefined if not found
      */
-    public async getIamDomain(): Promise<DataZoneCustomClient.Types.DomainSummary | undefined> {
+    public async getIamDomain(): Promise<DomainSummary | undefined> {
         const logger = getLogger('smus')
 
         try {
@@ -259,17 +270,16 @@ export class DataZoneCustomClientHelper {
      * @param domainId The ID of the domain to retrieve
      * @returns Promise resolving to the GetDomainOutput
      */
-    public async getDomain(domainId: string): Promise<DataZoneCustomClient.Types.GetDomainOutput> {
+    public async getDomain(domainId: string): Promise<GetDomainOutput> {
         try {
             this.logger.debug(`DataZoneCustomClientHelper: Getting domain with ID: ${domainId}`)
 
             const datazoneCustomClient = await this.getDataZoneCustomClient()
 
-            const response = await datazoneCustomClient
-                .getDomain({
-                    identifier: domainId,
-                })
-                .promise()
+            const command = new GetDomainCommand({
+                identifier: domainId,
+            })
+            const response = await datazoneCustomClient.send(command)
 
             this.logger.debug(`DataZoneCustomClientHelper: Successfully retrieved domain: ${domainId}`)
             return response
@@ -293,25 +303,23 @@ export class DataZoneCustomClientHelper {
             maxResults?: number
             nextToken?: string
         }
-    ): Promise<{ items: DataZoneCustomClient.Types.GroupProfileSummary[]; nextToken?: string }> {
+    ): Promise<{ items: GroupProfileSummary[]; nextToken?: string }> {
         try {
             this.logger.debug(
                 `DataZoneCustomClientHelper: Searching group profiles in domain ${domainIdentifier} with groupType: ${options?.groupType}, searchText: ${options?.searchText}`
             )
 
             const datazoneCustomClient = await this.getDataZoneCustomClient()
 
-            // Build the request parameters
-            const params: DataZoneCustomClient.Types.SearchGroupProfilesInput = {
+            // Call DataZone API to search group profiles
+            const command = new SearchGroupProfilesCommand({
                 domainIdentifier,
-                groupType: options?.groupType as DataZoneCustomClient.Types.GroupSearchType,
+                groupType: options?.groupType as GroupSearchType,
                 searchText: options?.searchText,
                 maxResults: options?.maxResults,
                 nextToken: options?.nextToken,
-            }
-
-            // Call DataZone API to search group profiles
-            const response = await datazoneCustomClient.searchGroupProfiles(params).promise()
+            })
+            const response = await datazoneCustomClient.send(command)
 
             const items = response.items || []
 
@@ -342,25 +350,23 @@ export class DataZoneCustomClientHelper {
             maxResults?: number
             nextToken?: string
         }
-    ): Promise<{ items: DataZoneCustomClient.Types.UserProfileSummary[]; nextToken?: string }> {
+    ): Promise<{ items: UserProfileSummary[]; nextToken?: string }> {
         try {
             this.logger.debug(
                 `DataZoneCustomClientHelper: Searching user profiles in domain ${domainIdentifier} with userType: ${options.userType}, searchText: ${options.searchText}`
             )
 
             const datazoneCustomClient = await this.getDataZoneCustomClient()
 
-            // Build the request parameters
-            const params: DataZoneCustomClient.Types.SearchUserProfilesInput = {
+            // Call DataZone API to search user profiles
+            const command = new SearchUserProfilesCommand({
                 domainIdentifier,
-                userType: options.userType as DataZoneCustomClient.Types.UserSearchType,
+                userType: options.userType as UserSearchType,
                 searchText: options.searchText,
                 maxResults: options.maxResults,
                 nextToken: options.nextToken,
-            }
-
-            // Call DataZone API to search user profiles
-            const response = await datazoneCustomClient.searchUserProfiles(params).promise()
+            })
+            const response = await datazoneCustomClient.send(command)
 
             const items = response.items || []