From 11026bc0e3bae90f77524f570fc3db15715268e8 Mon Sep 17 00:00:00 2001 From: ubaskota <19787410+ubaskota@users.noreply.github.com> Date: Wed, 29 Jul 2026 23:54:10 -0400 Subject: [PATCH 1/3] Add support for remaining config variables from the old to-be-deprecated Config interface --- .../python/codegen/ClientGenerator.java | 25 ++- .../smithy/python/codegen/CodegenUtils.java | 44 ++++ .../codegen/generators/ConfigGenerator.java | 189 +++++++++++++++++- .../src/smithy_aws_core/config/aws_config.py | 63 +++++- .../smithy_aws_core/config/merged_config.py | 44 ++++ .../src/smithy_aws_core/config/resolvers.py | 128 ++++++++++++ .../tests/unit/config/test_merged_config.py | 97 +++++++++ .../tests/unit/config/test_resolver.py | 174 ++++++++++++++++ .../src/smithy_core/aio/retries.py | 19 +- .../tests/unit/aio/test_retries.py | 59 ++++++ 10 files changed, 833 insertions(+), 9 deletions(-) diff --git a/codegen/core/src/main/java/software/amazon/smithy/python/codegen/ClientGenerator.java b/codegen/core/src/main/java/software/amazon/smithy/python/codegen/ClientGenerator.java index e9f5d7a35..ceda9dc55 100644 --- a/codegen/core/src/main/java/software/amazon/smithy/python/codegen/ClientGenerator.java +++ b/codegen/core/src/main/java/software/amazon/smithy/python/codegen/ClientGenerator.java @@ -72,10 +72,23 @@ private void generateService(PythonWriter writer) { } writer.addDependency(SmithyPythonDependency.SMITHY_CORE); + var asyncConfigSymbol = CodegenUtils.getAsyncConfigSymbol(context.settings(), context.model()); writer.write(""" - def __init__(self, config: $1T | None = None, plugins: list[$2T] | None = None): + def __init__( + self, + config: $1T | $6T | None = None, + plugins: list[$2T] | None = None, + ): $3C - self._config = config or $1T() + if isinstance(config, $6T): + self._config: $1T = config # type: ignore[assignment] + elif isinstance(config, $1T) or config is None: + self._config = config or $1T() + else: + raise $7T( + f"config must be $6L or $1L, got {type(config).__name__}. " + f"Use 'await $6L.resolve()' instead." + ) client_plugins: list[$2T] = [ $4C @@ -92,7 +105,9 @@ def __init__(self, config: $1T | None = None, plugins: list[$2T] | None = None): pluginSymbol, writer.consumer(w -> writeConstructorDocs(w, serviceSymbol.getName())), writer.consumer(w -> writeDefaultPlugins(w, defaultPlugins)), - RuntimeTypes.RETRY_STRATEGY_RESOLVER); + RuntimeTypes.RETRY_STRATEGY_RESOLVER, + asyncConfigSymbol, + RuntimeTypes.EXPECTATION_NOT_MET_ERROR); var topDownIndex = TopDownIndex.of(model); var eventStreamIndex = EventStreamIndex.of(model); @@ -249,7 +264,9 @@ private void writeSharedOperationInit( raise $2T("protocol and transport MUST be set on the config to make calls.") retry_strategy = await self._retry_strategy_resolver.resolve_retry_strategy( - retry_strategy=config.retry_strategy + retry_strategy=config.retry_strategy, + retry_mode=getattr(config, "retry_mode", None), + max_attempts=getattr(config, "max_attempts", None), ) pipeline = $3T( diff --git a/codegen/core/src/main/java/software/amazon/smithy/python/codegen/CodegenUtils.java b/codegen/core/src/main/java/software/amazon/smithy/python/codegen/CodegenUtils.java index a6def8968..09b73c189 100644 --- a/codegen/core/src/main/java/software/amazon/smithy/python/codegen/CodegenUtils.java +++ b/codegen/core/src/main/java/software/amazon/smithy/python/codegen/CodegenUtils.java @@ -26,6 +26,7 @@ import java.util.logging.Logger; import software.amazon.smithy.codegen.core.CodegenException; import software.amazon.smithy.codegen.core.Symbol; +import software.amazon.smithy.aws.traits.ServiceTrait; import software.amazon.smithy.model.Model; import software.amazon.smithy.model.knowledge.NullableIndex; import software.amazon.smithy.model.node.Node; @@ -87,6 +88,49 @@ public static Symbol getPluginSymbol(PythonSettings settings) { .build(); } + /** + * Gets the async configuration object symbol for the service. + * + *

This is the new async-resolved config class that inherits from AsyncAwsConfig. + * Derives the name from the SDK ID (e.g., "Bedrock Runtime" becomes + * "AsyncBedrockRuntimeConfig"). Falls back to "AsyncConfig" for non-AWS services. + * + * @param settings The client settings. + * @param model The model containing the service shape. + * @return Returns the async config symbol. + */ + public static Symbol getAsyncConfigSymbol(PythonSettings settings, Model model) { + var service = settings.service(model); + var name = service.getTrait(ServiceTrait.class) + .map(trait -> "Async" + StringUtils.capitalize(trait.getSdkId()).replace(" ", "") + "Config") + .orElse("AsyncConfig"); + return Symbol.builder() + .name(name) + .namespace(String.format("%s.config", settings.moduleName()), ".") + .definitionFile(String.format("./src/%s/config.py", settings.moduleName())) + .build(); + } + + /** + * Gets the async plugin type hint symbol for the service. + * + * @param settings The client settings. + * @param model The model containing the service shape. + * @return Returns the async plugin type hint symbol. + */ + public static Symbol getAsyncPluginSymbol(PythonSettings settings, Model model) { + var service = settings.service(model); + var name = service.getTrait(ServiceTrait.class) + .map(trait -> "Async" + StringUtils.capitalize(trait.getSdkId()).replace(" ", "") + "Plugin") + .orElse("AsyncPlugin"); + return Symbol.builder() + .name(name) + .namespace(String.format("%s.config", settings.moduleName()), ".") + .definitionFile(String.format("./src/%s/config.py", settings.moduleName())) + .build(); + } + + /** * Gets the service error symbol. * diff --git a/codegen/core/src/main/java/software/amazon/smithy/python/codegen/generators/ConfigGenerator.java b/codegen/core/src/main/java/software/amazon/smithy/python/codegen/generators/ConfigGenerator.java index 30d2b07ef..04c75c18d 100644 --- a/codegen/core/src/main/java/software/amazon/smithy/python/codegen/generators/ConfigGenerator.java +++ b/codegen/core/src/main/java/software/amazon/smithy/python/codegen/generators/ConfigGenerator.java @@ -275,6 +275,23 @@ public void run() { writer.write("$L: TypeAlias = Callable[[$T], None]", plugin.getName(), config); writer.writeDocs("A callable that allows customizing the config object on each request.", context); }); + + // Generate the async config subclass and its plugin type + var model = context.model(); + var asyncConfig = CodegenUtils.getAsyncConfigSymbol(context.settings(), model); + var asyncPlugin = CodegenUtils.getAsyncPluginSymbol(context.settings(), model); + context.writerDelegator().useFileWriter(asyncConfig.getDefinitionFile(), asyncConfig.getNamespace(), writer -> { + generateAsyncConfig(context, writer, asyncConfig); + + // Generate the async plugin type alias + writer.addStdlibImport("typing", "Callable"); + writer.addStdlibImport("typing", "TypeAlias"); + writer.write(""); + writer.write(""); + writer.write("$L: TypeAlias = Callable[[$L], None]", asyncPlugin.getName(), asyncConfig.getName()); + writer.writeDocs( + "A callable that allows customizing the async config object on each request.", context); + }); } private void writeInterceptorsType(PythonWriter writer) { @@ -340,10 +357,16 @@ private void generateConfig(GenerationContext context, PythonWriter writer) { writer.pushState(new ConfigSection(finalProperties)); writer.addLocallyDefinedSymbol(configSymbol); writer.addStdlibImport("dataclasses", "dataclass"); + writer.addStdlibImport("warnings"); + var asyncConfigName = CodegenUtils.getAsyncConfigSymbol(context.settings(), context.model()).getName(); writer.write(""" @dataclass(init=False) class $L: - \"""Configuration for $L.\""" + \"""Configuration for $L. + + .. deprecated:: + Use :class:`$L` with ``await $L.resolve()`` instead. + \""" ${C|} @@ -352,12 +375,22 @@ def __init__( *, ${C|} ): + warnings.warn( + "$L is deprecated, use $L.resolve() instead. " + "This class will be removed in a future version.", + DeprecationWarning, + stacklevel=2, + ) ${C|} """, configSymbol.getName(), serviceId, + asyncConfigName, + asyncConfigName, writer.consumer(w -> writePropertyDeclarations(w, finalProperties)), writer.consumer(w -> writeInitParams(w, finalProperties)), + configSymbol.getName(), + asyncConfigName, writer.consumer(w -> initializeProperties(w, finalProperties))); writer.popState(); } @@ -385,6 +418,160 @@ private void initializeProperties(PythonWriter writer, CollectionThis class uses the FieldSpec-based resolution pipeline and adds + * service-specific fields (endpoint_resolver, protocol, auth_schemes, + * auth_scheme_resolver) with their defaults derived from the Smithy model. + */ + private void generateAsyncConfig(GenerationContext context, PythonWriter writer, Symbol asyncConfigSymbol) { + var model = context.model(); + var service = context.settings().service(model); + final String serviceId = service.getTrait(ServiceTrait.class) + .map(ServiceTrait::getSdkId) + .orElse(context.settings().service().getName()); + + // Import AsyncAwsConfig base class + writer.addDependency(SmithyPythonDependency.SMITHY_AWS_CORE); + var asyncAwsConfigSymbol = Symbol.builder() + .name("AsyncAwsConfig") + .namespace("smithy_aws_core.config.aws_config", ".") + .addDependency(SmithyPythonDependency.SMITHY_AWS_CORE) + .build(); + + // Import FieldSpec and ClassVar + var fieldSpecSymbol = Symbol.builder() + .name("FieldSpec") + .namespace("smithy_aws_core.config.types", ".") + .addDependency(SmithyPythonDependency.SMITHY_AWS_CORE) + .build(); + writer.addStdlibImport("typing", "ClassVar"); + writer.addStdlibImport("typing", "Any"); + writer.addStdlibImport("dataclasses", "dataclass"); + + writer.write(""); + writer.write(""); + writer.write("@dataclass(kw_only=True)"); + writer.openBlock("class $L($T):", asyncConfigSymbol.getName(), asyncAwsConfigSymbol); + writer.write("\"\"\"$L configuration (async-resolved).\"\"\"", serviceId); + writer.write(""); + + // Write service-specific field declarations + writer.write("endpoint_resolver: $T | None = None", RuntimeTypes.ENDPOINT_RESOLVER); + writer.write("protocol: $T | None = None", Symbol.builder() + .name("ClientProtocol[Any, Any]") + .addReference(Symbol.builder() + .name("ClientProtocol") + .namespace("smithy_core.aio.interfaces", ".") + .addDependency(SmithyPythonDependency.SMITHY_CORE) + .build()) + .build()); + writer.write("auth_schemes: dict[$T, $T] | None = None", + RuntimeTypes.SHAPE_ID, + Symbol.builder() + .name("AuthScheme[Any, Any, Any, Any]") + .addReference(Symbol.builder() + .name("AuthScheme") + .namespace("smithy_core.aio.interfaces.auth", ".") + .addDependency(SmithyPythonDependency.SMITHY_CORE) + .build()) + .build()); + writer.write("auth_scheme_resolver: $T | None = None", + CodegenUtils.getHttpAuthSchemeResolverSymbol(context.settings())); + writer.write(""); + + // Write _FIELDS class variable with service-specific defaults + writer.openBlock("_FIELDS: ClassVar[dict[str, $T]] = {", fieldSpecSymbol); + writer.write("**$T._FIELDS,", asyncAwsConfigSymbol); + + // endpoint_uri FieldSpec — overrides base class with service-aware resolver + var makeEndpointResolverSymbol = Symbol.builder() + .name("make_endpoint_uri_resolver") + .namespace("smithy_aws_core.config.resolvers", ".") + .addDependency(SmithyPythonDependency.SMITHY_AWS_CORE) + .build(); + var snakeCaseServiceId = serviceId.replace(" ", "_").toLowerCase(); + writer.write("\"endpoint_uri\": $T(", fieldSpecSymbol); + writer.indent(); + writer.write("default=None,"); + writer.write("resolver=$T($S),", makeEndpointResolverSymbol, snakeCaseServiceId); + writer.dedent(); + writer.write("),"); + + // endpoint_resolver FieldSpec + var endpointPrefix = service.getTrait(ServiceTrait.class) + .map(ServiceTrait::getEndpointPrefix) + .orElse(context.settings().service().getName()); + var standardRegionalResolverSymbol = Symbol.builder() + .name("StandardRegionalEndpointsResolver") + .namespace("smithy_aws_core.endpoints.standard_regional", ".") + .addDependency(SmithyPythonDependency.SMITHY_AWS_CORE) + .build(); + writer.write("\"endpoint_resolver\": $T(", fieldSpecSymbol); + writer.indent(); + writer.write("default_factory=lambda: $T(endpoint_prefix=$S),", + standardRegionalResolverSymbol, endpointPrefix); + writer.dedent(); + writer.write("),"); + + // protocol FieldSpec + writer.write("\"protocol\": $T(", fieldSpecSymbol); + writer.indent(); + writer.write("default_factory=lambda: ${C|},", + writer.consumer(w -> context.protocolGenerator().initializeProtocol(context, w))); + writer.dedent(); + writer.write("),"); + + // auth_schemes FieldSpec + writer.write("\"auth_schemes\": $T(", fieldSpecSymbol); + writer.indent(); + writer.write("default_factory=lambda: ${C|},", + writer.consumer(w -> writeAsyncDefaultAuthSchemes(context, w))); + writer.dedent(); + writer.write("),"); + + // auth_scheme_resolver FieldSpec + writer.write("\"auth_scheme_resolver\": $T(", fieldSpecSymbol); + writer.indent(); + writer.write("default_factory=HTTPAuthSchemeResolver,"); + writer.dedent(); + writer.write("),"); + + // transport FieldSpec + writer.write("\"transport\": $T(", fieldSpecSymbol); + writer.indent(); + if (usesHttp2(context)) { + writer.addDependency(SmithyPythonDependency.SMITHY_HTTP.withOptionalDependencies("awscrt")); + writer.write("default_factory=lambda: $T(),", RuntimeTypes.AWS_CRT_HTTP_CLIENT); + } else { + writer.addDependency(SmithyPythonDependency.SMITHY_HTTP.withOptionalDependencies("aiohttp")); + writer.write("default_factory=lambda: $T(),", RuntimeTypes.AIOHTTP_CLIENT); + } + writer.dedent(); + writer.write("),"); + + writer.closeBlock("}"); + writer.closeBlock(""); + } + + private static void writeAsyncDefaultAuthSchemes(GenerationContext context, PythonWriter writer) { + var service = context.settings().service(context.model()); + writer.openBlock("{"); + for (PythonIntegration integration : context.integrations()) { + for (RuntimeClientPlugin plugin : integration.getClientPlugins(context)) { + if (plugin.matchesService(context.model(), service) && plugin.getAuthScheme().isPresent()) { + var scheme = plugin.getAuthScheme().get(); + writer.write("$T($S): ${C|},", + RuntimeTypes.SHAPE_ID, + scheme.getAuthTrait(), + writer.consumer(w -> scheme.initializeScheme(context, writer, service))); + } + } + } + writer.closeBlock("}"); + } + private static final class AddAuthHelper implements CodeInterceptor { @Override public Class sectionType() { diff --git a/packages/smithy-aws-core/src/smithy_aws_core/config/aws_config.py b/packages/smithy-aws-core/src/smithy_aws_core/config/aws_config.py index fc5b8f61c..ad16f7574 100644 --- a/packages/smithy-aws-core/src/smithy_aws_core/config/aws_config.py +++ b/packages/smithy-aws-core/src/smithy_aws_core/config/aws_config.py @@ -2,17 +2,29 @@ # SPDX-License-Identifier: Apache-2.0 from dataclasses import dataclass, field -from typing import Any, ClassVar, Self +from typing import TYPE_CHECKING, Any, ClassVar, Self from smithy_core.retries import RetryStrategyOptions +if TYPE_CHECKING: + from smithy_core.aio.interfaces import ClientTransport + from smithy_core.aio.interfaces.identity import IdentityResolver + from smithy_http.interfaces import HTTPRequestConfiguration + + from smithy_aws_core.identity import AWSCredentialsIdentity, AWSIdentityProperties + from .context import SharedConfigContext from .exceptions import ConfigError, ConfigValidationError from .filesystem import FileSystem from .resolvers import ( + resolve_aws_access_key_id, + resolve_aws_secret_access_key, + resolve_aws_session_token, + resolve_endpoint_uri, resolve_max_attempts, resolve_region, resolve_retry_mode, + resolve_sdk_ua_app_id, ) from .types import UNSET, ConfigSource, FieldSpec, Resolved from .validators import ( @@ -37,6 +49,17 @@ class AsyncAwsConfig: region: str | None = None retry_mode: str | None = None max_attempts: int | None = None + endpoint_uri: str | None = None + aws_access_key_id: str | None = None + aws_secret_access_key: str | None = None + aws_session_token: str | None = None + aws_credentials_identity_resolver: "IdentityResolver[AWSCredentialsIdentity, AWSIdentityProperties] | None" = None + sdk_ua_app_id: str | None = None + user_agent_extra: str | None = None + interceptors: list[Any] = field(default_factory=list) # type: ignore + http_request_config: "HTTPRequestConfiguration | None" = None + transport: "ClientTransport[Any, Any] | None" = None + retry_strategy: Any | None = None _ctx: SharedConfigContext | None = field(default=None, repr=False, compare=False) _sources: dict[str, ConfigSource] = field( # type: ignore[assignment] @@ -61,6 +84,44 @@ class AsyncAwsConfig: resolver=resolve_max_attempts, validator=validate_max_attempts, ), + "endpoint_uri": FieldSpec( + default=None, + resolver=resolve_endpoint_uri, + ), + "aws_access_key_id": FieldSpec( + default=None, + resolver=resolve_aws_access_key_id, + ), + "aws_secret_access_key": FieldSpec( + default=None, + resolver=resolve_aws_secret_access_key, + ), + "aws_session_token": FieldSpec( + default=None, + resolver=resolve_aws_session_token, + ), + "aws_credentials_identity_resolver": FieldSpec( + default=None, + ), + "sdk_ua_app_id": FieldSpec( + default=None, + resolver=resolve_sdk_ua_app_id, + ), + "user_agent_extra": FieldSpec( + default=None, + ), + "interceptors": FieldSpec( + default_factory=list, + ), + "http_request_config": FieldSpec( + default=None, + ), + "transport": FieldSpec( + default=None, + ), + "retry_strategy": FieldSpec( + default=None, + ), } def __post_init__(self) -> None: diff --git a/packages/smithy-aws-core/src/smithy_aws_core/config/merged_config.py b/packages/smithy-aws-core/src/smithy_aws_core/config/merged_config.py index 5e550732a..f496d9aca 100644 --- a/packages/smithy-aws-core/src/smithy_aws_core/config/merged_config.py +++ b/packages/smithy-aws-core/src/smithy_aws_core/config/merged_config.py @@ -113,3 +113,47 @@ def _merge_profiles( else: merged[name] = Section(properties=dict(section.properties)) return merged + + def get_service_config( + self, profile_name: str, service_id: str, key: str + ) -> str | None: + """Get a config value from the services section for a specific service. + + Looks up the services section referenced by the profile, then finds + the service-specific sub-property within it. + + For a config file like: + [profile default] + services = my-services + + [services my-services] + bedrock_runtime = + endpoint_url = http://localhost:5678 + + Usage: get_service_config("default", "bedrock_runtime", "endpoint_url") + + :param profile_name: The profile name to look up. + :param service_id: The service identifier (lowercase, underscored). + :param key: The property key within the service section. + + :returns: The value, or None if not found. + """ + # Get the services section name from the profile + profile = self._profiles.get(profile_name) + if profile is None: + return None + services_name = profile.properties.get("services") + if not services_name or not isinstance(services_name, str): + return None + + # Look up the services section + services_section = self._services.get(services_name) + if services_section is None: + return None + + # Get the service-specific sub-property + service_props = services_section.properties.get(service_id) + if not isinstance(service_props, dict): + return None + + return service_props.get(key.lower()) diff --git a/packages/smithy-aws-core/src/smithy_aws_core/config/resolvers.py b/packages/smithy-aws-core/src/smithy_aws_core/config/resolvers.py index 7f48ff99c..983420d7c 100644 --- a/packages/smithy-aws-core/src/smithy_aws_core/config/resolvers.py +++ b/packages/smithy-aws-core/src/smithy_aws_core/config/resolvers.py @@ -120,3 +120,131 @@ async def resolve_max_attempts(ctx: SharedConfigContext) -> Resolved[int | None] env_vars=("AWS_MAX_ATTEMPTS",), profile_keys=("max_attempts",), ) + + +async def resolve_endpoint_uri(ctx: SharedConfigContext) -> Resolved[str | None]: + """Resolve the endpoint URI from global environment or config file. + + This is the base resolver that only checks global sources. + For service-specific resolution, use make_endpoint_uri_resolver(). + + :param ctx: The shared resolution context. + :returns: Resolved endpoint URI value with source. + """ + return await _resolve_str( + ctx, + env_vars=("AWS_ENDPOINT_URL",), + profile_keys=("endpoint_url",), + ) + + +async def resolve_aws_access_key_id(ctx: SharedConfigContext) -> Resolved[str | None]: + """Resolve the AWS access key ID from environment or config file. + + :param ctx: The shared resolution context. + :returns: Resolved access key ID value with source. + """ + return await _resolve_str( + ctx, + env_vars=("AWS_ACCESS_KEY_ID",), + profile_keys=("aws_access_key_id",), + ) + + +async def resolve_aws_secret_access_key( + ctx: SharedConfigContext, +) -> Resolved[str | None]: + """Resolve the AWS secret access key from environment or config file. + + :param ctx: The shared resolution context. + :returns: Resolved secret access key value with source. + """ + return await _resolve_str( + ctx, + env_vars=("AWS_SECRET_ACCESS_KEY",), + profile_keys=("aws_secret_access_key",), + ) + + +async def resolve_aws_session_token(ctx: SharedConfigContext) -> Resolved[str | None]: + """Resolve the AWS session token from environment or config file. + + :param ctx: The shared resolution context. + :returns: Resolved session token value with source. + """ + return await _resolve_str( + ctx, + env_vars=("AWS_SESSION_TOKEN",), + profile_keys=("aws_session_token",), + ) + + +async def resolve_sdk_ua_app_id(ctx: SharedConfigContext) -> Resolved[str | None]: + """Resolve the SDK user-agent app ID from environment or config file. + + :param ctx: The shared resolution context. + :returns: Resolved app ID value with source. + """ + return await _resolve_str( + ctx, + env_vars=("AWS_SDK_UA_APP_ID",), + profile_keys=("sdk_ua_app_id",), + ) + + +class EndpointUriResolver: + """Service-aware endpoint URI resolver. + + Resolution order (first match wins): + 1. Service-specific env var (AWS_ENDPOINT_URL_) + 2. Global env var (AWS_ENDPOINT_URL) + 3. Service-specific config file (services section -> service_id -> endpoint_url) + 4. Global config file (profile -> endpoint_url) + """ + + def __init__(self, service_id: str): + """Initialize with a service identifier. + + :param service_id: The service identifier (e.g., "bedrock_runtime"). + Used to construct the service-specific env var and config lookup key. + """ + self._service_env_var = ( + f"AWS_ENDPOINT_URL_{service_id.replace('-', '_').upper()}" + ) + self._service_key = service_id.replace("-", "_").lower() + + async def __call__(self, ctx: SharedConfigContext) -> Resolved[str | None]: + """Resolve the endpoint URI from all sources. + + :param ctx: The shared resolution context. + :returns: Resolved endpoint URI value with source. + """ + value = os.environ.get(self._service_env_var) + if value: + return Resolved(value=value, source=ConfigSource.ENV) + + value = os.environ.get("AWS_ENDPOINT_URL") + if value: + return Resolved(value=value, source=ConfigSource.ENV) + + config_file = await ctx.parsed_profiles() + value = config_file.get_service_config( + ctx.profile_name, self._service_key, "endpoint_url" + ) + if value: + return Resolved(value=value, source=ConfigSource.PROFILE) + + value = config_file.get(ctx.profile_name, "endpoint_url") + if value: + return Resolved(value=value, source=ConfigSource.PROFILE) + + return Resolved(value=UNSET, source=ConfigSource.DEFAULT) # type: ignore[arg-type] + + +def make_endpoint_uri_resolver(service_id: str) -> EndpointUriResolver: + """Create a service-aware endpoint URI resolver. + + :param service_id: The service identifier (e.g., "bedrock_runtime"). + :returns: An EndpointUriResolver instance for use in FieldSpec. + """ + return EndpointUriResolver(service_id) diff --git a/packages/smithy-aws-core/tests/unit/config/test_merged_config.py b/packages/smithy-aws-core/tests/unit/config/test_merged_config.py index 53322d90c..eca880cdf 100644 --- a/packages/smithy-aws-core/tests/unit/config/test_merged_config.py +++ b/packages/smithy-aws-core/tests/unit/config/test_merged_config.py @@ -227,3 +227,100 @@ def test_services_property(self): ) assert "my-svc" in cf.services assert cf.services["my-svc"].properties == {"endpoint_url": "http://localhost"} + + +class TestGetServiceConfig: + """Tests for MergedConfig.get_service_config()""" + + def test_returns_service_specific_endpoint_url(self): + config_data = StandardizedOutput( + profiles={"default": Section(properties={"services": "my-services"})}, + services={ + "my-services": Section( + properties={ + "bedrock_runtime": {"endpoint_url": "https://custom.com"} + } + ) + }, + ) + cf = MergedConfig(config_data, StandardizedOutput()) + assert ( + cf.get_service_config("default", "bedrock_runtime", "endpoint_url") + == "https://custom.com" + ) + + def test_returns_none_when_profile_missing(self): + config_data = StandardizedOutput() + cf = MergedConfig(config_data, StandardizedOutput()) + assert ( + cf.get_service_config("default", "bedrock_runtime", "endpoint_url") is None + ) + + def test_returns_none_when_no_services_key_in_profile(self): + config_data = StandardizedOutput( + profiles={"default": Section(properties={"region": "us-east-1"})}, + ) + cf = MergedConfig(config_data, StandardizedOutput()) + assert ( + cf.get_service_config("default", "bedrock_runtime", "endpoint_url") is None + ) + + def test_returns_none_when_services_section_not_found(self): + config_data = StandardizedOutput( + profiles={"default": Section(properties={"services": "nonexistent"})}, + services={}, + ) + cf = MergedConfig(config_data, StandardizedOutput()) + assert ( + cf.get_service_config("default", "bedrock_runtime", "endpoint_url") is None + ) + + def test_returns_none_when_service_id_not_in_section(self): + config_data = StandardizedOutput( + profiles={"default": Section(properties={"services": "my-services"})}, + services={ + "my-services": Section( + properties={"dynamodb": {"endpoint_url": "https://dynamo.local"}} + ) + }, + ) + cf = MergedConfig(config_data, StandardizedOutput()) + assert ( + cf.get_service_config("default", "bedrock_runtime", "endpoint_url") is None + ) + + def test_returns_none_when_key_not_in_service(self): + config_data = StandardizedOutput( + profiles={"default": Section(properties={"services": "my-services"})}, + services={ + "my-services": Section( + properties={"bedrock_runtime": {"some_other_key": "value"}} + ) + }, + ) + cf = MergedConfig(config_data, StandardizedOutput()) + assert ( + cf.get_service_config("default", "bedrock_runtime", "endpoint_url") is None + ) + + def test_multiple_services_in_section(self): + config_data = StandardizedOutput( + profiles={"default": Section(properties={"services": "my-services"})}, + services={ + "my-services": Section( + properties={ + "bedrock_runtime": {"endpoint_url": "https://bedrock.local"}, + "dynamodb": {"endpoint_url": "https://dynamo.local"}, + } + ) + }, + ) + cf = MergedConfig(config_data, StandardizedOutput()) + assert ( + cf.get_service_config("default", "bedrock_runtime", "endpoint_url") + == "https://bedrock.local" + ) + assert ( + cf.get_service_config("default", "dynamodb", "endpoint_url") + == "https://dynamo.local" + ) diff --git a/packages/smithy-aws-core/tests/unit/config/test_resolver.py b/packages/smithy-aws-core/tests/unit/config/test_resolver.py index 5559cc35d..2879135f5 100644 --- a/packages/smithy-aws-core/tests/unit/config/test_resolver.py +++ b/packages/smithy-aws-core/tests/unit/config/test_resolver.py @@ -18,6 +18,8 @@ ProfileNotFoundError, ) from smithy_aws_core.config.resolvers import ( + EndpointUriResolver, + make_endpoint_uri_resolver, resolve_max_attempts, resolve_region, resolve_retry_mode, @@ -512,3 +514,175 @@ async def test_invalid_value_raises_error(self): ctx = SharedConfigContext() with pytest.raises(ConfigValidationError, match="Invalid integer value"): await resolve_max_attempts(ctx) + + +class TestMakeEndpointUriResolver: + """Tests for the service-aware endpoint URI resolver factory.""" + + @pytest.fixture + def resolver(self): + + return make_endpoint_uri_resolver("bedrock_runtime") + + @pytest.mark.asyncio + async def test_service_specific_env_var_takes_precedence( + self, resolver: EndpointUriResolver + ): + fs = FakeFileSystem( + { + "/fake/config": "[profile default]\nendpoint_url = https://global-profile.com\n" + } + ) + with patch.dict( + os.environ, + {"AWS_ENDPOINT_URL_BEDROCK_RUNTIME": "https://service-env.com"}, + clear=True, + ): + ctx = SharedConfigContext( + fs=fs, + config_file_path="/fake/config", + credentials_file_path="/fake/creds", + ) + result = await resolver(ctx) + assert result.value == "https://service-env.com" + assert result.source == ConfigSource.ENV + + @pytest.mark.asyncio + async def test_global_env_var_when_no_service_specific( + self, resolver: EndpointUriResolver + ): + with patch.dict( + os.environ, {"AWS_ENDPOINT_URL": "https://global-env.com"}, clear=True + ): + ctx = SharedConfigContext( + fs=NullFileSystem(), + config_file_path="/fake/config", + credentials_file_path="/fake/creds", + ) + result = await resolver(ctx) + assert result.value == "https://global-env.com" + assert result.source == ConfigSource.ENV + + @pytest.mark.asyncio + async def test_service_env_beats_global_env(self, resolver: EndpointUriResolver): + with patch.dict( + os.environ, + { + "AWS_ENDPOINT_URL_BEDROCK_RUNTIME": "https://service-env.com", + "AWS_ENDPOINT_URL": "https://global-env.com", + }, + clear=True, + ): + ctx = SharedConfigContext( + fs=NullFileSystem(), + config_file_path="/fake/config", + credentials_file_path="/fake/creds", + ) + result = await resolver(ctx) + assert result.value == "https://service-env.com" + + @pytest.mark.asyncio + async def test_service_specific_config_file(self, resolver: EndpointUriResolver): + fs = FakeFileSystem( + { + "/fake/config": ( + "[profile default]\n" + "services = my-services\n" + "\n" + "[services my-services]\n" + "bedrock_runtime =\n" + " endpoint_url = https://service-config.com\n" + ) + } + ) + with patch.dict(os.environ, {}, clear=True): + ctx = SharedConfigContext( + fs=fs, + config_file_path="/fake/config", + credentials_file_path="/fake/creds", + ) + result = await resolver(ctx) + assert result.value == "https://service-config.com" + assert result.source == ConfigSource.PROFILE + + @pytest.mark.asyncio + async def test_global_config_file_fallback(self, resolver: EndpointUriResolver): + fs = FakeFileSystem( + { + "/fake/config": "[profile default]\nendpoint_url = https://global-config.com\n" + } + ) + with patch.dict(os.environ, {}, clear=True): + ctx = SharedConfigContext( + fs=fs, + config_file_path="/fake/config", + credentials_file_path="/fake/creds", + ) + result = await resolver(ctx) + assert result.value == "https://global-config.com" + assert result.source == ConfigSource.PROFILE + + @pytest.mark.asyncio + async def test_service_config_beats_global_config( + self, resolver: EndpointUriResolver + ): + fs = FakeFileSystem( + { + "/fake/config": ( + "[profile default]\n" + "endpoint_url = https://global-config.com\n" + "services = my-services\n" + "\n" + "[services my-services]\n" + "bedrock_runtime =\n" + " endpoint_url = https://service-config.com\n" + ) + } + ) + with patch.dict(os.environ, {}, clear=True): + ctx = SharedConfigContext( + fs=fs, + config_file_path="/fake/config", + credentials_file_path="/fake/creds", + ) + result = await resolver(ctx) + assert result.value == "https://service-config.com" + + @pytest.mark.asyncio + async def test_env_beats_config_file(self, resolver: EndpointUriResolver): + fs = FakeFileSystem( + { + "/fake/config": ( + "[profile default]\n" + "endpoint_url = https://global-config.com\n" + "services = my-services\n" + "\n" + "[services my-services]\n" + "bedrock_runtime =\n" + " endpoint_url = https://service-config.com\n" + ) + } + ) + with patch.dict( + os.environ, {"AWS_ENDPOINT_URL": "https://global-env.com"}, clear=True + ): + ctx = SharedConfigContext( + fs=fs, + config_file_path="/fake/config", + credentials_file_path="/fake/creds", + ) + result = await resolver(ctx) + assert result.value == "https://global-env.com" + + @pytest.mark.asyncio + async def test_returns_unset_when_nothing_found( + self, resolver: EndpointUriResolver + ): + with patch.dict(os.environ, {}, clear=True): + ctx = SharedConfigContext( + fs=NullFileSystem(), + config_file_path="/fake/config", + credentials_file_path="/fake/creds", + ) + result = await resolver(ctx) + assert result.value is UNSET diff --git a/packages/smithy-core/src/smithy_core/aio/retries.py b/packages/smithy-core/src/smithy_core/aio/retries.py index e3fa6340e..1ebac9495 100644 --- a/packages/smithy-core/src/smithy_core/aio/retries.py +++ b/packages/smithy-core/src/smithy_core/aio/retries.py @@ -25,16 +25,29 @@ class RetryStrategyResolver: """ async def resolve_retry_strategy( - self, *, retry_strategy: RetryStrategy | RetryStrategyOptions | None + self, + *, + retry_strategy: RetryStrategy | RetryStrategyOptions | None, + retry_mode: RetryStrategyType | None = None, + max_attempts: int | None = None, ) -> RetryStrategy: """Resolve a retry strategy from the provided options, using cache when possible. - :param retry_strategy: An explicitly configured retry strategy or options for creating one. + :param retry_strategy: An explicitly configured retry strategy or options for + creating one. Takes precedence over ``retry_mode``/``max_attempts``. + :param retry_mode: Retry mode to fall back on when ``retry_strategy`` is None, + typically resolved from the ``AWS_RETRY_MODE`` env var or a config profile. + :param max_attempts: Maximum attempts to fall back on when ``retry_strategy`` is + None, typically resolved from ``AWS_MAX_ATTEMPTS`` or a config profile. """ if isinstance(retry_strategy, RetryStrategy): return retry_strategy elif retry_strategy is None: - retry_strategy = RetryStrategyOptions() + # Fall back to the separately-resolved config values. + retry_strategy = RetryStrategyOptions( + retry_mode=retry_mode if retry_mode is not None else "standard", + max_attempts=max_attempts, + ) elif not isinstance(retry_strategy, RetryStrategyOptions): # type: ignore[reportUnnecessaryIsInstance] raise TypeError( f"retry_strategy must be RetryStrategy, RetryStrategyOptions, or None, " diff --git a/packages/smithy-core/tests/unit/aio/test_retries.py b/packages/smithy-core/tests/unit/aio/test_retries.py index f35c50750..a9710f313 100644 --- a/packages/smithy-core/tests/unit/aio/test_retries.py +++ b/packages/smithy-core/tests/unit/aio/test_retries.py @@ -166,3 +166,62 @@ async def test_retry_strategy_resolver_rejects_invalid_type() -> None: match="retry_strategy must be RetryStrategy, RetryStrategyOptions, or None", ): await resolver.resolve_retry_strategy(retry_strategy="invalid") # type: ignore + + +async def test_retry_strategy_resolver_uses_max_attempts_fallback() -> None: + resolver = RetryStrategyResolver() + + strategy = await resolver.resolve_retry_strategy( + retry_strategy=None, max_attempts=9 + ) + + assert isinstance(strategy, StandardRetryStrategy) + assert strategy.max_attempts == 9 + + +async def test_retry_strategy_resolver_uses_retry_mode_fallback() -> None: + resolver = RetryStrategyResolver() + + strategy = await resolver.resolve_retry_strategy( + retry_strategy=None, retry_mode="simple", max_attempts=4 + ) + + assert isinstance(strategy, SimpleRetryStrategy) + assert strategy.max_attempts == 4 + + +async def test_retry_strategy_resolver_fallback_defaults_when_unset() -> None: + """Omitting both fallbacks must match the prior no-argument behavior.""" + resolver = RetryStrategyResolver() + + explicit = await resolver.resolve_retry_strategy( + retry_strategy=None, retry_mode=None, max_attempts=None + ) + baseline = await resolver.resolve_retry_strategy(retry_strategy=None) + + assert explicit is baseline + assert isinstance(explicit, StandardRetryStrategy) + assert explicit.max_attempts == 3 + + +async def test_explicit_retry_strategy_options_beat_fallbacks() -> None: + resolver = RetryStrategyResolver() + retry_strategy = RetryStrategyOptions(max_attempts=2) + + strategy = await resolver.resolve_retry_strategy( + retry_strategy=retry_strategy, max_attempts=9 + ) + + assert strategy.max_attempts == 2 + + +async def test_explicit_retry_strategy_instance_beats_fallbacks() -> None: + resolver = RetryStrategyResolver() + provided = SimpleRetryStrategy(max_attempts=7) + + strategy = await resolver.resolve_retry_strategy( + retry_strategy=provided, retry_mode="standard", max_attempts=9 + ) + + assert strategy is provided + assert strategy.max_attempts == 7 From ce550d57a86ae3d8b2265fb618316e2b4d7b5a73 Mon Sep 17 00:00:00 2001 From: ubaskota <19787410+ubaskota@users.noreply.github.com> Date: Mon, 3 Aug 2026 01:58:09 -0400 Subject: [PATCH 2/3] Update implementation and tests --- .../smithy/python/codegen/CodegenUtils.java | 3 +- .../codegen/generators/ConfigGenerator.java | 25 +- .../codegen/generators/EnumGenerator.java | 28 +- .../codegen/generators/IntEnumGenerator.java | 28 +- .../codegen/generators/UnionGenerator.java | 49 +-- .../src/smithy_aws_core/config/aws_config.py | 106 +++++- .../src/smithy_aws_core/config/resolvers.py | 57 +-- .../tests/unit/config/test_resolver.py | 332 +++++++++++++++++- 8 files changed, 498 insertions(+), 130 deletions(-) diff --git a/codegen/core/src/main/java/software/amazon/smithy/python/codegen/CodegenUtils.java b/codegen/core/src/main/java/software/amazon/smithy/python/codegen/CodegenUtils.java index 09b73c189..1e803f33d 100644 --- a/codegen/core/src/main/java/software/amazon/smithy/python/codegen/CodegenUtils.java +++ b/codegen/core/src/main/java/software/amazon/smithy/python/codegen/CodegenUtils.java @@ -24,9 +24,9 @@ import java.util.Optional; import java.util.Set; import java.util.logging.Logger; +import software.amazon.smithy.aws.traits.ServiceTrait; import software.amazon.smithy.codegen.core.CodegenException; import software.amazon.smithy.codegen.core.Symbol; -import software.amazon.smithy.aws.traits.ServiceTrait; import software.amazon.smithy.model.Model; import software.amazon.smithy.model.knowledge.NullableIndex; import software.amazon.smithy.model.node.Node; @@ -130,7 +130,6 @@ public static Symbol getAsyncPluginSymbol(PythonSettings settings, Model model) .build(); } - /** * Gets the service error symbol. * diff --git a/codegen/core/src/main/java/software/amazon/smithy/python/codegen/generators/ConfigGenerator.java b/codegen/core/src/main/java/software/amazon/smithy/python/codegen/generators/ConfigGenerator.java index 04c75c18d..359e77371 100644 --- a/codegen/core/src/main/java/software/amazon/smithy/python/codegen/generators/ConfigGenerator.java +++ b/codegen/core/src/main/java/software/amazon/smithy/python/codegen/generators/ConfigGenerator.java @@ -290,7 +290,8 @@ public void run() { writer.write(""); writer.write("$L: TypeAlias = Callable[[$L], None]", asyncPlugin.getName(), asyncConfig.getName()); writer.writeDocs( - "A callable that allows customizing the async config object on each request.", context); + "A callable that allows customizing the async config object on each request.", + context); }); } @@ -459,14 +460,15 @@ private void generateAsyncConfig(GenerationContext context, PythonWriter writer, // Write service-specific field declarations writer.write("endpoint_resolver: $T | None = None", RuntimeTypes.ENDPOINT_RESOLVER); - writer.write("protocol: $T | None = None", Symbol.builder() - .name("ClientProtocol[Any, Any]") - .addReference(Symbol.builder() - .name("ClientProtocol") - .namespace("smithy_core.aio.interfaces", ".") - .addDependency(SmithyPythonDependency.SMITHY_CORE) - .build()) - .build()); + writer.write("protocol: $T | None = None", + Symbol.builder() + .name("ClientProtocol[Any, Any]") + .addReference(Symbol.builder() + .name("ClientProtocol") + .namespace("smithy_core.aio.interfaces", ".") + .addDependency(SmithyPythonDependency.SMITHY_CORE) + .build()) + .build()); writer.write("auth_schemes: dict[$T, $T] | None = None", RuntimeTypes.SHAPE_ID, Symbol.builder() @@ -487,7 +489,7 @@ private void generateAsyncConfig(GenerationContext context, PythonWriter writer, // endpoint_uri FieldSpec — overrides base class with service-aware resolver var makeEndpointResolverSymbol = Symbol.builder() - .name("make_endpoint_uri_resolver") + .name("EndpointUriResolver") .namespace("smithy_aws_core.config.resolvers", ".") .addDependency(SmithyPythonDependency.SMITHY_AWS_CORE) .build(); @@ -511,7 +513,8 @@ private void generateAsyncConfig(GenerationContext context, PythonWriter writer, writer.write("\"endpoint_resolver\": $T(", fieldSpecSymbol); writer.indent(); writer.write("default_factory=lambda: $T(endpoint_prefix=$S),", - standardRegionalResolverSymbol, endpointPrefix); + standardRegionalResolverSymbol, + endpointPrefix); writer.dedent(); writer.write("),"); diff --git a/codegen/core/src/main/java/software/amazon/smithy/python/codegen/generators/EnumGenerator.java b/codegen/core/src/main/java/software/amazon/smithy/python/codegen/generators/EnumGenerator.java index b38106b96..85ce09be0 100644 --- a/codegen/core/src/main/java/software/amazon/smithy/python/codegen/generators/EnumGenerator.java +++ b/codegen/core/src/main/java/software/amazon/smithy/python/codegen/generators/EnumGenerator.java @@ -41,20 +41,24 @@ public void run() { writer.addStdlibImport("enum", "StrEnum"); writer.addDependency(SmithyPythonDependency.SMITHY_CORE); writer.addLocallyDefinedSymbol(enumSymbol); - writer.openBlock("class $L($T, StrEnum):", "", enumSymbol.getName(), RuntimeTypes.UNKNOWN_ENUM_MIXIN, () -> { - shape.getTrait(DocumentationTrait.class).ifPresent(trait -> { - writer.writeDocs(trait.getValue(), context); - }); + writer.openBlock("class $L($T, StrEnum):", + "", + enumSymbol.getName(), + RuntimeTypes.UNKNOWN_ENUM_MIXIN, + () -> { + shape.getTrait(DocumentationTrait.class).ifPresent(trait -> { + writer.writeDocs(trait.getValue(), context); + }); - for (MemberShape member : shape.members()) { - var name = context.symbolProvider().toMemberName(member); - var value = member.expectTrait(EnumValueTrait.class).expectStringValue(); - writer.write("$L = $S", name, value); - member.getTrait(DocumentationTrait.class).ifPresent(trait -> { - writer.writeDocs(trait.getValue(), context); + for (MemberShape member : shape.members()) { + var name = context.symbolProvider().toMemberName(member); + var value = member.expectTrait(EnumValueTrait.class).expectStringValue(); + writer.write("$L = $S", name, value); + member.getTrait(DocumentationTrait.class).ifPresent(trait -> { + writer.writeDocs(trait.getValue(), context); + }); + } }); - } - }); }); } } diff --git a/codegen/core/src/main/java/software/amazon/smithy/python/codegen/generators/IntEnumGenerator.java b/codegen/core/src/main/java/software/amazon/smithy/python/codegen/generators/IntEnumGenerator.java index e9fb98ecf..b29d17132 100644 --- a/codegen/core/src/main/java/software/amazon/smithy/python/codegen/generators/IntEnumGenerator.java +++ b/codegen/core/src/main/java/software/amazon/smithy/python/codegen/generators/IntEnumGenerator.java @@ -41,20 +41,24 @@ public void run() { writer.addStdlibImport("enum", "IntEnum"); writer.addDependency(SmithyPythonDependency.SMITHY_CORE); writer.addLocallyDefinedSymbol(enumSymbol); - writer.openBlock("class $L($T, IntEnum):", "", enumSymbol.getName(), RuntimeTypes.UNKNOWN_ENUM_MIXIN, () -> { - directive.shape().getTrait(DocumentationTrait.class).ifPresent(trait -> { - writer.writeDocs(trait.getValue(), directive.context()); - }); + writer.openBlock("class $L($T, IntEnum):", + "", + enumSymbol.getName(), + RuntimeTypes.UNKNOWN_ENUM_MIXIN, + () -> { + directive.shape().getTrait(DocumentationTrait.class).ifPresent(trait -> { + writer.writeDocs(trait.getValue(), directive.context()); + }); - for (MemberShape member : directive.shape().members()) { - var name = directive.symbolProvider().toMemberName(member); - var value = member.expectTrait(EnumValueTrait.class).expectIntValue(); - writer.write("$L = $L", name, value); - member.getTrait(DocumentationTrait.class).ifPresent(trait -> { - writer.writeDocs(trait.getValue(), directive.context()); + for (MemberShape member : directive.shape().members()) { + var name = directive.symbolProvider().toMemberName(member); + var value = member.expectTrait(EnumValueTrait.class).expectIntValue(); + writer.write("$L = $L", name, value); + member.getTrait(DocumentationTrait.class).ifPresent(trait -> { + writer.writeDocs(trait.getValue(), directive.context()); + }); + } }); - } - }); }); } } diff --git a/codegen/core/src/main/java/software/amazon/smithy/python/codegen/generators/UnionGenerator.java b/codegen/core/src/main/java/software/amazon/smithy/python/codegen/generators/UnionGenerator.java index badf2ea60..1cd2c62ea 100644 --- a/codegen/core/src/main/java/software/amazon/smithy/python/codegen/generators/UnionGenerator.java +++ b/codegen/core/src/main/java/software/amazon/smithy/python/codegen/generators/UnionGenerator.java @@ -157,30 +157,31 @@ private void generateDeserializer() { var schemaSymbol = symbol.expectProperty(SymbolProperties.SCHEMA); var unknownSymbol = symbol.expectProperty(SymbolProperties.UNION_UNKNOWN); writer.putContext("schema", schemaSymbol); - writer.write(""" - class $1L: - _result: $2T | None = None - - def deserialize(self, deserializer: ${shapeDeserializer:T}) -> $2T: - self._result = None - deserializer.read_struct($3T, self._consumer) - - if self._result is None: - raise ${serializationError:T}("Unions must have exactly one value, but found none.") - - return self._result - - def _consumer(self, schema: $4T, de: ${shapeDeserializer:T}) -> None: - match schema.expect_member_index(): - ${5C|} - case _: - self._set_result($6L(tag=schema.expect_member_name())) - - def _set_result(self, value: $2T) -> None: - if self._result is not None: - raise ${serializationError:T}("Unions must have exactly one value, but found more than one.") - self._result = value - """, + writer.write( + """ + class $1L: + _result: $2T | None = None + + def deserialize(self, deserializer: ${shapeDeserializer:T}) -> $2T: + self._result = None + deserializer.read_struct($3T, self._consumer) + + if self._result is None: + raise ${serializationError:T}("Unions must have exactly one value, but found none.") + + return self._result + + def _consumer(self, schema: $4T, de: ${shapeDeserializer:T}) -> None: + match schema.expect_member_index(): + ${5C|} + case _: + self._set_result($6L(tag=schema.expect_member_name())) + + def _set_result(self, value: $2T) -> None: + if self._result is not None: + raise ${serializationError:T}("Unions must have exactly one value, but found more than one.") + self._result = value + """, deserializerSymbol.getName(), symbol, schemaSymbol, diff --git a/packages/smithy-aws-core/src/smithy_aws_core/config/aws_config.py b/packages/smithy-aws-core/src/smithy_aws_core/config/aws_config.py index ad16f7574..84fcf144c 100644 --- a/packages/smithy-aws-core/src/smithy_aws_core/config/aws_config.py +++ b/packages/smithy-aws-core/src/smithy_aws_core/config/aws_config.py @@ -1,6 +1,7 @@ # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. # SPDX-License-Identifier: Apache-2.0 +import os from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, ClassVar, Self @@ -9,6 +10,7 @@ if TYPE_CHECKING: from smithy_core.aio.interfaces import ClientTransport from smithy_core.aio.interfaces.identity import IdentityResolver + from smithy_core.interfaces import URI from smithy_http.interfaces import HTTPRequestConfiguration from smithy_aws_core.identity import AWSCredentialsIdentity, AWSIdentityProperties @@ -17,9 +19,6 @@ from .exceptions import ConfigError, ConfigValidationError from .filesystem import FileSystem from .resolvers import ( - resolve_aws_access_key_id, - resolve_aws_secret_access_key, - resolve_aws_session_token, resolve_endpoint_uri, resolve_max_attempts, resolve_region, @@ -34,6 +33,8 @@ validate_retry_mode, ) +_CREDENTIAL_FIELDS = ("aws_access_key_id", "aws_secret_access_key", "aws_session_token") + @dataclass(kw_only=True) class AsyncAwsConfig: @@ -49,10 +50,10 @@ class AsyncAwsConfig: region: str | None = None retry_mode: str | None = None max_attempts: int | None = None - endpoint_uri: str | None = None - aws_access_key_id: str | None = None - aws_secret_access_key: str | None = None - aws_session_token: str | None = None + endpoint_uri: "str | URI | None" = None + aws_access_key_id: str | None = field(default=None, repr=False) + aws_secret_access_key: str | None = field(default=None, repr=False) + aws_session_token: str | None = field(default=None, repr=False) aws_credentials_identity_resolver: "IdentityResolver[AWSCredentialsIdentity, AWSIdentityProperties] | None" = None sdk_ua_app_id: str | None = None user_agent_extra: str | None = None @@ -90,15 +91,12 @@ class AsyncAwsConfig: ), "aws_access_key_id": FieldSpec( default=None, - resolver=resolve_aws_access_key_id, ), "aws_secret_access_key": FieldSpec( default=None, - resolver=resolve_aws_secret_access_key, ), "aws_session_token": FieldSpec( default=None, - resolver=resolve_aws_session_token, ), "aws_credentials_identity_resolver": FieldSpec( default=None, @@ -211,7 +209,15 @@ async def _resolve_fields(self, overrides: dict[str, Any]) -> None: f"Valid fields are: {sorted(self._FIELDS)}" ) + # Resolve credentials atomically before the field loop + await self._resolve_credentials(overrides) + for field_name, spec in self._FIELDS.items(): + # Skip credentials — already resolved atomically above + if field_name in _CREDENTIAL_FIELDS: + if field_name in self._sources: + continue + # check for overrides first if field_name in overrides: value = overrides[field_name] @@ -244,8 +250,88 @@ def _apply_default(self, field_name: str, spec: FieldSpec) -> None: setattr(self, field_name, value) self._sources[field_name] = ConfigSource.DEFAULT + async def _resolve_credentials(self, overrides: dict[str, Any]) -> None: + """Resolve credential fields atomically from a single source. + + Rules: + - If both aws_access_key_id and aws_secret_access_key are overridden, + resolve normally + - If only one credential is overridden, raise ConfigValidationError. + - Otherwise, resolve atomically: if both key and secret are present in + env, take all three from env. If both are in the profile, take all + three from profile. Token may be None in either case. + + This prevents mixing credentials from different sources. + """ + + required = {"aws_access_key_id", "aws_secret_access_key"} + + cred_overrides = {f for f in _CREDENTIAL_FIELDS if f in overrides} + if cred_overrides: + if required <= cred_overrides: + return + else: + raise ConfigValidationError( + f"Partial credential override: {sorted(cred_overrides)}. " + "Both 'aws_access_key_id' and 'aws_secret_access_key' must be " + "provided together when overriding credentials." + ) + + # Check env vars atomically + env_creds = ( + (os.environ.get("AWS_ACCESS_KEY_ID") or "").strip() or None, + (os.environ.get("AWS_SECRET_ACCESS_KEY") or "").strip() or None, + (os.environ.get("AWS_SESSION_TOKEN") or "").strip() or None, + ) + if env_creds[0] and env_creds[1]: + self._set_credentials(_CREDENTIAL_FIELDS, env_creds, ConfigSource.ENV) + return + + # Check profile atomically + ctx = self._ctx + if ctx is None: + raise ConfigError("Resolution context not initialized") + config_file = await ctx.parsed_profiles() + profile_creds = ( + config_file.get(ctx.profile_name, "aws_access_key_id"), + config_file.get(ctx.profile_name, "aws_secret_access_key"), + config_file.get(ctx.profile_name, "aws_session_token"), + ) + if profile_creds[0] and profile_creds[1]: + self._set_credentials( + _CREDENTIAL_FIELDS, profile_creds, ConfigSource.PROFILE + ) + + def _set_credentials( + self, + fields: tuple[str, ...], + values: tuple[str | None, ...], + source: ConfigSource, + ) -> None: + """Set credential fields atomically, bypassing __setattr__ tracking.""" + for field_name, value in zip(fields, values, strict=True): + object.__setattr__(self, field_name, value or None) + self._sources[field_name] = source + def __setattr__(self, name: str, value: Any) -> None: """Track provenance when fields are set with plugins after construction""" + # Reject unknown fields + if not name.startswith("_") and name not in self.__class__._FIELDS: + raise AttributeError( + f"'{type(self).__name__}' has no config field '{name}'" + ) + + # Block override for credentials after resolution + if ( + name in _CREDENTIAL_FIELDS + and hasattr(self, "_sources") + and name in self._sources + ): + raise AttributeError( + f"'{name}' cannot be modified after resolution. " + "Create a new config with the desired credentials instead." + ) + # Mark as override only if the field is in _FIELDS and was already resolved if ( name in self.__class__._FIELDS diff --git a/packages/smithy-aws-core/src/smithy_aws_core/config/resolvers.py b/packages/smithy-aws-core/src/smithy_aws_core/config/resolvers.py index 983420d7c..f09ecb54c 100644 --- a/packages/smithy-aws-core/src/smithy_aws_core/config/resolvers.py +++ b/packages/smithy-aws-core/src/smithy_aws_core/config/resolvers.py @@ -126,7 +126,7 @@ async def resolve_endpoint_uri(ctx: SharedConfigContext) -> Resolved[str | None] """Resolve the endpoint URI from global environment or config file. This is the base resolver that only checks global sources. - For service-specific resolution, use make_endpoint_uri_resolver(). + For service-specific resolution, use EndpointUriResolver(). :param ctx: The shared resolution context. :returns: Resolved endpoint URI value with source. @@ -138,47 +138,6 @@ async def resolve_endpoint_uri(ctx: SharedConfigContext) -> Resolved[str | None] ) -async def resolve_aws_access_key_id(ctx: SharedConfigContext) -> Resolved[str | None]: - """Resolve the AWS access key ID from environment or config file. - - :param ctx: The shared resolution context. - :returns: Resolved access key ID value with source. - """ - return await _resolve_str( - ctx, - env_vars=("AWS_ACCESS_KEY_ID",), - profile_keys=("aws_access_key_id",), - ) - - -async def resolve_aws_secret_access_key( - ctx: SharedConfigContext, -) -> Resolved[str | None]: - """Resolve the AWS secret access key from environment or config file. - - :param ctx: The shared resolution context. - :returns: Resolved secret access key value with source. - """ - return await _resolve_str( - ctx, - env_vars=("AWS_SECRET_ACCESS_KEY",), - profile_keys=("aws_secret_access_key",), - ) - - -async def resolve_aws_session_token(ctx: SharedConfigContext) -> Resolved[str | None]: - """Resolve the AWS session token from environment or config file. - - :param ctx: The shared resolution context. - :returns: Resolved session token value with source. - """ - return await _resolve_str( - ctx, - env_vars=("AWS_SESSION_TOKEN",), - profile_keys=("aws_session_token",), - ) - - async def resolve_sdk_ua_app_id(ctx: SharedConfigContext) -> Resolved[str | None]: """Resolve the SDK user-agent app ID from environment or config file. @@ -209,9 +168,10 @@ def __init__(self, service_id: str): Used to construct the service-specific env var and config lookup key. """ self._service_env_var = ( - f"AWS_ENDPOINT_URL_{service_id.replace('-', '_').upper()}" + f"AWS_ENDPOINT_URL_{service_id.replace(' ', '_').replace('-', '_').upper()}" ) - self._service_key = service_id.replace("-", "_").lower() + + self._service_key = service_id.replace(" ", "_").replace("-", "_").lower() async def __call__(self, ctx: SharedConfigContext) -> Resolved[str | None]: """Resolve the endpoint URI from all sources. @@ -239,12 +199,3 @@ async def __call__(self, ctx: SharedConfigContext) -> Resolved[str | None]: return Resolved(value=value, source=ConfigSource.PROFILE) return Resolved(value=UNSET, source=ConfigSource.DEFAULT) # type: ignore[arg-type] - - -def make_endpoint_uri_resolver(service_id: str) -> EndpointUriResolver: - """Create a service-aware endpoint URI resolver. - - :param service_id: The service identifier (e.g., "bedrock_runtime"). - :returns: An EndpointUriResolver instance for use in FieldSpec. - """ - return EndpointUriResolver(service_id) diff --git a/packages/smithy-aws-core/tests/unit/config/test_resolver.py b/packages/smithy-aws-core/tests/unit/config/test_resolver.py index 2879135f5..6e8df112f 100644 --- a/packages/smithy-aws-core/tests/unit/config/test_resolver.py +++ b/packages/smithy-aws-core/tests/unit/config/test_resolver.py @@ -19,7 +19,6 @@ ) from smithy_aws_core.config.resolvers import ( EndpointUriResolver, - make_endpoint_uri_resolver, resolve_max_attempts, resolve_region, resolve_retry_mode, @@ -179,7 +178,7 @@ async def test_invalid_override_triggers_validator(self): with pytest.raises( ConfigValidationError, match="Must be a valid AWS region" ): - await AsyncAwsConfig.resolve(region="bad-value!") + await AsyncAwsConfig.resolve(region="bad-value!", fs=NullFileSystem()) @pytest.mark.asyncio async def test_invalid_profile_raises_error(self): @@ -259,6 +258,30 @@ async def test_explicit_default_profile_is_validated(self): fs=NullFileSystem(), ) + @pytest.mark.asyncio + async def test_base_class_resolves_endpoint_uri_from_global_env(self): + with patch.dict( + os.environ, + {"AWS_REGION": "us-east-1", "AWS_ENDPOINT_URL": "https://localhost:4567"}, + clear=True, + ): + config = await AsyncAwsConfig.resolve(fs=NullFileSystem()) + assert config.endpoint_uri == "https://localhost:4567" + assert config.source_of("endpoint_uri") == ConfigSource.ENV + + @pytest.mark.asyncio + async def test_resolve_defaults_all_non_resolved_fields(self): + with patch.dict(os.environ, {"AWS_REGION": "us-east-1"}, clear=True): + config = await AsyncAwsConfig.resolve(fs=NullFileSystem()) + assert config.interceptors == [] + assert config.transport is None + assert config.retry_strategy is None + assert config.http_request_config is None + assert config.user_agent_extra is None + assert config.aws_credentials_identity_resolver is None + for name in ("interceptors", "transport", "user_agent_extra"): + assert config.source_of(name) == ConfigSource.DEFAULT + class TestProvenanceTracking: @pytest.mark.asyncio @@ -337,6 +360,13 @@ async def test_setattr_validates_during_override( with pytest.raises(ConfigValidationError, match=match): setattr(config, field_name, invalid_value) + @pytest.mark.asyncio + async def test_typo_in_field_name_raises_attribute_error(self): + with patch.dict(os.environ, {"AWS_REGION": "us-east-1"}, clear=True): + config = await AsyncAwsConfig.resolve(fs=NullFileSystem()) + with pytest.raises(AttributeError, match="has no config field 'regoin'"): + config.regoin = "us-west-2" + class TestSharedConfigContext: def test_default_profile_is_default(self): @@ -432,6 +462,7 @@ async def test_returns_unset_when_not_found(self): ctx = SharedConfigContext(fs=NullFileSystem()) result = await resolve_retry_mode(ctx) assert result.value is UNSET + assert result.source is ConfigSource.DEFAULT @pytest.mark.asyncio async def test_legacy_warns_and_maps_to_standard(self): @@ -516,13 +547,11 @@ async def test_invalid_value_raises_error(self): await resolve_max_attempts(ctx) -class TestMakeEndpointUriResolver: - """Tests for the service-aware endpoint URI resolver factory.""" - +class TestEndpointUriResolver: @pytest.fixture def resolver(self): - return make_endpoint_uri_resolver("bedrock_runtime") + return EndpointUriResolver("bedrock_runtime") @pytest.mark.asyncio async def test_service_specific_env_var_takes_precedence( @@ -686,3 +715,294 @@ async def test_returns_unset_when_nothing_found( ) result = await resolver(ctx) assert result.value is UNSET + + @pytest.mark.asyncio + async def test_spaced_sdk_id_produces_valid_env_var_name(self): + """Passing a raw SDK ID with spaces (e.g., 'Bedrock Runtime') should + still resolve from the correctly normalized env var.""" + resolver = EndpointUriResolver("Bedrock Runtime") + with patch.dict( + os.environ, + {"AWS_ENDPOINT_URL_BEDROCK_RUNTIME": "https://from-env.com"}, + clear=True, + ): + ctx = SharedConfigContext( + fs=NullFileSystem(), + config_file_path="/fake/config", + credentials_file_path="/fake/creds", + ) + result = await resolver(ctx) + assert result.value == "https://from-env.com" + assert result.source == ConfigSource.ENV + + +class TestReprDoesNotLeakSecrets: + @pytest.mark.asyncio + async def test_repr_does_not_leak_secrets(self): + with patch.dict( + os.environ, + { + "AWS_REGION": "us-east-1", + "AWS_ACCESS_KEY_ID": "AKIAIOSFODNN7EXAMPLE", + "AWS_SECRET_ACCESS_KEY": "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY", + "AWS_SESSION_TOKEN": "FwoGZXIvYXdzEBYaDHqa0AP", + }, + clear=True, + ): + config = await AsyncAwsConfig.resolve(fs=NullFileSystem()) + config_repr = repr(config) + + assert "AKIAIOSFODNN7EXAMPLE" not in config_repr + assert "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY" not in config_repr + assert "FwoGZXIvYXdzEBYaDHqa0AP" not in config_repr + + +class TestCredentialSetIsAtomic: + """Credentials must be resolved from a single source — never mixed.""" + + @pytest.mark.asyncio + async def test_env_credentials_do_not_mix_with_profile_token(self): + """If access_key and secret come from env, token must also come from env (or be None).""" + fs = FakeFileSystem( + {"/fake/credentials": "[default]\naws_session_token = TOKEN_FROM_PROFILE\n"} + ) + with patch.dict( + os.environ, + { + "AWS_REGION": "us-east-1", + "AWS_ACCESS_KEY_ID": "AKID_FROM_ENV", + "AWS_SECRET_ACCESS_KEY": "SECRET_FROM_ENV", + }, + clear=True, + ): + config = await AsyncAwsConfig.resolve( + fs=fs, + config_file_path="/fake/config", + credentials_file_path="/fake/credentials", + ) + assert config.aws_access_key_id == "AKID_FROM_ENV" + assert config.aws_secret_access_key == "SECRET_FROM_ENV" + assert config.aws_session_token is None # NOT from profile + assert config.source_of("aws_access_key_id") == ConfigSource.ENV + assert config.source_of("aws_session_token") == ConfigSource.ENV + + @pytest.mark.asyncio + async def test_all_three_from_env_when_all_set(self): + with patch.dict( + os.environ, + { + "AWS_REGION": "us-east-1", + "AWS_ACCESS_KEY_ID": "AKID", + "AWS_SECRET_ACCESS_KEY": "SECRET", + "AWS_SESSION_TOKEN": "TOKEN", + }, + clear=True, + ): + config = await AsyncAwsConfig.resolve(fs=NullFileSystem()) + assert config.aws_access_key_id == "AKID" + assert config.aws_secret_access_key == "SECRET" + assert config.aws_session_token == "TOKEN" + assert config.source_of("aws_access_key_id") == ConfigSource.ENV + assert config.source_of("aws_secret_access_key") == ConfigSource.ENV + assert config.source_of("aws_session_token") == ConfigSource.ENV + + @pytest.mark.asyncio + async def test_all_three_from_profile_when_no_env(self): + fs = FakeFileSystem( + { + "/fake/credentials": ( + "[default]\n" + "aws_access_key_id = AKID_PROFILE\n" + "aws_secret_access_key = SECRET_PROFILE\n" + "aws_session_token = TOKEN_PROFILE\n" + ) + } + ) + with patch.dict(os.environ, {"AWS_REGION": "us-east-1"}, clear=True): + config = await AsyncAwsConfig.resolve( + fs=fs, + config_file_path="/fake/config", + credentials_file_path="/fake/credentials", + ) + assert config.aws_access_key_id == "AKID_PROFILE" + assert config.aws_secret_access_key == "SECRET_PROFILE" + assert config.aws_session_token == "TOKEN_PROFILE" + assert config.source_of("aws_access_key_id") == ConfigSource.PROFILE + assert config.source_of("aws_secret_access_key") == ConfigSource.PROFILE + assert config.source_of("aws_session_token") == ConfigSource.PROFILE + + @pytest.mark.asyncio + async def test_profile_token_not_used_when_env_has_key_and_secret(self): + """Even if profile has all three, env key+secret means token comes from env too.""" + fs = FakeFileSystem( + { + "/fake/credentials": ( + "[default]\n" + "aws_access_key_id = AKID_PROFILE\n" + "aws_secret_access_key = SECRET_PROFILE\n" + "aws_session_token = TOKEN_PROFILE\n" + ) + } + ) + with patch.dict( + os.environ, + { + "AWS_REGION": "us-east-1", + "AWS_ACCESS_KEY_ID": "AKID_ENV", + "AWS_SECRET_ACCESS_KEY": "SECRET_ENV", + }, + clear=True, + ): + config = await AsyncAwsConfig.resolve( + fs=fs, + config_file_path="/fake/config", + credentials_file_path="/fake/credentials", + ) + # Env wins for all three — token is None because env doesn't have it + assert config.aws_access_key_id == "AKID_ENV" + assert config.aws_secret_access_key == "SECRET_ENV" + assert config.aws_session_token is None + assert config.source_of("aws_session_token") == ConfigSource.ENV + + @pytest.mark.asyncio + async def test_no_credentials_when_nothing_set(self): + with patch.dict(os.environ, {"AWS_REGION": "us-east-1"}, clear=True): + config = await AsyncAwsConfig.resolve(fs=NullFileSystem()) + assert config.aws_access_key_id is None + assert config.aws_secret_access_key is None + assert config.aws_session_token is None + assert config.source_of("aws_access_key_id") == ConfigSource.DEFAULT + + @pytest.mark.asyncio + async def test_partial_credential_override_raises_error(self): + """Overriding only one credential raises an error.""" + fs = FakeFileSystem( + { + "/fake/credentials": ( + "[default]\n" + "aws_access_key_id = AKID_PROFILE\n" + "aws_secret_access_key = SECRET_PROFILE\n" + ) + } + ) + with patch.dict(os.environ, {"AWS_REGION": "us-east-1"}, clear=True): + with pytest.raises( + ConfigValidationError, match="Partial credential override" + ): + await AsyncAwsConfig.resolve( + fs=fs, + config_file_path="/fake/config", + credentials_file_path="/fake/credentials", + aws_access_key_id="OVERRIDE_KEY", + ) + + @pytest.mark.asyncio + async def test_credentials_cannot_be_overridden_after_resolution(self): + with patch.dict( + os.environ, + { + "AWS_REGION": "us-east-1", + "AWS_ACCESS_KEY_ID": "AKID", + "AWS_SECRET_ACCESS_KEY": "SECRET", + }, + clear=True, + ): + config = await AsyncAwsConfig.resolve(fs=NullFileSystem()) + with pytest.raises( + AttributeError, match="cannot be modified after resolution" + ): + config.aws_access_key_id = "NEW_KEY" + + @pytest.mark.asyncio + async def test_session_token_only_override_raises_error(self): + with patch.dict(os.environ, {"AWS_REGION": "us-east-1"}, clear=True): + with pytest.raises( + ConfigValidationError, match="Partial credential override" + ): + await AsyncAwsConfig.resolve( + fs=NullFileSystem(), + aws_session_token="FRESH_TOKEN", + ) + + @pytest.mark.asyncio + async def test_env_session_token_only_falls_through_to_profile(self): + fs = FakeFileSystem( + { + "/fake/credentials": ( + "[default]\n" + "aws_access_key_id = AKID_PROFILE\n" + "aws_secret_access_key = SECRET_PROFILE\n" + "aws_session_token = TOKEN_PROFILE\n" + ) + } + ) + with patch.dict( + os.environ, + {"AWS_REGION": "us-east-1", "AWS_SESSION_TOKEN": "TOKEN_ENV"}, + clear=True, + ): + config = await AsyncAwsConfig.resolve( + fs=fs, + config_file_path="/fake/config", + credentials_file_path="/fake/credentials", + ) + # Token-only env doesn't trigger env path — all from profile + assert config.aws_access_key_id == "AKID_PROFILE" + assert config.aws_secret_access_key == "SECRET_PROFILE" + assert config.aws_session_token == "TOKEN_PROFILE" + assert config.source_of("aws_access_key_id") == ConfigSource.PROFILE + + @pytest.mark.asyncio + async def test_env_key_only_without_secret_falls_through_to_profile(self): + fs = FakeFileSystem( + { + "/fake/credentials": ( + "[default]\n" + "aws_access_key_id = AKID_PROFILE\n" + "aws_secret_access_key = SECRET_PROFILE\n" + ) + } + ) + with patch.dict( + os.environ, + {"AWS_REGION": "us-east-1", "AWS_ACCESS_KEY_ID": "AKID_ENV"}, + clear=True, + ): + config = await AsyncAwsConfig.resolve( + fs=fs, + config_file_path="/fake/config", + credentials_file_path="/fake/credentials", + ) + + assert config.aws_access_key_id == "AKID_PROFILE" + assert config.aws_secret_access_key == "SECRET_PROFILE" + assert config.source_of("aws_access_key_id") == ConfigSource.PROFILE + + @pytest.mark.asyncio + async def test_empty_string_env_credentials_fall_through_to_profile(self): + fs = FakeFileSystem( + { + "/fake/credentials": ( + "[default]\n" + "aws_access_key_id = AKID_PROFILE\n" + "aws_secret_access_key = SECRET_PROFILE\n" + ) + } + ) + with patch.dict( + os.environ, + { + "AWS_REGION": "us-east-1", + "AWS_ACCESS_KEY_ID": "", + "AWS_SECRET_ACCESS_KEY": "", + }, + clear=True, + ): + config = await AsyncAwsConfig.resolve( + fs=fs, + config_file_path="/fake/config", + credentials_file_path="/fake/credentials", + ) + assert config.aws_access_key_id == "AKID_PROFILE" + assert config.aws_secret_access_key == "SECRET_PROFILE" + assert config.source_of("aws_access_key_id") == ConfigSource.PROFILE From 7653aea4189f8cd8ac69c8b771e71483a52445f8 Mon Sep 17 00:00:00 2001 From: ubaskota <19787410+ubaskota@users.noreply.github.com> Date: Sun, 9 Aug 2026 18:46:18 -0400 Subject: [PATCH 3/3] Address comments --- .../codegen/AwsAsyncConfigIntegration.java | 320 ++++++++++++++ .../aws/codegen/AwsUserAgentIntegration.java | 36 +- ...hon.codegen.integrations.PythonIntegration | 1 + .../python/codegen/ClientGenerator.java | 149 ++++--- .../smithy/python/codegen/CodegenUtils.java | 67 +-- .../codegen/generators/ConfigGenerator.java | 283 ++++--------- .../codegen/sections/AsyncConfigSection.java | 17 + .../src/smithy_aws_core/config/aws_config.py | 176 +++++--- .../src/smithy_aws_core/config/context.py | 14 + .../src/smithy_aws_core/config/resolvers.py | 4 + .../tests/unit/config/test_resolver.py | 398 +++++++++--------- 11 files changed, 913 insertions(+), 552 deletions(-) create mode 100644 codegen/aws/core/src/main/java/software/amazon/smithy/python/aws/codegen/AwsAsyncConfigIntegration.java create mode 100644 codegen/core/src/main/java/software/amazon/smithy/python/codegen/sections/AsyncConfigSection.java diff --git a/codegen/aws/core/src/main/java/software/amazon/smithy/python/aws/codegen/AwsAsyncConfigIntegration.java b/codegen/aws/core/src/main/java/software/amazon/smithy/python/aws/codegen/AwsAsyncConfigIntegration.java new file mode 100644 index 000000000..2875cd3f6 --- /dev/null +++ b/codegen/aws/core/src/main/java/software/amazon/smithy/python/aws/codegen/AwsAsyncConfigIntegration.java @@ -0,0 +1,320 @@ +/* + * 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.LinkedHashSet; +import java.util.List; +import java.util.Locale; +import software.amazon.smithy.aws.traits.ServiceTrait; +import software.amazon.smithy.codegen.core.Symbol; +import software.amazon.smithy.model.knowledge.EventStreamIndex; +import software.amazon.smithy.model.knowledge.ServiceIndex; +import software.amazon.smithy.model.knowledge.TopDownIndex; +import software.amazon.smithy.model.node.ArrayNode; +import software.amazon.smithy.model.node.StringNode; +import software.amazon.smithy.model.shapes.OperationShape; +import software.amazon.smithy.python.codegen.CodegenUtils; +import software.amazon.smithy.python.codegen.ConfigProperty; +import software.amazon.smithy.python.codegen.GenerationContext; +import software.amazon.smithy.python.codegen.RuntimeTypes; +import software.amazon.smithy.python.codegen.SmithyPythonDependency; +import software.amazon.smithy.python.codegen.integrations.PythonIntegration; +import software.amazon.smithy.python.codegen.integrations.RuntimeClientPlugin; +import software.amazon.smithy.python.codegen.sections.AsyncConfigSection; +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; + +/** + * AWS integration that generates the async config subclass (e.g., AsyncBedrockRuntimeConfig) + * inheriting from AsyncAwsConfig with service-specific fields and defaults. + */ +@SmithyInternalApi +public class AwsAsyncConfigIntegration implements PythonIntegration { + + @Override + public List> interceptors( + GenerationContext context + ) { + return List.of(new AsyncConfigInterceptor(context)); + } + + private static final class AsyncConfigInterceptor + implements CodeInterceptor { + + private final GenerationContext context; + + AsyncConfigInterceptor(GenerationContext context) { + this.context = context; + } + + @Override + public Class sectionType() { + return AsyncConfigSection.class; + } + + @Override + public void write(PythonWriter writer, String previousText, AsyncConfigSection section) { + // Write any previous content first + writer.write(previousText); + + var model = context.model(); + var service = context.settings().service(model); + + // Gate on the same source of truth the core generators use to decide whether + // to emit references to these classes. If it says no symbol is generated, we + // must not define one, or the two would disagree. + var maybeAsyncConfigSymbol = CodegenUtils.getAsyncConfigSymbol(context.settings(), model); + var maybeAsyncPluginSymbol = CodegenUtils.getAsyncPluginSymbol(context.settings(), model); + if (maybeAsyncConfigSymbol.isEmpty() || maybeAsyncPluginSymbol.isEmpty()) { + return; + } + var asyncConfigSymbol = maybeAsyncConfigSymbol.get(); + var asyncPluginSymbol = maybeAsyncPluginSymbol.get(); + + final String serviceId = service.getTrait(ServiceTrait.class) + .map(ServiceTrait::getSdkId) + .orElse(context.settings().service().getName()); + + // Import AsyncAwsConfig base class + var asyncAwsConfigSymbol = Symbol.builder() + .name("AsyncAwsConfig") + .namespace("smithy_aws_core.config.aws_config", ".") + .addDependency(AwsPythonDependency.SMITHY_AWS_CORE) + .build(); + + // Import FieldSpec and ClassVar + var fieldSpecSymbol = Symbol.builder() + .name("FieldSpec") + .namespace("smithy_aws_core.config.types", ".") + .addDependency(AwsPythonDependency.SMITHY_AWS_CORE) + .build(); + writer.addStdlibImport("typing", "ClassVar"); + writer.addStdlibImport("typing", "Any"); + writer.addStdlibImport("dataclasses", "dataclass"); + + writer.write(""); + writer.write(""); + // repr=False is required: AsyncAwsConfig defines a __repr__ that filters out + // credential fields, and a generated __repr__ on this subclass would shadow it + // and leak secrets. + writer.write("@dataclass(kw_only=True, repr=False)"); + writer.openBlock("class $L($T):", asyncConfigSymbol.getName(), asyncAwsConfigSymbol); + writer.writeDocs(serviceId + " configuration (async-resolved).", context); + writer.write(""); + + // Write service-specific field declarations + writer.write("endpoint_resolver: $T | None = None", RuntimeTypes.ENDPOINT_RESOLVER); + writer.writeDocs("The endpoint resolver used to resolve the final endpoint per-operation " + + "based on the configuration.", context); + writer.write(""); + + writer.write("protocol: $T | None = None", + Symbol.builder() + .name("ClientProtocol[Any, Any]") + .addReference(Symbol.builder() + .name("ClientProtocol") + .namespace("smithy_core.aio.interfaces", ".") + .addDependency(SmithyPythonDependency.SMITHY_CORE) + .build()) + .build()); + writer.writeDocs("The protocol to serialize and deserialize requests with.", context); + writer.write(""); + + var serviceIndex = ServiceIndex.of(context.model()); + var hasAuth = !serviceIndex.getAuthSchemes(context.settings().service()).isEmpty(); + + if (hasAuth) { + writer.write("auth_schemes: dict[$T, $T] | None = None", + RuntimeTypes.SHAPE_ID, + Symbol.builder() + .name("AuthScheme[Any, Any, Any, Any]") + .addReference(Symbol.builder() + .name("AuthScheme") + .namespace("smithy_core.aio.interfaces.auth", ".") + .addDependency(SmithyPythonDependency.SMITHY_CORE) + .build()) + .build()); + writer.writeDocs("A map of auth scheme ids to auth schemes.", context); + writer.write(""); + + writer.write("auth_scheme_resolver: $T | None = None", + CodegenUtils.getHttpAuthSchemeResolverSymbol(context.settings())); + writer.writeDocs("An auth scheme resolver that determines the auth scheme " + + "for each operation.", context); + writer.write(""); + } + + // Plugin-contributed field declarations (e.g., api_key for @httpApiKeyAuth). + // + // More than one plugin can contribute the same property — region, for + // instance, comes from both the auth and regional-endpoints integrations + // — so track the names already written and emit each only once. + var writtenProperties = new LinkedHashSet(); + for (PythonIntegration integration : context.integrations()) { + for (RuntimeClientPlugin plugin : integration.getClientPlugins(context)) { + if (plugin.matchesService(model, service)) { + for (ConfigProperty property : plugin.getConfigProperties()) { + if (!writtenProperties.add(property.name())) { + continue; + } + writer.write("$L: $T | None = None", property.name(), property.type()); + writer.writeDocs(property.documentation(), context); + writer.write(""); + } + } + } + } + + // Write _FIELDS class variable with service-specific defaults + writer.openBlock("_FIELDS: ClassVar[dict[str, $T]] = {", fieldSpecSymbol); + + // Plugin-contributed FieldSpec entries. + // + // These are written *before* the base class spread on purpose. Plugins + // declare config properties for the legacy Config object, which has no + // base class, so some of them duplicate fields AsyncAwsConfig already + // owns and resolves (region, credentials, sdk_ua_app_id, ...). Emitting + // them first means the spread below wins for any such duplicate, so a + // bare FieldSpec(default=None) can never clobber a base spec that + // carries a resolver or validator. Properties the base doesn't declare + // (e.g. api_key for @httpApiKeyAuth) survive untouched. + // + // Reuse the set collected above so these entries stay in step with the + // field declarations and duplicate contributions are written once. + for (String propertyName : writtenProperties) { + writer.write("\"$L\": $T(default=None),", propertyName, fieldSpecSymbol); + } + + writer.write("**$T._FIELDS,", asyncAwsConfigSymbol); + + // Everything below deliberately overrides the base class and so must + // stay after the spread. + + // endpoint_uri FieldSpec — overrides base class with service-aware resolver + var endpointUriResolverSymbol = Symbol.builder() + .name("EndpointUriResolver") + .namespace("smithy_aws_core.config.resolvers", ".") + .addDependency(AwsPythonDependency.SMITHY_AWS_CORE) + .build(); + var snakeCaseServiceId = serviceId.replace(" ", "_").toLowerCase(); + writer.write("\"endpoint_uri\": $T(", fieldSpecSymbol); + writer.indent(); + writer.write("default=None,"); + writer.write("resolver=$T($S),", endpointUriResolverSymbol, snakeCaseServiceId); + writer.dedent(); + writer.write("),"); + + // endpoint_resolver FieldSpec + var endpointPrefix = service.getTrait(ServiceTrait.class) + .map(ServiceTrait::getEndpointPrefix) + .orElse(context.settings().service().getName()); + writer.write("\"endpoint_resolver\": $T(", fieldSpecSymbol); + writer.indent(); + writer.write("default_factory=lambda: $T(endpoint_prefix=$S),", + AwsRuntimeTypes.STANDARD_REGIONAL_ENDPOINTS_RESOLVER, + endpointPrefix); + writer.dedent(); + writer.write("),"); + + // protocol FieldSpec + writer.write("\"protocol\": $T(", fieldSpecSymbol); + writer.indent(); + writer.write("default_factory=lambda: ${C|},", + writer.consumer(w -> context.protocolGenerator().initializeProtocol(context, w))); + writer.dedent(); + writer.write("),"); + + // auth_schemes FieldSpec + if (hasAuth) { + writer.write("\"auth_schemes\": $T(", fieldSpecSymbol); + writer.indent(); + writer.write("default_factory=lambda: ${C|},", + writer.consumer(w -> writeAsyncDefaultAuthSchemes(context, w))); + writer.dedent(); + writer.write("),"); + + // auth_scheme_resolver FieldSpec + writer.write("\"auth_scheme_resolver\": $T(", fieldSpecSymbol); + writer.indent(); + writer.write("default_factory=$T,", + CodegenUtils.getHttpAuthSchemeResolverSymbol(context.settings())); + writer.dedent(); + writer.write("),"); + } + + // transport FieldSpec + writer.write("\"transport\": $T(", fieldSpecSymbol); + writer.indent(); + if (usesHttp2(context)) { + writer.addDependency(SmithyPythonDependency.SMITHY_HTTP.withOptionalDependencies("awscrt")); + writer.write("default_factory=lambda: $T(),", RuntimeTypes.AWS_CRT_HTTP_CLIENT); + } else { + writer.addDependency(SmithyPythonDependency.SMITHY_HTTP.withOptionalDependencies("aiohttp")); + writer.write("default_factory=lambda: $T(),", RuntimeTypes.AIOHTTP_CLIENT); + } + writer.dedent(); + writer.write("),"); + + writer.closeBlock("}"); + writer.closeBlock(""); + + // Generate the async plugin type alias + writer.addStdlibImport("typing", "Callable"); + writer.addStdlibImport("typing", "TypeAlias"); + writer.write(""); + writer.write(""); + writer.write("$L: TypeAlias = Callable[[$L], None]", + asyncPluginSymbol.getName(), + asyncConfigSymbol.getName()); + writer.writeDocs( + "A callable that allows customizing the async config object on each request.", + context); + } + + private static void writeAsyncDefaultAuthSchemes(GenerationContext context, PythonWriter writer) { + var service = context.settings().service(context.model()); + writer.openBlock("{"); + for (PythonIntegration integration : context.integrations()) { + for (RuntimeClientPlugin plugin : integration.getClientPlugins(context)) { + if (plugin.matchesService(context.model(), service) && plugin.getAuthScheme().isPresent()) { + var scheme = plugin.getAuthScheme().get(); + writer.write("$T($S): ${C|},", + RuntimeTypes.SHAPE_ID, + scheme.getAuthTrait(), + writer.consumer(w -> scheme.initializeScheme(context, writer, service))); + } + } + } + writer.closeBlock("}"); + } + + private static boolean usesHttp2(GenerationContext context) { + var configuration = context.applicationProtocol().configuration(); + var httpVersions = configuration.getArrayMember("http") + .orElse(ArrayNode.arrayNode()) + .getElementsAs(StringNode.class) + .stream() + .map(node -> node.getValue().toLowerCase(Locale.ENGLISH)) + .toList(); + + if (httpVersions.contains("h2")) { + return true; + } + + var eventIndex = EventStreamIndex.of(context.model()); + var topDownIndex = TopDownIndex.of(context.model()); + for (OperationShape operation : topDownIndex.getContainedOperations(context.settings().service())) { + if (eventIndex.getInputInfo(operation).isPresent() + || eventIndex.getOutputInfo(operation).isPresent()) { + return true; + } + } + + return false; + } + } +} diff --git a/codegen/aws/core/src/main/java/software/amazon/smithy/python/aws/codegen/AwsUserAgentIntegration.java b/codegen/aws/core/src/main/java/software/amazon/smithy/python/aws/codegen/AwsUserAgentIntegration.java index 423296913..7db2f728a 100644 --- a/codegen/aws/core/src/main/java/software/amazon/smithy/python/aws/codegen/AwsUserAgentIntegration.java +++ b/codegen/aws/core/src/main/java/software/amazon/smithy/python/aws/codegen/AwsUserAgentIntegration.java @@ -22,6 +22,19 @@ public class AwsUserAgentIntegration implements PythonIntegration { public static final String USER_AGENT_PLUGIN = """ + def aws_user_agent_plugin(config: $1T | $5T): + config.interceptors.append( + $2T( + ua_suffix=config.user_agent_extra, + ua_app_id=config.sdk_ua_app_id, + sdk_version=$3T, + service_id=$4S, + ) + ) + """; + + // Variant for services without a generated async config, which must not be referenced. + private static final String USER_AGENT_PLUGIN_SYNC_ONLY = """ def aws_user_agent_plugin(config: $1T): config.interceptors.append( $2T( @@ -96,12 +109,23 @@ public List 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, + CodegenUtils.getConfigSymbol(c.settings()), + userAgentInterceptor, + versionSymbol, + serviceId, + asyncConfig.get()); + } else { + writer.write(USER_AGENT_PLUGIN_SYNC_ONLY, + CodegenUtils.getConfigSymbol(c.settings()), + userAgentInterceptor, + versionSymbol, + serviceId); + } }); return List.of(filename); }) diff --git a/codegen/aws/core/src/main/resources/META-INF/services/software.amazon.smithy.python.codegen.integrations.PythonIntegration b/codegen/aws/core/src/main/resources/META-INF/services/software.amazon.smithy.python.codegen.integrations.PythonIntegration index a338df30c..8fd7115d5 100644 --- a/codegen/aws/core/src/main/resources/META-INF/services/software.amazon.smithy.python.codegen.integrations.PythonIntegration +++ b/codegen/aws/core/src/main/resources/META-INF/services/software.amazon.smithy.python.codegen.integrations.PythonIntegration @@ -8,3 +8,4 @@ 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.AwsAsyncConfigIntegration diff --git a/codegen/core/src/main/java/software/amazon/smithy/python/codegen/ClientGenerator.java b/codegen/core/src/main/java/software/amazon/smithy/python/codegen/ClientGenerator.java index ceda9dc55..08f04775b 100644 --- a/codegen/core/src/main/java/software/amazon/smithy/python/codegen/ClientGenerator.java +++ b/codegen/core/src/main/java/software/amazon/smithy/python/codegen/ClientGenerator.java @@ -61,53 +61,70 @@ private void generateService(PythonWriter writer) { .orElse("Client for " + service.getId().getName()); writer.writeDocs(docs, context); - var defaultPlugins = new LinkedHashSet(); - - for (PythonIntegration integration : context.integrations()) { - for (RuntimeClientPlugin runtimeClientPlugin : integration.getClientPlugins(context)) { - if (runtimeClientPlugin.matchesService(model, service)) { - runtimeClientPlugin.getPythonPlugin().ifPresent(defaultPlugins::add); - } - } - } - writer.addDependency(SmithyPythonDependency.SMITHY_CORE); + // Services with a generated async config accept either config type and resolve + // lazily on first use; the rest keep the synchronous constructor. var asyncConfigSymbol = CodegenUtils.getAsyncConfigSymbol(context.settings(), context.model()); - writer.write(""" - def __init__( - self, - config: $1T | $6T | None = None, - plugins: list[$2T] | None = None, - ): - $3C - if isinstance(config, $6T): - self._config: $1T = config # type: ignore[assignment] - elif isinstance(config, $1T) or config is None: + if (asyncConfigSymbol.isPresent()) { + writer.addStdlibImport("asyncio"); + writer.write(""" + def __init__( + self, + config: $1T | $5T | None = None, + plugins: list[$2T] | None = None, + ): + $3C + if isinstance(config, $5T): + self._config: $1T | $5T | None = config + elif isinstance(config, $1T): + self._config = config + elif config is None: + self._config = None + else: + raise $6T( + f"config must be $5L or $1L, got {type(config).__name__}. " + f"Use 'await $5L.resolve()' instead." + ) + + self._plugins = plugins + self._derive_lock = asyncio.Lock() + self._retry_strategy_resolver = $4T() + + async def _ensure_config(self) -> $1T | $5T: + if self._config is not None: + return self._config + async with self._derive_lock: + if self._config is not None: + return self._config + self._config = await $5T.resolve() + return self._config + """, + configSymbol, + pluginSymbol, + writer.consumer(w -> writeConstructorDocs(w, serviceSymbol.getName())), + RuntimeTypes.RETRY_STRATEGY_RESOLVER, + asyncConfigSymbol.get(), + RuntimeTypes.EXPECTATION_NOT_MET_ERROR); + } else { + writer.write(""" + def __init__( + self, + config: $1T | None = None, + plugins: list[$2T] | None = None, + ): + $3C self._config = config or $1T() - else: - raise $7T( - f"config must be $6L or $1L, got {type(config).__name__}. " - f"Use 'await $6L.resolve()' instead." - ) - - client_plugins: list[$2T] = [ - $4C - ] - if plugins: - client_plugins.extend(plugins) - - for plugin in client_plugins: - plugin(self._config) + self._plugins = plugins + self._retry_strategy_resolver = $4T() - self._retry_strategy_resolver = $5T() - """, - configSymbol, - pluginSymbol, - writer.consumer(w -> writeConstructorDocs(w, serviceSymbol.getName())), - writer.consumer(w -> writeDefaultPlugins(w, defaultPlugins)), - RuntimeTypes.RETRY_STRATEGY_RESOLVER, - asyncConfigSymbol, - RuntimeTypes.EXPECTATION_NOT_MET_ERROR); + async def _ensure_config(self) -> $1T: + return self._config + """, + configSymbol, + pluginSymbol, + writer.consumer(w -> writeConstructorDocs(w, serviceSymbol.getName())), + RuntimeTypes.RETRY_STRATEGY_RESOLVER); + } var topDownIndex = TopDownIndex.of(model); var eventStreamIndex = EventStreamIndex.of(model); @@ -239,9 +256,18 @@ private void writeSharedOperationInit( """, operationDocs, inputDocs, outputDocs); }); + // Service-scoped and operation-scoped plugins are collected separately because a + // RuntimeClientPlugin is always one or the other, never both: setting either + // predicate forces the other to always-false. Service-scoped plugins used to be + // applied once in the constructor, but the config may not exist until the first + // call, so they are applied here instead. + var servicePlugins = new LinkedHashSet(); var defaultPlugins = new LinkedHashSet(); for (PythonIntegration integration : context.integrations()) { for (RuntimeClientPlugin runtimeClientPlugin : integration.getClientPlugins(context)) { + if (runtimeClientPlugin.matchesService(model, service)) { + runtimeClientPlugin.getPythonPlugin().ifPresent(servicePlugins::add); + } if (runtimeClientPlugin.matchesOperation(model, service, operation)) { runtimeClientPlugin.getPythonPlugin().ifPresent(defaultPlugins::add); } @@ -252,16 +278,36 @@ private void writeSharedOperationInit( writer.addStdlibImport("copy", "deepcopy"); writer.write(""" - operation_plugins: list[Plugin] = [ + client_plugins: list[Plugin] = [ $1C ] + operation_plugins: list[Plugin] = [ + $2C + ] if plugins: operation_plugins.extend(plugins) - config = deepcopy(self._config) + # deepcopy keeps plugin mutations (e.g. appending interceptors) scoped to + # this call, so applying client_plugins per-call cannot accumulate on the + # shared config. + config = deepcopy(await self._ensure_config()) + for plugin in client_plugins: + plugin(config) + if self._plugins: + for plugin in self._plugins: + plugin(config) for plugin in operation_plugins: plugin(config) - if config.protocol is None or config.transport is None: - raise $2T("protocol and transport MUST be set on the config to make calls.") + if ( + config.protocol is None + or config.transport is None + or config.endpoint_resolver is None + or config.auth_scheme_resolver is None + or config.auth_schemes is None + ): + raise $3T( + "protocol, transport, endpoint_resolver, auth_scheme_resolver," + " and auth_schemes MUST be set on the config to make calls." + ) retry_strategy = await self._retry_strategy_resolver.resolve_retry_strategy( retry_strategy=config.retry_strategy, @@ -269,21 +315,22 @@ private void writeSharedOperationInit( max_attempts=getattr(config, "max_attempts", None), ) - pipeline = $3T( + pipeline = $4T( protocol=config.protocol, transport=config.transport ) - call = $4T( + call = $5T( input=input, operation=${operation:T}, - context=$5T({"config": config}), - interceptor=$6T(config.interceptors), + context=$6T({"config": config}), + interceptor=$7T(config.interceptors), auth_scheme_resolver=config.auth_scheme_resolver, supported_auth_schemes=config.auth_schemes, endpoint_resolver=config.endpoint_resolver, retry_strategy=retry_strategy, ) """, + writer.consumer(w -> writeDefaultPlugins(w, servicePlugins)), writer.consumer(w -> writeDefaultPlugins(w, defaultPlugins)), RuntimeTypes.EXPECTATION_NOT_MET_ERROR, RuntimeTypes.REQUEST_PIPELINE, diff --git a/codegen/core/src/main/java/software/amazon/smithy/python/codegen/CodegenUtils.java b/codegen/core/src/main/java/software/amazon/smithy/python/codegen/CodegenUtils.java index 1e803f33d..e1f6b5f98 100644 --- a/codegen/core/src/main/java/software/amazon/smithy/python/codegen/CodegenUtils.java +++ b/codegen/core/src/main/java/software/amazon/smithy/python/codegen/CodegenUtils.java @@ -89,45 +89,54 @@ public static Symbol getPluginSymbol(PythonSettings settings) { } /** - * Gets the async configuration object symbol for the service. + * Gets the async configuration object symbol for the service, if one is generated. * *

This is the new async-resolved config class that inherits from AsyncAwsConfig. * Derives the name from the SDK ID (e.g., "Bedrock Runtime" becomes - * "AsyncBedrockRuntimeConfig"). Falls back to "AsyncConfig" for non-AWS services. + * "AsyncBedrockRuntimeConfig"). + * + *

The async config class lives in {@code smithy-aws-core} and is only generated + * for AWS services, so this returns an empty {@code Optional} otherwise. This is the + * single source of truth for whether the class exists: generators must not emit + * references to it when this is empty, and the integration that defines it gates + * itself on this same result. Callers that need the name unconditionally would + * reintroduce references to a class nobody defines. * * @param settings The client settings. * @param model The model containing the service shape. - * @return Returns the async config symbol. + * @return Returns the async config symbol, or empty if none is generated. */ - public static Symbol getAsyncConfigSymbol(PythonSettings settings, Model model) { - var service = settings.service(model); - var name = service.getTrait(ServiceTrait.class) - .map(trait -> "Async" + StringUtils.capitalize(trait.getSdkId()).replace(" ", "") + "Config") - .orElse("AsyncConfig"); - return Symbol.builder() - .name(name) - .namespace(String.format("%s.config", settings.moduleName()), ".") - .definitionFile(String.format("./src/%s/config.py", settings.moduleName())) - .build(); + public static Optional getAsyncConfigSymbol(PythonSettings settings, Model model) { + return asyncConfigSymbolName(settings, model, "Config"); } /** - * Gets the async plugin type hint symbol for the service. + * Gets the async plugin type hint symbol for the service, if one is generated. * * @param settings The client settings. * @param model The model containing the service shape. - * @return Returns the async plugin type hint symbol. + * @return Returns the async plugin symbol, or empty if none is generated. + * @see #getAsyncConfigSymbol(PythonSettings, Model) */ - public static Symbol getAsyncPluginSymbol(PythonSettings settings, Model model) { - var service = settings.service(model); - var name = service.getTrait(ServiceTrait.class) - .map(trait -> "Async" + StringUtils.capitalize(trait.getSdkId()).replace(" ", "") + "Plugin") - .orElse("AsyncPlugin"); - return Symbol.builder() + public static Optional getAsyncPluginSymbol(PythonSettings settings, Model model) { + return asyncConfigSymbolName(settings, model, "Plugin"); + } + + private static Optional asyncConfigSymbolName( + PythonSettings settings, + Model model, + String suffix + ) { + if (!isAwsService(settings, model)) { + return Optional.empty(); + } + var sdkId = settings.service(model).expectTrait(ServiceTrait.class).getSdkId(); + var name = "Async" + StringUtils.capitalize(sdkId).replace(" ", "") + suffix; + return Optional.of(Symbol.builder() .name(name) .namespace(String.format("%s.config", settings.moduleName()), ".") .definitionFile(String.format("./src/%s/config.py", settings.moduleName())) - .build(); + .build()); } /** @@ -343,8 +352,18 @@ private static ZonedDateTime parseHttpDate(Node value) { * @return Returns true if the service is an AWS service, false otherwise. */ public static boolean isAwsService(GenerationContext context) { - var service = context.model().expectShape(context.settings().service()); - return service.hasTrait(software.amazon.smithy.aws.traits.ServiceTrait.class); + return isAwsService(context.settings(), context.model()); + } + + /** + * Determines whether the service being generated is an AWS service. + * + * @param settings The client settings. + * @param model The model containing the service shape. + * @return Returns true if the service is an AWS service, false otherwise. + */ + public static boolean isAwsService(PythonSettings settings, Model model) { + return model.expectShape(settings.service()).hasTrait(ServiceTrait.class); } /** diff --git a/codegen/core/src/main/java/software/amazon/smithy/python/codegen/generators/ConfigGenerator.java b/codegen/core/src/main/java/software/amazon/smithy/python/codegen/generators/ConfigGenerator.java index 359e77371..248f2fa98 100644 --- a/codegen/core/src/main/java/software/amazon/smithy/python/codegen/generators/ConfigGenerator.java +++ b/codegen/core/src/main/java/software/amazon/smithy/python/codegen/generators/ConfigGenerator.java @@ -27,6 +27,7 @@ import software.amazon.smithy.python.codegen.SymbolProperties; import software.amazon.smithy.python.codegen.integrations.PythonIntegration; import software.amazon.smithy.python.codegen.integrations.RuntimeClientPlugin; +import software.amazon.smithy.python.codegen.sections.AsyncConfigSection; import software.amazon.smithy.python.codegen.sections.ConfigSection; import software.amazon.smithy.python.codegen.sections.InitDefaultEndpointResolverSection; import software.amazon.smithy.python.codegen.writer.PythonWriter; @@ -263,36 +264,33 @@ public void run() { context.writerDelegator().useFileWriter(config.getDefinitionFile(), config.getNamespace(), writer -> { writeInterceptorsType(writer); generateConfig(context, writer); + + // Emit the async config section — AWS integrations intercept this + // to generate the service-specific async config subclass. + writer.pushState(new AsyncConfigSection()); + writer.popState(); }); // Generate the plugin symbol. This is just a callable. We could do something // like have a class to implement, but that seems unnecessarily burdensome for // a single function. + // + // For AWS services that have an async config, the Plugin type accepts both + // Config and the async config. For non-AWS services without an async config, + // the Plugin type accepts only Config. var plugin = CodegenUtils.getPluginSymbol(context.settings()); + var asyncConfigForPlugin = CodegenUtils.getAsyncConfigSymbol(context.settings(), context.model()); context.writerDelegator().useFileWriter(plugin.getDefinitionFile(), plugin.getNamespace(), writer -> { writer.addStdlibImport("typing", "Callable"); writer.addStdlibImport("typing", "TypeAlias"); - writer.write("$L: TypeAlias = Callable[[$T], None]", plugin.getName(), config); + if (asyncConfigForPlugin.isPresent()) { + writer.write("$L: TypeAlias = Callable[[$T | $T], None]", + plugin.getName(), config, asyncConfigForPlugin.get()); + } else { + writer.write("$L: TypeAlias = Callable[[$T], None]", plugin.getName(), config); + } writer.writeDocs("A callable that allows customizing the config object on each request.", context); }); - - // Generate the async config subclass and its plugin type - var model = context.model(); - var asyncConfig = CodegenUtils.getAsyncConfigSymbol(context.settings(), model); - var asyncPlugin = CodegenUtils.getAsyncPluginSymbol(context.settings(), model); - context.writerDelegator().useFileWriter(asyncConfig.getDefinitionFile(), asyncConfig.getNamespace(), writer -> { - generateAsyncConfig(context, writer, asyncConfig); - - // Generate the async plugin type alias - writer.addStdlibImport("typing", "Callable"); - writer.addStdlibImport("typing", "TypeAlias"); - writer.write(""); - writer.write(""); - writer.write("$L: TypeAlias = Callable[[$L], None]", asyncPlugin.getName(), asyncConfig.getName()); - writer.writeDocs( - "A callable that allows customizing the async config object on each request.", - context); - }); } private void writeInterceptorsType(PythonWriter writer) { @@ -358,41 +356,66 @@ private void generateConfig(GenerationContext context, PythonWriter writer) { writer.pushState(new ConfigSection(finalProperties)); writer.addLocallyDefinedSymbol(configSymbol); writer.addStdlibImport("dataclasses", "dataclass"); - writer.addStdlibImport("warnings"); - var asyncConfigName = CodegenUtils.getAsyncConfigSymbol(context.settings(), context.model()).getName(); - writer.write(""" - @dataclass(init=False) - class $L: - \"""Configuration for $L. - - .. deprecated:: - Use :class:`$L` with ``await $L.resolve()`` instead. - \""" - - ${C|} - - def __init__( - self, - *, + // This class is only deprecated where an async replacement is generated to point + // at. For services without one it remains the supported config class. + var asyncConfigSymbol = CodegenUtils.getAsyncConfigSymbol(context.settings(), context.model()); + if (asyncConfigSymbol.isPresent()) { + var asyncConfigName = asyncConfigSymbol.get().getName(); + writer.addStdlibImport("warnings"); + writer.write(""" + @dataclass(init=False) + class $L: + \"""Configuration for $L. + + .. deprecated:: + Use :class:`$L` with ``await $L.resolve()`` instead. + \""" + ${C|} - ): - warnings.warn( - "$L is deprecated, use $L.resolve() instead. " - "This class will be removed in a future version.", - DeprecationWarning, - stacklevel=2, - ) + + def __init__( + self, + *, + ${C|} + ): + warnings.warn( + "$L is deprecated, use $L.resolve() instead. " + "This class will be removed in a future version.", + DeprecationWarning, + stacklevel=2, + ) + ${C|} + """, + configSymbol.getName(), + serviceId, + asyncConfigName, + asyncConfigName, + writer.consumer(w -> writePropertyDeclarations(w, finalProperties)), + writer.consumer(w -> writeInitParams(w, finalProperties)), + configSymbol.getName(), + asyncConfigName, + writer.consumer(w -> initializeProperties(w, finalProperties))); + } else { + writer.write(""" + @dataclass(init=False) + class $L: + \"""Configuration for $L.\""" + ${C|} - """, - configSymbol.getName(), - serviceId, - asyncConfigName, - asyncConfigName, - writer.consumer(w -> writePropertyDeclarations(w, finalProperties)), - writer.consumer(w -> writeInitParams(w, finalProperties)), - configSymbol.getName(), - asyncConfigName, - writer.consumer(w -> initializeProperties(w, finalProperties))); + + def __init__( + self, + *, + ${C|} + ): + ${C|} + """, + configSymbol.getName(), + serviceId, + writer.consumer(w -> writePropertyDeclarations(w, finalProperties)), + writer.consumer(w -> writeInitParams(w, finalProperties)), + writer.consumer(w -> initializeProperties(w, finalProperties))); + } writer.popState(); } @@ -419,162 +442,6 @@ private void initializeProperties(PythonWriter writer, CollectionThis class uses the FieldSpec-based resolution pipeline and adds - * service-specific fields (endpoint_resolver, protocol, auth_schemes, - * auth_scheme_resolver) with their defaults derived from the Smithy model. - */ - private void generateAsyncConfig(GenerationContext context, PythonWriter writer, Symbol asyncConfigSymbol) { - var model = context.model(); - var service = context.settings().service(model); - final String serviceId = service.getTrait(ServiceTrait.class) - .map(ServiceTrait::getSdkId) - .orElse(context.settings().service().getName()); - - // Import AsyncAwsConfig base class - writer.addDependency(SmithyPythonDependency.SMITHY_AWS_CORE); - var asyncAwsConfigSymbol = Symbol.builder() - .name("AsyncAwsConfig") - .namespace("smithy_aws_core.config.aws_config", ".") - .addDependency(SmithyPythonDependency.SMITHY_AWS_CORE) - .build(); - - // Import FieldSpec and ClassVar - var fieldSpecSymbol = Symbol.builder() - .name("FieldSpec") - .namespace("smithy_aws_core.config.types", ".") - .addDependency(SmithyPythonDependency.SMITHY_AWS_CORE) - .build(); - writer.addStdlibImport("typing", "ClassVar"); - writer.addStdlibImport("typing", "Any"); - writer.addStdlibImport("dataclasses", "dataclass"); - - writer.write(""); - writer.write(""); - writer.write("@dataclass(kw_only=True)"); - writer.openBlock("class $L($T):", asyncConfigSymbol.getName(), asyncAwsConfigSymbol); - writer.write("\"\"\"$L configuration (async-resolved).\"\"\"", serviceId); - writer.write(""); - - // Write service-specific field declarations - writer.write("endpoint_resolver: $T | None = None", RuntimeTypes.ENDPOINT_RESOLVER); - writer.write("protocol: $T | None = None", - Symbol.builder() - .name("ClientProtocol[Any, Any]") - .addReference(Symbol.builder() - .name("ClientProtocol") - .namespace("smithy_core.aio.interfaces", ".") - .addDependency(SmithyPythonDependency.SMITHY_CORE) - .build()) - .build()); - writer.write("auth_schemes: dict[$T, $T] | None = None", - RuntimeTypes.SHAPE_ID, - Symbol.builder() - .name("AuthScheme[Any, Any, Any, Any]") - .addReference(Symbol.builder() - .name("AuthScheme") - .namespace("smithy_core.aio.interfaces.auth", ".") - .addDependency(SmithyPythonDependency.SMITHY_CORE) - .build()) - .build()); - writer.write("auth_scheme_resolver: $T | None = None", - CodegenUtils.getHttpAuthSchemeResolverSymbol(context.settings())); - writer.write(""); - - // Write _FIELDS class variable with service-specific defaults - writer.openBlock("_FIELDS: ClassVar[dict[str, $T]] = {", fieldSpecSymbol); - writer.write("**$T._FIELDS,", asyncAwsConfigSymbol); - - // endpoint_uri FieldSpec — overrides base class with service-aware resolver - var makeEndpointResolverSymbol = Symbol.builder() - .name("EndpointUriResolver") - .namespace("smithy_aws_core.config.resolvers", ".") - .addDependency(SmithyPythonDependency.SMITHY_AWS_CORE) - .build(); - var snakeCaseServiceId = serviceId.replace(" ", "_").toLowerCase(); - writer.write("\"endpoint_uri\": $T(", fieldSpecSymbol); - writer.indent(); - writer.write("default=None,"); - writer.write("resolver=$T($S),", makeEndpointResolverSymbol, snakeCaseServiceId); - writer.dedent(); - writer.write("),"); - - // endpoint_resolver FieldSpec - var endpointPrefix = service.getTrait(ServiceTrait.class) - .map(ServiceTrait::getEndpointPrefix) - .orElse(context.settings().service().getName()); - var standardRegionalResolverSymbol = Symbol.builder() - .name("StandardRegionalEndpointsResolver") - .namespace("smithy_aws_core.endpoints.standard_regional", ".") - .addDependency(SmithyPythonDependency.SMITHY_AWS_CORE) - .build(); - writer.write("\"endpoint_resolver\": $T(", fieldSpecSymbol); - writer.indent(); - writer.write("default_factory=lambda: $T(endpoint_prefix=$S),", - standardRegionalResolverSymbol, - endpointPrefix); - writer.dedent(); - writer.write("),"); - - // protocol FieldSpec - writer.write("\"protocol\": $T(", fieldSpecSymbol); - writer.indent(); - writer.write("default_factory=lambda: ${C|},", - writer.consumer(w -> context.protocolGenerator().initializeProtocol(context, w))); - writer.dedent(); - writer.write("),"); - - // auth_schemes FieldSpec - writer.write("\"auth_schemes\": $T(", fieldSpecSymbol); - writer.indent(); - writer.write("default_factory=lambda: ${C|},", - writer.consumer(w -> writeAsyncDefaultAuthSchemes(context, w))); - writer.dedent(); - writer.write("),"); - - // auth_scheme_resolver FieldSpec - writer.write("\"auth_scheme_resolver\": $T(", fieldSpecSymbol); - writer.indent(); - writer.write("default_factory=HTTPAuthSchemeResolver,"); - writer.dedent(); - writer.write("),"); - - // transport FieldSpec - writer.write("\"transport\": $T(", fieldSpecSymbol); - writer.indent(); - if (usesHttp2(context)) { - writer.addDependency(SmithyPythonDependency.SMITHY_HTTP.withOptionalDependencies("awscrt")); - writer.write("default_factory=lambda: $T(),", RuntimeTypes.AWS_CRT_HTTP_CLIENT); - } else { - writer.addDependency(SmithyPythonDependency.SMITHY_HTTP.withOptionalDependencies("aiohttp")); - writer.write("default_factory=lambda: $T(),", RuntimeTypes.AIOHTTP_CLIENT); - } - writer.dedent(); - writer.write("),"); - - writer.closeBlock("}"); - writer.closeBlock(""); - } - - private static void writeAsyncDefaultAuthSchemes(GenerationContext context, PythonWriter writer) { - var service = context.settings().service(context.model()); - writer.openBlock("{"); - for (PythonIntegration integration : context.integrations()) { - for (RuntimeClientPlugin plugin : integration.getClientPlugins(context)) { - if (plugin.matchesService(context.model(), service) && plugin.getAuthScheme().isPresent()) { - var scheme = plugin.getAuthScheme().get(); - writer.write("$T($S): ${C|},", - RuntimeTypes.SHAPE_ID, - scheme.getAuthTrait(), - writer.consumer(w -> scheme.initializeScheme(context, writer, service))); - } - } - } - writer.closeBlock("}"); - } - private static final class AddAuthHelper implements CodeInterceptor { @Override public Class sectionType() { diff --git a/codegen/core/src/main/java/software/amazon/smithy/python/codegen/sections/AsyncConfigSection.java b/codegen/core/src/main/java/software/amazon/smithy/python/codegen/sections/AsyncConfigSection.java new file mode 100644 index 000000000..6f5eb6378 --- /dev/null +++ b/codegen/core/src/main/java/software/amazon/smithy/python/codegen/sections/AsyncConfigSection.java @@ -0,0 +1,17 @@ +/* + * 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; +import software.amazon.smithy.utils.SmithyInternalApi; + +/** + * Section marker emitted after the legacy Config class in config.py. + * + *

AWS integrations intercept this section to generate the async config + * subclass (e.g., AsyncBedrockRuntimeConfig) that inherits from AsyncAwsConfig. + */ +@SmithyInternalApi +public record AsyncConfigSection() implements CodeSection {} diff --git a/packages/smithy-aws-core/src/smithy_aws_core/config/aws_config.py b/packages/smithy-aws-core/src/smithy_aws_core/config/aws_config.py index 84fcf144c..21841a1b1 100644 --- a/packages/smithy-aws-core/src/smithy_aws_core/config/aws_config.py +++ b/packages/smithy-aws-core/src/smithy_aws_core/config/aws_config.py @@ -1,8 +1,7 @@ # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. # SPDX-License-Identifier: Apache-2.0 -import os -from dataclasses import dataclass, field +from dataclasses import dataclass, field, fields from typing import TYPE_CHECKING, Any, ClassVar, Self from smithy_core.retries import RetryStrategyOptions @@ -13,7 +12,10 @@ from smithy_core.interfaces import URI from smithy_http.interfaces import HTTPRequestConfiguration - from smithy_aws_core.identity import AWSCredentialsIdentity, AWSIdentityProperties + from smithy_aws_core.identity.components import ( + AWSCredentialsIdentity, + AWSIdentityProperties, + ) from .context import SharedConfigContext from .exceptions import ConfigError, ConfigValidationError @@ -48,19 +50,67 @@ class AsyncAwsConfig: """ region: str | None = None + """The AWS region to connect to. + """ + retry_mode: str | None = None + """The retry mode to use. ``standard`` is the only accepted override. + + ``legacy`` and ``adaptive`` are rejected when set here; when they come from + the environment or a config file they warn and fall back to ``standard``. + """ + max_attempts: int | None = None + """The maximum number of attempts to make per request, including the initial + attempt. Must be an integer of at least 1.""" + endpoint_uri: "str | URI | None" = None + """A static URI to route requests to.""" + aws_access_key_id: str | None = field(default=None, repr=False) + """The identifier for a secret access key. + + Set this together with ``aws_secret_access_key`` to supply credentials in + code. Cannot be modified after resolution; see + ``aws_credentials_identity_resolver`` to supply credentials dynamically. + """ + aws_secret_access_key: str | None = field(default=None, repr=False) + """A secret access key that can be used to sign requests. + + Must be set together with ``aws_access_key_id``. + """ + aws_session_token: str | None = field(default=None, repr=False) + """An access key ID that identifies temporary security credentials.""" + aws_credentials_identity_resolver: "IdentityResolver[AWSCredentialsIdentity, AWSIdentityProperties] | None" = None + """Resolves AWS Credentials. + + Set automatically to a ``StaticCredentialsResolver`` when + ``aws_access_key_id`` and ``aws_secret_access_key`` are supplied in code. + """ + sdk_ua_app_id: str | None = None + """A unique and opaque application ID that is appended to the User-Agent + header.""" + user_agent_extra: str | None = None + """Additional suffix to be added to the User-Agent header.""" + interceptors: list[Any] = field(default_factory=list) # type: ignore + """The list of interceptors, which are hooks that are called during the + execution of a request.""" + http_request_config: "HTTPRequestConfiguration | None" = None + """Configuration for individual HTTP requests.""" + transport: "ClientTransport[Any, Any] | None" = None + """The transport to use to send requests""" + retry_strategy: Any | None = None + """The retry strategy or options for configuring retry behavior. + """ _ctx: SharedConfigContext | None = field(default=None, repr=False, compare=False) _sources: dict[str, ConfigSource] = field( # type: ignore[assignment] @@ -122,6 +172,22 @@ class AsyncAwsConfig: ), } + def __repr__(self) -> str: + """Render the config without exposing credential material. + + Defined on the base class so that every subclass inherits the + filtering, rather than relying on each subclass to mark its own + credential fields ``repr=False``. Subclasses must be declared with + ``@dataclass(repr=False)`` so they inherit this instead of generating + their own ``__repr__``. + """ + rendered = ", ".join( + f"{f.name}={getattr(self, f.name)!r}" + for f in fields(self) + if f.repr and f.name not in _CREDENTIAL_FIELDS + ) + return f"{type(self).__name__}({rendered})" + def __post_init__(self) -> None: """Block direct construction. Use resolve() instead.""" raise ConfigError( @@ -209,15 +275,11 @@ async def _resolve_fields(self, overrides: dict[str, Any]) -> None: f"Valid fields are: {sorted(self._FIELDS)}" ) - # Resolve credentials atomically before the field loop - await self._resolve_credentials(overrides) + # Validate credential overrides and auto-wire the identity resolver + # before the field loop, so the loop sees the resolver as an override. + self._resolve_credentials(overrides) for field_name, spec in self._FIELDS.items(): - # Skip credentials — already resolved atomically above - if field_name in _CREDENTIAL_FIELDS: - if field_name in self._sources: - continue - # check for overrides first if field_name in overrides: value = overrides[field_name] @@ -250,71 +312,57 @@ def _apply_default(self, field_name: str, spec: FieldSpec) -> None: setattr(self, field_name, value) self._sources[field_name] = ConfigSource.DEFAULT - async def _resolve_credentials(self, overrides: dict[str, Any]) -> None: - """Resolve credential fields atomically from a single source. + def _resolve_credentials(self, overrides: dict[str, Any]) -> None: + """Validate in-code credentials and auto-wire StaticCredentialsResolver. Rules: - If both aws_access_key_id and aws_secret_access_key are overridden, - resolve normally - - If only one credential is overridden, raise ConfigValidationError. - - Otherwise, resolve atomically: if both key and secret are present in - env, take all three from env. If both are in the profile, take all - three from profile. Token may be None in either case. - - This prevents mixing credentials from different sources. + auto-set aws_credentials_identity_resolver to a + StaticCredentialsResolver (unless the caller already provided one). + Only the overridden values are used, so a session token present in a + profile is not picked up here. + - If credentials are overridden but the key/secret pair is incomplete, + raise ConfigValidationError. + - If no credential is overridden, credentials are resolved from the + remaining sources. """ required = {"aws_access_key_id", "aws_secret_access_key"} cred_overrides = {f for f in _CREDENTIAL_FIELDS if f in overrides} - if cred_overrides: - if required <= cred_overrides: - return - else: - raise ConfigValidationError( - f"Partial credential override: {sorted(cred_overrides)}. " - "Both 'aws_access_key_id' and 'aws_secret_access_key' must be " - "provided together when overriding credentials." - ) - - # Check env vars atomically - env_creds = ( - (os.environ.get("AWS_ACCESS_KEY_ID") or "").strip() or None, - (os.environ.get("AWS_SECRET_ACCESS_KEY") or "").strip() or None, - (os.environ.get("AWS_SESSION_TOKEN") or "").strip() or None, - ) - if env_creds[0] and env_creds[1]: - self._set_credentials(_CREDENTIAL_FIELDS, env_creds, ConfigSource.ENV) + + if not cred_overrides: return - # Check profile atomically - ctx = self._ctx - if ctx is None: - raise ConfigError("Resolution context not initialized") - config_file = await ctx.parsed_profiles() - profile_creds = ( - config_file.get(ctx.profile_name, "aws_access_key_id"), - config_file.get(ctx.profile_name, "aws_secret_access_key"), - config_file.get(ctx.profile_name, "aws_session_token"), - ) - if profile_creds[0] and profile_creds[1]: - self._set_credentials( - _CREDENTIAL_FIELDS, profile_creds, ConfigSource.PROFILE + if not required <= cred_overrides: + raise ConfigValidationError( + f"Partial credential override: {sorted(cred_overrides)}. " + "Both 'aws_access_key_id' and 'aws_secret_access_key' must be " + "provided together when overriding credentials." ) - def _set_credentials( - self, - fields: tuple[str, ...], - values: tuple[str | None, ...], - source: ConfigSource, - ) -> None: - """Set credential fields atomically, bypassing __setattr__ tracking.""" - for field_name, value in zip(fields, values, strict=True): - object.__setattr__(self, field_name, value or None) - self._sources[field_name] = source + # Auto-wire StaticCredentialsResolver if user didn't provide one + if overrides.get("aws_credentials_identity_resolver") is None: + # Lazy import to avoid circular dependency + from smithy_aws_core.identity.components import AWSCredentialsIdentity + from smithy_aws_core.identity.static import StaticCredentialsResolver + + identity = AWSCredentialsIdentity( + access_key_id=overrides["aws_access_key_id"], + secret_access_key=overrides["aws_secret_access_key"], + session_token=overrides.get("aws_session_token"), + ) + overrides["aws_credentials_identity_resolver"] = StaticCredentialsResolver( + identity=identity + ) def __setattr__(self, name: str, value: Any) -> None: - """Track provenance when fields are set with plugins after construction""" + """Guard and track config fields set after resolution. + + Rejects unknown field names, blocks credential mutation, validates the + new value, and records the field as an override so ``source_of()`` + stays accurate when plugins customize a config per request. + """ # Reject unknown fields if not name.startswith("_") and name not in self.__class__._FIELDS: raise AttributeError( @@ -328,8 +376,10 @@ def __setattr__(self, name: str, value: Any) -> None: and name in self._sources ): raise AttributeError( - f"'{name}' cannot be modified after resolution. " - "Create a new config with the desired credentials instead." + f"'{name}' cannot be modified after resolution. Pass credentials " + f"to `await {type(self).__name__}.resolve(...)`, or set " + "'aws_credentials_identity_resolver' to supply credentials " + "dynamically." ) # Mark as override only if the field is in _FIELDS and was already resolved diff --git a/packages/smithy-aws-core/src/smithy_aws_core/config/context.py b/packages/smithy-aws-core/src/smithy_aws_core/config/context.py index ae38ff400..0df9c4971 100644 --- a/packages/smithy-aws-core/src/smithy_aws_core/config/context.py +++ b/packages/smithy-aws-core/src/smithy_aws_core/config/context.py @@ -167,6 +167,20 @@ def http_client(self) -> Any | None: """HTTP client for network-based resolvers.""" return self._http_client + def __deepcopy__(self, memo: Any) -> "SharedConfigContext": + """Return self rather than a copy. + + The context is read-only once resolution finishes: resolvers have + already pulled their values onto the config's fields, and nothing on + the request path reads it again. Generated clients deep-copy the + config on every operation call to keep plugin mutations scoped to + that call, which would otherwise rebuild the whole parsed profile + tree per request — work proportional to the size of the caller's + shared config files. Sharing this instead keeps that cost flat + without weakening the isolation of the fields plugins actually write. + """ + return self + async def parsed_profiles(self) -> MergedConfig: """Get the parsed and merged config/credentials file data. diff --git a/packages/smithy-aws-core/src/smithy_aws_core/config/resolvers.py b/packages/smithy-aws-core/src/smithy_aws_core/config/resolvers.py index f09ecb54c..3f6a66631 100644 --- a/packages/smithy-aws-core/src/smithy_aws_core/config/resolvers.py +++ b/packages/smithy-aws-core/src/smithy_aws_core/config/resolvers.py @@ -92,6 +92,10 @@ async def resolve_retry_mode(ctx: SharedConfigContext) -> Resolved[str | None]: env_vars=("AWS_RETRY_MODE",), profile_keys=("retry_mode",), ) + + if result.value is not UNSET: + result = Resolved(value=result.value.lower(), source=result.source) + if result.value == "legacy": warnings.warn( "'legacy' retry mode is not supported, using 'standard' instead.", diff --git a/packages/smithy-aws-core/tests/unit/config/test_resolver.py b/packages/smithy-aws-core/tests/unit/config/test_resolver.py index 6e8df112f..6008632c5 100644 --- a/packages/smithy-aws-core/tests/unit/config/test_resolver.py +++ b/packages/smithy-aws-core/tests/unit/config/test_resolver.py @@ -7,6 +7,8 @@ """ import os +from copy import deepcopy +from dataclasses import dataclass from unittest.mock import patch import pytest @@ -19,11 +21,14 @@ ) from smithy_aws_core.config.resolvers import ( EndpointUriResolver, + resolve_endpoint_uri, resolve_max_attempts, resolve_region, resolve_retry_mode, + resolve_sdk_ua_app_id, ) from smithy_aws_core.config.types import UNSET, ConfigSource +from smithy_aws_core.identity.static import StaticCredentialsResolver class NullFileSystem: @@ -397,6 +402,58 @@ async def test_parsed_profiles_caches_result(self): result2 = await ctx.parsed_profiles() assert result1 is result2 + def test_deepcopy_returns_same_instance(self): + with patch.dict(os.environ, {}, clear=True): + ctx = SharedConfigContext(fs=NullFileSystem()) + assert deepcopy(ctx) is ctx + + +class TestConfigDeepCopy: + """Generated clients deep-copy the config on every operation call. + + The copy exists to keep plugin mutations scoped to a single call, so the + fields plugins write must be independent per copy. The resolution context + is read-only afterwards and is shared instead, which keeps the per-request + cost from scaling with the size of the caller's shared config files. + """ + + @pytest.mark.asyncio + async def test_resolution_context_is_shared(self): + fs = FakeFileSystem({"/fake/config": "[profile default]\nregion = us-east-1\n"}) + with patch.dict(os.environ, {}, clear=True): + config = await AsyncAwsConfig.resolve( + fs=fs, + config_file_path="/fake/config", + credentials_file_path="/fake/credentials", + ) + # Sanity check: there is a context to share, so the assertion below + # is meaningful. + assert config.resolution_context() is not None + assert deepcopy(config).resolution_context() is config.resolution_context() + + @pytest.mark.asyncio + async def test_mutable_fields_are_isolated(self): + with patch.dict(os.environ, {"AWS_REGION": "us-east-1"}, clear=True): + config = await AsyncAwsConfig.resolve(fs=NullFileSystem()) + + first = deepcopy(config) + second = deepcopy(config) + + # A plugin appending an interceptor must not affect the shared config + # or any other in-flight call. + first.interceptors.append("first-only") + second.interceptors.append("second-only") + assert first.interceptors == ["first-only"] + assert second.interceptors == ["second-only"] + assert config.interceptors == [] + + # Scalar overrides and their provenance stay per-copy too. + first.region = "eu-west-2" + assert first.region == "eu-west-2" + assert config.region == "us-east-1" + assert first.source_of("region") is ConfigSource.OVERRIDE + assert config.source_of("region") is ConfigSource.ENV + class TestResolveRetryMode: @pytest.mark.asyncio @@ -488,6 +545,24 @@ async def test_adaptive_warns_and_maps_to_standard(self): assert result.value == "standard" assert result.source == ConfigSource.ENV + @pytest.mark.asyncio + @pytest.mark.parametrize( + "env_value,expected", + [ + ("STANDARD", "standard"), + ("Standard", "standard"), + ("LEGACY", "standard"), + ("Legacy", "standard"), + ("ADAPTIVE", "standard"), + ("Adaptive", "standard"), + ], + ) + async def test_retry_mode_is_case_insensitive(self, env_value: str, expected: str): + with patch.dict(os.environ, {"AWS_RETRY_MODE": env_value}, clear=True): + ctx = SharedConfigContext(fs=NullFileSystem()) + result = await resolve_retry_mode(ctx) + assert result.value == expected + class TestResolveMaxAttempts: @pytest.mark.asyncio @@ -739,130 +814,54 @@ async def test_spaced_sdk_id_produces_valid_env_var_name(self): class TestReprDoesNotLeakSecrets: @pytest.mark.asyncio async def test_repr_does_not_leak_secrets(self): - with patch.dict( - os.environ, - { - "AWS_REGION": "us-east-1", - "AWS_ACCESS_KEY_ID": "AKIAIOSFODNN7EXAMPLE", - "AWS_SECRET_ACCESS_KEY": "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY", - "AWS_SESSION_TOKEN": "FwoGZXIvYXdzEBYaDHqa0AP", - }, - clear=True, - ): - config = await AsyncAwsConfig.resolve(fs=NullFileSystem()) - config_repr = repr(config) + with patch.dict(os.environ, {"AWS_REGION": "us-east-1"}, clear=True): + config = await AsyncAwsConfig.resolve( + fs=NullFileSystem(), + aws_access_key_id="AKIAIOSFODNN7EXAMPLE", + aws_secret_access_key="wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY", + aws_session_token="FwoGZXIvYXdzEBYaDHqa0AP", + ) + + # Sanity check: the credentials really are populated, so the + # assertions below are meaningful. + assert config.aws_access_key_id == "AKIAIOSFODNN7EXAMPLE" + config_repr = repr(config) assert "AKIAIOSFODNN7EXAMPLE" not in config_repr assert "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY" not in config_repr assert "FwoGZXIvYXdzEBYaDHqa0AP" not in config_repr - - -class TestCredentialSetIsAtomic: - """Credentials must be resolved from a single source — never mixed.""" + assert "region='us-east-1'" in config_repr @pytest.mark.asyncio - async def test_env_credentials_do_not_mix_with_profile_token(self): - """If access_key and secret come from env, token must also come from env (or be None).""" - fs = FakeFileSystem( - {"/fake/credentials": "[default]\naws_session_token = TOKEN_FROM_PROFILE\n"} - ) - with patch.dict( - os.environ, - { - "AWS_REGION": "us-east-1", - "AWS_ACCESS_KEY_ID": "AKID_FROM_ENV", - "AWS_SECRET_ACCESS_KEY": "SECRET_FROM_ENV", - }, - clear=True, - ): - config = await AsyncAwsConfig.resolve( - fs=fs, - config_file_path="/fake/config", - credentials_file_path="/fake/credentials", - ) - assert config.aws_access_key_id == "AKID_FROM_ENV" - assert config.aws_secret_access_key == "SECRET_FROM_ENV" - assert config.aws_session_token is None # NOT from profile - assert config.source_of("aws_access_key_id") == ConfigSource.ENV - assert config.source_of("aws_session_token") == ConfigSource.ENV + async def test_subclass_repr_does_not_leak_secrets(self): + """Subclasses declared with repr=False inherit the filtered __repr__. - @pytest.mark.asyncio - async def test_all_three_from_env_when_all_set(self): - with patch.dict( - os.environ, - { - "AWS_REGION": "us-east-1", - "AWS_ACCESS_KEY_ID": "AKID", - "AWS_SECRET_ACCESS_KEY": "SECRET", - "AWS_SESSION_TOKEN": "TOKEN", - }, - clear=True, - ): - config = await AsyncAwsConfig.resolve(fs=NullFileSystem()) - assert config.aws_access_key_id == "AKID" - assert config.aws_secret_access_key == "SECRET" - assert config.aws_session_token == "TOKEN" - assert config.source_of("aws_access_key_id") == ConfigSource.ENV - assert config.source_of("aws_secret_access_key") == ConfigSource.ENV - assert config.source_of("aws_session_token") == ConfigSource.ENV + This mirrors what codegen emits for service-specific async configs. + """ + + @dataclass(kw_only=True, repr=False) + class ServiceConfig(AsyncAwsConfig): + aws_access_key_id: str | None = None + aws_secret_access_key: str | None = None + aws_session_token: str | None = None - @pytest.mark.asyncio - async def test_all_three_from_profile_when_no_env(self): - fs = FakeFileSystem( - { - "/fake/credentials": ( - "[default]\n" - "aws_access_key_id = AKID_PROFILE\n" - "aws_secret_access_key = SECRET_PROFILE\n" - "aws_session_token = TOKEN_PROFILE\n" - ) - } - ) with patch.dict(os.environ, {"AWS_REGION": "us-east-1"}, clear=True): - config = await AsyncAwsConfig.resolve( - fs=fs, - config_file_path="/fake/config", - credentials_file_path="/fake/credentials", + config = await ServiceConfig.resolve( + fs=NullFileSystem(), + aws_access_key_id="AKIAIOSFODNN7EXAMPLE", + aws_secret_access_key="wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY", ) - assert config.aws_access_key_id == "AKID_PROFILE" - assert config.aws_secret_access_key == "SECRET_PROFILE" - assert config.aws_session_token == "TOKEN_PROFILE" - assert config.source_of("aws_access_key_id") == ConfigSource.PROFILE - assert config.source_of("aws_secret_access_key") == ConfigSource.PROFILE - assert config.source_of("aws_session_token") == ConfigSource.PROFILE - @pytest.mark.asyncio - async def test_profile_token_not_used_when_env_has_key_and_secret(self): - """Even if profile has all three, env key+secret means token comes from env too.""" - fs = FakeFileSystem( - { - "/fake/credentials": ( - "[default]\n" - "aws_access_key_id = AKID_PROFILE\n" - "aws_secret_access_key = SECRET_PROFILE\n" - "aws_session_token = TOKEN_PROFILE\n" - ) - } - ) - with patch.dict( - os.environ, - { - "AWS_REGION": "us-east-1", - "AWS_ACCESS_KEY_ID": "AKID_ENV", - "AWS_SECRET_ACCESS_KEY": "SECRET_ENV", - }, - clear=True, - ): - config = await AsyncAwsConfig.resolve( - fs=fs, - config_file_path="/fake/config", - credentials_file_path="/fake/credentials", - ) - # Env wins for all three — token is None because env doesn't have it - assert config.aws_access_key_id == "AKID_ENV" - assert config.aws_secret_access_key == "SECRET_ENV" - assert config.aws_session_token is None - assert config.source_of("aws_session_token") == ConfigSource.ENV + assert config.aws_access_key_id == "AKIAIOSFODNN7EXAMPLE" + + config_repr = repr(config) + assert config_repr.startswith("ServiceConfig(") + assert "AKIAIOSFODNN7EXAMPLE" not in config_repr + assert "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY" not in config_repr + + +class TestIncodeStaticCredentialResolution: + """Credentials must be resolved from a single source — never mixed.""" @pytest.mark.asyncio async def test_no_credentials_when_nothing_set(self): @@ -876,21 +875,12 @@ async def test_no_credentials_when_nothing_set(self): @pytest.mark.asyncio async def test_partial_credential_override_raises_error(self): """Overriding only one credential raises an error.""" - fs = FakeFileSystem( - { - "/fake/credentials": ( - "[default]\n" - "aws_access_key_id = AKID_PROFILE\n" - "aws_secret_access_key = SECRET_PROFILE\n" - ) - } - ) with patch.dict(os.environ, {"AWS_REGION": "us-east-1"}, clear=True): with pytest.raises( ConfigValidationError, match="Partial credential override" ): await AsyncAwsConfig.resolve( - fs=fs, + fs=NullFileSystem(), config_file_path="/fake/config", credentials_file_path="/fake/credentials", aws_access_key_id="OVERRIDE_KEY", @@ -898,16 +888,12 @@ async def test_partial_credential_override_raises_error(self): @pytest.mark.asyncio async def test_credentials_cannot_be_overridden_after_resolution(self): - with patch.dict( - os.environ, - { - "AWS_REGION": "us-east-1", - "AWS_ACCESS_KEY_ID": "AKID", - "AWS_SECRET_ACCESS_KEY": "SECRET", - }, - clear=True, - ): - config = await AsyncAwsConfig.resolve(fs=NullFileSystem()) + with patch.dict(os.environ, {"AWS_REGION": "us-east-1"}, clear=True): + config = await AsyncAwsConfig.resolve( + fs=NullFileSystem(), + aws_access_key_id="AKID", + aws_secret_access_key="SECRET", + ) with pytest.raises( AttributeError, match="cannot be modified after resolution" ): @@ -925,84 +911,96 @@ async def test_session_token_only_override_raises_error(self): ) @pytest.mark.asyncio - async def test_env_session_token_only_falls_through_to_profile(self): - fs = FakeFileSystem( - { - "/fake/credentials": ( - "[default]\n" - "aws_access_key_id = AKID_PROFILE\n" - "aws_secret_access_key = SECRET_PROFILE\n" - "aws_session_token = TOKEN_PROFILE\n" - ) - } - ) - with patch.dict( - os.environ, - {"AWS_REGION": "us-east-1", "AWS_SESSION_TOKEN": "TOKEN_ENV"}, - clear=True, - ): + async def test_key_and_secret_auto_wires_static_resolver(self): + with patch.dict(os.environ, {"AWS_REGION": "us-east-1"}, clear=True): config = await AsyncAwsConfig.resolve( - fs=fs, - config_file_path="/fake/config", - credentials_file_path="/fake/credentials", + fs=NullFileSystem(), + aws_access_key_id="AKID", + aws_secret_access_key="SECRET", ) - # Token-only env doesn't trigger env path — all from profile - assert config.aws_access_key_id == "AKID_PROFILE" - assert config.aws_secret_access_key == "SECRET_PROFILE" - assert config.aws_session_token == "TOKEN_PROFILE" - assert config.source_of("aws_access_key_id") == ConfigSource.PROFILE + assert config.aws_access_key_id == "AKID" + assert config.aws_secret_access_key == "SECRET" + assert config.aws_credentials_identity_resolver is not None + identity = await config.aws_credentials_identity_resolver.get_identity( + properties={} + ) + assert identity.access_key_id == "AKID" + assert identity.secret_access_key == "SECRET" + + @pytest.mark.asyncio + async def test_explicit_resolver_not_overwritten(self): + custom_resolver = StaticCredentialsResolver() + with patch.dict(os.environ, {"AWS_REGION": "us-east-1"}, clear=True): + config = await AsyncAwsConfig.resolve( + fs=NullFileSystem(), + aws_access_key_id="AKID", + aws_secret_access_key="SECRET", + aws_credentials_identity_resolver=custom_resolver, + ) + assert config.aws_credentials_identity_resolver is custom_resolver + + @pytest.mark.asyncio + async def test_no_credentials_leaves_resolver_none(self): + with patch.dict(os.environ, {"AWS_REGION": "us-east-1"}, clear=True): + config = await AsyncAwsConfig.resolve(fs=NullFileSystem()) + assert config.aws_credentials_identity_resolver is None + +class TestResolveSdkUaAppId: @pytest.mark.asyncio - async def test_env_key_only_without_secret_falls_through_to_profile(self): + async def test_resolves_from_env(self): + with patch.dict(os.environ, {"AWS_SDK_UA_APP_ID": "my-app"}, clear=True): + ctx = SharedConfigContext(fs=NullFileSystem()) + result = await resolve_sdk_ua_app_id(ctx) + assert result.value == "my-app" + assert result.source == ConfigSource.ENV + + @pytest.mark.asyncio + async def test_resolves_from_profile(self): fs = FakeFileSystem( - { - "/fake/credentials": ( - "[default]\n" - "aws_access_key_id = AKID_PROFILE\n" - "aws_secret_access_key = SECRET_PROFILE\n" - ) - } + {"/fake/config": "[profile default]\nsdk_ua_app_id = profile-app\n"} ) + with patch.dict(os.environ, {}, clear=True): + ctx = SharedConfigContext(fs=fs, config_file_path="/fake/config") + result = await resolve_sdk_ua_app_id(ctx) + assert result.value == "profile-app" + assert result.source == ConfigSource.PROFILE + + @pytest.mark.asyncio + async def test_returns_unset_when_not_configured(self): + with patch.dict(os.environ, {}, clear=True): + ctx = SharedConfigContext(fs=NullFileSystem()) + result = await resolve_sdk_ua_app_id(ctx) + assert result.value is UNSET + + +class TestResolveEndpointUri: + @pytest.mark.asyncio + async def test_resolves_from_env(self): with patch.dict( - os.environ, - {"AWS_REGION": "us-east-1", "AWS_ACCESS_KEY_ID": "AKID_ENV"}, - clear=True, + os.environ, {"AWS_ENDPOINT_URL": "https://custom.endpoint"}, clear=True ): - config = await AsyncAwsConfig.resolve( - fs=fs, - config_file_path="/fake/config", - credentials_file_path="/fake/credentials", - ) - - assert config.aws_access_key_id == "AKID_PROFILE" - assert config.aws_secret_access_key == "SECRET_PROFILE" - assert config.source_of("aws_access_key_id") == ConfigSource.PROFILE + ctx = SharedConfigContext(fs=NullFileSystem()) + result = await resolve_endpoint_uri(ctx) + assert result.value == "https://custom.endpoint" + assert result.source == ConfigSource.ENV @pytest.mark.asyncio - async def test_empty_string_env_credentials_fall_through_to_profile(self): + async def test_resolves_from_profile(self): fs = FakeFileSystem( { - "/fake/credentials": ( - "[default]\n" - "aws_access_key_id = AKID_PROFILE\n" - "aws_secret_access_key = SECRET_PROFILE\n" - ) + "/fake/config": "[profile default]\nendpoint_url = https://profile.endpoint\n" } ) - with patch.dict( - os.environ, - { - "AWS_REGION": "us-east-1", - "AWS_ACCESS_KEY_ID": "", - "AWS_SECRET_ACCESS_KEY": "", - }, - clear=True, - ): - config = await AsyncAwsConfig.resolve( - fs=fs, - config_file_path="/fake/config", - credentials_file_path="/fake/credentials", - ) - assert config.aws_access_key_id == "AKID_PROFILE" - assert config.aws_secret_access_key == "SECRET_PROFILE" - assert config.source_of("aws_access_key_id") == ConfigSource.PROFILE + with patch.dict(os.environ, {}, clear=True): + ctx = SharedConfigContext(fs=fs, config_file_path="/fake/config") + result = await resolve_endpoint_uri(ctx) + assert result.value == "https://profile.endpoint" + assert result.source == ConfigSource.PROFILE + + @pytest.mark.asyncio + async def test_returns_unset_when_not_configured(self): + with patch.dict(os.environ, {}, clear=True): + ctx = SharedConfigContext(fs=NullFileSystem()) + result = await resolve_endpoint_uri(ctx) + assert result.value is UNSET