fix: overhaul endpoint discovery to favor endpointUrl when provided, make endpoint discoverers interfaces so custom implementations can be provided, and to respect environmental config (#1506)

ianbotsf · web-flow · commit bd6fff1e9a09 · 2025-01-17T11:37:41.000-08:00
diff --git a/.changes/49af01b8-6fed-4add-ace0-9f027e83425a.json b/.changes/49af01b8-6fed-4add-ace0-9f027e83425a.json
@@ -0,0 +1,9 @@
+{
+    "id": "49af01b8-6fed-4add-ace0-9f027e83425a",
+    "type": "feature",
+    "description": "⚠️ **IMPORTANT**: Refactor endpoint discoverer classes into interfaces so custom implementations may be provided",
+    "issues": [
+        "awslabs/aws-sdk-kotlin#1413"
+    ],
+    "requiresMinorVersionBump": true
+}
diff --git a/.changes/929f0e2a-3af9-4f73-9f1b-b4e97f91f0db.json b/.changes/929f0e2a-3af9-4f73-9f1b-b4e97f91f0db.json
@@ -0,0 +1,9 @@
+{
+    "id": "929f0e2a-3af9-4f73-9f1b-b4e97f91f0db",
+    "type": "feature",
+    "description": "⚠️ **IMPORTANT**: Add support for enabling/disabling endpoint discovery via [standard cross-SDK config mechanisms](https://docs.aws.amazon.com/sdkref/latest/guide/feature-endpoint-discovery.html)",
+    "issues": [
+        "awslabs/aws-sdk-kotlin#1413"
+    ],
+    "requiresMinorVersionBump": true
+}
diff --git a/.changes/e6515649-dab5-4be9-b4b4-b289369960d5.json b/.changes/e6515649-dab5-4be9-b4b4-b289369960d5.json
@@ -0,0 +1,8 @@
+{
+    "id": "e6515649-dab5-4be9-b4b4-b289369960d5",
+    "type": "bugfix",
+    "description": "Favor `endpointUrl` instead of endpoint discovery if both are provided",
+    "issues": [
+        "awslabs/aws-sdk-kotlin#1413"
+    ]
+}
diff --git a/aws-runtime/aws-config/api/aws-config.api b/aws-runtime/aws-config/api/aws-config.api
@@ -259,6 +259,7 @@ public final class aws/sdk/kotlin/runtime/config/AwsSdkSetting {
 	public final fun getAwsEc2MetadataDisabled ()Laws/smithy/kotlin/runtime/config/EnvironmentSetting;
 	public final fun getAwsEc2MetadataServiceEndpoint ()Laws/smithy/kotlin/runtime/config/EnvironmentSetting;
 	public final fun getAwsEc2MetadataServiceEndpointMode ()Laws/smithy/kotlin/runtime/config/EnvironmentSetting;
+	public final fun getAwsEndpointDiscoveryEnabled ()Laws/smithy/kotlin/runtime/config/EnvironmentSetting;
 	public final fun getAwsEndpointUrl ()Laws/smithy/kotlin/runtime/config/EnvironmentSetting;
 	public final fun getAwsExecutionEnv ()Laws/smithy/kotlin/runtime/config/EnvironmentSetting;
 	public final fun getAwsIgnoreEndpointUrls ()Laws/smithy/kotlin/runtime/config/EnvironmentSetting;
@@ -311,6 +312,8 @@ public final class aws/sdk/kotlin/runtime/config/endpoints/ResolversKt {
 	public static final fun resolveAccountId (Laws/sdk/kotlin/runtime/config/endpoints/AccountIdEndpointMode;Laws/smithy/kotlin/runtime/collections/Attributes;)Ljava/lang/String;
 	public static final fun resolveAccountIdEndpointMode (Laws/smithy/kotlin/runtime/util/PlatformProvider;Laws/smithy/kotlin/runtime/util/LazyAsyncValue;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
 	public static synthetic fun resolveAccountIdEndpointMode$default (Laws/smithy/kotlin/runtime/util/PlatformProvider;Laws/smithy/kotlin/runtime/util/LazyAsyncValue;Lkotlin/coroutines/Continuation;ILjava/lang/Object;)Ljava/lang/Object;
+	public static final fun resolveEndpointDiscoveryEnabled (Laws/smithy/kotlin/runtime/util/PlatformProvider;Laws/smithy/kotlin/runtime/util/LazyAsyncValue;ZLkotlin/coroutines/Continuation;)Ljava/lang/Object;
+	public static synthetic fun resolveEndpointDiscoveryEnabled$default (Laws/smithy/kotlin/runtime/util/PlatformProvider;Laws/smithy/kotlin/runtime/util/LazyAsyncValue;ZLkotlin/coroutines/Continuation;ILjava/lang/Object;)Ljava/lang/Object;
 	public static final fun resolveEndpointUrl (Laws/smithy/kotlin/runtime/util/LazyAsyncValue;Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;Laws/smithy/kotlin/runtime/util/PlatformProvider;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
 	public static synthetic fun resolveEndpointUrl$default (Laws/smithy/kotlin/runtime/util/LazyAsyncValue;Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;Laws/smithy/kotlin/runtime/util/PlatformProvider;Lkotlin/coroutines/Continuation;ILjava/lang/Object;)Ljava/lang/Object;
 	public static final fun resolveUseDualStack (Laws/smithy/kotlin/runtime/util/PlatformProvider;Laws/smithy/kotlin/runtime/util/LazyAsyncValue;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
@@ -490,6 +493,7 @@ public final class aws/sdk/kotlin/runtime/config/profile/AwsProfileKt {
 	public static synthetic fun getBooleanOrNull$default (Laws/sdk/kotlin/runtime/config/profile/ConfigSection;Ljava/lang/String;Ljava/lang/String;ILjava/lang/Object;)Ljava/lang/Boolean;
 	public static final fun getCredentialProcess (Laws/sdk/kotlin/runtime/config/profile/ConfigSection;)Ljava/lang/String;
 	public static final fun getDisableRequestCompression (Laws/sdk/kotlin/runtime/config/profile/ConfigSection;)Ljava/lang/Boolean;
+	public static final fun getEndpointDiscoveryEnabled (Laws/sdk/kotlin/runtime/config/profile/ConfigSection;)Ljava/lang/Boolean;
 	public static final fun getEndpointUrl (Laws/sdk/kotlin/runtime/config/profile/ConfigSection;)Laws/smithy/kotlin/runtime/net/url/Url;
 	public static final fun getIgnoreEndpointUrls (Laws/sdk/kotlin/runtime/config/profile/ConfigSection;)Ljava/lang/Boolean;
 	public static final fun getIntOrNull (Laws/sdk/kotlin/runtime/config/profile/ConfigSection;Ljava/lang/String;Ljava/lang/String;)Ljava/lang/Integer;
diff --git a/aws-runtime/aws-config/common/src/aws/sdk/kotlin/runtime/config/AwsSdkSetting.kt b/aws-runtime/aws-config/common/src/aws/sdk/kotlin/runtime/config/AwsSdkSetting.kt
@@ -211,6 +211,15 @@ public object AwsSdkSetting {
     public val AwsSigV4aSigningRegionSet: EnvironmentSetting<String> =
         strEnvSetting("aws.sigV4aSigningRegionSet", "AWS_SIGV4A_SIGNING_REGION_SET")
 
+    /**
+     * A flag indicating whether endpoint discovery is enabled for AWS services that support it. The implicit default
+     * value for this setting is:
+     * * `true` for services which _require_ EP discovery (e.g., Timestream)
+     * * `false` for services which _allow but do not require_ EP discovery (e.g., DynamoDB)
+     */
+    public val AwsEndpointDiscoveryEnabled: EnvironmentSetting<Boolean> =
+        boolEnvSetting("aws.endpointDiscoveryEnabled", "AWS_ENABLE_ENDPOINT_DISCOVERY")
+
     /**
      * Configures request checksum calculation
      */
diff --git a/aws-runtime/aws-config/common/src/aws/sdk/kotlin/runtime/config/endpoints/Resolvers.kt b/aws-runtime/aws-config/common/src/aws/sdk/kotlin/runtime/config/endpoints/Resolvers.kt
@@ -101,3 +101,16 @@ public fun resolveAccountId(endpointMode: AccountIdEndpointMode, attributes: Att
     AccountIdEndpointMode.DISABLED -> null
     AccountIdEndpointMode.REQUIRED -> attributes.getOrNull(AwsClientOption.AccountId) ?: throw ConfigurationException("AccountIdEndpointMode is set to required but no AWS account ID found")
 }
+
+/**
+ * Resolve the endpoint discovery mode
+ */
+@InternalSdkApi
+public suspend fun resolveEndpointDiscoveryEnabled(
+    provider: PlatformProvider = PlatformProvider.System,
+    profile: LazyAsyncValue<AwsProfile> = asyncLazy { loadAwsSharedConfig(provider).activeProfile },
+    serviceRequiresEpDiscovery: Boolean,
+): Boolean =
+    AwsSdkSetting.AwsEndpointDiscoveryEnabled.resolve(provider)
+        ?: profile.get().endpointDiscoveryEnabled
+        ?: serviceRequiresEpDiscovery
diff --git a/aws-runtime/aws-config/common/src/aws/sdk/kotlin/runtime/config/profile/AwsProfile.kt b/aws-runtime/aws-config/common/src/aws/sdk/kotlin/runtime/config/profile/AwsProfile.kt
@@ -169,6 +169,14 @@ public val AwsProfile.requestMinCompressionSizeBytes: Long?
 public val AwsProfile.sigV4aSigningRegionSet: String?
     get() = getOrNull("sigv4a_signing_region_set")
 
+/**
+ * A flag indicating whether endpoint discovery should be enabled for a service that supports it. This setting has no
+ * effect for services which _do not_ support endpoint discovery.
+ */
+@InternalSdkApi
+public val AwsProfile.endpointDiscoveryEnabled: Boolean?
+    get() = getBooleanOrNull("endpoint_discovery_enabled")
+
 /**
  * Configures request checksum calculation
  */
diff --git a/aws-runtime/aws-config/common/test/aws/sdk/kotlin/runtime/config/endpoints/ResolveEndpointDiscoveryTest.kt b/aws-runtime/aws-config/common/test/aws/sdk/kotlin/runtime/config/endpoints/ResolveEndpointDiscoveryTest.kt
@@ -0,0 +1,77 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+package aws.sdk.kotlin.runtime.config.endpoints
+
+import aws.sdk.kotlin.runtime.config.profile.*
+import aws.sdk.kotlin.runtime.config.profile.FileType
+import aws.sdk.kotlin.runtime.config.profile.parse
+import aws.sdk.kotlin.runtime.config.profile.toSharedConfig
+import aws.smithy.kotlin.runtime.telemetry.logging.Logger
+import aws.smithy.kotlin.runtime.util.TestPlatformProvider
+import aws.smithy.kotlin.runtime.util.asyncLazy
+import kotlinx.coroutines.test.runTest
+import kotlin.test.Test
+import kotlin.test.assertEquals
+
+class ResolveEndpointDiscoveryTest {
+    @Test
+    fun testPrecedenceSysProps() = assertEpDiscovery(
+        sysProps = mapOf("aws.endpointDiscoveryEnabled" to "true"),
+        env = mapOf("AWS_ENABLE_ENDPOINT_DISCOVERY" to "false"),
+        config = """
+            [${Literals.DEFAULT_PROFILE}]
+            endpoint_discovery_enabled = false
+        """.trimIndent(),
+        serviceRequiresEpDiscovery = false,
+        expected = true,
+    )
+
+    @Test
+    fun testPrecedenceEnvVars() = assertEpDiscovery(
+        env = mapOf("AWS_ENABLE_ENDPOINT_DISCOVERY" to "true"),
+        config = """
+            [${Literals.DEFAULT_PROFILE}]
+            endpoint_discovery_enabled = false
+        """.trimIndent(),
+        serviceRequiresEpDiscovery = false,
+        expected = true,
+    )
+
+    @Test
+    fun testPrecedenceConfig() = assertEpDiscovery(
+        config = """
+            [${Literals.DEFAULT_PROFILE}]
+            endpoint_discovery_enabled = true
+        """.trimIndent(),
+        serviceRequiresEpDiscovery = false,
+        expected = true,
+    )
+
+    @Test
+    fun testPrecedenceDefault() = assertEpDiscovery(
+        serviceRequiresEpDiscovery = true,
+        expected = true,
+    )
+}
+
+fun assertEpDiscovery(
+    sysProps: Map<String, String> = mapOf(),
+    env: Map<String, String> = mapOf(),
+    config: String = "",
+    serviceRequiresEpDiscovery: Boolean,
+    expected: Boolean,
+) = runTest {
+    val provider = TestPlatformProvider(env, sysProps)
+    val source = AwsConfigurationSource(Literals.DEFAULT_PROFILE, "", "")
+
+    val profile = asyncLazy {
+        parse(Logger.None, FileType.CONFIGURATION, config)
+            .toSharedConfig(source)
+            .activeProfile
+    }
+
+    val actual = resolveEndpointDiscoveryEnabled(provider, profile, serviceRequiresEpDiscovery)
+    assertEquals(expected, actual)
+}
diff --git a/codegen/aws-sdk-codegen/src/main/kotlin/aws/sdk/kotlin/codegen/AwsRuntimeTypes.kt b/codegen/aws-sdk-codegen/src/main/kotlin/aws/sdk/kotlin/codegen/AwsRuntimeTypes.kt
@@ -38,6 +38,7 @@ object AwsRuntimeTypes {
 
         object Endpoints : RuntimeTypePackage(AwsKotlinDependency.AWS_CONFIG, "config.endpoints") {
             val AccountIdEndpointMode = symbol("AccountIdEndpointMode")
+            val resolveEndpointDiscoveryEnabled = symbol("resolveEndpointDiscoveryEnabled")
             val resolveEndpointUrl = symbol("resolveEndpointUrl")
             val resolveAccountId = symbol("resolveAccountId")
             val resolveAccountIdEndpointMode = symbol("resolveAccountIdEndpointMode")
diff --git a/codegen/aws-sdk-codegen/src/main/kotlin/aws/sdk/kotlin/codegen/PresignerGenerator.kt b/codegen/aws-sdk-codegen/src/main/kotlin/aws/sdk/kotlin/codegen/PresignerGenerator.kt
@@ -44,7 +44,6 @@ class PresignerGenerator : KotlinIntegration {
      * signing.
      */
     object UnsignedRequestCustomizationSection : SectionId {
-        val CodegenContext: SectionKey<CodegenContext> = SectionKey("CodegenContext")
         val OperationId: SectionKey<String> = SectionKey("OperationId")
         val HttpBindingResolver: SectionKey<HttpBindingResolver> = SectionKey("HttpBindingResolver")
         val DefaultTimestampFormat: SectionKey<TimestampFormatTrait.Format> = SectionKey("DefaultTimestampFormat")
@@ -126,7 +125,6 @@ class PresignerGenerator : KotlinIntegration {
 
             val contextMap: Map<SectionKey<*>, Any> = mapOf(
                 UnsignedRequestCustomizationSection.OperationId to op.id.toString(),
-                UnsignedRequestCustomizationSection.CodegenContext to ctx,
                 UnsignedRequestCustomizationSection.HttpBindingResolver to httpBindingResolver,
                 UnsignedRequestCustomizationSection.DefaultTimestampFormat to defaultTimestampFormat,
             )
diff --git a/codegen/aws-sdk-codegen/src/main/kotlin/aws/sdk/kotlin/codegen/customization/ReplaceServiceExceptionBase.kt b/codegen/aws-sdk-codegen/src/main/kotlin/aws/sdk/kotlin/codegen/customization/ReplaceServiceExceptionBase.kt
@@ -27,7 +27,7 @@ class ReplaceServiceExceptionBase : KotlinIntegration {
         get() = listOf(SectionWriterBinding(ExceptionBaseClassGenerator.ExceptionBaseClassSection, exceptionSectionWriter))
 
     private val exceptionSectionWriter = SectionWriter { writer, _ ->
-        val ctx = writer.getContextValue(ExceptionBaseClassGenerator.ExceptionBaseClassSection.CodegenContext)
+        val ctx = writer.getContextValue(CodegenContext.Key)
         AwsServiceExceptionBaseClassGenerator().render(ctx, writer)
     }
 }
diff --git a/codegen/aws-sdk-codegen/src/main/kotlin/aws/sdk/kotlin/codegen/customization/polly/PollyPresigner.kt b/codegen/aws-sdk-codegen/src/main/kotlin/aws/sdk/kotlin/codegen/customization/polly/PollyPresigner.kt
@@ -36,7 +36,7 @@ class PollyPresigner : KotlinIntegration {
     }
 
     private val customizeUnsignedRequest = SectionWriter { writer, _ ->
-        val ctx = writer.getContextValue(PresignerGenerator.UnsignedRequestCustomizationSection.CodegenContext)
+        val ctx = writer.getContextValue(CodegenContext.Key)
         val operation = ctx.model.expectShape<OperationShape>(writer.getContextValue(PresignerGenerator.UnsignedRequestCustomizationSection.OperationId))
         val resolver = writer.getContextValue(PresignerGenerator.UnsignedRequestCustomizationSection.HttpBindingResolver)
         val defaultTimestampFormat = writer.getContextValue(PresignerGenerator.UnsignedRequestCustomizationSection.DefaultTimestampFormat)
diff --git a/codegen/aws-sdk-codegen/src/main/kotlin/aws/sdk/kotlin/codegen/endpoints/AwsEndpointDiscoveryIntegration.kt b/codegen/aws-sdk-codegen/src/main/kotlin/aws/sdk/kotlin/codegen/endpoints/AwsEndpointDiscoveryIntegration.kt
@@ -0,0 +1,95 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+package aws.sdk.kotlin.codegen.endpoints
+
+import aws.sdk.kotlin.codegen.AwsRuntimeTypes
+import aws.sdk.kotlin.codegen.ServiceClientCompanionObjectWriter
+import software.amazon.smithy.kotlin.codegen.KotlinSettings
+import software.amazon.smithy.kotlin.codegen.core.CodegenContext
+import software.amazon.smithy.kotlin.codegen.core.KotlinDelegator
+import software.amazon.smithy.kotlin.codegen.core.getContextValue
+import software.amazon.smithy.kotlin.codegen.integration.AppendingSectionWriter
+import software.amazon.smithy.kotlin.codegen.integration.KotlinIntegration
+import software.amazon.smithy.kotlin.codegen.integration.SectionWriterBinding
+import software.amazon.smithy.kotlin.codegen.model.asNullable
+import software.amazon.smithy.kotlin.codegen.rendering.endpoints.discovery.DefaultEndpointDiscovererGenerator
+import software.amazon.smithy.kotlin.codegen.rendering.endpoints.discovery.EndpointDiscovererInterfaceGenerator
+import software.amazon.smithy.kotlin.codegen.rendering.endpoints.discovery.EndpointDiscoveryIntegration
+import software.amazon.smithy.kotlin.codegen.rendering.util.ConfigProperty
+import software.amazon.smithy.kotlin.codegen.rendering.util.ConfigPropertyType
+import software.amazon.smithy.model.Model
+
+class AwsEndpointDiscoveryIntegration : KotlinIntegration {
+    override val order: Byte = (EndpointDiscoveryIntegration.ORDER + 1).toByte() // after EndpointDiscoveryIntegration
+
+    override fun additionalServiceConfigProps(ctx: CodegenContext): List<ConfigProperty> {
+        val endpointDiscoveryOptional = EndpointDiscoveryIntegration.isOptionalFor(ctx)
+        val interfaceSymbol = EndpointDiscovererInterfaceGenerator.symbolFor(ctx.settings)
+        return listOf(
+            ConfigProperty {
+                name = EndpointDiscoveryIntegration.CLIENT_CONFIG_NAME
+                symbol = interfaceSymbol.asNullable()
+
+                if (endpointDiscoveryOptional) {
+                    documentation = """
+                        The endpoint discoverer for this client, if applicable. By default, no endpoint discovery is
+                        provided. To use endpoint discovery, set this to a valid [${interfaceSymbol.name}] instance.
+                    """.trimIndent()
+                    propertyType = ConfigPropertyType.SymbolDefault
+                } else {
+                    val defaultImplSymbol = DefaultEndpointDiscovererGenerator.symbolFor(ctx.settings)
+
+                    documentation = """
+                        The endpoint discoverer for this client, [${defaultImplSymbol.name}] by default.
+                    """.trimIndent()
+                    propertyType = ConfigPropertyType.Custom(
+                        render = { prop, writer ->
+                            writer.write(
+                                "#1L val #2L: #3T = builder.#2L ?: #4T()",
+                                ctx.settings.api.visibility,
+                                prop.propertyName,
+                                prop.symbol,
+                                defaultImplSymbol,
+                            )
+                        },
+                    )
+                }
+            },
+        )
+    }
+
+    override fun enabledForService(model: Model, settings: KotlinSettings): Boolean =
+        EndpointDiscoveryIntegration.isEnabledFor(model, settings)
+
+    private val resolveEndpointDiscoverer = AppendingSectionWriter { writer ->
+        val ctx = writer.getContextValue(CodegenContext.Key)
+        val endpointDiscoveryOptional = EndpointDiscoveryIntegration.isOptionalFor(ctx)
+
+        writer.write(
+            "val epDiscoveryEnabled = #T(profile = activeProfile, serviceRequiresEpDiscovery = #L)",
+            AwsRuntimeTypes.Config.Endpoints.resolveEndpointDiscoveryEnabled,
+            !endpointDiscoveryOptional,
+        )
+
+        writer.write(
+            "builder.config.#1L = builder.config.#1L ?: if (epDiscoveryEnabled) #2T() else null",
+            EndpointDiscoveryIntegration.CLIENT_CONFIG_NAME,
+            DefaultEndpointDiscovererGenerator.symbolFor(ctx.settings),
+        )
+    }
+
+    override val sectionWriters = listOf(
+        SectionWriterBinding(ServiceClientCompanionObjectWriter.FinalizeEnvironmentalConfig, resolveEndpointDiscoverer),
+    )
+
+    override fun writeAdditionalFiles(ctx: CodegenContext, delegator: KotlinDelegator) {
+        // EndpointDiscoveryIntegration already renders the default endpoint discoverer for services that _require_ EP
+        // discovery. So we only need to render it for services which _do not require_ EP discovery in order to support
+        // enabling discovery via environmental config.
+        if (EndpointDiscoveryIntegration.isOptionalFor(ctx)) {
+            DefaultEndpointDiscovererGenerator(ctx, delegator).render()
+        }
+    }
+}
diff --git a/codegen/aws-sdk-codegen/src/main/resources/META-INF/services/software.amazon.smithy.kotlin.codegen.integration.KotlinIntegration b/codegen/aws-sdk-codegen/src/main/resources/META-INF/services/software.amazon.smithy.kotlin.codegen.integration.KotlinIntegration
@@ -7,6 +7,7 @@ aws.sdk.kotlin.codegen.customization.DefaultMiddleware
 aws.sdk.kotlin.codegen.customization.AccountIdEndpointBuiltinCustomization
 aws.sdk.kotlin.codegen.customization.PresignableModelIntegration
 aws.sdk.kotlin.codegen.customization.BackfillOptionalAuth
+aws.sdk.kotlin.codegen.endpoints.AwsEndpointDiscoveryIntegration
 aws.sdk.kotlin.codegen.endpoints.BindAwsEndpointBuiltins
 aws.sdk.kotlin.codegen.endpoints.AddAwsEndpointFunctions
 aws.sdk.kotlin.codegen.PresignerGenerator
