Skip to content

feat(smus): Upgrade sdk v2 to sdk v3 - SQLWorkbench #8407

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
Open

feat(smus): Upgrade sdk v2 to sdk v3 - SQLWorkbench #8407

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,095 changes: 1,095 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 @@ -581,6 +581,7 @@
"@amzn/amazon-q-developer-streaming-client": "file:../../src.gen/@amzn/amazon-q-developer-streaming-client",
"@amzn/codewhisperer-streaming": "file:../../src.gen/@amzn/codewhisperer-streaming",
"@amzn/sagemaker-client": "file:../../src.gen/@amzn/sagemaker-client/1.0.0.tgz",
"@amzn/sql-workbench-client": "file:../../src.gen/@amzn/sql-workbench-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/gluecatalogapi.json',
serviceName: 'GlueCatalogApi',
},
{
serviceJsonPath: 'src/sagemakerunifiedstudio/shared/client/sqlworkbench.json',
serviceName: 'SQLWorkbench',
},
{
serviceJsonPath: 'src/sagemakerunifiedstudio/shared/client/datazonecustomclient.json',
serviceName: 'DataZoneCustomClient',
Expand Down
2 changes: 1 addition & 1 deletion packages/core/src/sagemakerunifiedstudio/explorer/nodes/redshiftStrategy.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -458,7 +458,7 @@ function createSchemaNode(schemaName: string, connectionConfig: ConnectionConfig
parentType: ResourceType.SCHEMA,
},
{
parentId: schemaConnectionConfig.database,
parentId: schemaConnectionConfig.database || '',
Copy link
Contributor

@liuzulin liuzulin Dec 11, 2025

Choose a reason for hiding this comment

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

Why we need this change? let's avoid default to empty unless it's necessary

Copy link
Contributor Author

@PotatoWKY PotatoWKY Dec 12, 2025

Choose a reason for hiding this comment

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

Yeah this is needed.

SDK v3 Type Definitions:

From the new @amzn/sql-workbench-client SDK v3 (3.0.0.tgz):

// ParentResource interface - parentId is required string
export interface ParentResource {
    parentId: string;  // Required, non-nullable
    parentType: string;
}

// DatabaseConnectionConfiguration - database is optional
export interface DatabaseConnectionConfiguration {
    // ...
    database?: string | undefined;  // Optional
}

The Problem:

schemaConnectionConfig.database has type string | undefined (because database is optional in DatabaseConnectionConfiguration)
parentId in ParentResource requires type string (non-nullable)
TypeScript won't allow assigning string | undefined to string
The Fix:

parentId: schemaConnectionConfig.database || ''

parentType: ResourceType.DATABASE,
},
],
Expand Down
154 changes: 66 additions & 88 deletions packages/core/src/sagemakerunifiedstudio/shared/client/sqlWorkbenchClient.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,40 +3,28 @@
* SPDX-License-Identifier: Apache-2.0
*/

import { Service } from 'aws-sdk'
import { ServiceConfigurationOptions } from 'aws-sdk/lib/service'
import globals from '../../../shared/extensionGlobals'
import { getLogger } from '../../../shared/logger/logger'
import * as SQLWorkbench from './sqlworkbench'
import apiConfig = require('./sqlworkbench.json')
import {
SQLWorkbench,
GetResourcesCommand,
ExecuteQueryCommand,
GetResourcesRequest,
GetResourcesResponse,
ExecuteQueryRequest,
DatabaseConnectionConfiguration,
ParentResource,
DatabaseIntegrationConnectionAuthenticationTypes,
} from '@amzn/sql-workbench-client'
import { v4 as uuidv4 } from 'uuid'
import { getRedshiftTypeFromHost } from '../../explorer/nodes/utils'
import { DatabaseIntegrationConnectionAuthenticationTypes, RedshiftType } from '../../explorer/nodes/types'
import { RedshiftType } from '../../explorer/nodes/types'
import { ConnectionCredentialsProvider } from '../../auth/providers/connectionCredentialsProvider'
import { adaptConnectionCredentialsProvider } from './credentialsAdapter'

/**
* Connection configuration for SQL Workbench
* This is an alias for the SDK's DatabaseConnectionConfiguration type
*/
export interface ConnectionConfig {
id: string
type: string
databaseType: string
connectableResourceIdentifier: string
connectableResourceType: string
database: string
auth?: {
secretArn?: string
}
}

/**
* Resource parent information
*/
export interface ParentResource {
parentId: string
parentType: string
}
export type ConnectionConfig = DatabaseConnectionConfiguration

/**
* Gets a SQL Workbench ARN
Expand Down Expand Up @@ -89,7 +77,7 @@ export async function createRedshiftConnectionConfig(
}

// Determine auth type based on the provided parameters
let authType: string
let authType: DatabaseIntegrationConnectionAuthenticationTypes

if (secretArn) {
authType = DatabaseIntegrationConnectionAuthenticationTypes.SECRET
Expand Down Expand Up @@ -118,11 +106,7 @@ export async function createRedshiftConnectionConfig(
}

// Add auth object for SECRET authentication type
if (
(authType as DatabaseIntegrationConnectionAuthenticationTypes) ===
DatabaseIntegrationConnectionAuthenticationTypes.SECRET &&
secretArn
) {
if (authType === DatabaseIntegrationConnectionAuthenticationTypes.SECRET && secretArn) {
connectionConfig.auth = { secretArn }
}

Expand Down Expand Up @@ -177,7 +161,7 @@ export class SQLWorkbenchClient {
/**
* Gets resources from SQL Workbench
* @param params Request parameters
* @returns Raw response from getResources API
* @returns Response containing resources and optional next token
*/
public async getResources(params: {
connection: ConnectionConfig
Expand All @@ -187,13 +171,13 @@ export class SQLWorkbenchClient {
parents?: ParentResource[]
pageToken?: string
forceRefresh?: boolean
}): Promise<SQLWorkbench.GetResourcesResponse> {
}): Promise<GetResourcesResponse> {
try {
this.logger.info(`SQLWorkbenchClient: Getting resources in region ${this.region}`)

const sqlClient = await this.getSQLClient()

const requestParams = {
const requestParams: GetResourcesRequest = {
connection: params.connection,
type: params.resourceType,
maxItems: params.maxItems || 100,
Expand All @@ -203,13 +187,9 @@ export class SQLWorkbenchClient {
accountSettings: {},
}

// Call the GetResources API
const response = await sqlClient.getResources(requestParams).promise()

return {
resources: response.resources || [],
nextToken: response.nextToken,
}
// Call the GetResources API using SDK v3 Command pattern
const command = new GetResourcesCommand(requestParams)
return await sqlClient.send(command)
} catch (err) {
this.logger.error('SQLWorkbenchClient: Failed to get resources: %s', err as Error)
throw err
Expand All @@ -228,26 +208,27 @@ export class SQLWorkbenchClient {

const sqlClient = await this.getSQLClient()

// Call the ExecuteQuery API
const response = await sqlClient
.executeQuery({
connection: connectionConfig as any,
databaseType: 'REDSHIFT',
accountSettings: {},
executionContext: [
{
parentType: 'DATABASE',
parentId: connectionConfig.database || '',
},
],
query,
queryExecutionType: 'NO_SESSION',
queryResponseDeliveryType: 'ASYNC',
maxItems: 100,
ignoreHistory: true,
tabId: 'data_explorer',
})
.promise()
const requestParams: ExecuteQueryRequest = {
connection: connectionConfig,
databaseType: 'REDSHIFT',
accountSettings: {},
executionContext: [
{
parentType: 'DATABASE',
parentId: connectionConfig.database || '',
},
],
query,
queryExecutionType: 'NO_SESSION',
queryResponseDeliveryType: 'ASYNC',
maxItems: 100,
ignoreHistory: true,
tabId: 'data_explorer',
}

// Call the ExecuteQuery API using SDK v3 Command pattern
const command = new ExecuteQueryCommand(requestParams)
const response = await sqlClient.send(command)

// Log the response
this.logger.info(
Expand All @@ -261,9 +242,6 @@ export class SQLWorkbenchClient {
}
}

/**
* Gets the SQL client, initializing it if necessary
*/
/**
* Gets the SQL Workbench endpoint URL for the given region
* @param region AWS region
Expand All @@ -273,6 +251,9 @@ export class SQLWorkbenchClient {
return `https://api-v2.sqlworkbench.${region}.amazonaws.com`
}

/**
* Gets the SQL client, initializing it if necessary
*/
private async getSQLClient(): Promise<SQLWorkbench> {
if (!this.sqlClient) {
try {
Expand All @@ -281,30 +262,27 @@ export class SQLWorkbenchClient {
this.logger.info(`Using SQL Workbench endpoint: ${endpoint}`)

if (this.connectionCredentialsProvider) {
// Create client with provided credentials
this.sqlClient = (await globals.sdkClientBuilder.createAwsService(
Service,
{
apiConfig: apiConfig,
region: this.region,
endpoint: endpoint,
credentialProvider: adaptConnectionCredentialsProvider(this.connectionCredentialsProvider),
} as ServiceConfigurationOptions,
undefined,
false
)) as SQLWorkbench
// Create client with credential provider function for auto-refresh
const awsCredentialProvider = async () => {
const credentials = await this.connectionCredentialsProvider!.getCredentials()
return {
accessKeyId: credentials.accessKeyId,
secretAccessKey: credentials.secretAccessKey,
sessionToken: credentials.sessionToken,
expiration: credentials.expiration,
}
}
this.sqlClient = new SQLWorkbench({
region: this.region,
endpoint: endpoint,
credentials: awsCredentialProvider,
})
} else {
// Use the SDK client builder for default credentials
this.sqlClient = (await globals.sdkClientBuilder.createAwsService(
Service,
{
apiConfig: apiConfig,
region: this.region,
endpoint: endpoint,
} as ServiceConfigurationOptions,
undefined,
false
)) as SQLWorkbench
// Use default credentials
this.sqlClient = new SQLWorkbench({
region: this.region,
endpoint: endpoint,
})
}

this.logger.debug('SQLWorkbenchClient: Successfully created SQL client')
Expand Down
Loading
Loading