Merge master into feature/mcp-registry-dev

Author: aws-toolkit-automation

packages/core/src/sagemakerunifiedstudio/shared/telemetry.ts

@@ -21,6 +21,7 @@ import { ConnectionCredentialsProvider } from '../auth/providers/connectionCrede
 import { DataZoneConnection } from './client/datazoneClient'
 import { createDZClientBaseOnDomainMode } from '../explorer/nodes/utils'
 
+const notSet = 'not-set'
 /**
  * Records space telemetry
  */
@@ -30,37 +31,46 @@ export async function recordSpaceTelemetry(
 ) {
     const logger = getLogger('smus')
 
+    const parent = node.resource.getParent() as SageMakerUnifiedStudioSpacesParentNode
+    const authProvider = SmusAuthenticationProvider.fromContext()
+    const projectId = parent?.getProjectId()
+    const domainId = parent?.getAuthProvider()?.getDomainId()
+
+    span.record({
+        smusAuthMode: authProvider.activeConnection?.type,
+        smusSpaceKey: node.resource.DomainSpaceKey,
+        smusDomainRegion: node.resource.regionCode,
+        smusDomainId: domainId,
+        smusProjectId: projectId,
+    })
+
     try {
-        const parent = node.resource.getParent() as SageMakerUnifiedStudioSpacesParentNode
-        const authProvider = SmusAuthenticationProvider.fromContext()
         const accountId = await authProvider.getDomainAccountId()
-        const projectId = parent?.getProjectId()
-
-        // Get project account ID and region
-        let projectAccountId: string | undefined
-        let projectRegion: string | undefined
+        span.record({ smusDomainAccountId: accountId })
+    } catch (err) {
+        span.record({ smusDomainAccountId: notSet })
+        logger.warn(`Failed to record domain account Id for telemetry in domain ${domainId}: ${(err as Error).message}`)
+    }
 
-        if (projectId) {
-            projectAccountId = await authProvider.getProjectAccountId(projectId)
+    if (projectId) {
+        try {
+            const projectAccountId = await authProvider.getProjectAccountId(projectId)
+            span.record({ smusProjectAccountId: projectAccountId })
+        } catch (err) {
+            span.record({ smusProjectAccountId: notSet })
+            logger.warn(
+                `Failed to record project account Id for telemetry in domain ${domainId}: ${(err as Error).message}`
+            )
+        }
 
-            // Get project region from tooling environment
+        try {
             const dzClient = await createDZClientBaseOnDomainMode(authProvider)
             const toolingEnv = await dzClient.getToolingEnvironment(projectId)
-            projectRegion = toolingEnv.awsAccountRegion
+            span.record({ smusProjectRegion: toolingEnv.awsAccountRegion })
+        } catch (err) {
+            span.record({ smusProjectRegion: notSet })
+            logger.warn(`Failed to get project region for telemetry: ${(err as Error).message}`)
         }
-
-        span.record({
-            smusAuthMode: authProvider.activeConnection?.type,
-            smusSpaceKey: node.resource.DomainSpaceKey,
-            smusDomainRegion: node.resource.regionCode,
-            smusDomainId: parent?.getAuthProvider()?.getDomainId(),
-            smusDomainAccountId: accountId,
-            smusProjectId: projectId,
-            smusProjectAccountId: projectAccountId,
-            smusProjectRegion: projectRegion,
-        })
-    } catch (err) {
-        logger.error(`Failed to record space telemetry: ${(err as Error).message}`)
     }
 }
 
@@ -90,6 +100,7 @@ export async function recordAuthTelemetry(
             smusDomainAccountId: accountId,
         })
     } catch (err) {
+        span.record({ smusDomainAccountId: notSet })
         logger.error(
             `Failed to record Domain AccountId in data connection telemetry for domain ${domainId} in region ${region}: ${err}`
         )
@@ -106,23 +117,25 @@ export async function recordDataConnectionTelemetry(
 ) {
     const logger = getLogger('smus')
 
+    const isInSmusSpace = getContext('aws.smus.inSmusSpaceEnvironment')
+    const authProvider = SmusAuthenticationProvider.fromContext()
+
+    span.record({
+        smusAuthMode: authProvider.activeConnection?.type,
+        smusToolkitEnv: isInSmusSpace ? 'smus_space' : 'local',
+        smusDomainId: connection.domainId,
+        smusProjectId: connection.projectId,
+        smusConnectionId: connection.connectionId,
+        smusConnectionType: connection.type,
+        smusProjectRegion: connection.location?.awsRegion,
+        smusProjectAccountId: connection.location?.awsAccountId,
+    })
+
     try {
-        const isInSmusSpace = getContext('aws.smus.inSmusSpaceEnvironment')
-        const authProvider = SmusAuthenticationProvider.fromContext()
         const accountId = await connectionCredentialsProvider.getDomainAccountId()
-
-        span.record({
-            smusAuthMode: authProvider.activeConnection?.type,
-            smusToolkitEnv: isInSmusSpace ? 'smus_space' : 'local',
-            smusDomainId: connection.domainId,
-            smusDomainAccountId: accountId,
-            smusProjectId: connection.projectId,
-            smusConnectionId: connection.connectionId,
-            smusConnectionType: connection.type,
-            smusProjectRegion: connection.location?.awsRegion,
-            smusProjectAccountId: connection.location?.awsAccountId,
-        })
+        span.record({ smusDomainAccountId: accountId })
     } catch (err) {
-        logger.error(`Failed to record data connection telemetry: ${(err as Error).message}`)
+        span.record({ smusDomainAccountId: notSet })
+        logger.warn(`Failed to record domain account ID for data connection telemetry: ${(err as Error).message}`)
     }
 }

packages/core/src/test/sagemakerunifiedstudio/testUtils.ts

@@ -65,25 +65,22 @@ export function createMockUnauthenticatedAuthProvider(): any {
  Creates a mock space node for SageMaker Unified Studio tests
  */
 export function createMockSpaceNode(): any {
+    const mockParent = {
+        getAuthProvider: sinon.stub().returns({
+            activeConnection: { domainId: 'test-domain' },
+            getDomainAccountId: sinon.stub().resolves('123456789012'),
+            getDomainId: sinon.stub().returns('test-domain'),
+        }),
+        getProjectId: sinon.stub().returns('test-project'),
+    }
+
     return {
         resource: {
             sageMakerClient: {},
             DomainSpaceKey: 'test-space-key',
             regionCode: 'us-east-1',
-            getParent: sinon.stub().returns({
-                getAuthProvider: sinon.stub().returns({
-                    activeConnection: { domainId: 'test-domain' },
-                    getDomainAccountId: sinon.stub().resolves('123456789012'),
-                }),
-                getProjectId: sinon.stub().returns('test-project'),
-            }),
+            getParent: sinon.stub().returns(mockParent),
         },
-        getParent: sinon.stub().returns({
-            getAuthProvider: sinon.stub().returns({
-                activeConnection: { domainId: 'test-domain' },
-                getDomainAccountId: sinon.stub().resolves('123456789012'),
-            }),
-            getProjectId: sinon.stub().returns('test-project'),
-        }),
+        getParent: sinon.stub().returns(mockParent),
     }
 }