Skip to content
Open
Show file tree
Hide file tree
Changes from 3 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 3 additions & 1 deletion src/envars/aws_cloudformation.py
Original file line number Diff line number Diff line change
@@ -1,9 +1,11 @@
import boto3

from .aws_config import AWS_CLIENT_CONFIG


class CloudFormationExports:
def __init__(self, region_name: str | None = None):
self.client = boto3.client("cloudformation", region_name=region_name)
self.client = boto3.client("cloudformation", region_name=region_name, config=AWS_CLIENT_CONFIG)
self._exports_cache: dict[str, str] | None = None

def _populate_exports_cache(self):
Expand Down
45 changes: 45 additions & 0 deletions src/envars/aws_config.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,45 @@
"""Shared botocore client configuration for envars' AWS calls.

Without an explicit ``Config``, ``boto3.client(...)`` inherits the botocore defaults:
``connect_timeout=60s``, ``read_timeout=60s`` and legacy retries (up to 5 attempts).
A single stalled STS/KMS/SSM/CloudFormation endpoint can therefore hang for up to
``5 x 60 = ~300s``, and a full resolve makes several such calls (STS for location
auto-detect, one KMS decrypt per secret, plus any ``parameter_store:`` /
``cloudformation_export:`` lookups). These bounds turn a silent multi-minute hang into
a fast, legible failure while still tolerating a transient blip.

Override per-environment via the ``ENVARS_AWS_*`` variables if the defaults are too tight.
"""

import os

from botocore.config import Config


def _int_env(name: str, default: int) -> int:
"""Reads a positive int from the environment, falling back to ``default``."""
try:
value = int(os.environ[name])
except (KeyError, ValueError):
return default
return value if value > 0 else default


def _build_config() -> Config:
"""Builds the bounded AWS client config, reading ``ENVARS_AWS_*`` overrides at call time.

Comment thread
kthhrv marked this conversation as resolved.
Outdated
Bounded so a stalled endpoint fails in seconds, not minutes. botocore treats
``retries.max_attempts`` as the RETRY count, so ``max_attempts=2`` resolves to 3 total
attempts (1 initial + 2 retries, i.e. ``total_max_attempts=3``); with ``read_timeout=5s``
the worst case is ~17s, versus botocore's default of ``read_timeout(60) x legacy 5
attempts = ~300s``. A fully-unreachable endpoint fails the whole resolve regardless, so we
fail fast rather than wait it out.
"""
return Config(
connect_timeout=_int_env("ENVARS_AWS_CONNECT_TIMEOUT", 3),
read_timeout=_int_env("ENVARS_AWS_READ_TIMEOUT", 5),
retries={"max_attempts": _int_env("ENVARS_AWS_MAX_ATTEMPTS", 2), "mode": "standard"},
)


AWS_CLIENT_CONFIG = _build_config()
4 changes: 3 additions & 1 deletion src/envars/aws_kms.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,13 +3,15 @@
import boto3
from botocore.exceptions import ClientError

from .aws_config import AWS_CLIENT_CONFIG


class AWSKMSAgent:
"""A class to handle AWS KMS operations."""

def __init__(self, region_name: str | None = None):
"""Initializes the KMS client."""
self.kms_client = boto3.client("kms", region_name=region_name)
self.kms_client = boto3.client("kms", region_name=region_name, config=AWS_CLIENT_CONFIG)

def encrypt(self, data: str, key_id: str, encryption_context: dict[str, str]) -> str:
"""Encrypts data using the specified KMS key."""
Expand Down
4 changes: 3 additions & 1 deletion src/envars/aws_ssm.py
Original file line number Diff line number Diff line change
@@ -1,9 +1,11 @@
import boto3

from .aws_config import AWS_CLIENT_CONFIG


class SSMParameterStore:
def __init__(self, region_name: str | None = None):
self.client = boto3.client("ssm", region_name=region_name)
self.client = boto3.client("ssm", region_name=region_name, config=AWS_CLIENT_CONFIG)

def get_parameter(self, name: str, with_decryption: bool = True) -> str | None:
try:
Expand Down
13 changes: 9 additions & 4 deletions src/envars/cloud_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,10 +2,12 @@
import sys

import boto3
from botocore.exceptions import NoCredentialsError
from botocore.exceptions import BotoCoreError, ClientError
from google.auth import default as google_auth_default
from google.auth.exceptions import DefaultCredentialsError

from .aws_config import AWS_CLIENT_CONFIG


def _debug(message):
"""Prints a debug message to stderr if ENVARS_DEBUG is set."""
Expand All @@ -16,11 +18,14 @@ def _debug(message):
def get_aws_account_id() -> str | None:
"""Retrieves the AWS account ID from the current credentials."""
try:
account_id = boto3.client("sts").get_caller_identity().get("Account")
account_id = boto3.client("sts", config=AWS_CLIENT_CONFIG).get_caller_identity().get("Account")
_debug(f"Found AWS Account ID: {account_id}")
return account_id
except NoCredentialsError:
_debug("No AWS credentials found.")
except (BotoCoreError, ClientError) as e:
# NoCredentialsError, connect/read timeouts and API errors all land here. This is
# best-effort location auto-detection, so degrade to "unknown" (the caller then asks
# for --loc) instead of hanging/crashing when STS is slow or unreachable.
_debug(f"Could not determine AWS account ID: {type(e).__name__}: {e}")
return None


Expand Down
97 changes: 97 additions & 0 deletions tests/test_aws_config.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,97 @@
"""Tests for the shared bounded AWS client configuration (TOPS-2500)."""

from unittest.mock import MagicMock, patch

import boto3
from botocore.exceptions import ReadTimeoutError

from src.envars import aws_config, cloud_utils
from src.envars.aws_cloudformation import CloudFormationExports
from src.envars.aws_kms import AWSKMSAgent
from src.envars.aws_ssm import SSMParameterStore


def test_default_config_is_bounded():
"""The shared config caps timeouts and retries well below botocore's defaults (60s/60s/5).

Asserts individual keys rather than full-dict equality on ``retries``: botocore rewrites
that dict in place when a client is built from the config (``max_attempts`` becomes
``total_max_attempts``), so an equality assertion would be order- and version-fragile.
"""
cfg = aws_config.AWS_CLIENT_CONFIG
assert cfg.connect_timeout == 3
assert cfg.read_timeout == 5
assert cfg.retries["mode"] == "standard"


def test_default_makes_three_total_attempts():
"""max_attempts=2 resolves to 3 total HTTP attempts (1 initial + 2 retries).

botocore treats retries.max_attempts as the retry count, so total_max_attempts = N + 1
(verified against botocore 1.39.4). Pinning the resolved value keeps the ~300s -> ~17s
worst case from silently regressing if that mapping ever changes, and makes the PR's
"3 total attempts" claim executable. Built from a fresh _build_config() so botocore's
in-place rewrite of retries doesn't leak into the shared AWS_CLIENT_CONFIG other tests read.
"""
client = boto3.client("sts", region_name="eu-west-1", config=aws_config._build_config())
assert client.meta.config.retries["total_max_attempts"] == 3


def test_env_overrides_flow_through(monkeypatch):
"""All three ENVARS_AWS_* overrides wire through to the built config, not just read_timeout."""
monkeypatch.setenv("ENVARS_AWS_CONNECT_TIMEOUT", "7")
monkeypatch.setenv("ENVARS_AWS_READ_TIMEOUT", "11")
monkeypatch.setenv("ENVARS_AWS_MAX_ATTEMPTS", "4")
client = boto3.client("sts", region_name="eu-west-1", config=aws_config._build_config())
assert client.meta.config.connect_timeout == 7
assert client.meta.config.read_timeout == 11
assert client.meta.config.retries["total_max_attempts"] == 5 # input 4 -> 5 total (N+1)


def test_int_env_parses_positive_override(monkeypatch):
"""_int_env returns a valid positive override from the environment."""
monkeypatch.setenv("ENVARS_AWS_READ_TIMEOUT", "42")
assert aws_config._int_env("ENVARS_AWS_READ_TIMEOUT", 10) == 42

Comment thread
kthhrv marked this conversation as resolved.

def test_int_env_falls_back_when_unset_invalid_or_non_positive(monkeypatch):
"""Unset, non-numeric, and zero/negative values all fall back to the default."""
monkeypatch.delenv("ENVARS_AWS_READ_TIMEOUT", raising=False)
assert aws_config._int_env("ENVARS_AWS_READ_TIMEOUT", 10) == 10 # unset

monkeypatch.setenv("ENVARS_AWS_READ_TIMEOUT", "not-an-int")
assert aws_config._int_env("ENVARS_AWS_READ_TIMEOUT", 10) == 10 # invalid

monkeypatch.setenv("ENVARS_AWS_READ_TIMEOUT", "0")
assert aws_config._int_env("ENVARS_AWS_READ_TIMEOUT", 10) == 10 # non-positive


def test_kms_client_uses_bounded_config():
"""AWSKMSAgent constructs its client with the shared bounded config."""
agent = AWSKMSAgent(region_name="eu-west-1")
assert agent.kms_client.meta.config.connect_timeout == 3
assert agent.kms_client.meta.config.read_timeout == 5


def test_ssm_client_uses_bounded_config():
"""SSMParameterStore constructs its client with the shared bounded config."""
store = SSMParameterStore(region_name="eu-west-1")
assert store.client.meta.config.read_timeout == 5


def test_cloudformation_client_uses_bounded_config():
"""CloudFormationExports constructs its client with the shared bounded config."""
exports = CloudFormationExports(region_name="eu-west-1")
assert exports.client.meta.config.read_timeout == 5


def test_get_aws_account_id_returns_none_on_timeout():
"""A stalled STS endpoint degrades to None instead of raising ReadTimeoutError.

Regression for the DSS-3428 crash: get_aws_account_id only caught NoCredentialsError,
so a stalled endpoint surfaced as an uncaught traceback out of `envars exec`.
"""
stalled = MagicMock()
stalled.get_caller_identity.side_effect = ReadTimeoutError(endpoint_url="https://sts.eu-west-1.amazonaws.com/")
with patch.object(boto3, "client", return_value=stalled):
assert cloud_utils.get_aws_account_id() is None