aws-lambda-powertools 3.31.0__tar.gz → 3.31.1__tar.gz

{aws_lambda_powertools-3.31.0 → aws_lambda_powertools-3.31.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: aws_lambda_powertools
-Version: 3.31.0
+Version: 3.31.1
 Summary: Powertools for AWS Lambda (Python) is a developer toolkit to implement Serverless best practices and increase developer velocity.
 License: MIT
 Keywords: aws_lambda_powertools,aws,tracing,logging,lambda,powertools,feature_flags,idempotency,middleware
@@ -29,7 +29,7 @@ Provides-Extra: tracer
 Provides-Extra: validation
 Provides-Extra: valkey
 Requires-Dist: avro (>=1.12.0,<2.0.0) ; extra == "kafka-consumer-avro"
-Requires-Dist: aws-encryption-sdk (>=3.3.1,!=4.0.0,!=4.0.1,!=4.0.2,!=4.0.3,!=4.0.4,<5.0.0) ; extra == "all" or extra == "datamasking"
+Requires-Dist: aws-encryption-sdk (>=4.0.5,<5.0.0) ; extra == "all" or extra == "datamasking"
 Requires-Dist: aws-xray-sdk (>=2.8.0,<3.0.0) ; extra == "tracer" or extra == "all"
 Requires-Dist: boto3 (>=1.34.32,<2.0.0) ; extra == "aws-sdk"
 Requires-Dist: datadog-lambda (>=8.114.0,<9.0.0) ; extra == "datadog"

{aws_lambda_powertools-3.31.0 → aws_lambda_powertools-3.31.1}/aws_lambda_powertools/event_handler/events_appsync/appsync_events.py

@@ -321,7 +321,7 @@ class AppSyncEventsResolver(Router):
                     if isinstance(ret, Exception)
                     else {"id": e.get("id"), "payload": ret}
                 )
-                for e, ret in zip(self.current_event.events, results)
+                for e, ret in zip(self.current_event.events, results, strict=True)
             ],
         )
 

{aws_lambda_powertools-3.31.0 → aws_lambda_powertools-3.31.1}/aws_lambda_powertools/event_handler/openapi/dependant.py

@@ -156,6 +156,7 @@ def get_dependant(
     call: Callable[..., Any],
     name: str | None = None,
     responses: dict[int, OpenAPIResponse] | None = None,
+    is_dependency: bool = False,
 ) -> Dependant:
     """
     Returns a dependant model for a handler function. A dependant model is a model that contains
@@ -200,6 +201,7 @@ def get_dependant(
             sub_dependant = get_dependant(
                 path=path,
                 call=depends_instance.dependency,
+                is_dependency=True,
             )
             dependant.dependencies.append(
                 DependencyParam(
@@ -229,7 +231,12 @@ def get_dependant(
         else:
             add_param_to_fields(field=param_field, dependant=dependant)
 
-    _add_return_annotation(dependant, endpoint_signature)
+    # A dependency's return value is injected directly into the handler, never serialized as a
+    # response body nor validated as an OpenAPI parameter — so building a Pydantic schema for it is
+    # both unused (no reader consumes a sub-dependency's return_param) and actively harmful: it crashes
+    # for arbitrary return types (e.g. boto3/botocore clients). Only the top-level handler needs it.
+    if not is_dependency:
+        _add_return_annotation(dependant, endpoint_signature)
     _add_extra_responses(dependant, responses)
 
     return dependant

{aws_lambda_powertools-3.31.0 → aws_lambda_powertools-3.31.1}/aws_lambda_powertools/shared/version.py

@@ -1,3 +1,3 @@
 """Exposes version constant to avoid circular dependencies."""
 
-VERSION = "3.31.0"
+VERSION = "3.31.1"

{aws_lambda_powertools-3.31.0 → aws_lambda_powertools-3.31.1}/aws_lambda_powertools/utilities/circuit_breaker_alpha/base.py

@@ -10,6 +10,8 @@ from __future__ import annotations
 
 import datetime
 import logging
+import os
+import threading
 import uuid
 from typing import TYPE_CHECKING, Any
 
@@ -37,8 +39,31 @@ _LOCAL_SUCCESSES: dict[str, int] = {}
 # recovered), so stale local failure streaks can be invalidated.
 _LAST_OBSERVED_STATE: dict[str, CircuitState] = {}
 
-# Stable per-environment identifier used to claim the half-open probe lock.
-_ENVIRONMENT_ID = uuid.uuid4().hex
+# Guards the three dicts above. Increments are read-modify-write and a threshold
+# crossing must be observed by exactly one thread, so every access goes through this
+# lock. Held only while mutating the dicts, never across persistence writes or user
+# callbacks.
+_COUNTERS_LOCK = threading.Lock()
+
+# Identifier used to claim the half-open probe lock, unique per thread so the store's
+# conditional election picks a single prober across threads as well as processes.
+_PROBE_OWNER = threading.local()
+
+
+def _probe_owner_id() -> str:
+    """
+    Return this thread's stable probe-owner identifier, minting it on first use.
+
+    A uuid in thread-local storage rather than ``threading.get_ident()``: the OS reuses
+    thread ids, and a recycled id would let an unrelated thread pass the owner check and
+    probe alongside the real owner. The pid check re-mints the id in forked children,
+    which inherit the forking thread's local storage.
+    """
+    pid = os.getpid()
+    if getattr(_PROBE_OWNER, "pid", None) != pid:
+        _PROBE_OWNER.id = f"{pid}#{uuid.uuid4().hex}"
+        _PROBE_OWNER.pid = pid
+    return _PROBE_OWNER.id
 
 
 class CircuitBreakerHandler:
@@ -111,32 +136,35 @@ class CircuitBreakerHandler:
             # If we previously observed a non-CLOSED state and the circuit is now back to
             # CLOSED, another environment completed the recovery cycle. Reset local counters
             # so a stale partial failure streak doesn't immediately re-trip the circuit.
-            prev = _LAST_OBSERVED_STATE.get(self.name)
-            if prev is not None and prev != CircuitState.CLOSED:
-                _LOCAL_FAILURES[self.name] = 0
-            _LAST_OBSERVED_STATE[self.name] = CircuitState.CLOSED
+            with _COUNTERS_LOCK:
+                prev = _LAST_OBSERVED_STATE.get(self.name)
+                if prev is not None and prev != CircuitState.CLOSED:
+                    _LOCAL_FAILURES[self.name] = 0
+                _LAST_OBSERVED_STATE[self.name] = CircuitState.CLOSED
             return self._call_closed()
 
         if record.state == CircuitState.OPEN:
-            _LAST_OBSERVED_STATE[self.name] = CircuitState.OPEN
+            with _COUNTERS_LOCK:
+                _LAST_OBSERVED_STATE[self.name] = CircuitState.OPEN
             # ``opened_at`` may legitimately be 0 (epoch); treat only None as missing.
             opened_at = record.opened_at if record.opened_at is not None else self._now()
             if self._now() >= opened_at + self.config.recovery_timeout:
                 # Recovery window elapsed: try to become the single prober.
-                if self.persistence_store.try_acquire_half_open(self.name, _ENVIRONMENT_ID, opened_at):
+                if self.persistence_store.try_acquire_half_open(self.name, _probe_owner_id(), opened_at):
                     self._notify(CircuitState.OPEN, CircuitState.HALF_OPEN, opened_at=opened_at)
                     return self._call_probe()
             return self._open_response(record.to_circuit_info())
 
-        # HALF_OPEN: only the environment that owns the probe lock runs.
-        _LAST_OBSERVED_STATE[self.name] = CircuitState.HALF_OPEN
-        if record.half_open_owner == _ENVIRONMENT_ID:
+        # HALF_OPEN: only the thread that owns the probe lock runs.
+        with _COUNTERS_LOCK:
+            _LAST_OBSERVED_STATE[self.name] = CircuitState.HALF_OPEN
+        if record.half_open_owner == _probe_owner_id():
             return self._call_probe()
 
         # If the probe lease has expired (owner recycled mid-probe), take over.
         if record.probe_lease_expiry is not None and self._now() >= record.probe_lease_expiry:
             logger.debug("Circuit '%s' probe lease expired; attempting takeover.", self.name)
-            if self.persistence_store.try_acquire_half_open(self.name, _ENVIRONMENT_ID, record.opened_at or 0):
+            if self.persistence_store.try_acquire_half_open(self.name, _probe_owner_id(), record.opened_at or 0):
                 return self._call_probe()
 
         return self._open_response(record.to_circuit_info())
@@ -148,9 +176,14 @@ class CircuitBreakerHandler:
         except Exception as exc:
             if not self.config.counts_as_failure(exc):
                 raise
-            failures = _LOCAL_FAILURES.get(self.name, 0) + 1
-            _LOCAL_FAILURES[self.name] = failures
-            if failures >= self.config.failure_threshold:
+            # Increment and reset atomically so exactly one thread observes the threshold
+            # crossing; racing threads would otherwise lose increments (tripping late) or
+            # each persist the same transition.
+            with _COUNTERS_LOCK:
+                failures = _LOCAL_FAILURES.get(self.name, 0) + 1
+                tripped = failures >= self.config.failure_threshold
+                _LOCAL_FAILURES[self.name] = 0 if tripped else failures
+            if tripped:
                 logger.debug("Circuit '%s' tripping CLOSED to OPEN after %d failures.", self.name, failures)
                 opened_at = self._now()
                 self._safe_persist(
@@ -159,11 +192,11 @@ class CircuitBreakerHandler:
                     failure_count=failures,
                     opened_at=opened_at,
                 )
-                _LOCAL_FAILURES[self.name] = 0
                 self._notify(CircuitState.CLOSED, CircuitState.OPEN, opened_at=opened_at)
             raise
         else:
-            _LOCAL_FAILURES[self.name] = 0
+            with _COUNTERS_LOCK:
+                _LOCAL_FAILURES[self.name] = 0
             return result
 
     def _call_probe(self) -> Any:
@@ -176,17 +209,20 @@ class CircuitBreakerHandler:
             logger.debug("Circuit '%s' probe failed; reopening.", self.name)
             opened_at = self._now()
             self._safe_persist(self.persistence_store.save_reopen, self.name, opened_at=opened_at)
-            _LOCAL_SUCCESSES[self.name] = 0
+            with _COUNTERS_LOCK:
+                _LOCAL_SUCCESSES[self.name] = 0
             self._notify(CircuitState.HALF_OPEN, CircuitState.OPEN, opened_at=opened_at)
             raise
         else:
-            successes = _LOCAL_SUCCESSES.get(self.name, 0) + 1
-            _LOCAL_SUCCESSES[self.name] = successes
-            if successes >= self.config.success_threshold:
+            with _COUNTERS_LOCK:
+                successes = _LOCAL_SUCCESSES.get(self.name, 0) + 1
+                closed = successes >= self.config.success_threshold
+                _LOCAL_SUCCESSES[self.name] = 0 if closed else successes
+                if closed:
+                    _LOCAL_FAILURES[self.name] = 0
+            if closed:
                 logger.debug("Circuit '%s' closing after %d probe successes.", self.name, successes)
                 self._safe_persist(self.persistence_store.save_closed, self.name)
-                _LOCAL_SUCCESSES[self.name] = 0
-                _LOCAL_FAILURES[self.name] = 0
                 self._notify(CircuitState.HALF_OPEN, CircuitState.CLOSED)
             return result
 

{aws_lambda_powertools-3.31.0 → aws_lambda_powertools-3.31.1}/aws_lambda_powertools/utilities/circuit_breaker_alpha/config.py

@@ -4,8 +4,13 @@ Configuration for the Circuit Breaker utility.
 
 from __future__ import annotations
 
+from typing import TYPE_CHECKING
+
 from aws_lambda_powertools.utilities.circuit_breaker_alpha.exceptions import CircuitBreakerConfigError
 
+if TYPE_CHECKING:
+    from collections.abc import Iterable
+
 
 class CircuitBreakerConfig:
     """
@@ -24,12 +29,14 @@ class CircuitBreakerConfig:
     success_threshold : int
         Number of *consecutive* probe successes required to close a half-open circuit.
         Defaults to 3.
-    handled_exceptions : tuple[type[Exception], ...] | None
+    handled_exceptions : type[Exception] | Iterable[type[Exception]] | None
         Allowlist: only these exception types count as failures; anything else
-        propagates without affecting the circuit. Mutually exclusive with
+        propagates without affecting the circuit. Accepts a single exception type or
+        an iterable of them (normalized to a tuple). Mutually exclusive with
         ``ignored_exceptions``. Defaults to ``None`` (treated as ``(Exception,)``).
-    ignored_exceptions : tuple[type[Exception], ...] | None
-        Denylist: every exception counts as a failure *except* these. Mutually
+    ignored_exceptions : type[Exception] | Iterable[type[Exception]] | None
+        Denylist: every exception counts as a failure *except* these. Accepts a single
+        exception type or an iterable of them (normalized to a tuple). Mutually
         exclusive with ``handled_exceptions``. Defaults to ``None``.
     local_cache_max_age : int
         Seconds a circuit's state is cached in the execution environment before a
@@ -38,8 +45,9 @@ class CircuitBreakerConfig:
     Raises
     ------
     CircuitBreakerConfigError
-        If both ``handled_exceptions`` and ``ignored_exceptions`` are provided, or a
-        numeric tunable is not a positive integer.
+        If both ``handled_exceptions`` and ``ignored_exceptions`` are provided, a
+        numeric tunable is not a positive integer, or an exception allowlist/denylist
+        is empty or contains a value that is not an exception type.
 
     Example
     -------
@@ -57,10 +65,16 @@ class CircuitBreakerConfig:
         failure_threshold: int = 5,
         recovery_timeout: int = 30,
         success_threshold: int = 3,
-        handled_exceptions: tuple[type[Exception], ...] | None = None,
-        ignored_exceptions: tuple[type[Exception], ...] | None = None,
+        handled_exceptions: type[Exception] | Iterable[type[Exception]] | None = None,
+        ignored_exceptions: type[Exception] | Iterable[type[Exception]] | None = None,
         local_cache_max_age: int = 5,
     ):
+        # Normalize first: a single exception type or any iterable becomes a tuple, and a
+        # bad value fails here (at construction) rather than as a cryptic isinstance
+        # TypeError later, the first time the circuit evaluates a failure.
+        handled_exceptions = self._normalize_exceptions(handled_exceptions, "handled_exceptions")
+        ignored_exceptions = self._normalize_exceptions(ignored_exceptions, "ignored_exceptions")
+
         self._validate(
             failure_threshold=failure_threshold,
             recovery_timeout=recovery_timeout,
@@ -105,6 +119,48 @@ class CircuitBreakerConfig:
                 f"local_cache_max_age must be a non-negative integer, got {local_cache_max_age!r}.",
             )
 
+    @classmethod
+    def _normalize_exceptions(
+        cls,
+        value: type[Exception] | Iterable[type[Exception]] | None,
+        field: str,
+    ) -> tuple[type[Exception], ...] | None:
+        """Coerce a single exception type or an iterable of them into a validated, non-empty tuple.
+
+        Runs at construction so a bad value fails immediately with a clear error, rather
+        than as a cryptic ``isinstance`` ``TypeError`` from ``counts_as_failure`` the
+        first time the circuit evaluates a failure (i.e. only once the dependency is
+        already unhealthy).
+        """
+        if value is None:
+            return None
+
+        invalid = f"{field} must be an exception type or an iterable of exception types, got {value!r}."
+        # A str is iterable; reject it rather than iterate it as a sequence of characters.
+        if isinstance(value, str):
+            raise CircuitBreakerConfigError(invalid)
+
+        if isinstance(value, type):
+            # ty (unlike mypy) does not narrow the union here, so it needs the ignore.
+            exceptions: tuple[type[Exception], ...] = (value,)  # ty: ignore[invalid-assignment]
+        else:
+            try:
+                exceptions = tuple(value)
+            except TypeError:
+                raise CircuitBreakerConfigError(invalid) from None
+
+        cls._validate_exception_types(exceptions, field)
+        return exceptions
+
+    @staticmethod
+    def _validate_exception_types(exceptions: tuple[type[Exception], ...], field: str) -> None:
+        """Require a non-empty tuple whose every element is an exception type."""
+        if not exceptions:
+            raise CircuitBreakerConfigError(f"{field} must contain at least one exception type.")
+        for exception in exceptions:
+            if not (isinstance(exception, type) and issubclass(exception, Exception)):
+                raise CircuitBreakerConfigError(f"{field} must contain only exception types, got {exception!r}.")
+
     def counts_as_failure(self, exception: Exception) -> bool:
         """
         Decide whether an exception raised by the protected call counts as a circuit failure.

{aws_lambda_powertools-3.31.0 → aws_lambda_powertools-3.31.1}/aws_lambda_powertools/utilities/circuit_breaker_alpha/persistence/base.py

@@ -10,8 +10,9 @@ from __future__ import annotations
 
 import datetime
 import logging
+import threading
 from abc import ABC, abstractmethod
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, NamedTuple
 
 from aws_lambda_powertools.shared.cache_dict import LRUDict
 from aws_lambda_powertools.utilities.circuit_breaker_alpha.persistence.record import CircuitStateRecord
@@ -32,6 +33,17 @@ LOCAL_CACHE_MAX_ITEMS = 1024
 PERSISTED_STATE_TTL_BUFFER = 3600
 
 
+class _CircuitSettings(NamedTuple):
+    """Per-circuit tunables the layer captures from :meth:`configure`."""
+
+    local_cache_max_age: int
+    recovery_timeout: int
+
+
+# Fallback for direct layer use before configure() has run; mirrors CircuitBreakerConfig defaults.
+_DEFAULT_SETTINGS = _CircuitSettings(local_cache_max_age=5, recovery_timeout=30)
+
+
 class CircuitBreakerExistingLockError(Exception):
     """Internal signal that a conditional half-open probe write lost the race."""
 
@@ -50,19 +62,26 @@ class CircuitBreakerPersistenceLayer(ABC):
 
     def __init__(self) -> None:
         """Initialize defaults; real configuration happens in :meth:`configure`."""
-        self.circuit_name: str = ""
-        self.local_cache_max_age: int = 5
-        self.recovery_timeout: int = 30
+        # Per-circuit tunables, keyed by circuit name. One persistence instance is shared
+        # by every circuit (and thread) using the same store, so these must never live in
+        # plain instance attributes: circuits with different configs would stamp each
+        # other's TTL and probe lease. A plain dict, not an LRUDict: evicting a live
+        # circuit's settings would silently swap in the defaults (wrong lease and TTL),
+        # whereas evicting a cache entry below only costs a store re-read.
+        self._settings: dict[str, _CircuitSettings] = {}
         # Maps circuit name -> the unix timestamp the locally cached record goes stale.
         # Kept separate from the record's durable ``expiry_timestamp`` (the store TTL) so
         # the short in-memory freshness window is never mistaken for the long store TTL.
         self._cache: LRUDict = LRUDict(max_items=LOCAL_CACHE_MAX_ITEMS)
+        # One lock for both maps: LRUDict reorders entries even on reads, so unguarded
+        # concurrent access can corrupt it or raise.
+        self._lock = threading.Lock()
 
     def configure(self, config: CircuitBreakerConfig, circuit_name: str) -> None:
         """
-        Bind the layer to a circuit and its configuration.
+        Bind a circuit's configuration to the layer.
 
-        Called once per invocation by the handler; the assignments are cheap and the
+        Called once per invocation by the handler; the assignment is cheap and the
         same persistence instance is reused across invocations within an environment.
 
         Parameters
@@ -70,43 +89,59 @@ class CircuitBreakerPersistenceLayer(ABC):
         config : CircuitBreakerConfig
             Configuration providing the local cache TTL and recovery timeout.
         circuit_name : str
-            The circuit this layer instance serves.
+            The circuit these settings apply to.
         """
-        self.circuit_name = circuit_name
-        self.local_cache_max_age = config.local_cache_max_age
-        self.recovery_timeout = config.recovery_timeout
+        with self._lock:
+            self._settings[circuit_name] = _CircuitSettings(
+                local_cache_max_age=config.local_cache_max_age,
+                recovery_timeout=config.recovery_timeout,
+            )
+
+    def _settings_for(self, name: str) -> _CircuitSettings:
+        """Return a circuit's configured settings, or the defaults if never configured."""
+        with self._lock:
+            return self._settings.get(name, _DEFAULT_SETTINGS)
 
     # ------------------------------------------------------------------ cache
 
     def _cache_key(self, name: str) -> str:
         return name
 
-    def _durable_ttl(self) -> int:
+    def _durable_ttl(self, name: str) -> int:
         """
         Compute the store TTL stamped on a persisted record.
 
         Sized to outlive a full recovery window so a live circuit is never reaped
         mid-cycle, while an abandoned circuit (no further writes) self-cleans soon after.
         """
-        return int(datetime.datetime.now().timestamp()) + self.recovery_timeout + PERSISTED_STATE_TTL_BUFFER
+        now = int(datetime.datetime.now().timestamp())
+        return now + self._settings_for(name).recovery_timeout + PERSISTED_STATE_TTL_BUFFER
 
     def _save_to_cache(self, record: CircuitStateRecord) -> None:
         """Cache a record locally with a short in-memory freshness window."""
-        local_expiry = int(datetime.datetime.now().timestamp()) + self.local_cache_max_age
-        self._cache[self._cache_key(record.name)] = (local_expiry, record)
+        local_expiry = int(datetime.datetime.now().timestamp()) + self._settings_for(record.name).local_cache_max_age
+        with self._lock:
+            self._cache[self._cache_key(record.name)] = (local_expiry, record)
 
     def _retrieve_from_cache(self, name: str) -> CircuitStateRecord | None:
         """Return a cached record if present and still within its local freshness window."""
-        cached = self._cache.get(self._cache_key(name))
-        if cached is None:
-            return None
-
-        local_expiry, record = cached
-        if int(datetime.datetime.now().timestamp()) >= local_expiry:
-            del self._cache[self._cache_key(name)]
-            return None
-
-        return record
+        with self._lock:
+            cached = self._cache.get(self._cache_key(name))
+            if cached is None:
+                return None
+
+            local_expiry, record = cached
+            if int(datetime.datetime.now().timestamp()) >= local_expiry:
+                # Guarded del, not pop: on Python 3.10 OrderedDict.pop re-enters the
+                # subclass __getitem__ after detaching the node, so LRUDict.pop raises
+                # KeyError for a *present* key and corrupts the dict (fixed in 3.11).
+                try:
+                    del self._cache[self._cache_key(name)]
+                except KeyError:
+                    pass
+                return None
+
+            return record
 
     # ------------------------------------------------------------- public API
 
@@ -175,19 +210,19 @@ class CircuitBreakerPersistenceLayer(ABC):
             state=CircuitState.OPEN,
             failure_count=failure_count,
             opened_at=opened_at,
-            expiry_timestamp=self._durable_ttl(),
+            expiry_timestamp=self._durable_ttl(name),
         )
         self._put_record(record)
         self._save_to_cache(record)
 
     def try_acquire_half_open(self, name: str, owner: str, opened_at: int) -> bool:
         """
-        Atomically elect a single environment to run the half-open probe.
+        Atomically elect a single worker to run the half-open probe.
 
         The conditional write succeeds only when the circuit is OPEN with no existing
         lock owner AND the ``opened_at`` matches what the caller observed (guards against
         stale eventually-consistent reads). A lease expiry is stamped so that if the
-        winning environment is recycled before completing the probe, others can take over
+        winning worker is recycled before completing the probe, others can take over
         once the lease lapses.
 
         Parameters
@@ -195,25 +230,26 @@ class CircuitBreakerPersistenceLayer(ABC):
         name : str
             Circuit name.
         owner : str
-            Identifier of the environment attempting the probe.
+            Identifier of the worker (one thread in one execution environment)
+            attempting the probe.
         opened_at : int
             The ``opened_at`` the caller observed, kept stable across the transition.
 
         Returns
         -------
         bool
-            ``True`` if this environment won the probe lock, ``False`` if another
-            environment already holds it.
+            ``True`` if this worker won the probe lock, ``False`` if another
+            worker already holds it.
         """
         # Lease = recovery_timeout gives the probe a full cycle to complete.
-        probe_lease_expiry = int(datetime.datetime.now().timestamp()) + self.recovery_timeout
+        probe_lease_expiry = int(datetime.datetime.now().timestamp()) + self._settings_for(name).recovery_timeout
         record = CircuitStateRecord(
             name=name,
             state=CircuitState.HALF_OPEN,
             opened_at=opened_at,
             half_open_owner=owner,
             probe_lease_expiry=probe_lease_expiry,
-            expiry_timestamp=self._durable_ttl(),
+            expiry_timestamp=self._durable_ttl(name),
         )
         try:
             self._put_record(record, condition="half_open", expected_opened_at=opened_at)
@@ -228,7 +264,7 @@ class CircuitBreakerPersistenceLayer(ABC):
             name=name,
             state=CircuitState.CLOSED,
             failure_count=0,
-            expiry_timestamp=self._durable_ttl(),
+            expiry_timestamp=self._durable_ttl(name),
         )
         self._update_record(record)
         self._save_to_cache(record)
@@ -245,7 +281,7 @@ class CircuitBreakerPersistenceLayer(ABC):
             name=name,
             state=CircuitState.OPEN,
             opened_at=opened_at,
-            expiry_timestamp=self._durable_ttl(),
+            expiry_timestamp=self._durable_ttl(name),
         )
         self._update_record(record)
         self._save_to_cache(record)

{aws_lambda_powertools-3.31.0 → aws_lambda_powertools-3.31.1}/aws_lambda_powertools/utilities/circuit_breaker_alpha/persistence/dynamodb.py

@@ -6,19 +6,23 @@ from __future__ import annotations
 
 import datetime
 import logging
+import os
+import warnings
 from typing import TYPE_CHECKING
 
 import boto3
 from boto3.dynamodb.types import TypeDeserializer
 from botocore.exceptions import ClientError
 
-from aws_lambda_powertools.shared import user_agent
+from aws_lambda_powertools.shared import constants, user_agent
+from aws_lambda_powertools.utilities.circuit_breaker_alpha.exceptions import CircuitBreakerConfigError
 from aws_lambda_powertools.utilities.circuit_breaker_alpha.persistence.base import (
     CircuitBreakerExistingLockError,
     CircuitBreakerPersistenceLayer,
 )
 from aws_lambda_powertools.utilities.circuit_breaker_alpha.persistence.record import CircuitStateRecord
 from aws_lambda_powertools.utilities.circuit_breaker_alpha.states import CircuitState
+from aws_lambda_powertools.warnings import PowertoolsUserWarning
 
 if TYPE_CHECKING:
     from botocore.config import Config
@@ -41,6 +45,13 @@ class CircuitBreakerDynamoDBPersistence(CircuitBreakerPersistenceLayer):
         Name of the DynamoDB table that stores circuit state.
     key_attr : str
         Partition key attribute holding the circuit name. Defaults to ``"id"``.
+    static_pk_value : str, optional
+        Partition key value used when ``sort_key_attr`` is set, so the circuit name
+        moves to the sort key. Defaults to ``"circuit_breaker#<function-name>"``.
+    sort_key_attr : str, optional
+        Sort key attribute holding the circuit name. When set, the table is treated as
+        composite: ``static_pk_value`` is written to the partition key and the circuit
+        name to the sort key. Omit it for the default partition-key-only behavior.
     state_attr : str
         Attribute holding the circuit state. Defaults to ``"state"``.
     failure_count_attr : str
@@ -75,6 +86,8 @@ class CircuitBreakerDynamoDBPersistence(CircuitBreakerPersistenceLayer):
         self,
         table_name: str,
         key_attr: str = "id",
+        static_pk_value: str | None = None,
+        sort_key_attr: str | None = None,
         state_attr: str = "state",
         failure_count_attr: str = "failure_count",
         opened_at_attr: str = "opened_at",
@@ -92,8 +105,25 @@ class CircuitBreakerDynamoDBPersistence(CircuitBreakerPersistenceLayer):
 
         user_agent.register_feature_to_client(client=self.client, feature="circuit_breaker")
 
+        if sort_key_attr == key_attr:
+            raise CircuitBreakerConfigError(
+                f"key_attr [{key_attr}] and sort_key_attr [{sort_key_attr}] cannot be the same!",
+            )
+
+        if static_pk_value is not None and sort_key_attr is None:
+            warnings.warn(
+                "static_pk_value is ignored unless sort_key_attr is also set.",
+                category=PowertoolsUserWarning,
+                stacklevel=2,
+            )
+
+        if static_pk_value is None:
+            static_pk_value = f"circuit_breaker#{os.getenv(constants.LAMBDA_FUNCTION_NAME_ENV, '')}"
+
         self.table_name = table_name
         self.key_attr = key_attr
+        self.static_pk_value = static_pk_value
+        self.sort_key_attr = sort_key_attr
         self.state_attr = state_attr
         self.failure_count_attr = failure_count_attr
         self.opened_at_attr = opened_at_attr
@@ -111,7 +141,7 @@ class CircuitBreakerDynamoDBPersistence(CircuitBreakerPersistenceLayer):
         opened_at = data.get(self.opened_at_attr)
         probe_lease_expiry = data.get(self.probe_lease_expiry_attr)
         return CircuitStateRecord(
-            name=data[self.key_attr],
+            name=data[self.sort_key_attr] if self.sort_key_attr else data[self.key_attr],
             state=CircuitState(data[self.state_attr]),
             failure_count=int(data.get(self.failure_count_attr, 0)),
             opened_at=int(opened_at) if opened_at is not None else None,
@@ -130,10 +160,21 @@ class CircuitBreakerDynamoDBPersistence(CircuitBreakerPersistenceLayer):
         """Wrap a value as a DynamoDB string attribute."""
         return {"S": value}
 
+    def _get_key(self, name: str) -> dict:
+        """Build the simple or composite primary key for a circuit.
+
+        When ``sort_key_attr`` is set the key is composite: ``static_pk_value`` in the
+        partition and the circuit name in the sort key. Otherwise the name is the
+        partition key.
+        """
+        if self.sort_key_attr:
+            return {self.key_attr: self._s(self.static_pk_value), self.sort_key_attr: self._s(name)}
+        return {self.key_attr: self._s(name)}
+
     def _record_to_item(self, record: CircuitStateRecord) -> dict:
         """Translate a :class:`CircuitStateRecord` into a DynamoDB item."""
         item: dict = {
-            self.key_attr: self._s(record.name),
+            **self._get_key(record.name),
             self.state_attr: self._s(str(record.state)),
             self.failure_count_attr: self._n(record.failure_count),
         }
@@ -152,7 +193,7 @@ class CircuitBreakerDynamoDBPersistence(CircuitBreakerPersistenceLayer):
         # and halves the read cost on the hot path.
         response = self.client.get_item(
             TableName=self.table_name,
-            Key={self.key_attr: self._s(name)},
+            Key=self._get_key(name),
             ConsistentRead=False,
         )
         item = response.get("Item")
@@ -242,7 +283,7 @@ class CircuitBreakerDynamoDBPersistence(CircuitBreakerPersistenceLayer):
 
         return {
             "TableName": self.table_name,
-            "Key": {self.key_attr: self._s(record.name)},
+            "Key": self._get_key(record.name),
             "UpdateExpression": update_expression,
             "ExpressionAttributeNames": expression_attr_names,
             "ExpressionAttributeValues": expression_attr_values,