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 1cac89621..95c6e0d5a 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 @@ -83,6 +83,8 @@ private void generateService(PythonWriter writer) { } } + writer.addDependency(SmithyPythonDependency.SMITHY_CORE); + writer.addImport("smithy_core.retries", "RetryStrategyResolver"); writer.write(""" def __init__(self, config: $1T | None = None, plugins: list[$2T] | None = None): self._config = config or $1T() @@ -95,6 +97,8 @@ def __init__(self, config: $1T | None = None, plugins: list[$2T] | None = None): for plugin in client_plugins: plugin(self._config) + + self._retry_strategy_resolver = RetryStrategyResolver() """, configSymbol, pluginSymbol, writer.consumer(w -> writeDefaultPlugins(w, defaultPlugins))); var topDownIndex = TopDownIndex.of(model); @@ -187,6 +191,8 @@ private void writeSharedOperationInit(PythonWriter writer, OperationShape operat writer.addImport("smithy_core.types", "TypedProperties"); writer.addImport("smithy_core.aio.client", "RequestPipeline"); writer.addImport("smithy_core.exceptions", "ExpectationNotMetError"); + writer.addImport("smithy_core.retries", "RetryStrategyOptions"); + writer.addImport("smithy_core.interfaces.retries", "RetryStrategy"); writer.addStdlibImport("copy", "deepcopy"); writer.write(""" @@ -200,6 +206,24 @@ private void writeSharedOperationInit(PythonWriter writer, OperationShape operat plugin(config) if config.protocol is None or config.transport is None: raise ExpectationNotMetError("protocol and transport MUST be set on the config to make calls.") + + # Resolve retry strategy from config + if isinstance(config.retry_strategy, RetryStrategy): + retry_strategy = config.retry_strategy + elif isinstance(config.retry_strategy, RetryStrategyOptions): + retry_strategy = await self._retry_strategy_resolver.resolve_retry_strategy( + options=config.retry_strategy + ) + elif config.retry_strategy is None: + retry_strategy = await self._retry_strategy_resolver.resolve_retry_strategy( + options=RetryStrategyOptions() + ) + else: + raise TypeError( + f"retry_strategy must be RetryStrategy, RetryStrategyOptions, or None, " + f"got {type(config.retry_strategy).__name__}" + ) + pipeline = RequestPipeline( protocol=config.protocol, transport=config.transport @@ -212,7 +236,7 @@ raise ExpectationNotMetError("protocol and transport MUST be set on the config t auth_scheme_resolver=config.auth_scheme_resolver, supported_auth_schemes=config.auth_schemes, endpoint_resolver=config.endpoint_resolver, - retry_strategy=config.retry_strategy, + retry_strategy=retry_strategy, ) """, writer.consumer(w -> writeDefaultPlugins(w, defaultPlugins))); 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 45b6324d7..2b788e47f 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 @@ -55,17 +55,20 @@ public final class ConfigGenerator implements Runnable { ConfigProperty.builder() .name("retry_strategy") .type(Symbol.builder() - .name("RetryStrategy") - .namespace("smithy_core.interfaces.retries", ".") - .addDependency(SmithyPythonDependency.SMITHY_CORE) + .name("RetryStrategy | RetryStrategyOptions") + .addReference(Symbol.builder() + .name("RetryStrategy") + .namespace("smithy_core.interfaces.retries", ".") + .addDependency(SmithyPythonDependency.SMITHY_CORE) + .build()) + .addReference(Symbol.builder() + .name("RetryStrategyOptions") + .namespace("smithy_core.retries", ".") + .addDependency(SmithyPythonDependency.SMITHY_CORE) + .build()) .build()) - .documentation("The retry strategy for issuing retry tokens and computing retry delays.") - .nullable(false) - .initialize(writer -> { - writer.addDependency(SmithyPythonDependency.SMITHY_CORE); - writer.addImport("smithy_core.retries", "SimpleRetryStrategy"); - writer.write("self.retry_strategy = retry_strategy or SimpleRetryStrategy()"); - }) + .documentation( + "The retry strategy or options for configuring retry behavior. Can be either a fully configured RetryStrategy or RetryStrategyOptions to create one.") .build(), ConfigProperty.builder() .name("endpoint_uri") @@ -379,7 +382,7 @@ private void writeInitParams(PythonWriter writer, Collection pro } private void documentProperties(PythonWriter writer, Collection properties) { - writer.writeDocs(() ->{ + writer.writeDocs(() -> { var iter = properties.iterator(); writer.write("\nConstructor.\n"); while (iter.hasNext()) { diff --git a/packages/smithy-core/src/smithy_core/interfaces/retries.py b/packages/smithy-core/src/smithy_core/interfaces/retries.py index a5c9d428b..19e96362a 100644 --- a/packages/smithy-core/src/smithy_core/interfaces/retries.py +++ b/packages/smithy-core/src/smithy_core/interfaces/retries.py @@ -52,6 +52,7 @@ class RetryToken(Protocol): """Delay in seconds to wait before the retry attempt.""" +@runtime_checkable class RetryStrategy(Protocol): """Issuer of :py:class:`RetryToken`s.""" diff --git a/packages/smithy-core/src/smithy_core/retries.py b/packages/smithy-core/src/smithy_core/retries.py index 06bf6f988..78d139d90 100644 --- a/packages/smithy-core/src/smithy_core/retries.py +++ b/packages/smithy-core/src/smithy_core/retries.py @@ -4,9 +4,52 @@ from collections.abc import Callable from dataclasses import dataclass from enum import Enum +from functools import lru_cache +from typing import Literal from .exceptions import RetryError from .interfaces import retries as retries_interface +from .interfaces.retries import RetryStrategy + +RetryStrategyType = Literal["simple"] + + +@dataclass(kw_only=True, frozen=True) +class RetryStrategyOptions: + """Options for configuring retry behavior.""" + + retry_mode: RetryStrategyType = "simple" + """The retry mode to use.""" + + max_attempts: int = 3 + """Maximum number of attempts (initial attempt plus retries).""" + + +class RetryStrategyResolver: + """Retry strategy resolver that caches retry strategies based on configuration options. + + This resolver caches retry strategy instances based on their configuration to reuse existing + instances of RetryStrategy with the same settings. Uses LRU cache for thread-safe caching. + """ + + async def resolve_retry_strategy( + self, *, options: RetryStrategyOptions + ) -> RetryStrategy: + """Resolve a retry strategy from the provided options, using cache when possible. + + :param options: The retry strategy options to use for creating the strategy. + """ + return self._create_retry_strategy(options.retry_mode, options.max_attempts) + + @lru_cache + def _create_retry_strategy( + self, retry_mode: RetryStrategyType, max_attempts: int + ) -> RetryStrategy: + match retry_mode: + case "simple": + return SimpleRetryStrategy(max_attempts=max_attempts) + case _: + raise ValueError(f"Unknown retry mode: {retry_mode}") class ExponentialBackoffJitterType(Enum): diff --git a/packages/smithy-core/tests/unit/test_retries.py b/packages/smithy-core/tests/unit/test_retries.py index 0b3c23be4..fd87a789b 100644 --- a/packages/smithy-core/tests/unit/test_retries.py +++ b/packages/smithy-core/tests/unit/test_retries.py @@ -1,10 +1,16 @@ # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. # SPDX-License-Identifier: Apache-2.0 - import pytest from smithy_core.exceptions import CallError, RetryError -from smithy_core.retries import ExponentialBackoffJitterType as EBJT -from smithy_core.retries import ExponentialRetryBackoffStrategy, SimpleRetryStrategy +from smithy_core.retries import ( + ExponentialBackoffJitterType as EBJT, +) +from smithy_core.retries import ( + ExponentialRetryBackoffStrategy, + RetryStrategyOptions, + RetryStrategyResolver, + SimpleRetryStrategy, +) @pytest.mark.parametrize( @@ -100,3 +106,36 @@ def test_simple_retry_does_not_retry_unsafe() -> None: token = strategy.acquire_initial_retry_token() with pytest.raises(RetryError): strategy.refresh_retry_token_for_retry(token_to_renew=token, error=error) + + +async def test_caching_retry_strategy_default_resolution() -> None: + resolver = RetryStrategyResolver() + options = RetryStrategyOptions() + + strategy = await resolver.resolve_retry_strategy(options=options) + + assert isinstance(strategy, SimpleRetryStrategy) + assert strategy.max_attempts == 3 + + +async def test_caching_retry_strategy_resolver_creates_strategies_by_options() -> None: + resolver = RetryStrategyResolver() + + options1 = RetryStrategyOptions(max_attempts=3) + options2 = RetryStrategyOptions(max_attempts=5) + + strategy1 = await resolver.resolve_retry_strategy(options=options1) + strategy2 = await resolver.resolve_retry_strategy(options=options2) + + assert strategy1.max_attempts == 3 + assert strategy2.max_attempts == 5 + + +async def test_caching_retry_strategy_resolver_caches_strategies() -> None: + resolver = RetryStrategyResolver() + + options = RetryStrategyOptions(max_attempts=5) + strategy1 = await resolver.resolve_retry_strategy(options=options) + strategy2 = await resolver.resolve_retry_strategy(options=options) + + assert strategy1 is strategy2