deps(sdkv3): implement logging and endpoint override for client builder. (#6441)

Hweinstock · web-flow · commit 66f050ccdc4f · 2025-02-11T12:32:43.000-05:00
## Problem The V3 client builder currently only adds telemetry. We want some additional functionality on each request/response. ## Solution - Allow user to override the endpoint used by the sdk. - Add debugging messages for each request and response. --- - 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.
diff --git a/packages/core/src/shared/awsClientBuilderV3.ts b/packages/core/src/shared/awsClientBuilderV3.ts
@@ -9,9 +9,12 @@ import { AwsCredentialIdentityProvider, RetryStrategyV2 } from '@smithy/types'
 import { getUserAgent } from './telemetry/util'
 import { DevSettings } from './settings'
 import {
+    BuildHandler,
+    BuildMiddleware,
     DeserializeHandler,
-    DeserializeHandlerOptions,
     DeserializeMiddleware,
+    FinalizeHandler,
+    FinalizeRequestMiddleware,
     HandlerExecutionContext,
     MetadataBearer,
     MiddlewareStack,
@@ -21,13 +24,14 @@ import {
     RetryStrategy,
     UserAgent,
 } from '@aws-sdk/types'
-import { HttpResponse } from '@aws-sdk/protocol-http'
+import { HttpResponse, HttpRequest } from '@aws-sdk/protocol-http'
 import { ConfiguredRetryStrategy } from '@smithy/util-retry'
-import { telemetry } from './telemetry'
+import { telemetry } from './telemetry/telemetry'
 import { getRequestId, getTelemetryReason, getTelemetryReasonDesc, getTelemetryResult } from './errors'
-import { extensionVersion } from '.'
-import { getLogger } from './logger'
+import { extensionVersion } from './vscode/env'
+import { getLogger } from './logger/logger'
 import { partialClone } from './utilities/collectionUtils'
+import { selectFrom } from './utilities/tsUtils'
 
 export type AwsClientConstructor<C> = new (o: AwsClientOptions) => C
 
@@ -96,8 +100,9 @@ export class AWSClientBuilderV3 {
         }
 
         const service = new type(opt)
-        // TODO: add middleware for logging, telemetry, endpoints.
-        service.middlewareStack.add(telemetryMiddleware, { step: 'deserialize' } as DeserializeHandlerOptions)
+        service.middlewareStack.add(telemetryMiddleware, { step: 'deserialize' })
+        service.middlewareStack.add(loggingMiddleware, { step: 'finalizeRequest' })
+        service.middlewareStack.add(getEndpointMiddleware(settings), { step: 'build' })
         return service
     }
 }
@@ -123,29 +128,67 @@ export function recordErrorTelemetry(err: Error, serviceName?: string) {
 function logAndThrow(e: any, serviceId: string, errorMessageAppend: string): never {
     if (e instanceof Error) {
         recordErrorTelemetry(e, serviceId)
-        const err = { ...e }
-        delete err['stack']
-        getLogger().error('API Response %s: %O', errorMessageAppend, err)
+        getLogger().error('API Response %s: %O', errorMessageAppend, e)
     }
     throw e
 }
-/**
- * Telemetry logic to be added to all created clients. Adds logging and emitting metric on errors.
- */
+
 const telemetryMiddleware: DeserializeMiddleware<any, any> =
-    (next: DeserializeHandler<any, any>, context: HandlerExecutionContext) => async (args: any) => {
-        if (!HttpResponse.isInstance(args.request)) {
-            return next(args)
-        }
-        const serviceId = getServiceId(context as object)
-        const { hostname, path } = args.request
-        const logTail = `(${hostname} ${path})`
-        const result = await next(args).catch((e: any) => logAndThrow(e, serviceId, logTail))
+    (next: DeserializeHandler<any, any>, context: HandlerExecutionContext) => async (args: any) =>
+        emitOnRequest(next, context, args)
+
+const loggingMiddleware: FinalizeRequestMiddleware<any, any> = (next: FinalizeHandler<any, any>) => async (args: any) =>
+    logOnRequest(next, args)
+
+function getEndpointMiddleware(settings: DevSettings = DevSettings.instance): BuildMiddleware<any, any> {
+    return (next: BuildHandler<any, any>, context: HandlerExecutionContext) => async (args: any) =>
+        overwriteEndpoint(next, context, settings, args)
+}
+
+export async function emitOnRequest(next: DeserializeHandler<any, any>, context: HandlerExecutionContext, args: any) {
+    if (!HttpResponse.isInstance(args.request)) {
+        return next(args)
+    }
+    const serviceId = getServiceId(context as object)
+    const { hostname, path } = args.request
+    const logTail = `(${hostname} ${path})`
+    try {
+        const result = await next(args)
         if (HttpResponse.isInstance(result.response)) {
-            // TODO: omit credentials / sensitive info from the logs / telemetry.
+            // TODO: omit credentials / sensitive info from the telemetry.
             const output = partialClone(result.output, 3)
-            getLogger().debug('API Response %s: %O', logTail, output)
+            getLogger().debug(`API Response %s: %O`, logTail, output)
         }
-
         return result
+    } catch (e: any) {
+        logAndThrow(e, serviceId, logTail)
+    }
+}
+
+export async function logOnRequest(next: FinalizeHandler<any, any>, args: any) {
+    if (HttpRequest.isInstance(args.request)) {
+        const { hostname, path } = args.request
+        // TODO: omit credentials / sensitive info from the logs.
+        const input = partialClone(args.input, 3)
+        getLogger().debug(`API Request (%s %s): %O`, hostname, path, input)
     }
+    return next(args)
+}
+
+export function overwriteEndpoint(
+    next: BuildHandler<any, any>,
+    context: HandlerExecutionContext,
+    settings: DevSettings,
+    args: any
+) {
+    if (HttpRequest.isInstance(args.request)) {
+        const serviceId = getServiceId(context as object)
+        const endpoint = serviceId ? settings.get('endpoints', {})[serviceId] : undefined
+        if (endpoint) {
+            const url = new URL(endpoint)
+            Object.assign(args.request, selectFrom(url, 'hostname', 'port', 'protocol', 'pathname'))
+            args.request.path = args.request.pathname
+        }
+    }
+    return next(args)
+}
diff --git a/packages/core/src/test/globalSetup.test.ts b/packages/core/src/test/globalSetup.test.ts
@@ -167,6 +167,10 @@ async function writeLogsToFile(testName: string) {
     await fs.appendFile(testLogOutput, entries?.join('\n') ?? '')
 }
 
+export function assertLogsContainAllOf(keywords: string[], exactMatch: boolean, severity: LogLevel) {
+    return keywords.map((k) => assertLogsContain(k, exactMatch, severity))
+}
+
 // TODO: merge this with `toolkitLogger.test.ts:checkFile`
 export function assertLogsContain(text: string, exactMatch: boolean, severity: LogLevel) {
     const logs = getTestLogger().getLoggedEntries(severity)
diff --git a/packages/core/src/test/shared/awsClientBuilderV3.test.ts b/packages/core/src/test/shared/awsClientBuilderV3.test.ts
@@ -2,18 +2,28 @@
  * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
-
+import sinon from 'sinon'
 import assert from 'assert'
 import { version } from 'vscode'
 import { getClientId } from '../../shared/telemetry/util'
 import { FakeMemento } from '../fakeExtensionContext'
 import { FakeAwsContext } from '../utilities/fakeAwsContext'
 import { GlobalState } from '../../shared/globalState'
-import { AWSClientBuilderV3, getServiceId, recordErrorTelemetry } from '../../shared/awsClientBuilderV3'
+import {
+    AWSClientBuilderV3,
+    emitOnRequest,
+    getServiceId,
+    logOnRequest,
+    overwriteEndpoint,
+    recordErrorTelemetry,
+} from '../../shared/awsClientBuilderV3'
 import { Client } from '@aws-sdk/smithy-client'
-import { extensionVersion } from '../../shared'
+import { DevSettings, extensionVersion } from '../../shared'
 import { assertTelemetry } from '../testUtil'
 import { telemetry } from '../../shared/telemetry'
+import { HttpRequest, HttpResponse } from '@aws-sdk/protocol-http'
+import { assertLogsContain, assertLogsContainAllOf } from '../globalSetup.test'
+import { TestSettings } from '../utilities/testSettingsConfiguration'
 import { CredentialsShim } from '../../auth/deprecated/loginManager'
 import { Credentials } from '@aws-sdk/types'
 import { oneDay } from '../../shared/datetime'
@@ -25,39 +35,126 @@ describe('AwsClientBuilderV3', function () {
         builder = new AWSClientBuilderV3(new FakeAwsContext())
     })
 
-    describe('createAndConfigureSdkClient', function () {
-        it('includes Toolkit user-agent if no options are specified', async function () {
-            const service = await builder.createAwsService(Client)
-            const clientId = getClientId(new GlobalState(new FakeMemento()))
-
-            assert.ok(service.config.userAgent)
-            assert.strictEqual(
-                service.config.userAgent![0][0].replace('---Insiders', ''),
-                `AWS-Toolkit-For-VSCode/testPluginVersion Visual-Studio-Code/${version} ClientId/${clientId}`
-            )
-            assert.strictEqual(service.config.userAgent![0][1], extensionVersion)
+    it('includes Toolkit user-agent if no options are specified', async function () {
+        const service = await builder.createAwsService(Client)
+        const clientId = getClientId(new GlobalState(new FakeMemento()))
+
+        assert.ok(service.config.userAgent)
+        assert.strictEqual(
+            service.config.userAgent![0][0].replace('---Insiders', ''),
+            `AWS-Toolkit-For-VSCode/testPluginVersion Visual-Studio-Code/${version} ClientId/${clientId}`
+        )
+        assert.strictEqual(service.config.userAgent![0][1], extensionVersion)
+    })
+
+    it('adds region to client', async function () {
+        const service = await builder.createAwsService(Client, { region: 'us-west-2' })
+
+        assert.ok(service.config.region)
+        assert.strictEqual(service.config.region, 'us-west-2')
+    })
+
+    it('adds Client-Id to user agent', async function () {
+        const service = await builder.createAwsService(Client)
+        const clientId = getClientId(new GlobalState(new FakeMemento()))
+        const regex = new RegExp(`ClientId/${clientId}`)
+        assert.ok(service.config.userAgent![0][0].match(regex))
+    })
+
+    it('does not override custom user-agent if specified in options', async function () {
+        const service = await builder.createAwsService(Client, {
+            userAgent: [['CUSTOM USER AGENT']],
         })
 
-        it('adds region to client', async function () {
-            const service = await builder.createAwsService(Client, { region: 'us-west-2' })
+        assert.strictEqual(service.config.userAgent[0][0], 'CUSTOM USER AGENT')
+    })
 
-            assert.ok(service.config.region)
-            assert.strictEqual(service.config.region, 'us-west-2')
+    describe('middlewareStack', function () {
+        let args: { request: { hostname: string; path: string }; input: any }
+        let context: { clientName?: string; commandName?: string }
+        let response: { response: { statusCode: number }; output: { message: string } }
+        let httpRequestStub: sinon.SinonStub
+        let httpResponseStub: sinon.SinonStub
+
+        before(function () {
+            httpRequestStub = sinon.stub(HttpRequest, 'isInstance')
+            httpResponseStub = sinon.stub(HttpResponse, 'isInstance')
+            httpRequestStub.callsFake(() => true)
+            httpResponseStub.callsFake(() => true)
         })
 
-        it('adds Client-Id to user agent', async function () {
-            const service = await builder.createAwsService(Client)
-            const clientId = getClientId(new GlobalState(new FakeMemento()))
-            const regex = new RegExp(`ClientId/${clientId}`)
-            assert.ok(service.config.userAgent![0][0].match(regex))
+        beforeEach(function () {
+            args = {
+                request: {
+                    hostname: 'testHost',
+                    path: 'testPath',
+                },
+                input: {
+                    testKey: 'testValue',
+                },
+            }
+            context = {
+                clientName: 'fooClient',
+            }
+            response = {
+                response: {
+                    statusCode: 200,
+                },
+                output: {
+                    message: 'test output',
+                },
+            }
+        })
+        after(function () {
+            sinon.restore()
+        })
+
+        it('logs messages on request', async function () {
+            await logOnRequest((_: any) => _, args as any)
+            assertLogsContainAllOf(['testHost', 'testPath'], false, 'debug')
+        })
+
+        it('adds telemetry metadata and logs on error failure', async function () {
+            const next = (_: any) => {
+                throw new Error('test error')
+            }
+            await telemetry.vscode_executeCommand.run(async (span) => {
+                await assert.rejects(emitOnRequest(next, context, args))
+            })
+            assertLogsContain('test error', false, 'error')
+            assertTelemetry('vscode_executeCommand', { requestServiceType: 'foo' })
         })
 
-        it('does not override custom user-agent if specified in options', async function () {
-            const service = await builder.createAwsService(Client, {
-                userAgent: [['CUSTOM USER AGENT']],
+        it('does not emit telemetry, but still logs on successes', async function () {
+            const next = async (_: any) => {
+                return response
+            }
+            await telemetry.vscode_executeCommand.run(async (span) => {
+                assert.deepStrictEqual(await emitOnRequest(next, context, args), response)
             })
+            assertLogsContainAllOf(['testHost', 'testPath'], false, 'debug')
+            assert.throws(() => assertTelemetry('vscode_executeCommand', { requestServiceType: 'foo' }))
+        })
+
+        it('custom endpoints overwrite request url', async function () {
+            const settings = new TestSettings()
+            await settings.update('aws.dev.endpoints', { foo: 'http://example.com:3000/path' })
+            const next = async (args: any) => args
+            const newArgs: any = await overwriteEndpoint(next, context, new DevSettings(settings), args)
+
+            assert.strictEqual(newArgs.request.hostname, 'example.com')
+            assert.strictEqual(newArgs.request.protocol, 'http:')
+            assert.strictEqual(newArgs.request.port, '3000')
+            assert.strictEqual(newArgs.request.pathname, '/path')
+        })
+
+        it('custom endpoints are not overwritten if not specified', async function () {
+            const settings = new TestSettings()
+            const next = async (args: any) => args
+            const newArgs: any = await overwriteEndpoint(next, context, new DevSettings(settings), args)
 
-            assert.strictEqual(service.config.userAgent[0][0], 'CUSTOM USER AGENT')
+            assert.strictEqual(newArgs.request.hostname, 'testHost')
+            assert.strictEqual(newArgs.request.path, 'testPath')
         })
     })
 
