Skip to content

Commit 454d669

Browse files
authored
chore: move env names to constant file (#264)
* chore: shift vars to constant to ease changing * improv: update tests to include non-truthy choice var
1 parent ea90421 commit 454d669

File tree

10 files changed

+141
-42
lines changed

10 files changed

+141
-42
lines changed

aws_lambda_powertools/logging/formatter.py

+3-1
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,8 @@
33
import os
44
from typing import Dict, Iterable, Optional, Union
55

6+
from ..shared import constants
7+
68
STD_LOGGING_KEYS = (
79
"name",
810
"msg",
@@ -73,7 +75,7 @@ def _build_root_keys(**kwargs):
7375

7476
@staticmethod
7577
def _get_latest_trace_id():
76-
xray_trace_id = os.getenv("_X_AMZN_TRACE_ID")
78+
xray_trace_id = os.getenv(constants.XRAY_TRACE_ID_ENV)
7779
return xray_trace_id.split(";")[0].replace("Root=", "") if xray_trace_id else None
7880

7981
def update_formatter(self, **kwargs):

aws_lambda_powertools/logging/logger.py

+12-6
Original file line numberDiff line numberDiff line change
@@ -4,9 +4,10 @@
44
import os
55
import random
66
import sys
7-
from distutils.util import strtobool
87
from typing import Any, Callable, Dict, Union
98

9+
from ..shared import constants
10+
from ..shared.functions import resolve_env_var_choice, resolve_truthy_env_var_choice
1011
from .exceptions import InvalidLoggerSamplingRateError
1112
from .filters import SuppressFilter
1213
from .formatter import JsonFormatter
@@ -122,8 +123,12 @@ def __init__(
122123
stream: sys.stdout = None,
123124
**kwargs,
124125
):
125-
self.service = service or os.getenv("POWERTOOLS_SERVICE_NAME") or "service_undefined"
126-
self.sampling_rate = sampling_rate or os.getenv("POWERTOOLS_LOGGER_SAMPLE_RATE") or 0.0
126+
self.service = resolve_env_var_choice(
127+
choice=service, env=os.getenv(constants.SERVICE_NAME_ENV, "service_undefined")
128+
)
129+
self.sampling_rate = resolve_env_var_choice(
130+
choice=sampling_rate, env=os.getenv(constants.LOGGER_LOG_SAMPLING_RATE, 0.0)
131+
)
127132
self.log_level = self._get_log_level(level)
128133
self.child = child
129134
self._handler = logging.StreamHandler(stream) if stream is not None else logging.StreamHandler(sys.stdout)
@@ -193,7 +198,7 @@ def _configure_sampling(self):
193198
f"Please review POWERTOOLS_LOGGER_SAMPLE_RATE environment variable."
194199
)
195200

196-
def inject_lambda_context(self, lambda_handler: Callable[[Dict, Any], Any] = None, log_event: bool = False):
201+
def inject_lambda_context(self, lambda_handler: Callable[[Dict, Any], Any] = None, log_event: bool = None):
197202
"""Decorator to capture Lambda contextual info and inject into logger
198203
199204
Parameters
@@ -242,8 +247,9 @@ def handler(event, context):
242247
logger.debug("Decorator called with parameters")
243248
return functools.partial(self.inject_lambda_context, log_event=log_event)
244249

245-
log_event_env_option = str(os.getenv("POWERTOOLS_LOGGER_LOG_EVENT", "false"))
246-
log_event = strtobool(log_event_env_option) or log_event
250+
log_event = resolve_truthy_env_var_choice(
251+
choice=log_event, env=os.getenv(constants.LOGGER_LOG_EVENT_ENV, "false")
252+
)
247253

248254
@functools.wraps(lambda_handler)
249255
def decorate(event, context):

aws_lambda_powertools/metrics/base.py

+4-2
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,8 @@
1010

1111
import fastjsonschema
1212

13+
from ..shared import constants
14+
from ..shared.functions import resolve_env_var_choice
1315
from .exceptions import MetricUnitError, MetricValueError, SchemaValidationError
1416

1517
logger = logging.getLogger(__name__)
@@ -88,8 +90,8 @@ def __init__(
8890
):
8991
self.metric_set = metric_set if metric_set is not None else {}
9092
self.dimension_set = dimension_set if dimension_set is not None else {}
91-
self.namespace = namespace or os.getenv("POWERTOOLS_METRICS_NAMESPACE")
92-
self.service = service or os.environ.get("POWERTOOLS_SERVICE_NAME")
93+
self.namespace = resolve_env_var_choice(choice=namespace, env=os.getenv(constants.METRICS_NAMESPACE_ENV))
94+
self.service = resolve_env_var_choice(choice=service, env=os.getenv(constants.SERVICE_NAME_ENV))
9395
self._metric_units = [unit.value for unit in MetricUnit]
9496
self._metric_unit_options = list(MetricUnit.__members__)
9597
self.metadata_set = self.metadata_set if metadata_set is not None else {}

aws_lambda_powertools/middleware_factory/factory.py

+6-3
Original file line numberDiff line numberDiff line change
@@ -2,16 +2,17 @@
22
import inspect
33
import logging
44
import os
5-
from distutils.util import strtobool
65
from typing import Callable
76

7+
from ..shared import constants
8+
from ..shared.functions import resolve_truthy_env_var_choice
89
from ..tracing import Tracer
910
from .exceptions import MiddlewareInvalidArgumentError
1011

1112
logger = logging.getLogger(__name__)
1213

1314

14-
def lambda_handler_decorator(decorator: Callable = None, trace_execution=False):
15+
def lambda_handler_decorator(decorator: Callable = None, trace_execution: bool = None):
1516
"""Decorator factory for decorating Lambda handlers.
1617
1718
You can use lambda_handler_decorator to create your own middlewares,
@@ -104,7 +105,9 @@ def lambda_handler(event, context):
104105
if decorator is None:
105106
return functools.partial(lambda_handler_decorator, trace_execution=trace_execution)
106107

107-
trace_execution = trace_execution or strtobool(str(os.getenv("POWERTOOLS_TRACE_MIDDLEWARES", False)))
108+
trace_execution = resolve_truthy_env_var_choice(
109+
choice=trace_execution, env=os.getenv(constants.MIDDLEWARE_FACTORY_TRACE_ENV, "false")
110+
)
108111

109112
@functools.wraps(decorator)
110113
def final_decorator(func: Callable = None, **kwargs):
+14-3
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,15 @@
1-
import os
1+
TRACER_CAPTURE_RESPONSE_ENV: str = "POWERTOOLS_TRACER_CAPTURE_RESPONSE"
2+
TRACER_CAPTURE_ERROR_ENV: str = "POWERTOOLS_TRACER_CAPTURE_ERROR"
3+
TRACER_DISABLED_ENV: str = "POWERTOOLS_TRACE_DISABLED"
24

3-
TRACER_CAPTURE_RESPONSE_ENV: str = os.getenv("POWERTOOLS_TRACER_CAPTURE_RESPONSE", "true")
4-
TRACER_CAPTURE_ERROR_ENV: str = os.getenv("POWERTOOLS_TRACER_CAPTURE_ERROR", "true")
5+
LOGGER_LOG_SAMPLING_RATE: str = "POWERTOOLS_LOGGER_SAMPLE_RATE"
6+
LOGGER_LOG_EVENT_ENV: str = "POWERTOOLS_LOGGER_LOG_EVENT"
7+
8+
MIDDLEWARE_FACTORY_TRACE_ENV: str = "POWERTOOLS_TRACE_MIDDLEWARES"
9+
10+
METRICS_NAMESPACE_ENV: str = "POWERTOOLS_METRICS_NAMESPACE"
11+
12+
SAM_LOCAL_ENV: str = "AWS_SAM_LOCAL"
13+
CHALICE_LOCAL_ENV: str = "AWS_CHALICE_CLI_MODE"
14+
SERVICE_NAME_ENV: str = "POWERTOOLS_SERVICE_NAME"
15+
XRAY_TRACE_ID_ENV: str = "_X_AMZN_TRACE_ID"
+38-1
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,42 @@
11
from distutils.util import strtobool
2+
from typing import Any, Union
23

34

4-
def resolve_env_var_choice(env: str, choice: bool = None) -> bool:
5+
def resolve_truthy_env_var_choice(env: Any, choice: bool = None) -> bool:
6+
""" Pick explicit choice over truthy env value, if available, otherwise return truthy env value
7+
8+
NOTE: Environment variable should be resolved by the caller.
9+
10+
Parameters
11+
----------
12+
env : Any
13+
environment variable actual value
14+
choice : bool
15+
explicit choice
16+
17+
Returns
18+
-------
19+
choice : str
20+
resolved choice as either bool or environment value
21+
"""
522
return choice if choice is not None else strtobool(env)
23+
24+
25+
def resolve_env_var_choice(env: Any, choice: bool = None) -> Union[bool, Any]:
26+
""" Pick explicit choice over env, if available, otherwise return env value received
27+
28+
NOTE: Environment variable should be resolved by the caller.
29+
30+
Parameters
31+
----------
32+
env : Any
33+
environment variable actual value
34+
choice : bool
35+
explicit choice
36+
37+
Returns
38+
-------
39+
choice : str
40+
resolved choice as either bool or environment value
41+
"""
42+
return choice if choice is not None else env

aws_lambda_powertools/tracing/tracer.py

+20-15
Original file line numberDiff line numberDiff line change
@@ -4,14 +4,13 @@
44
import inspect
55
import logging
66
import os
7-
from distutils.util import strtobool
87
from typing import Any, Callable, Dict, List, Optional, Tuple
98

109
import aws_xray_sdk
1110
import aws_xray_sdk.core
1211

13-
from aws_lambda_powertools.shared.constants import TRACER_CAPTURE_ERROR_ENV, TRACER_CAPTURE_RESPONSE_ENV
14-
from aws_lambda_powertools.shared.functions import resolve_env_var_choice
12+
from ..shared import constants
13+
from ..shared.functions import resolve_truthy_env_var_choice
1514

1615
is_cold_start = True
1716
logger = logging.getLogger(__name__)
@@ -283,9 +282,12 @@ def handler(event, context):
283282
)
284283

285284
lambda_handler_name = lambda_handler.__name__
286-
287-
capture_response = resolve_env_var_choice(env=TRACER_CAPTURE_RESPONSE_ENV, choice=capture_response)
288-
capture_error = resolve_env_var_choice(env=TRACER_CAPTURE_ERROR_ENV, choice=capture_error)
285+
capture_response = resolve_truthy_env_var_choice(
286+
env=os.getenv(constants.TRACER_CAPTURE_RESPONSE_ENV, "true"), choice=capture_response
287+
)
288+
capture_error = resolve_truthy_env_var_choice(
289+
env=os.getenv(constants.TRACER_CAPTURE_ERROR_ENV, "true"), choice=capture_error
290+
)
289291

290292
@functools.wraps(lambda_handler)
291293
def decorate(event, context):
@@ -478,8 +480,12 @@ async def async_tasks():
478480

479481
method_name = f"{method.__name__}"
480482

481-
capture_response = resolve_env_var_choice(env=TRACER_CAPTURE_RESPONSE_ENV, choice=capture_response)
482-
capture_error = resolve_env_var_choice(env=TRACER_CAPTURE_ERROR_ENV, choice=capture_error)
483+
capture_response = resolve_truthy_env_var_choice(
484+
env=os.getenv(constants.TRACER_CAPTURE_RESPONSE_ENV, "true"), choice=capture_response
485+
)
486+
capture_error = resolve_truthy_env_var_choice(
487+
env=os.getenv(constants.TRACER_CAPTURE_ERROR_ENV, "true"), choice=capture_error
488+
)
483489

484490
if inspect.iscoroutinefunction(method):
485491
return self._decorate_async_function(
@@ -681,14 +687,13 @@ def _is_tracer_disabled() -> bool:
681687
bool
682688
"""
683689
logger.debug("Verifying whether Tracing has been disabled")
684-
is_lambda_sam_cli = os.getenv("AWS_SAM_LOCAL")
685-
is_chalice_cli = os.getenv("AWS_CHALICE_CLI_MODE")
686-
env_option = str(os.getenv("POWERTOOLS_TRACE_DISABLED", "false"))
687-
disabled_env = strtobool(env_option)
690+
is_lambda_sam_cli = os.getenv(constants.SAM_LOCAL_ENV)
691+
is_chalice_cli = os.getenv(constants.CHALICE_LOCAL_ENV)
692+
is_disabled = resolve_truthy_env_var_choice(env=os.getenv(constants.TRACER_DISABLED_ENV, "false"))
688693

689-
if disabled_env:
694+
if is_disabled:
690695
logger.debug("Tracing has been disabled via env var POWERTOOLS_TRACE_DISABLED")
691-
return disabled_env
696+
return is_disabled
692697

693698
if is_lambda_sam_cli or is_chalice_cli:
694699
logger.debug("Running under SAM CLI env or not in Lambda env; disabling Tracing")
@@ -706,7 +711,7 @@ def __build_config(
706711
):
707712
""" Populates Tracer config for new and existing initializations """
708713
is_disabled = disabled if disabled is not None else self._is_tracer_disabled()
709-
is_service = service if service is not None else os.getenv("POWERTOOLS_SERVICE_NAME")
714+
is_service = service if service is not None else os.getenv(constants.SERVICE_NAME_ENV)
710715

711716
self._config["provider"] = provider if provider is not None else self._config["provider"]
712717
self._config["auto_patch"] = auto_patch if auto_patch is not None else self._config["auto_patch"]

tests/functional/test_logger.py

+3-3
Original file line numberDiff line numberDiff line change
@@ -373,9 +373,9 @@ def test_logger_do_not_log_twice_when_root_logger_is_setup(stdout, service_name)
373373

374374
# WHEN we create a new Logger and child Logger
375375
logger = Logger(service=service_name, stream=stdout)
376-
child_logger = Logger(child=True, stream=stdout)
377-
logger.info("hello")
378-
child_logger.info("hello again")
376+
child_logger = Logger(service=service_name, child=True, stream=stdout)
377+
logger.info("PARENT")
378+
child_logger.info("CHILD")
379379

380380
# THEN it should only contain only two log entries
381381
# since child's log records propagated to root logger should be rejected
+5-3
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,11 @@
1-
from aws_lambda_powertools.shared.functions import resolve_env_var_choice
1+
from aws_lambda_powertools.shared.functions import resolve_env_var_choice, resolve_truthy_env_var_choice
22

33

44
def test_resolve_env_var_choice_explicit_wins_over_env_var():
5-
assert resolve_env_var_choice(env="true", choice=False) is False
5+
assert resolve_truthy_env_var_choice(env="true", choice=False) is False
6+
assert resolve_env_var_choice(env="something", choice=False) is False
67

78

89
def test_resolve_env_var_choice_env_wins_over_absent_explicit():
9-
assert resolve_env_var_choice(env="true") == 1
10+
assert resolve_truthy_env_var_choice(env="true") == 1
11+
assert resolve_env_var_choice(env="something") == "something"

tests/unit/test_tracing.py

+36-5
Original file line numberDiff line numberDiff line change
@@ -506,11 +506,9 @@ def test_tracer_lambda_handler_override_response_as_metadata(mocker, provider_st
506506
# GIVEN tracer is initialized
507507
provider = provider_stub(in_subsegment=in_subsegment_mock.in_subsegment)
508508

509-
mocker.patch("aws_lambda_powertools.tracing.tracer.TRACER_CAPTURE_RESPONSE_ENV", return_value=True)
510509
tracer = Tracer(provider=provider, auto_patch=False)
511510

512-
# WHEN capture_lambda_handler decorator is used
513-
# with capture_response set to False
511+
# WHEN capture_lambda_handler decorator is used with capture_response set to False
514512
@tracer.capture_lambda_handler(capture_response=False)
515513
def handler(event, context):
516514
return "response"
@@ -526,8 +524,7 @@ def test_tracer_method_override_response_as_metadata(provider_stub, in_subsegmen
526524
provider = provider_stub(in_subsegment=in_subsegment_mock.in_subsegment)
527525
tracer = Tracer(provider=provider, auto_patch=False)
528526

529-
# WHEN capture_method decorator is used
530-
# and the method response is empty
527+
# WHEN capture_method decorator is used with capture_response set to False
531528
@tracer.capture_method(capture_response=False)
532529
def greeting(name, message):
533530
return "response"
@@ -536,3 +533,37 @@ def greeting(name, message):
536533

537534
# THEN we should not add any metadata
538535
assert in_subsegment_mock.put_metadata.call_count == 0
536+
537+
538+
def test_tracer_lambda_handler_override_error_as_metadata(mocker, provider_stub, in_subsegment_mock):
539+
# GIVEN tracer is initialized
540+
provider = provider_stub(in_subsegment=in_subsegment_mock.in_subsegment)
541+
tracer = Tracer(provider=provider, auto_patch=False)
542+
543+
# WHEN capture_lambda_handler decorator is used with capture_error set to False
544+
@tracer.capture_lambda_handler(capture_error=False)
545+
def handler(event, context):
546+
raise ValueError("error")
547+
548+
with pytest.raises(ValueError):
549+
handler({}, mocker.MagicMock())
550+
551+
# THEN we should not add any metadata
552+
assert in_subsegment_mock.put_metadata.call_count == 0
553+
554+
555+
def test_tracer_method_override_error_as_metadata(provider_stub, in_subsegment_mock):
556+
# GIVEN tracer is initialized
557+
provider = provider_stub(in_subsegment=in_subsegment_mock.in_subsegment)
558+
tracer = Tracer(provider=provider, auto_patch=False)
559+
560+
# WHEN capture_method decorator is used with capture_error set to False
561+
@tracer.capture_method(capture_error=False)
562+
def greeting(name, message):
563+
raise ValueError("error")
564+
565+
with pytest.raises(ValueError):
566+
greeting(name="Foo", message="Bar")
567+
568+
# THEN we should not add any metadata
569+
assert in_subsegment_mock.put_metadata.call_count == 0

0 commit comments

Comments
 (0)