Merge feat/async-config to develop

codegen/aws/core/src/it/java/software/amazon/smithy/python/codegen/test/AwsCodegenTest.java

@@ -4,6 +4,11 @@
  */
 package software.amazon.smithy.python.codegen.test;
 
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+import java.io.IOException;
+import java.nio.file.Files;
 import java.nio.file.Path;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.io.TempDir;
@@ -19,7 +24,7 @@
 public class AwsCodegenTest {
 
     @Test
-    public void testCodegen(@TempDir Path tempDir) {
+    public void testCodegen(@TempDir Path tempDir) throws IOException {
         PythonClientCodegenPlugin plugin = new PythonClientCodegenPlugin();
         Model model = Model.assembler(AwsCodegenTest.class.getClassLoader())
                 .discoverModels(AwsCodegenTest.class.getClassLoader())
@@ -36,6 +41,44 @@ public void testCodegen(@TempDir Path tempDir) {
                 .model(model)
                 .build();
         plugin.execute(context);
+
+        var config = Files.readString(tempDir.resolve("src/restjson/config.py"));
+        assertTrue(config.contains("Overrides(AwsConfigOverrides, total=False):"));
+        assertTrue(config.contains("api_key: str | None"));
+        assertTrue(config.contains("@dataclass(kw_only=True, repr=False, init=False)"));
+        assertTrue(config.contains("**overrides: Unpack["));
+        assertTrue(config.contains(
+                "interceptors: list[_ServiceInterceptor] = field(default_factory=lambda: [])"));
+        assertTrue(config.contains(
+                "def set_auth_scheme(self, scheme: AuthScheme[Any, Any, Any, Any]) -> None:"));
+        assertTrue(config.contains("auth_schemes = dict(self.auth_schemes or {})"));
+        assertTrue(config.contains("auth_schemes[scheme.scheme_id] = scheme"));
+        assertTrue(config.contains("self.auth_schemes = auth_schemes"));
+
+        var client = Files.readString(tempDir.resolve("src/restjson/client.py"));
+        assertInOrder(
+                client,
+                "config = await AsyncRESTJSONConfig.resolve()",
+                "for plugin in self._client_plugins:",
+                "for plugin in self._plugins:",
+                "self._config = config",
+                "if operation_plugins:",
+                "config = deepcopy(self._config)",
+                "for plugin in operation_plugins:",
+                "config = self._config");
+        assertFalse(client.contains("plugin(self._config)"));
+        assertTrue(client.contains("retry_mode=config.retry_mode"));
+        assertTrue(client.contains("max_attempts=config.max_attempts"));
+        assertFalse(client.contains("getattr(config, \"retry_mode\""));
+        assertFalse(client.contains("getattr(config, \"max_attempts\""));
     }
 
+    private static void assertInOrder(String value, String... fragments) {
+        var index = 0;
+        for (String fragment : fragments) {
+            index = value.indexOf(fragment, index);
+            assertTrue(index >= 0, "Missing or out-of-order fragment: " + fragment);
+            index += fragment.length();
+        }
+    }
 }

codegen/aws/core/src/main/java/software/amazon/smithy/python/aws/codegen/AwsAuthIntegration.java

@@ -79,7 +79,7 @@ public List<RuntimeClientPlugin> getClientPlugins(GenerationContext context) {
                         .addConfigProperty(ConfigProperty.builder()
                                 .name("aws_session_token")
                                 .type(Symbol.builder().name("str").build())
-                                .documentation("An access key ID that identifies temporary security credentials.")
+                                .documentation("The session token used with temporary AWS credentials.")
                                 .nullable(true)
                                 .build())
                         .authScheme(new Sigv4AuthScheme())

codegen/aws/core/src/main/java/software/amazon/smithy/python/aws/codegen/AwsUserAgentIntegration.java

@@ -33,6 +33,19 @@ def aws_user_agent_plugin(config: $1T):
                 )
             """;
 
+    // Variant for services without a generated async config, which uses old Config.
+    private static final String USER_AGENT_PLUGIN_SYNC_ONLY = """
+            def aws_user_agent_plugin(config: $1T):
+                config.interceptors.append(
+                    $2T(
+                        ua_suffix=config.user_agent_extra,
+                        ua_app_id=config.sdk_ua_app_id,
+                        sdk_version=$3T,
+                        service_id=$4S,
+                    )
+                )
+            """;
+
     @Override
     public List<RuntimeClientPlugin> getClientPlugins(GenerationContext context) {
         if (context.applicationProtocol().isHttpProtocol()) {
@@ -96,12 +109,22 @@ public List<RuntimeClientPlugin> getClientPlugins(GenerationContext context) {
                                                 filename,
                                                 moduleName + ".",
                                                 writer -> {
-                                                    writer.write(USER_AGENT_PLUGIN,
-                                                            CodegenUtils.getConfigSymbol(c.settings()),
-                                                            userAgentInterceptor,
-                                                            versionSymbol,
-                                                            serviceId);
-
+                                                    var asyncConfig = CodegenUtils.getAsyncConfigSymbol(
+                                                            c.settings(),
+                                                            c.model());
+                                                    if (asyncConfig.isPresent()) {
+                                                        writer.write(USER_AGENT_PLUGIN,
+                                                                asyncConfig.get(),
+                                                                userAgentInterceptor,
+                                                                versionSymbol,
+                                                                serviceId);
+                                                    } else {
+                                                        writer.write(USER_AGENT_PLUGIN_SYNC_ONLY,
+                                                                CodegenUtils.getConfigSymbol(c.settings()),
+                                                                userAgentInterceptor,
+                                                                versionSymbol,
+                                                                serviceId);
+                                                    }
                                                 });
                                 return List.of(filename);
                             })

codegen/aws/core/src/main/java/software/amazon/smithy/python/aws/codegen/customizations/apigateway/ApiGatewayIntegration.java

@@ -100,7 +100,8 @@ public List<RuntimeClientPlugin> getClientPlugins(GenerationContext context) {
                                                         httpRequest,
                                                         requestContext,
                                                         field,
-                                                        CodegenUtils.getConfigSymbol(c.settings()));
+                                                        CodegenUtils.getAsyncConfigSymbol(c.settings(), c.model())
+                                                                .orElseThrow());
                                             });
                             return List.of(filename);
                         })

codegen/aws/core/src/main/java/software/amazon/smithy/python/aws/codegen/AwsDynamoDbRetryIntegration.java → codegen/aws/core/src/main/java/software/amazon/smithy/python/aws/codegen/customizations/dynamodb/AwsDynamoDbRetryIntegration.java

@@ -2,13 +2,15 @@
  * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
-package software.amazon.smithy.python.aws.codegen;
+package software.amazon.smithy.python.aws.codegen.customizations.dynamodb;
 
 import java.util.List;
 import java.util.Set;
 import software.amazon.smithy.aws.traits.ServiceTrait;
 import software.amazon.smithy.codegen.core.Symbol;
 import software.amazon.smithy.codegen.core.SymbolReference;
+import software.amazon.smithy.python.aws.codegen.AwsPythonDependency;
+import software.amazon.smithy.python.codegen.CodegenUtils;
 import software.amazon.smithy.python.codegen.GenerationContext;
 import software.amazon.smithy.python.codegen.SmithyPythonDependency;
 import software.amazon.smithy.python.codegen.integrations.PythonIntegration;
@@ -28,11 +30,7 @@ public final class AwsDynamoDbRetryIntegration implements PythonIntegration {
             _DYNAMODB_DEFAULT_MAX_BACKOFF = 20
 
 
-            class _RetryConfig(Protocol):
-                retry_strategy: $1T | $2T | None
-
-
-            def dynamodb_retry_plugin(config: _RetryConfig) -> None:
+            def dynamodb_retry_plugin(config: $1T) -> None:
                 \"\"\"Apply DynamoDB's standard-mode retry defaults for any option left unset.\"\"\"
                 retry_strategy = config.retry_strategy
                 if retry_strategy is not None and not isinstance(
@@ -45,15 +43,10 @@ def dynamodb_retry_plugin(config: _RetryConfig) -> None:
                     retry_mode = retry_strategy.retry_mode
                     max_attempts = retry_strategy.max_attempts
                 else:
-                    # Read independently resolved AsyncConfig fields when available. A legacy
-                    # Config has no scalar retry fields, so None represents an unset value.
-                    retry_mode = getattr(config, "retry_mode", None) or "standard"
-                    max_attempts = getattr(config, "max_attempts", None)
-                    source_of = getattr(config, "source_of", None)
-                    if (
-                        source_of is not None
-                        and source_of("max_attempts") == $4T.DEFAULT
-                    ):
+                    # Fall back to the config's independently resolved retry fields.
+                    retry_mode = config.retry_mode or "standard"
+                    max_attempts = config.max_attempts
+                    if config.source_of("max_attempts") == $4T.DEFAULT:
                         max_attempts = None
 
                 if retry_mode != "standard":
@@ -85,10 +78,6 @@ public List<RuntimeClientPlugin> getClientPlugins(GenerationContext context) {
                         .name("dynamodb_retry_plugin")
                         .build())
                 .build();
-        final Symbol retryStrategy = Symbol.builder()
-                .namespace("smithy_core.aio.interfaces.retries", ".")
-                .name("RetryStrategy")
-                .build();
         final Symbol retryStrategyOptions = Symbol.builder()
                 .namespace("smithy_core.retries", ".")
                 .name("RetryStrategyOptions")
@@ -125,10 +114,10 @@ public List<RuntimeClientPlugin> getClientPlugins(GenerationContext context) {
                                             writer -> {
                                                 writer.addDependency(SmithyPythonDependency.SMITHY_CORE);
                                                 writer.addDependency(AwsPythonDependency.SMITHY_AWS_CORE);
-                                                writer.addStdlibImport("typing", "Protocol");
                                                 writer.write(
                                                         DYNAMODB_RETRY_MODULE,
-                                                        retryStrategy,
+                                                        CodegenUtils.getAsyncConfigSymbol(c.settings(), c.model())
+                                                                .orElseThrow(),
                                                         retryStrategyOptions,
                                                         standardRetryStrategy,
                                                         configSource,

codegen/aws/core/src/main/resources/META-INF/services/software.amazon.smithy.python.codegen.integrations.PythonIntegration

@@ -9,4 +9,5 @@ software.amazon.smithy.python.aws.codegen.AwsProtocolsIntegration
 software.amazon.smithy.python.aws.codegen.AwsServiceIdIntegration
 software.amazon.smithy.python.aws.codegen.AwsUserAgentIntegration
 software.amazon.smithy.python.aws.codegen.AwsStandardRegionalEndpointsIntegration
-software.amazon.smithy.python.aws.codegen.AwsDynamoDbRetryIntegration
+software.amazon.smithy.python.aws.codegen.customizations.dynamodb.AwsDynamoDbRetryIntegration
+software.amazon.smithy.python.aws.codegen.AwsAsyncConfigIntegration

codegen/core/src/it/java/software/amazon/smithy/python/codegen/test/PythonCodegenTest.java

@@ -4,6 +4,10 @@
  */
 package software.amazon.smithy.python.codegen.test;
 
+import static org.junit.jupiter.api.Assertions.assertFalse;
+
+import java.io.IOException;
+import java.nio.file.Files;
 import java.nio.file.Path;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.io.TempDir;
@@ -20,7 +24,7 @@
 public class PythonCodegenTest {
 
     @Test
-    public void testCodegen(@TempDir Path tempDir) {
+    public void testCodegen(@TempDir Path tempDir) throws IOException {
         // TODO: Move this to its own package once client codegen is in its own package
         PythonClientCodegenPlugin plugin = new PythonClientCodegenPlugin();
         Model model = Model.assembler(PythonCodegenTest.class.getClassLoader())
@@ -38,5 +42,9 @@ public void testCodegen(@TempDir Path tempDir) {
                 .model(model)
                 .build();
         plugin.execute(context);
+
+        var client = Files.readString(tempDir.resolve("src/weather/client.py"));
+        assertFalse(client.contains("retry_mode="));
+        assertFalse(client.contains("max_attempts="));
     }
 }