Integrate identity chain into SDK client construction (#771)

Author: arandito

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

@@ -0,0 +1,89 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+package software.amazon.smithy.python.aws.codegen;
+
+import java.util.List;
+import software.amazon.smithy.aws.traits.auth.SigV4Trait;
+import software.amazon.smithy.codegen.core.Symbol;
+import software.amazon.smithy.python.codegen.GenerationContext;
+import software.amazon.smithy.python.codegen.SmithyPythonDependency;
+import software.amazon.smithy.python.codegen.integrations.PythonIntegration;
+import software.amazon.smithy.python.codegen.sections.ClientSetupSection;
+import software.amazon.smithy.python.codegen.writer.PythonWriter;
+import software.amazon.smithy.utils.CodeInterceptor;
+import software.amazon.smithy.utils.CodeSection;
+import software.amazon.smithy.utils.SmithyInternalApi;
+
+/**
+ * Sets up the default AWS credentials identity chain during client setup.
+ */
+@SmithyInternalApi
+public class AwsIdentityIntegration implements PythonIntegration {
+
+    @Override
+    public List<? extends CodeInterceptor<? extends CodeSection, PythonWriter>> interceptors(
+            GenerationContext context
+    ) {
+        var service = context.settings().service(context.model());
+        if (!service.hasTrait(SigV4Trait.class)) {
+            return List.of();
+        }
+        return List.of(new CredentialsIdentitySetupInterceptor());
+    }
+
+    /**
+     * Initializes the default AWS credentials identity chain during client setup.
+     */
+    private static final class CredentialsIdentitySetupInterceptor
+            implements CodeInterceptor<ClientSetupSection, PythonWriter> {
+
+        @Override
+        public Class<ClientSetupSection> sectionType() {
+            return ClientSetupSection.class;
+        }
+
+        @Override
+        public void write(PythonWriter writer, String previousText, ClientSetupSection section) {
+            writer.write(previousText);
+            writer.addStdlibImport("typing", "cast");
+            writer.write("""
+                    if self._config.aws_credentials_identity_resolver is None:
+                        config_context = self._config.resolution_context()
+                        config_file = None
+                        profile_name = None
+                        if config_context is not None:
+                            config_file = await config_context.parsed_profiles()
+                            if config_context.profile_source is $4T.OVERRIDE:
+                                profile_name = config_context.profile_name
+                        self._config.aws_credentials_identity_resolver = await $1T.create(
+                            $2T,
+                            config_file=config_file,
+                            profile_name=profile_name,
+                            region_override=self._config.region,
+                            http_client=cast($3T | None, self._config.transport),
+                        )""",
+                    Symbol.builder()
+                            .name("IdentityChain")
+                            .namespace("smithy_aws_core.identity.chain", ".")
+                            .addDependency(AwsPythonDependency.SMITHY_AWS_CORE)
+                            .build(),
+                    Symbol.builder()
+                            .name("AWSCredentialsIdentity")
+                            .namespace("smithy_aws_core.identity", ".")
+                            .addDependency(AwsPythonDependency.SMITHY_AWS_CORE)
+                            .build(),
+                    Symbol.builder()
+                            .name("HTTPClient")
+                            .namespace("smithy_http.aio.interfaces", ".")
+                            .addDependency(SmithyPythonDependency.SMITHY_HTTP)
+                            .build(),
+                    Symbol.builder()
+                            .name("ConfigSource")
+                            .namespace("smithy_aws_core.config", ".")
+                            .addDependency(AwsPythonDependency.SMITHY_AWS_CORE)
+                            .build());
+        }
+    }
+}

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

@@ -5,6 +5,7 @@
 
 software.amazon.smithy.python.aws.codegen.customizations.apigateway.ApiGatewayIntegration
 software.amazon.smithy.python.aws.codegen.AwsAuthIntegration
+software.amazon.smithy.python.aws.codegen.AwsIdentityIntegration
 software.amazon.smithy.python.aws.codegen.AwsProtocolsIntegration
 software.amazon.smithy.python.aws.codegen.AwsServiceIdIntegration
 software.amazon.smithy.python.aws.codegen.AwsUserAgentIntegration

codegen/core/src/main/java/software/amazon/smithy/python/codegen/ClientGenerator.java

@@ -22,6 +22,7 @@
 import software.amazon.smithy.model.traits.StringTrait;
 import software.amazon.smithy.python.codegen.integrations.PythonIntegration;
 import software.amazon.smithy.python.codegen.integrations.RuntimeClientPlugin;
+import software.amazon.smithy.python.codegen.sections.ClientSetupSection;
 import software.amazon.smithy.python.codegen.writer.PythonWriter;
 import software.amazon.smithy.utils.SmithyInternalApi;
 
@@ -113,6 +114,7 @@ async def _ensure_setup(self) -> None:
                                         for plugin in self._plugins:
                                             plugin(config)
                                     self._config = config
+                                    ${7C|}
                                     self._setup_done = True
                     """,
                     configSym,
@@ -126,6 +128,10 @@ async def _ensure_setup(self) -> None:
                         } else {
                             w.write("config = $T()", configSym);
                         }
+                    }),
+                    writer.consumer(w -> {
+                        w.pushState(new ClientSetupSection());
+                        w.popState();
                     }));
 
             var topDownIndex = TopDownIndex.of(model);

codegen/core/src/main/java/software/amazon/smithy/python/codegen/sections/ClientSetupSection.java

@@ -0,0 +1,12 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+package software.amazon.smithy.python.codegen.sections;
+
+import software.amazon.smithy.utils.CodeSection;
+
+/**
+ * Section for service-specific lazy client setup.
+ */
+public record ClientSetupSection() implements CodeSection {}

packages/smithy-aws-core/.changes/next-release/smithy-aws-core-enhancement-aeae35f3537641a9b52ba87ce11cafb2.json

@@ -0,0 +1,4 @@
+{
+  "type": "enhancement",
+  "description": "Updated `SharedConfigContext` to track the source of the active profile as a `ConfigSource`."
+}

packages/smithy-aws-core/src/smithy_aws_core/config/aws_config.py

@@ -2,6 +2,7 @@
 # SPDX-License-Identifier: Apache-2.0
 
 from collections.abc import Mapping
+from copy import deepcopy
 from dataclasses import dataclass, field, fields
 from typing import TYPE_CHECKING, Any, ClassVar, Self, TypedDict, Unpack
 
@@ -272,11 +273,11 @@ async def _resolve(
             credentials_file_path=credentials_file_path,
         )
 
-        # Fail fast on a bad profile
-        profile_origin = ctx.profile_origin
-        if profile_origin is not None:
+        # Fail fast on a bad profile when one was provided (not the default)
+        profile_source = ctx.profile_source
+        if profile_source is not ConfigSource.DEFAULT:
             config_file = await ctx.parsed_profiles()
-            validate_profile(ctx.profile_name, config_file.profiles, profile_origin)
+            validate_profile(ctx.profile_name, config_file.profiles, profile_source)
 
         # Create the instance without calling the blocked constructor
         instance = cls._create_instance()
@@ -439,3 +440,18 @@ def __setattr__(self, name: str, value: Any) -> None:
                 spec.validator(value)
             self._sources[name] = ConfigSource.OVERRIDE
         super().__setattr__(name, value)
+
+    def __deepcopy__(self, memo: dict[int, Any]) -> Self:
+        """Deep-copy the config while sharing resources that must not be duplicated."""
+        for shared in (
+            self.aws_credentials_identity_resolver,
+            self.transport,
+            self.retry_strategy,
+        ):
+            if shared is not None:
+                memo[id(shared)] = shared
+        new = self._create_instance()
+        memo[id(self)] = new
+        for f in fields(self):
+            object.__setattr__(new, f.name, deepcopy(getattr(self, f.name), memo))
+        return new

packages/smithy-aws-core/src/smithy_aws_core/config/context.py

@@ -13,6 +13,7 @@
 )
 from .filesystem import DefaultFileSystem, FileSystem
 from .merged_config import MergedConfig
+from .types import ConfigSource
 
 logger = logging.getLogger(__name__)
 
@@ -136,7 +137,7 @@ def __init__(
         """
         self._fs: FileSystem = fs if fs is not None else DefaultFileSystem()
         self._http_client: Any | None = http_client
-        self._profile_name, self._profile_origin = self._resolve_profile_name(
+        self._profile_name, self._profile_source = self._resolve_profile_name(
             profile_name
         )
         self._config_file_path: Path | None = (
@@ -153,9 +154,9 @@ def profile_name(self) -> str:
         return self._profile_name
 
     @property
-    def profile_origin(self) -> str | None:
-        """Where the active profile name came from, or None if it defaulted."""
-        return self._profile_origin
+    def profile_source(self) -> ConfigSource:
+        """The source of the active profile name."""
+        return self._profile_source
 
     @property
     def fs(self) -> FileSystem:
@@ -197,19 +198,20 @@ async def parsed_profiles(self) -> MergedConfig:
 
     def _resolve_profile_name(
         self, explicit_profile: str | None
-    ) -> tuple[str, str | None]:
+    ) -> tuple[str, ConfigSource]:
         """Determine the active profile name and where it came from.
 
         Priority: explicit argument > AWS_PROFILE env var > "default"
 
-        :returns: Tuple of (profile_name, origin), where origin describes the
-            source for error messages and is None when the name was defaulted.
+        :returns: Tuple of (profile_name, source), where source is the
+            provenance of the name: ``OVERRIDE`` for the explicit argument,
+            ``ENV`` for ``AWS_PROFILE``, and ``DEFAULT`` for the fallback.
         """
         if explicit_profile is not None:
-            return explicit_profile, "the profile argument"
+            return explicit_profile, ConfigSource.OVERRIDE
 
         env_profile = os.environ.get(_PROFILE_ENV_VAR)
-        if env_profile is not None:
-            return env_profile, _PROFILE_ENV_VAR
+        if env_profile:
+            return env_profile, ConfigSource.ENV
 
-        return _DEFAULT_PROFILE, None
+        return _DEFAULT_PROFILE, ConfigSource.DEFAULT

packages/smithy-aws-core/src/smithy_aws_core/config/validators.py

@@ -8,6 +8,7 @@
 from smithy_core.retries import RetryStrategyType
 
 from .exceptions import ConfigValidationError, ProfileNotFoundError
+from .types import ConfigSource
 
 _REGION_PATTERN = re.compile(r"^(?![0-9]+$)(?!-)[a-zA-Z0-9-]{1,63}(?<!-)$")
 
@@ -84,16 +85,22 @@ def validate_max_attempts(
 def validate_profile(
     profile_name: str,
     available_profiles: Collection[str],
-    origin: str,
+    source: ConfigSource,
 ) -> None:
     """Validate that a requested profile exists in the config files.
 
     :param profile_name: The active profile name to check.
     :param available_profiles: Profile names defined in the config files.
-    :param origin: Where the profile name came from, used in the error message.
+    :param source: Where the profile name came from, used to format the error.
     :raises ProfileNotFoundError: If the profile is not defined.
     """
     if profile_name not in available_profiles:
+        # Only OVERRIDE and ENV reach here; a DEFAULT profile is never validated.
+        source_str = (
+            "profile argument"
+            if source is ConfigSource.OVERRIDE
+            else "AWS_PROFILE environment variable"
+        )
         raise ProfileNotFoundError(
-            f"Profile {profile_name!r} (from {origin}) not found in config file."
+            f"Profile {profile_name!r} from {source_str} was not found in config file."
         )

packages/smithy-aws-core/tests/unit/config/test_resolver.py

@@ -29,7 +29,11 @@
     resolve_sdk_ua_app_id,
 )
 from smithy_aws_core.config.types import UNSET, ConfigSource
+from smithy_aws_core.identity.environment import EnvironmentCredentialsResolver
 from smithy_aws_core.identity.static import StaticCredentialsResolver
+from smithy_core.aio.retries import StandardRetryStrategy
+from smithy_http.interfaces import HTTPRequestConfiguration
+from smithy_http.testing import MockHTTPClient
 
 
 class NullFileSystem:
@@ -190,7 +194,8 @@ async def test_invalid_override_triggers_validator(self):
     async def test_invalid_profile_raises_error(self):
         with patch.dict(os.environ, {"AWS_REGION": "us-east-1"}, clear=True):
             with pytest.raises(
-                ProfileNotFoundError, match="'FOOBAR' \\(from the profile argument\\)"
+                ProfileNotFoundError,
+                match="Profile 'FOOBAR' from profile argument was not found in config file",
             ):
                 await AsyncAwsConfig.resolve(
                     profile="FOOBAR",
@@ -205,7 +210,8 @@ async def test_invalid_profile_set_via_env_var_raises_error(self):
             clear=True,
         ):
             with pytest.raises(
-                ProfileNotFoundError, match="'FOOBAR' \\(from AWS_PROFILE\\)"
+                ProfileNotFoundError,
+                match="Profile 'FOOBAR' from AWS_PROFILE environment variable was not found in config file",
             ):
                 await AsyncAwsConfig.resolve(
                     fs=NullFileSystem(),
@@ -216,7 +222,8 @@ async def test_unknown_profile_raises_when_config_file_has_others(self):
         fs = FakeFileSystem({"/fake/config": "[profile work]\nregion = eu-west-1\n"})
         with patch.dict(os.environ, {}, clear=True):
             with pytest.raises(
-                ProfileNotFoundError, match="'other' \\(from the profile argument\\)"
+                ProfileNotFoundError,
+                match="Profile 'other' from profile argument was not found in config file",
             ):
                 await AsyncAwsConfig.resolve(
                     profile="other",
@@ -230,7 +237,8 @@ async def test_invalid_aws_profile_env_raises_profile_error(self):
         fs = FakeFileSystem({"/fake/config": "[profile work]\nregion = eu-west-1\n"})
         with patch.dict(os.environ, {"AWS_PROFILE": "wrok"}, clear=True):
             with pytest.raises(
-                ProfileNotFoundError, match="'wrok' \\(from AWS_PROFILE\\)"
+                ProfileNotFoundError,
+                match="Profile 'wrok' from AWS_PROFILE environment variable was not found in config file",
             ):
                 await AsyncAwsConfig.resolve(
                     fs=fs,
@@ -396,16 +404,25 @@ def test_default_profile_is_default(self):
         with patch.dict(os.environ, {}, clear=True):
             ctx = SharedConfigContext()
             assert ctx.profile_name == "default"
+            assert ctx.profile_source is ConfigSource.DEFAULT
 
     def test_profile_from_aws_profile_env(self):
         with patch.dict(os.environ, {"AWS_PROFILE": "work"}, clear=True):
             ctx = SharedConfigContext()
             assert ctx.profile_name == "work"
+            assert ctx.profile_source is ConfigSource.ENV
+
+    def test_empty_aws_profile_env_treated_as_absent(self):
+        with patch.dict(os.environ, {"AWS_PROFILE": ""}, clear=True):
+            ctx = SharedConfigContext()
+            assert ctx.profile_name == "default"
+            assert ctx.profile_source is ConfigSource.DEFAULT
 
     def test_explicit_profile_overrides_env(self):
         with patch.dict(os.environ, {"AWS_PROFILE": "work"}, clear=True):
             ctx = SharedConfigContext(profile_name="custom")
             assert ctx.profile_name == "custom"
+            assert ctx.profile_source is ConfigSource.OVERRIDE
 
     @pytest.mark.asyncio
     async def test_parsed_profiles_caches_result(self):
@@ -472,6 +489,43 @@ async def test_mutable_fields_are_isolated(self):
         assert first.source_of("region") is ConfigSource.OVERRIDE
         assert config.source_of("region") is ConfigSource.ENV
 
+    @pytest.mark.asyncio
+    async def test_shared_resources_are_shared_by_identity(self):
+        with patch.dict(os.environ, {"AWS_REGION": "us-east-1"}, clear=True):
+            config = await AsyncAwsConfig.resolve(fs=NullFileSystem())
+
+        # The transport, credentials resolver, and retry strategy hold network
+        # clients, locks, and shared retry quotas that must not be duplicated
+        transport = MockHTTPClient()
+        resolver = EnvironmentCredentialsResolver()
+        retry_strategy = StandardRetryStrategy()
+        config.transport = transport
+        config.aws_credentials_identity_resolver = resolver
+        config.retry_strategy = retry_strategy
+
+        copy = deepcopy(config)
+        assert copy is not config
+        assert copy.transport is transport
+        assert copy.aws_credentials_identity_resolver is resolver
+        assert copy.retry_strategy is retry_strategy
+
+    @pytest.mark.asyncio
+    async def test_deepcopy_with_no_shared_resources(self):
+        with patch.dict(os.environ, {"AWS_REGION": "us-east-1"}, clear=True):
+            config = await AsyncAwsConfig.resolve(fs=NullFileSystem())
+        config.http_request_config = HTTPRequestConfiguration(read_timeout=1.0)
+
+        copy = deepcopy(config)
+        assert copy is not config
+        # None resources are skipped by the identity-sharing shortcut.
+        assert copy.transport is None
+        assert copy.aws_credentials_identity_resolver is None
+        assert copy.retry_strategy is None
+
+        assert copy.region == "us-east-1"
+        assert copy.http_request_config == config.http_request_config
+        assert copy.http_request_config is not config.http_request_config
+
 
 class TestResolveRetryMode:
     @pytest.mark.asyncio