Merge master into feature/console-session-profile

Author: aws-toolkit-automation

package-lock.json

@@ -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/glue-catalog-client": "file:../../src.gen/@amzn/glue-catalog-client/0.0.1.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

@@ -241,10 +241,6 @@ void (async () => {
             serviceJsonPath: 'src/codewhisperer/client/user-service-2.json',
             serviceName: 'CodeWhispererUserClient',
         },
-        {
-            serviceJsonPath: 'src/sagemakerunifiedstudio/shared/client/gluecatalogapi.json',
-            serviceName: 'GlueCatalogApi',
-        },
         {
             serviceJsonPath: 'src/sagemakerunifiedstudio/shared/client/sqlworkbench.json',
             serviceName: 'SQLWorkbench',

packages/core/scripts/build/generateServiceClient.ts

@@ -7,7 +7,7 @@ import * as vscode from 'vscode'
 import { TreeNode } from '../../../shared/treeview/resourceTreeDataProvider'
 import { getLogger } from '../../../shared/logger/logger'
 import { DataZoneConnection } from '../../shared/client/datazoneClient'
-import { GlueCatalog, GlueCatalogClient } from '../../shared/client/glueCatalogClient'
+import { Catalog } from '@amzn/glue-catalog-client'
 import { GlueClient } from '../../shared/client/glueClient'
 import { ConnectionClientStore } from '../../shared/client/connectionClientStore'
 import {
@@ -38,6 +38,7 @@ import { Column, Database, Table } from '@aws-sdk/client-glue'
 import { ConnectionCredentialsProvider } from '../../auth/providers/connectionCredentialsProvider'
 import { telemetry } from '../../../shared/telemetry/telemetry'
 import { recordDataConnectionTelemetry } from '../../shared/telemetry'
+import { GlueCatalogClient } from '../../shared/client/glueCatalogClient'
 
 /**
  * Lakehouse data node for SageMaker Unified Studio
@@ -249,8 +250,8 @@ function createAwsDataCatalogNode(parent: LakehouseNode, glueClient: GlueClient)
 }
 
 export interface CatalogTree {
-    parent: GlueCatalog
-    children?: GlueCatalog[]
+    parent: Catalog
+    children?: Catalog[]
 }
 
 /**
@@ -265,7 +266,7 @@ export interface CatalogTree {
  *
  * Without the first pass, we'd need O(n²) time to find parent catalogs for each child catalog.
  */
-function buildCatalogTree(catalogs: GlueCatalog[]): CatalogTree[] {
+function buildCatalogTree(catalogs: Catalog[]): CatalogTree[] {
     const catalogMap: Record<string, CatalogTree> = {}
     const rootCatalogs: CatalogTree[] = []
 
@@ -372,7 +373,7 @@ async function getCatalogs(
  */
 function createCatalogNode(
     catalogId: string,
-    catalog: GlueCatalog,
+    catalog: Catalog,
     glueClient: GlueClient,
     parent: LakehouseNode,
     isParent: boolean = false

packages/core/src/sagemakerunifiedstudio/explorer/nodes/lakehouseStrategy.ts

@@ -22,7 +22,7 @@ import {
 } from './utils'
 import { createPlaceholderItem } from '../../../shared/treeview/utils'
 import { ConnectionCredentialsProvider } from '../../auth/providers/connectionCredentialsProvider'
-import { GlueCatalog } from '../../shared/client/glueCatalogClient'
+import { Catalog } from '@amzn/glue-catalog-client'
 import { handleCredExpiredError } from '../../shared/credentialExpiryHandler'
 import { telemetry } from '../../../shared/telemetry/telemetry'
 import { recordDataConnectionTelemetry } from '../../shared/telemetry'
@@ -169,7 +169,7 @@ export function createRedshiftConnectionNode(
 
                 // Fetch Glue catalogs for filtering purposes only
                 // This will help determine which catalogs are accessible within the project
-                let glueCatalogs: GlueCatalog[] = []
+                let glueCatalogs: Catalog[] = []
                 try {
                     glueCatalogs = await listGlueCatalogs(
                         connection.connectionId,
@@ -771,7 +771,7 @@ async function listGlueCatalogs(
     connectionId: string,
     region: string,
     connectionCredentialsProvider: ConnectionCredentialsProvider
-): Promise<GlueCatalog[]> {
+): Promise<Catalog[]> {
     const clientStore = ConnectionClientStore.getInstance()
     const glueCatalogClient = clientStore.getGlueCatalogClient(connectionId, region, connectionCredentialsProvider)
 

packages/core/src/sagemakerunifiedstudio/explorer/nodes/redshiftStrategy.ts

@@ -24,7 +24,7 @@ import { getContext } from '../../../shared/vscode/setContext'
 import { SmusAuthenticationProvider } from '../../auth/providers/smusAuthenticationProvider'
 import { SmusIamConnection } from '../../auth/model'
 import { ConnectionStatus } from '@aws-sdk/client-datazone'
-import { GlueCatalog } from '../../shared/client/glueCatalogClient'
+import { Catalog } from '@amzn/glue-catalog-client'
 
 /**
  * Polling interval in milliseconds for checking space status updates
@@ -35,7 +35,7 @@ export const PENDING_NODE_POLLING_INTERVAL_MS = 5000
 /**
  * Check if a catalog is a RedLake catalog
  */
-export const isRedLakeCatalog = (catalog?: GlueCatalog) => {
+export const isRedLakeCatalog = (catalog?: Catalog) => {
     return (
         catalog?.FederatedCatalog?.ConnectionName === 'aws:redshift' ||
         catalog?.CatalogProperties?.DataLakeAccessProperties?.CatalogType === 'aws:redshift'
@@ -45,7 +45,7 @@ export const isRedLakeCatalog = (catalog?: GlueCatalog) => {
 /**
  * Check if a catalog is a S3 table catalog
  */
-export const isS3TablesCatalog = (catalog?: GlueCatalog) => {
+export const isS3TablesCatalog = (catalog?: Catalog) => {
     return catalog?.FederatedCatalog?.ConnectionName === 'aws:s3tables'
 }
 

packages/core/src/sagemakerunifiedstudio/explorer/nodes/utils.ts

@@ -3,25 +3,15 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { Service } from 'aws-sdk'
-import globals from '../../../shared/extensionGlobals'
 import { getLogger } from '../../../shared/logger/logger'
-import * as GlueCatalogApi from './gluecatalogapi'
-import apiConfig = require('./gluecatalogapi.json')
-import { ServiceConfigurationOptions } from 'aws-sdk/lib/service'
+import { GlueCatalog, Catalog } from '@amzn/glue-catalog-client'
 import { ConnectionCredentialsProvider } from '../../auth/providers/connectionCredentialsProvider'
-import { adaptConnectionCredentialsProvider } from './credentialsAdapter'
-
-/**
- * Represents a Glue catalog
- */
-export type GlueCatalog = GlueCatalogApi.Types.Catalog
 
 /**
  * Client for interacting with Glue Catalog API
  */
 export class GlueCatalogClient {
-    private glueClient: GlueCatalogApi | undefined
+    private glueClient: GlueCatalog | undefined
     private static instance: GlueCatalogClient | undefined
     private readonly logger = getLogger('smus')
 
@@ -67,21 +57,19 @@ export class GlueCatalogClient {
      * @param nextToken Optional pagination token
      * @returns Object containing catalogs and nextToken
      */
-    public async getCatalogs(nextToken?: string): Promise<{ catalogs: GlueCatalog[]; nextToken?: string }> {
+    public async getCatalogs(nextToken?: string): Promise<{ catalogs: Catalog[]; nextToken?: string }> {
         try {
             this.logger.info(`GlueCatalogClient: Getting catalogs in region ${this.region}`)
 
             const glueClient = await this.getGlueCatalogClient()
 
             // Call the GetCatalogs API with pagination
-            const response = await glueClient
-                .getCatalogs({
-                    Recursive: true,
-                    NextToken: nextToken,
-                })
-                .promise()
+            const response = await glueClient.getCatalogs({
+                Recursive: true,
+                NextToken: nextToken,
+            })
 
-            const catalogs: GlueCatalog[] = response.CatalogList || []
+            const catalogs: Catalog[] = response.CatalogList || []
 
             this.logger.info(`GlueCatalogClient: Found ${catalogs.length} catalogs in this page`)
             return {
@@ -97,32 +85,30 @@ export class GlueCatalogClient {
     /**
      * Gets the Glue client, initializing it if necessary
      */
-    private async getGlueCatalogClient(): Promise<GlueCatalogApi> {
+    private async getGlueCatalogClient(): Promise<GlueCatalog> {
         if (!this.glueClient) {
             try {
                 if (this.connectionCredentialsProvider) {
-                    // Create client with provided credentials
-                    this.glueClient = (await globals.sdkClientBuilder.createAwsService(
-                        Service,
-                        {
-                            apiConfig: apiConfig,
-                            region: this.region,
-                            credentialProvider: adaptConnectionCredentialsProvider(this.connectionCredentialsProvider),
-                        } as ServiceConfigurationOptions,
-                        undefined,
-                        false
-                    )) as GlueCatalogApi
+                    // 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.glueClient = new GlueCatalog({
+                        region: this.region,
+                        credentials: awsCredentialProvider,
+                    })
                 } else {
-                    // Use the SDK client builder for default credentials
-                    this.glueClient = (await globals.sdkClientBuilder.createAwsService(
-                        Service,
-                        {
-                            apiConfig: apiConfig,
-                            region: this.region,
-                        } as ServiceConfigurationOptions,
-                        undefined,
-                        false
-                    )) as GlueCatalogApi
+                    // Use default credentials
+                    this.glueClient = new GlueCatalog({
+                        region: this.region,
+                    })
                 }
 
                 this.logger.debug('GlueCatalogClient: Successfully created Glue client')