Self review & refactoring

Author: 0marperez

aws-runtime/aws-config/common/src/aws/sdk/kotlin/runtime/auth/credentials/LazilyInitializedCredentialsProvider.kt

@@ -22,11 +22,12 @@ public class LazilyInitializedCredentialsProvider(
 ) : CredentialsProvider {
     private val provider = lazy(initializer)
 
-    override suspend fun resolve(attributes: Attributes): Credentials {
-        val credentials = provider.value.resolve(attributes)
-        if (businessMetric == null) return credentials
-        return credentials.emitBusinessMetric(businessMetric)
-    }
+    override suspend fun resolve(attributes: Attributes): Credentials =
+        if (businessMetric == null) {
+            provider.value.resolve(attributes)
+        } else {
+            provider.value.resolve(attributes).emitBusinessMetric(businessMetric)
+        }
 
     override fun toString(): String = providerName
 }

aws-runtime/aws-config/common/test/aws/sdk/kotlin/runtime/util/AwsBusinessMetricsTestUtils.kt

@@ -37,5 +37,8 @@ internal fun testAttributes(vararg metrics: BusinessMetric): Attributes {
 internal fun Credentials.withBusinessMetrics(vararg metrics: BusinessMetric): Credentials =
     emitBusinessMetrics(metrics.toSet())
 
+/**
+ * Converts a [String] into an [AwsBusinessMetric.Credentials] if the identifier matches
+ */
 internal fun String.toAwsCredentialsBusinessMetric(): BusinessMetric =
     AwsBusinessMetric.Credentials.entries.find { it.identifier == this } ?: throw Exception("String '$this' is not an AWS business metric")

aws-runtime/aws-config/jvm/test/aws/sdk/kotlin/runtime/auth/credentials/DefaultChainCredentialsProviderTest.kt

@@ -11,7 +11,6 @@ import aws.sdk.kotlin.runtime.util.toAwsCredentialsBusinessMetric
 import aws.smithy.kotlin.runtime.auth.awscredentials.Credentials
 import aws.smithy.kotlin.runtime.auth.awscredentials.copy
 import aws.smithy.kotlin.runtime.httptest.TestConnection
-import aws.smithy.kotlin.runtime.operation.ExecutionContext
 import aws.smithy.kotlin.runtime.time.Instant
 import aws.smithy.kotlin.runtime.util.Filesystem
 import aws.smithy.kotlin.runtime.util.OperatingSystem
@@ -169,8 +168,7 @@ class DefaultChainCredentialsProviderTest {
     fun executeTest(name: String) = runTest {
         val test = makeTest(name)
         val provider = DefaultChainCredentialsProvider(platformProvider = test.testPlatform, httpClient = test.testEngine)
-        val attributes = ExecutionContext()
-        val actual = runCatching { provider.resolve(attributes) }
+        val actual = runCatching { provider.resolve() }
         val expected = test.expected
         when {
             expected is TestResult.Ok && actual.isFailure -> error("expected success, got error: $actual")

codegen/aws-sdk-codegen/src/main/kotlin/aws/sdk/kotlin/codegen/BusinessMetricsIntegration.kt

@@ -0,0 +1,32 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+package aws.sdk.kotlin.codegen.businessmetrics
+
+import aws.sdk.kotlin.codegen.AwsRuntimeTypes
+import software.amazon.smithy.kotlin.codegen.core.KotlinWriter
+import software.amazon.smithy.kotlin.codegen.integration.KotlinIntegration
+import software.amazon.smithy.kotlin.codegen.rendering.protocol.ProtocolGenerator
+import software.amazon.smithy.kotlin.codegen.rendering.protocol.ProtocolMiddleware
+import software.amazon.smithy.model.shapes.OperationShape
+
+/**
+ * Renders the addition of the [aws.sdk.kotlin.runtime.http.interceptors.businessmetrics.BusinessMetricsInterceptor]
+ */
+class BusinessMetricsInterceptorIntegration : KotlinIntegration {
+    override fun customizeMiddleware(
+        ctx: ProtocolGenerator.GenerationContext,
+        resolved: List<ProtocolMiddleware>,
+    ): List<ProtocolMiddleware> = resolved + userAgentBusinessMetricsMiddleware
+
+    private val userAgentBusinessMetricsMiddleware = object : ProtocolMiddleware {
+        override val name: String = "UserAgentBusinessMetrics"
+        override fun render(ctx: ProtocolGenerator.GenerationContext, op: OperationShape, writer: KotlinWriter) {
+            writer.write(
+                "op.interceptors.add(#T())",
+                AwsRuntimeTypes.Http.Interceptors.BusinessMetrics.BusinessMetricsInterceptor,
+            )
+        }
+    }
+}

codegen/aws-sdk-codegen/src/main/kotlin/aws/sdk/kotlin/codegen/businessmetrics/BusinessMetricsIntegration.kt

@@ -0,0 +1,37 @@
+package aws.sdk.kotlin.codegen.businessmetrics
+
+import aws.sdk.kotlin.codegen.AwsRuntimeTypes
+import software.amazon.smithy.kotlin.codegen.core.KotlinWriter
+import software.amazon.smithy.kotlin.codegen.core.RuntimeTypes
+import software.amazon.smithy.kotlin.codegen.core.withBlock
+import software.amazon.smithy.kotlin.codegen.integration.KotlinIntegration
+import software.amazon.smithy.kotlin.codegen.rendering.protocol.ProtocolGenerator
+import software.amazon.smithy.kotlin.codegen.rendering.protocol.ProtocolMiddleware
+import software.amazon.smithy.model.shapes.OperationShape
+
+/**
+ * Renders the addition of some of the credentials related business metrics.
+ */
+class CredentialsBusinessMetricsIntegration : KotlinIntegration {
+    override fun customizeMiddleware(
+        ctx: ProtocolGenerator.GenerationContext,
+        resolved: List<ProtocolMiddleware>,
+    ): List<ProtocolMiddleware> = resolved + credentialsBusinessMetricsMiddleware
+
+    private val credentialsBusinessMetricsMiddleware = object : ProtocolMiddleware {
+        override val name: String = "credentialsOverrideBusinessMetricsMiddleware"
+        override fun render(ctx: ProtocolGenerator.GenerationContext, op: OperationShape, writer: KotlinWriter) {
+            writer.withBlock(
+                "if (config.credentialsProvider is #T) {",
+                "}",
+                AwsRuntimeTypes.Config.Credentials.StaticCredentialsProvider,
+            ) {
+                write(
+                    "op.context.#T(#T.Credentials.CREDENTIALS_CODE)",
+                    RuntimeTypes.Core.BusinessMetrics.emitBusinessMetric,
+                    AwsRuntimeTypes.Http.Interceptors.BusinessMetrics.AwsBusinessMetric,
+                )
+            }
+        }
+    }
+}

codegen/aws-sdk-codegen/src/main/kotlin/aws/sdk/kotlin/codegen/businessmetrics/CredentialsBusinessMetricsIntegration.kt

@@ -0,0 +1,43 @@
+package aws.sdk.kotlin.codegen.businessmetrics
+
+import aws.sdk.kotlin.codegen.AwsRuntimeTypes
+import software.amazon.smithy.kotlin.codegen.core.RuntimeTypes
+import software.amazon.smithy.kotlin.codegen.core.RuntimeTypes.Auth.Signing.AwsSigningCommon.AwsSigningAttributes
+import software.amazon.smithy.kotlin.codegen.integration.KotlinIntegration
+import software.amazon.smithy.kotlin.codegen.integration.SectionWriter
+import software.amazon.smithy.kotlin.codegen.integration.SectionWriterBinding
+import software.amazon.smithy.kotlin.codegen.rendering.endpoints.EndpointBusinessMetrics
+
+/**
+ * Renders the addition of endpoint & endpoint adjacent business metrics.
+ */
+class EndpointBusinessMetricsIntegration : KotlinIntegration {
+    override val sectionWriters: List<SectionWriterBinding>
+        get() = listOf(
+            SectionWriterBinding(EndpointBusinessMetrics, endpointBusinessMetricsSectionWriter),
+        )
+
+    private val endpointBusinessMetricsSectionWriter = SectionWriter { writer, _ ->
+        writer.write("")
+        writer.write(
+            "if (endpoint.attributes.contains(#T)) request.context.#T(#T.SERVICE_ENDPOINT_OVERRIDE)",
+            RuntimeTypes.Core.BusinessMetrics.ServiceEndpointOverride,
+            RuntimeTypes.Core.BusinessMetrics.emitBusinessMetric,
+            RuntimeTypes.Core.BusinessMetrics.SmithyBusinessMetric,
+        )
+        writer.write(
+            "if (endpoint.attributes.contains(#T)) request.context.#T(#T.ACCOUNT_ID_BASED_ENDPOINT)",
+            RuntimeTypes.Core.BusinessMetrics.AccountIdBasedEndpointAccountId,
+            RuntimeTypes.Core.BusinessMetrics.emitBusinessMetric,
+            RuntimeTypes.Core.BusinessMetrics.SmithyBusinessMetric,
+        )
+        writer.write(
+            "if (endpoint.attributes.contains(#T.SigningService) && endpoint.attributes[#T.SigningService] == \"s3express\") request.context.#T(#T.S3_EXPRESS_BUCKET)",
+            AwsSigningAttributes,
+            AwsSigningAttributes,
+            RuntimeTypes.Core.BusinessMetrics.emitBusinessMetric,
+            AwsRuntimeTypes.Http.Interceptors.BusinessMetrics.AwsBusinessMetric,
+        )
+        writer.write("")
+    }
+}

codegen/aws-sdk-codegen/src/main/kotlin/aws/sdk/kotlin/codegen/businessmetrics/EndpointBusinessMetricsIntegration.kt

@@ -42,5 +42,7 @@ aws.sdk.kotlin.codegen.customization.cloudfrontkeyvaluestore.BackfillSigV4ACusto
 aws.sdk.kotlin.codegen.customization.s3.express.SigV4S3ExpressAuthSchemeIntegration
 aws.sdk.kotlin.codegen.customization.s3.express.S3ExpressIntegration
 aws.sdk.kotlin.codegen.customization.s3.S3ExpiresIntegration
-aws.sdk.kotlin.codegen.BusinessMetricsIntegration
+aws.sdk.kotlin.codegen.businessmetrics.BusinessMetricsInterceptorIntegration
+aws.sdk.kotlin.codegen.businessmetrics.CredentialsBusinessMetricsIntegration
+aws.sdk.kotlin.codegen.businessmetrics.EndpointBusinessMetricsIntegration
 aws.sdk.kotlin.codegen.smoketests.SmokeTestsDenyListIntegration