feat: add codegen support for extra metadata (#846)

Author: ianbotsf

aws-runtime/aws-http/common/src/aws/sdk/kotlin/runtime/http/interceptors/AddUserAgentMetadataInterceptor.kt

@@ -0,0 +1,26 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+package aws.sdk.kotlin.runtime.http.interceptors
+
+import aws.sdk.kotlin.runtime.InternalSdkApi
+import aws.sdk.kotlin.runtime.http.operation.CustomUserAgentMetadata
+import aws.smithy.kotlin.runtime.client.RequestInterceptorContext
+import aws.smithy.kotlin.runtime.http.interceptors.HttpInterceptor
+
+/**
+ * Adds metadata to the user agent sent in requests.
+ * @param metadata A map of keys/values to add to existing metadata.
+ */
+@InternalSdkApi
+public class AddUserAgentMetadataInterceptor(private val metadata: Map<String, String>) : HttpInterceptor {
+    override fun readBeforeExecution(context: RequestInterceptorContext<Any>) {
+        val existing = context.executionContext.getOrNull(CustomUserAgentMetadata.ContextKey)
+
+        val new = existing ?: CustomUserAgentMetadata()
+        metadata.forEach { (k, v) -> new.add(k, v) }
+
+        context.executionContext[CustomUserAgentMetadata.ContextKey] = new
+    }
+}

aws-runtime/aws-http/common/test/aws/sdk/kotlin/runtime/http/interceptors/AddUserAgentMetadataInterceptorTest.kt

@@ -0,0 +1,49 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+package aws.sdk.kotlin.runtime.http.interceptors
+
+import aws.sdk.kotlin.runtime.http.operation.CustomUserAgentMetadata
+import aws.smithy.kotlin.runtime.client.RequestInterceptorContext
+import aws.smithy.kotlin.runtime.operation.ExecutionContext
+import aws.smithy.kotlin.runtime.util.get
+import kotlin.test.Test
+import kotlin.test.assertEquals
+
+class AddUserAgentMetadataInterceptorTest {
+    @Test
+    fun testAddNew() {
+        val ctx = interceptorContext()
+
+        val interceptor = AddUserAgentMetadataInterceptor(metadata)
+        interceptor.readBeforeExecution(ctx)
+
+        val actual = ctx.executionContext[CustomUserAgentMetadata.ContextKey].extras
+        assertEquals("bar", actual["foo"])
+    }
+
+    @Test
+    fun testMerge() {
+        val existingMetadata = mapOf("foo" to "This value will be replaced", "baz" to "qux")
+        val executionContext = ExecutionContext().apply {
+            set(CustomUserAgentMetadata.ContextKey, CustomUserAgentMetadata(extras = existingMetadata))
+        }
+        val ctx = interceptorContext(executionContext)
+
+        val interceptor = AddUserAgentMetadataInterceptor(metadata)
+        interceptor.readBeforeExecution(ctx)
+
+        val actual = ctx.executionContext[CustomUserAgentMetadata.ContextKey].extras
+        assertEquals("bar", actual["foo"])
+        assertEquals("qux", actual["baz"])
+    }
+}
+
+private val metadata = mapOf("foo" to "bar")
+
+private fun interceptorContext(executionContext: ExecutionContext = ExecutionContext()) =
+    object : RequestInterceptorContext<Any> {
+        override val request: Any = Unit
+        override val executionContext: ExecutionContext = executionContext
+    }

codegen/sdk/build.gradle.kts

@@ -307,13 +307,14 @@ fun forwardProperty(name: String) {
 val generateSmithyProjections = tasks.named<SmithyBuild>("generateSmithyProjections") {
     doFirst {
         forwardProperty("aws.partitions_file")
+        forwardProperty("aws.user_agent.add_metadata")
     }
 }
 
 val stageSdks = tasks.register("stageSdks") {
     group = "codegen"
     description = "relocate generated SDK(s) from build directory to services/ dir"
-    dependsOn(tasks.named("generateSmithyProjections"))
+    dependsOn(generateSmithyProjections)
     doLast {
         println("discoveredServices = ${discoveredServices.joinToString { it.sdkId }}")
         discoveredServices.forEach {
@@ -335,7 +336,7 @@ tasks.register("bootstrap") {
     group = "codegen"
     description = "Generate AWS SDK's and register them with the build"
 
-    dependsOn(tasks.named("generateSmithyProjections"))
+    dependsOn(generateSmithyProjections)
     finalizedBy(stageSdks)
 }
 

codegen/smithy-aws-kotlin-codegen/src/main/kotlin/aws/sdk/kotlin/codegen/AddUserAgentMetadataIntegration.kt

@@ -0,0 +1,66 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+package aws.sdk.kotlin.codegen
+
+import software.amazon.smithy.kotlin.codegen.KotlinSettings
+import software.amazon.smithy.kotlin.codegen.core.KotlinWriter
+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.Model
+import software.amazon.smithy.model.node.Node
+import software.amazon.smithy.model.shapes.OperationShape
+import java.io.File
+
+internal const val ADD_USER_AGENT_METADATA_ENV_VAR = "ADD_USER_AGENT_METADATA"
+internal const val ADD_USER_AGENT_METADATA_SYS_PROP = "aws.user_agent.add_metadata"
+
+/**
+ * Adds additional metadata to the user agent sent in requests. This additional metadata is loaded from a JSON file
+ * during the **:codegen:sdk:bootstrap** task. The file name is detected by one of the following means:
+ *
+ * * The `ADD_USER_AGENT_METADATA` environment variable
+ * * The `aws.user_agent.add_metadata` project property. This may be set via:
+ *     * The Gradle command line argument `-Paws.user_agent.add_metadata=<path>` –or–
+ *     * A line in the **local.properties** file
+ *
+ * If neither option is configured, no additional metadata is added to the user agent.
+ */
+class AddUserAgentMetadataIntegration : KotlinIntegration {
+    private val metadataFile: String? by lazy {
+        System.getenv(ADD_USER_AGENT_METADATA_ENV_VAR) ?: System.getProperty(ADD_USER_AGENT_METADATA_SYS_PROP)
+    }
+
+    private val metadata by lazy {
+        val json = metadataFile?.let { File(it).readText() } ?: "{}" // Default to empty JSON object (no extra metadata)
+        val node = Node.parse(json).expectObjectNode()
+        node.members.entries.associate { (k, v) -> k.value to v.expectStringNode().value }
+    }
+
+    override fun customizeMiddleware(
+        ctx: ProtocolGenerator.GenerationContext,
+        resolved: List<ProtocolMiddleware>,
+    ): List<ProtocolMiddleware> = resolved + addUserAgentMetadataMiddleware
+
+    override fun enabledForService(model: Model, settings: KotlinSettings): Boolean = metadataFile != null
+
+    private val addUserAgentMetadataMiddleware = object : ProtocolMiddleware {
+        override val name: String = "AddUserAgentMetadata"
+
+        override fun render(ctx: ProtocolGenerator.GenerationContext, op: OperationShape, writer: KotlinWriter) {
+            writer.write(
+                "op.interceptors.add(#T(extraMetadata))",
+                AwsRuntimeTypes.Http.Interceptors.AddUserAgentMetadataInterceptor,
+            )
+        }
+
+        override fun renderProperties(writer: KotlinWriter) {
+            writer.withBlock("private val extraMetadata: Map<String, String> = mapOf(", ")") {
+                metadata.forEach { (k, v) -> writer.write("#S to #S,", k, v) }
+            }
+        }
+    }
+}

codegen/smithy-aws-kotlin-codegen/src/main/kotlin/aws/sdk/kotlin/codegen/AwsRuntimeTypes.kt

@@ -50,9 +50,14 @@ object AwsRuntimeTypes {
     }
 
     object Http : RuntimeTypePackage(AwsKotlinDependency.AWS_HTTP) {
+        object Interceptors : RuntimeTypePackage(AwsKotlinDependency.AWS_HTTP, "interceptors") {
+            val AddUserAgentMetadataInterceptor = symbol("AddUserAgentMetadataInterceptor")
+        }
+
         object Retries {
             val AwsDefaultRetryPolicy = symbol("AwsDefaultRetryPolicy", "retries")
         }
+
         object Middleware {
             val AwsRetryHeaderMiddleware = symbol("AwsRetryHeaderMiddleware", "middleware")
         }

codegen/smithy-aws-kotlin-codegen/src/main/resources/META-INF/services/software.amazon.smithy.kotlin.codegen.integration.KotlinIntegration

@@ -1,4 +1,5 @@
 aws.sdk.kotlin.codegen.SdkProtocolGeneratorSupplier
+aws.sdk.kotlin.codegen.AddUserAgentMetadataIntegration
 aws.sdk.kotlin.codegen.AwsRetryHeaderIntegration
 aws.sdk.kotlin.codegen.customization.s3.S3GeneratorSupplier
 aws.sdk.kotlin.codegen.GradleGenerator

codegen/smithy-aws-kotlin-codegen/src/test/kotlin/aws/sdk/kotlin/codegen/AddUserAgentMetadataIntegrationTest.kt

@@ -0,0 +1,87 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+package aws.sdk.kotlin.codegen
+
+import aws.sdk.kotlin.codegen.testutil.lines
+import software.amazon.smithy.kotlin.codegen.test.*
+import java.nio.file.Files
+import kotlin.io.path.deleteIfExists
+import kotlin.io.path.writeText
+import kotlin.test.Test
+import kotlin.test.assertFalse
+import kotlin.test.assertTrue
+
+class AddUserAgentMetadataIntegrationTest {
+    private val model = """
+        @http(method: "GET", uri: "/foo")
+        operation Foo { }
+    """
+        .trimIndent()
+        .prependNamespaceAndService(operations = listOf("Foo"))
+        .toSmithyModel()
+    private val settings = model.defaultSettings()
+
+    @Test
+    fun testDisabled() {
+        assertFalse(AddUserAgentMetadataIntegration().enabledForService(model, settings))
+    }
+
+    @Test
+    fun testEnabledBySystemProperty() {
+        withSysProp {
+            assertTrue(AddUserAgentMetadataIntegration().enabledForService(model, settings))
+        }
+    }
+
+    @Test
+    fun testRenderProperty() {
+        withSysProp {
+            val ctx = model.newTestContext(integrations = listOf(AddUserAgentMetadataIntegration()))
+            val generator = MockHttpProtocolGenerator()
+            generator.generateProtocolClient(ctx.generationCtx)
+
+            ctx.generationCtx.delegator.finalize()
+            ctx.generationCtx.delegator.flushWriters()
+
+            val actual = ctx.manifest.expectFileString("/src/main/kotlin/com/test/DefaultTestClient.kt")
+            val expected = """
+                private val extraMetadata: Map<String, String> = mapOf(
+                    "foo" to "bar",
+                )
+            """.replaceIndent("    ")
+            actual.shouldContainOnlyOnceWithDiff(expected)
+        }
+    }
+
+    @Test
+    fun testRenderInterceptor() {
+        withSysProp {
+            val ctx = model.newTestContext(integrations = listOf(AddUserAgentMetadataIntegration()))
+            val generator = MockHttpProtocolGenerator()
+            generator.generateProtocolClient(ctx.generationCtx)
+
+            ctx.generationCtx.delegator.finalize()
+            ctx.generationCtx.delegator.flushWriters()
+
+            val actual = ctx.manifest.expectFileString("/src/main/kotlin/com/test/DefaultTestClient.kt")
+
+            val fooMethod = actual.lines("    override suspend fun foo(input: FooRequest): FooResponse {", "    }")
+            val expectedInterceptor = "op.interceptors.add(AddUserAgentMetadataInterceptor(extraMetadata))"
+            fooMethod.shouldContainOnlyOnceWithDiff(expectedInterceptor)
+        }
+    }
+}
+
+private inline fun withSysProp(block: () -> Unit) {
+    val path = Files.createTempFile(null, null)
+    try {
+        path.writeText(""" { "foo" : "bar" } """)
+        System.setProperty(ADD_USER_AGENT_METADATA_SYS_PROP, path.toString())
+        block()
+    } finally {
+        System.clearProperty(ADD_USER_AGENT_METADATA_SYS_PROP)
+        path.deleteIfExists()
+    }
+}

codegen/smithy-aws-kotlin-codegen/src/test/kotlin/aws/sdk/kotlin/codegen/customization/s3/ContinueIntegrationTest.kt

@@ -1,5 +1,6 @@
 package aws.sdk.kotlin.codegen.customization.s3
 
+import aws.sdk.kotlin.codegen.testutil.lines
 import org.junit.jupiter.api.Test
 import software.amazon.smithy.codegen.core.SymbolProvider
 import software.amazon.smithy.kotlin.codegen.KotlinSettings
@@ -120,16 +121,3 @@ object FooMiddleware : ProtocolMiddleware {
     override fun render(ctx: ProtocolGenerator.GenerationContext, op: OperationShape, writer: KotlinWriter) =
         fail("Unexpected call to `FooMiddleware.render")
 }
-
-private fun String.lines(fromLine: String, toLine: String): String {
-    val allLines = lines()
-
-    val fromIdx = allLines.indexOf(fromLine)
-    assertNotEquals(-1, fromIdx, """Could not find from line "$fromLine" in all lines""")
-
-    val toIdxOffset = allLines.drop(fromIdx + 1).indexOf(toLine)
-    assertNotEquals(-1, toIdxOffset, """Could not find to line "$toLine" in all lines""")
-
-    val toIdx = toIdxOffset + fromIdx + 1
-    return allLines.subList(fromIdx, toIdx + 1).joinToString("\n")
-}

codegen/smithy-aws-kotlin-codegen/src/test/kotlin/aws/sdk/kotlin/codegen/testutil/StringExt.kt

@@ -0,0 +1,20 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+package aws.sdk.kotlin.codegen.testutil
+
+import kotlin.test.assertNotEquals
+
+public fun String.lines(fromLine: String, toLine: String): String {
+    val allLines = lines()
+
+    val fromIdx = allLines.indexOf(fromLine)
+    assertNotEquals(-1, fromIdx, """Could not find from line "$fromLine" in all lines""")
+
+    val toIdxOffset = allLines.drop(fromIdx + 1).indexOf(toLine)
+    assertNotEquals(-1, toIdxOffset, """Could not find to line "$toLine" in all lines""")
+
+    val toIdx = toIdxOffset + fromIdx + 1
+    return allLines.subList(fromIdx, toIdx + 1).joinToString("\n")
+}