feat(sagemakerunifiedstudio): Add telemetry for login, signout, access project, start/stop space and data explorer (aws#2207)

## Problem
Telemetry is missing from the feature.

## Solution

- update packages/core/src/shared/telemetry/vscodeTelemetry.json with
new types, metrics, and metadata
- run npm run generateTelemetry in core package. this will create
required obj that are accessible via telemetry.{metric name}
- then emit metric in the code

---

- Treat all work as PUBLIC. Private `feature/x` branches will not be
squash-merged at release time.
- Your code changes must meet the guidelines in
[CONTRIBUTING.md](https://github.com/aws/aws-toolkit-vscode/blob/master/CONTRIBUTING.md#guidelines).
- License: I confirm that my contribution is made under the terms of the
Apache 2.0 license.

---------

Co-authored-by: Keyvan Zare Rami <[email protected]>
Co-authored-by: Zulin Liu <[email protected]>

Author: 3 people

packages/core/src/awsService/sagemaker/commands.ts

@@ -158,9 +158,13 @@ export async function stopSpace(
         if (error.name === 'AccessDeniedException') {
             throw new ToolkitError('You do not have permission to stop spaces. Please contact your administrator', {
                 cause: error,
+                code: error.name,
             })
         } else {
-            throw err
+            throw new ToolkitError(`Failed to stop space: ${spaceName}`, {
+                cause: error,
+                code: error.name,
+            })
         }
     }
     await tryRefreshNode(node)
@@ -185,14 +189,19 @@ export async function openRemoteConnect(
             await tryRefreshNode(node)
             const appType = node.spaceApp.SpaceSettingsSummary?.AppType
             if (!appType) {
-                throw new ToolkitError('AppType is undefined for the selected space. Cannot start remote connection.')
+                throw new ToolkitError('AppType is undefined for the selected space. Cannot start remote connection.', {
+                    code: 'undefinedAppType',
+                })
             }
             await client.waitForAppInService(node.spaceApp.DomainId!, node.spaceApp.SpaceName!, appType)
             await tryRemoteConnection(node, ctx)
         } catch (err: any) {
             // Ignore InstanceTypeError since it means the user decided not to use an instanceType with more memory
             if (err.code !== InstanceTypeError) {
-                throw err
+                throw new ToolkitError('Remote connection failed.', {
+                    cause: err as Error,
+                    code: err.code,
+                })
             }
         }
     } else if (node.getStatus() === 'Running') {

packages/core/src/sagemakerunifiedstudio/auth/providers/smusAuthenticationProvider.ts

@@ -12,7 +12,7 @@ import { SsoConnection } from '../../../auth/connection'
 import { showReauthenticateMessage } from '../../../shared/utilities/messages'
 import * as localizedText from '../../../shared/localizedText'
 import { ToolkitPromptSettings } from '../../../shared/settings'
-import { setContext } from '../../../shared/vscode/setContext'
+import { setContext, getContext } from '../../../shared/vscode/setContext'
 import { SmusUtils, SmusErrorCodes } from '../../shared/smusUtils'
 import { createSmusProfile, isValidSmusConnection, SmusConnection } from '../model'
 import { getLogger } from '../../../shared/logger/logger'
@@ -22,6 +22,7 @@ import { ConnectionCredentialsProvider } from './connectionCredentialsProvider'
 import { ConnectionClientStore } from '../../shared/client/connectionClientStore'
 import { getResourceMetadata } from '../../shared/utils/resourceMetadataUtils'
 import { fromIni } from '@aws-sdk/credential-providers'
+import { randomUUID } from '../../../shared/crypto'
 
 /**
  * Sets the context variable for SageMaker Unified Studio connection state
@@ -86,16 +87,17 @@ export class SmusAuthenticationProvider {
      * Gets the active connection
      */
     public get activeConnection() {
-        if (SmusUtils.isInSmusSpaceEnvironment()) {
+        if (getContext('aws.smus.inSmusSpaceEnvironment')) {
             const resourceMetadata = getResourceMetadata()!
             if (resourceMetadata.AdditionalMetadata!.DataZoneDomainRegion) {
                 return {
                     domainId: resourceMetadata.AdditionalMetadata!.DataZoneDomainId!,
                     ssoRegion: resourceMetadata.AdditionalMetadata!.DataZoneDomainRegion!,
                     // The following fields won't be needed in SMUS space environment
-                    // Set them to be empty string for type checks only
-                    domainUrl: '',
-                    id: '',
+                    // Craft the domain url with known information
+                    // Use randome id as placeholder
+                    domainUrl: `https://${resourceMetadata.AdditionalMetadata!.DataZoneDomainId!}.sagemaker.${resourceMetadata.AdditionalMetadata!.DataZoneDomainRegion!}.on.aws/`,
+                    id: randomUUID(),
                 }
             } else {
                 throw new ToolkitError('Domain region not found in metadata file.')
@@ -115,7 +117,9 @@ export class SmusAuthenticationProvider {
      * Checks if the connection is valid
      */
     public isConnectionValid(): boolean {
-        if (SmusUtils.isInSmusSpaceEnvironment()) {
+        // When in SMUS space, the extension is already running in projet context and sign in is not needed
+        // Set isConnectionValid to always true
+        if (getContext('aws.smus.inSmusSpaceEnvironment')) {
             return true
         }
         return this.activeConnection !== undefined && !this.secondaryAuth.isConnectionExpired
@@ -125,7 +129,9 @@ export class SmusAuthenticationProvider {
      * Checks if connected to SMUS
      */
     public isConnected(): boolean {
-        if (SmusUtils.isInSmusSpaceEnvironment()) {
+        // When in SMUS space, the extension is already running in projet context and sign in is not needed
+        // Set isConnected to always true
+        if (getContext('aws.smus.inSmusSpaceEnvironment')) {
             return true
         }
         return this.activeConnection !== undefined
@@ -365,7 +371,7 @@ export class SmusAuthenticationProvider {
      * @returns Domain ID
      */
     public getDomainId(): string {
-        if (SmusUtils.isInSmusSpaceEnvironment()) {
+        if (getContext('aws.smus.inSmusSpaceEnvironment')) {
             return getResourceMetadata()!.AdditionalMetadata!.DataZoneDomainId!
         }
 
@@ -387,7 +393,7 @@ export class SmusAuthenticationProvider {
     }
 
     public getDomainRegion(): string {
-        if (SmusUtils.isInSmusSpaceEnvironment()) {
+        if (getContext('aws.smus.inSmusSpaceEnvironment')) {
             const resourceMetadata = getResourceMetadata()!
             if (resourceMetadata.AdditionalMetadata!.DataZoneDomainRegion) {
                 return resourceMetadata.AdditionalMetadata!.DataZoneDomainRegion
@@ -409,7 +415,9 @@ export class SmusAuthenticationProvider {
     public async getDerCredentialsProvider(): Promise<any> {
         const logger = getLogger()
 
-        if (SmusUtils.isInSmusSpaceEnvironment()) {
+        if (getContext('aws.smus.inSmusSpaceEnvironment')) {
+            // When in SMUS space, DomainExecutionRoleCreds can be found in config file
+            // Read the credentials from credential profile DomainExecutionRoleCreds
             const credentials = fromIni({ profile: 'DomainExecutionRoleCreds' })
             return {
                 getCredentials: async () => await credentials(),

packages/core/src/sagemakerunifiedstudio/explorer/activation.ts

@@ -19,6 +19,8 @@ import { SageMakerUnifiedStudioProjectNode } from './nodes/sageMakerUnifiedStudi
 import { getLogger } from '../../shared/logger/logger'
 import { setSmusConnectedContext, SmusAuthenticationProvider } from '../auth/providers/smusAuthenticationProvider'
 import { setupUserActivityMonitoring } from '../../awsService/sagemaker/sagemakerSpace'
+import { telemetry } from '../../shared/telemetry/telemetry'
+import { SageMakerUnifiedStudioSpacesParentNode } from './nodes/sageMakerUnifiedStudioSpacesParentNode'
 
 export async function activate(extensionContext: vscode.ExtensionContext): Promise<void> {
     // Initialize the SMUS authentication provider
@@ -71,7 +73,19 @@ export async function activate(extensionContext: vscode.ExtensionContext): Promi
             if (!validateNode(node)) {
                 return
             }
-            await stopSpace(node.resource, extensionContext, node.resource.sageMakerClient)
+            await telemetry.smus_stopSpace.run(async (span) => {
+                span.record({
+                    smusSpaceKey: node.resource.DomainSpaceKey,
+                    smusDomainRegion: node.resource.regionCode,
+                    smusDomainId: (
+                        node.resource.getParent() as SageMakerUnifiedStudioSpacesParentNode
+                    )?.getAuthProvider()?.activeConnection?.domainId,
+                    smusProjectId: (
+                        node.resource.getParent() as SageMakerUnifiedStudioSpacesParentNode
+                    )?.getProjectId(),
+                })
+                await stopSpace(node.resource, extensionContext, node.resource.sageMakerClient)
+            })
         }),
 
         vscode.commands.registerCommand(
@@ -80,7 +94,19 @@ export async function activate(extensionContext: vscode.ExtensionContext): Promi
                 if (!validateNode(node)) {
                     return
                 }
-                await openRemoteConnect(node.resource, extensionContext, node.resource.sageMakerClient)
+                await telemetry.smus_startSpace.run(async (span) => {
+                    span.record({
+                        smusSpaceKey: node.resource.DomainSpaceKey,
+                        smusDomainRegion: node.resource.regionCode,
+                        smusDomainId: (
+                            node.resource.getParent() as SageMakerUnifiedStudioSpacesParentNode
+                        )?.getAuthProvider()?.activeConnection?.domainId,
+                        smusProjectId: (
+                            node.resource.getParent() as SageMakerUnifiedStudioSpacesParentNode
+                        )?.getProjectId(),
+                    })
+                    await openRemoteConnect(node.resource, extensionContext, node.resource.sageMakerClient)
+                })
             }
         ),
 
@@ -112,7 +138,7 @@ export async function activate(extensionContext: vscode.ExtensionContext): Promi
 /**
  * Checks if a node  is undefined and shows a warning message if so.
  */
-function validateNode(node: unknown): boolean {
+function validateNode(node: SagemakerUnifiedStudioSpaceNode): boolean {
     if (!node) {
         void vscode.window.showWarningMessage('Space information is being refreshed. Please try again shortly.')
         return false

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

@@ -33,6 +33,8 @@ import {
 import { createPlaceholderItem } from '../../../shared/treeview/utils'
 import { Column, Database, Table } from '@aws-sdk/client-glue'
 import { ConnectionCredentialsProvider } from '../../auth/providers/connectionCredentialsProvider'
+import { telemetry } from '../../../shared/telemetry/telemetry'
+import { getContext } from '../../../shared/vscode/setContext'
 
 /**
  * Lakehouse data node for SageMaker Unified Studio
@@ -149,49 +151,63 @@ export function createLakehouseConnectionNode(
             },
         },
         async (node) => {
-            try {
-                logger.info(`Loading Lakehouse catalogs for connection ${connection.name}`)
-
-                // Check if this is a default connection
-                const isDefaultConnection =
-                    DATA_DEFAULT_IAM_CONNECTION_NAME_REGEXP.test(connection.name) ||
-                    DATA_DEFAULT_LAKEHOUSE_CONNECTION_NAME_REGEXP.test(connection.name) ||
-                    DATA_DEFAULT_ATHENA_CONNECTION_NAME_REGEXP.test(connection.name)
-
-                // Follow the reference pattern with Promise.allSettled
-                const [awsDataCatalogResult, catalogsResult] = await Promise.allSettled([
-                    // AWS Data Catalog node (only for default connections)
-                    isDefaultConnection
-                        ? Promise.resolve([createAwsDataCatalogNode(node, glueClient)])
-                        : Promise.resolve([]),
-                    // Get catalogs by calling Glue API
-                    getCatalogs(glueCatalogClient, glueClient, node),
-                ])
-
-                const awsDataCatalog = awsDataCatalogResult.status === 'fulfilled' ? awsDataCatalogResult.value : []
-                const apiCatalogs = catalogsResult.status === 'fulfilled' ? catalogsResult.value : []
-                const errors: LakehouseNode[] = []
-
-                if (awsDataCatalogResult.status === 'rejected') {
-                    const errorMessage = (awsDataCatalogResult.reason as Error).message
-                    void vscode.window.showErrorMessage(errorMessage)
-                    errors.push(createErrorItem(errorMessage, 'aws-data-catalog', node.id) as LakehouseNode)
-                }
+            return telemetry.smus_renderLakehouseNode.run(async (span) => {
+                const isInSmusSpace = getContext('aws.smus.inSmusSpaceEnvironment')
+
+                span.record({
+                    smusToolkitEnv: isInSmusSpace ? 'smus_space' : 'local',
+                    smusDomainId: connection.domainId,
+                    smusProjectId: connection.projectId,
+                    smusConnectionId: connection.connectionId,
+                    smusConnectionType: connection.type,
+                    smusProjectRegion: connection.location?.awsRegion,
+                })
+                try {
+                    logger.info(`Loading Lakehouse catalogs for connection ${connection.name}`)
+
+                    // Check if this is a default connection
+                    const isDefaultConnection =
+                        DATA_DEFAULT_IAM_CONNECTION_NAME_REGEXP.test(connection.name) ||
+                        DATA_DEFAULT_LAKEHOUSE_CONNECTION_NAME_REGEXP.test(connection.name) ||
+                        DATA_DEFAULT_ATHENA_CONNECTION_NAME_REGEXP.test(connection.name)
+
+                    // Follow the reference pattern with Promise.allSettled
+                    const [awsDataCatalogResult, catalogsResult] = await Promise.allSettled([
+                        // AWS Data Catalog node (only for default connections)
+                        isDefaultConnection
+                            ? Promise.resolve([createAwsDataCatalogNode(node, glueClient)])
+                            : Promise.resolve([]),
+                        // Get catalogs by calling Glue API
+                        getCatalogs(glueCatalogClient, glueClient, node),
+                    ])
+
+                    const awsDataCatalog = awsDataCatalogResult.status === 'fulfilled' ? awsDataCatalogResult.value : []
+                    const apiCatalogs = catalogsResult.status === 'fulfilled' ? catalogsResult.value : []
+                    const errors: LakehouseNode[] = []
+
+                    if (awsDataCatalogResult.status === 'rejected') {
+                        const errorMessage = (awsDataCatalogResult.reason as Error).message
+                        void vscode.window.showErrorMessage(errorMessage)
+                        errors.push(createErrorItem(errorMessage, 'aws-data-catalog', node.id) as LakehouseNode)
+                    }
+
+                    if (catalogsResult.status === 'rejected') {
+                        const errorMessage = (catalogsResult.reason as Error).message
+                        void vscode.window.showErrorMessage(errorMessage)
+                        errors.push(createErrorItem(errorMessage, 'catalogs', node.id) as LakehouseNode)
+                    }
 
-                if (catalogsResult.status === 'rejected') {
-                    const errorMessage = (catalogsResult.reason as Error).message
+                    const allNodes = [...awsDataCatalog, ...apiCatalogs, ...errors]
+                    return allNodes.length > 0
+                        ? allNodes
+                        : [createPlaceholderItem(NO_DATA_FOUND_MESSAGE) as LakehouseNode]
+                } catch (err) {
+                    logger.error(`Failed to get Lakehouse catalogs: ${(err as Error).message}`)
+                    const errorMessage = (err as Error).message
                     void vscode.window.showErrorMessage(errorMessage)
-                    errors.push(createErrorItem(errorMessage, 'catalogs', node.id) as LakehouseNode)
+                    return [createErrorItem(errorMessage, 'lakehouse-catalogs', node.id) as LakehouseNode]
                 }
-
-                const allNodes = [...awsDataCatalog, ...apiCatalogs, ...errors]
-                return allNodes.length > 0 ? allNodes : [createPlaceholderItem(NO_DATA_FOUND_MESSAGE) as LakehouseNode]
-            } catch (err) {
-                logger.error(`Failed to get Lakehouse catalogs: ${(err as Error).message}`)
-                const errorMessage = (err as Error).message
-                void vscode.window.showErrorMessage(errorMessage)
-                return [createErrorItem(errorMessage, 'lakehouse-catalogs', node.id) as LakehouseNode]
-            }
+            })
         }
     )
 }
@@ -229,7 +245,9 @@ function createAwsDataCatalogNode(parent: LakehouseNode, glueClient: GlueClient)
                 nextToken = token
             } while (nextToken)
 
-            return allDatabases.map((database) => createDatabaseNode(database.Name || '', database, glueClient, node))
+            return allDatabases.length > 0
+                ? allDatabases.map((database) => createDatabaseNode(database.Name || '', database, glueClient, node))
+                : [createPlaceholderItem(NO_DATA_FOUND_MESSAGE) as LakehouseNode]
         }
     )
 }
@@ -392,9 +410,11 @@ function createCatalogNode(
                           nextToken = token
                       } while (nextToken)
 
-                      return allDatabases.map((database) =>
-                          createDatabaseNode(database.Name || '', database, glueClient, node)
-                      )
+                      return allDatabases.length > 0
+                          ? allDatabases.map((database) =>
+                                createDatabaseNode(database.Name || '', database, glueClient, node)
+                            )
+                          : [createPlaceholderItem(NO_DATA_FOUND_MESSAGE) as LakehouseNode]
                   } catch (err) {
                       logger.error(`Failed to get databases for catalog ${catalogId}: ${(err as Error).message}`)
                       const errorMessage = (err as Error).message
@@ -513,7 +533,10 @@ function createTableNode(
                 const columns = tableDetails?.StorageDescriptor?.Columns || []
                 const partitions = tableDetails?.PartitionKeys || []
 
-                return [...columns, ...partitions].map((column) => createColumnNode(column.Name || '', column, node))
+                const allColumns = [...columns, ...partitions]
+                return allColumns.length > 0
+                    ? allColumns.map((column) => createColumnNode(column.Name || '', column, node))
+                    : [createPlaceholderItem(NO_DATA_FOUND_MESSAGE) as LakehouseNode]
             } catch (err) {
                 logger.error(`Failed to get columns for table ${tableName}: ${(err as Error).message}`)
                 return []
@@ -565,7 +588,9 @@ function createContainerNode(
         },
         async (node) => {
             // Map items to nodes
-            return items.map((item) => createTableNode(item.Name || '', item, glueClient, node))
+            return items.length > 0
+                ? items.map((item) => createTableNode(item.Name || '', item, glueClient, node))
+                : [createPlaceholderItem(NO_DATA_FOUND_MESSAGE) as LakehouseNode]
         }
     )
 }