Skip to content

feat(smus): Upgrade sdk v2 to sdk v3 - DataZoneCustomClient #8411

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
PotatoWKY wants to merge 1 commit into
base: master
Choose a base branch
from
+1,208 −1,051
Open

feat(smus): Upgrade sdk v2 to sdk v3 - DataZoneCustomClient #8411

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1,087 changes: 1,087 additions & 0 deletions package-lock.json
View file
Open in desktop

Large diffs are not rendered by default.

1 change: 1 addition & 0 deletions packages/core/package.json
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
4 changes: 0 additions & 4 deletions packages/core/scripts/build/generateServiceClient.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
})()
118 changes: 62 additions & 56 deletions packages/core/src/sagemakerunifiedstudio/shared/client/datazoneCustomClientHelper.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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'
Expand All @@ -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')

Expand Down Expand Up @@ -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,
}
}
Comment on lines +107 to +115
Copy link
Contributor

@vpbhargav vpbhargav Dec 11, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This wrapper needed? Unable to pass this.credentialProvider directly to the clientConfig? Seems like it but just want to confirm.

Copy link
Contributor Author

@PotatoWKY PotatoWKY Dec 11, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah it is needed. I am not Typescript pro, but here is my understanding:

  • passing this.credentialProvider would only pass the method,
  • with the wrapper it also passes the instance

we would need the correct instance to get the credentials
Inspired by @liuzulin:
https://github.com/aws/aws-toolkit-vscode/blame/20a5e850bc483e0d649eddf1880014e2873ca565/packages/core/src/sagemakerunifiedstudio/shared/client/s3Client.ts#L159


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) {
Expand All @@ -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 || []

Expand All @@ -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
Expand All @@ -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 {
Expand Down Expand Up @@ -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
Expand All @@ -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 || []

Expand Down Expand Up @@ -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 || []

Expand Down
Loading
Loading