feat: enable machinelearning endpoint customization (#378)

Author: ianbotsf

codegen/protocol-tests/build.gradle.kts

@@ -37,7 +37,8 @@ val enabledProtocols = listOf(
 
     // service specific tests
     ProtocolTest("apigateway", "com.amazonaws.apigateway#BackplaneControlService"),
-    ProtocolTest("glacier", "com.amazonaws.glacier#Glacier")
+    ProtocolTest("glacier", "com.amazonaws.glacier#Glacier"),
+    ProtocolTest("machinelearning", "com.amazonaws.machinelearning#AmazonML_20141212", sdkId = "Machine Learning"),
 )
 
 // This project doesn't produce a JAR.
@@ -79,7 +80,7 @@ tasks.create<SmithyBuild>("generateSdk") {
 // force rebuild every time while developing
 tasks["generateSdk"].outputs.upToDateWhen { false }
 
-data class ProtocolTest(val projectionName: String, val serviceShapeId: String) {
+data class ProtocolTest(val projectionName: String, val serviceShapeId: String, val sdkId: String? = null) {
     val packageName: String
         get() = projectionName.toLowerCase().filter { it.isLetterOrDigit() }
 }
@@ -90,6 +91,7 @@ data class ProtocolTest(val projectionName: String, val serviceShapeId: String)
 // it's rebuilt each time codegen is performed.
 fun generateSmithyBuild(tests: List<ProtocolTest>): String {
     val projections = tests.joinToString(",") { test ->
+        val sdkIdEntry = test.sdkId?.let { """"sdkId": "$it",""" } ?: ""
         """
         "${test.projectionName}": {
           "transforms": [
@@ -109,6 +111,7 @@ fun generateSmithyBuild(tests: List<ProtocolTest>): String {
                 "name": "aws.sdk.kotlin.services.${test.packageName}",
                 "version": "1.0"
               },
+              $sdkIdEntry
               "build": {
                 "rootProject": true,
                 "optInAnnotations": [
@@ -131,7 +134,6 @@ fun generateSmithyBuild(tests: List<ProtocolTest>): String {
     """.trimIndent()
 }
 
-
 open class ProtocolTestTask : DefaultTask() {
     /**
      * The protocol name
@@ -194,7 +196,6 @@ enabledProtocols.forEach {
     }
 }
 
-
 tasks.register("testAllProtocols") {
     group = "Verification"
     val allTests = tasks.withType<ProtocolTestTask>()

codegen/smithy-aws-kotlin-codegen/src/main/kotlin/aws/sdk/kotlin/codegen/customization/machinelearning/MachineLearningEndpointCustomization.kt

@@ -0,0 +1,47 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0.
+ */
+
+package aws.sdk.kotlin.codegen.customization.machinelearning
+
+import aws.sdk.kotlin.codegen.protocols.middleware.ResolveAwsEndpointMiddleware
+import aws.sdk.kotlin.codegen.sdkId
+import software.amazon.smithy.kotlin.codegen.KotlinSettings
+import software.amazon.smithy.kotlin.codegen.core.KotlinWriter
+import software.amazon.smithy.kotlin.codegen.integration.KotlinIntegration
+import software.amazon.smithy.kotlin.codegen.model.buildSymbol
+import software.amazon.smithy.kotlin.codegen.model.expectShape
+import software.amazon.smithy.kotlin.codegen.rendering.protocol.*
+import software.amazon.smithy.model.Model
+import software.amazon.smithy.model.shapes.OperationShape
+import software.amazon.smithy.model.shapes.ServiceShape
+
+class MachineLearningEndpointCustomization : KotlinIntegration {
+    override fun customizeMiddleware(
+        ctx: ProtocolGenerator.GenerationContext,
+        resolved: List<ProtocolMiddleware>
+    ): List<ProtocolMiddleware> =
+        super
+            .customizeMiddleware(ctx, resolved)
+            .replace(endpointResolverMiddleware) { it is ResolveAwsEndpointMiddleware }
+
+    private val endpointResolverMiddleware = object : HttpFeatureMiddleware() {
+        override val name: String = "ResolvePredictEndpoint"
+
+        override fun isEnabledFor(ctx: ProtocolGenerator.GenerationContext, op: OperationShape): Boolean =
+            op.id.name == "Predict"
+
+        override fun renderConfigure(writer: KotlinWriter) {
+            writer.addImport(machineLearningSymbol("ResolvePredictEndpoint"))
+        }
+
+        private fun machineLearningSymbol(name: String) = buildSymbol {
+            this.name = name
+            namespace = "aws.sdk.kotlin.services.machinelearning.internal"
+        }
+    }
+
+    override fun enabledForService(model: Model, settings: KotlinSettings): Boolean =
+        model.expectShape<ServiceShape>(settings.service).sdkId.equals("Machine Learning", ignoreCase = true)
+}

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

@@ -15,3 +15,4 @@ aws.sdk.kotlin.codegen.customization.glacier.GlacierAccountIdDefault
 aws.sdk.kotlin.codegen.customization.polly.PollyPresigner
 aws.sdk.kotlin.codegen.customization.BoxServices
 aws.sdk.kotlin.codegen.customization.glacier.GlacierBodyChecksum
+aws.sdk.kotlin.codegen.customization.machinelearning.MachineLearningEndpointCustomization

services/machinelearning/common/src/aws/sdk/kotlin/services/machinelearning/internal/PredictEndpointKey.kt

@@ -0,0 +1,16 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0.
+ */
+
+package aws.sdk.kotlin.services.machinelearning.internal
+
+import aws.sdk.kotlin.runtime.endpoint.AwsEndpoint
+import aws.smithy.kotlin.runtime.client.ExecutionContext
+import aws.smithy.kotlin.runtime.util.AttributeKey
+
+internal val predictEndpointKey = AttributeKey<AwsEndpoint>("PredictEndpointKey")
+
+internal var ExecutionContext.predictEndpoint: AwsEndpoint?
+    get() = getOrNull(predictEndpointKey)
+    set(value) = set(predictEndpointKey, value!!)

services/machinelearning/common/src/aws/sdk/kotlin/services/machinelearning/internal/ResolvePredictEndpoint.kt

@@ -0,0 +1,44 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0.
+ */
+
+package aws.sdk.kotlin.services.machinelearning.internal
+
+import aws.sdk.kotlin.runtime.endpoint.AwsEndpoint
+import aws.sdk.kotlin.services.machinelearning.model.MachineLearningException
+import aws.sdk.kotlin.services.machinelearning.model.PredictRequest
+import aws.smithy.kotlin.runtime.http.Feature
+import aws.smithy.kotlin.runtime.http.FeatureKey
+import aws.smithy.kotlin.runtime.http.HttpClientFeatureFactory
+import aws.smithy.kotlin.runtime.http.middleware.setRequestEndpoint
+import aws.smithy.kotlin.runtime.http.operation.SdkHttpOperation
+
+internal class ResolvePredictEndpoint : Feature {
+    companion object Feature : HttpClientFeatureFactory<Unit, ResolvePredictEndpoint> {
+        override val key: FeatureKey<ResolvePredictEndpoint> = FeatureKey("ResolvePredictEndpoint")
+        override fun create(block: Unit.() -> Unit): ResolvePredictEndpoint = ResolvePredictEndpoint()
+    }
+
+    override fun <I, O> install(operation: SdkHttpOperation<I, O>) {
+        operation.execution.initialize.intercept { req, next ->
+            val input = req.subject as PredictRequest
+            if (input.predictEndpoint == null || input.predictEndpoint.isBlank()) {
+                throw MachineLearningException("Predict requires predictEnpoint to be set to a non-empty value")
+            }
+            // Stash the endpoint for later use by the mutate interceptor
+            req.context.predictEndpoint = AwsEndpoint(input.predictEndpoint)
+
+            next.call(req)
+        }
+
+        operation.execution.mutate.intercept { req, next ->
+            // This should've been set by the initialize interceptor
+            val endpoint = req.context.predictEndpoint
+            requireNotNull(endpoint) { "Predict endpoint wasn't set by middleware." }
+            setRequestEndpoint(req, endpoint.endpoint)
+
+            next.call(req)
+        }
+    }
+}