feat(user-agent): add custom header User-Agent to AWS SDK requests

aws_lambda_powertools/__init__.py

@@ -7,6 +7,7 @@
 from .logging import Logger
 from .metrics import Metrics, single_metric
 from .package_logger import set_package_logger_handler
+from .shared import user_agent
 from .tracing import Tracer
 
 __author__ = """Amazon Web Services"""
@@ -20,3 +21,5 @@
 PACKAGE_PATH = Path(__file__).parent
 
 set_package_logger_handler()
+
+user_agent.inject_user_agent()

aws_lambda_powertools/shared/user_agent.py

@@ -0,0 +1,70 @@
+import logging
+import os
+
+from importlib_metadata import version
+
+pt_version = version("aws-lambda-powertools")
+
+try:
+    from botocore import handlers
+except ImportError:
+    # if botocore failed to import, user might be using custom runtime. We can ignore here
+    handlers = None
+
+logger = logging.getLogger(__name__)
+
+ENV_KEY = "AWS_EXECUTION_ENV"
+EXEC_ENV = os.environ.get(ENV_KEY, "NA")
+target_sdk_event = "request-created"
+
+
+# add user-agent like powertools/2.14.1 EXEC-ENV/AWS_Lambda_python3.9
+def _add_pt_version(request, **kwargs):
+    headers = request.headers
+    # assume we always have a User-Agent.
+    # Q: do we need to check if field: User-Agent exist in request.headers?
+    # Q: do we care if User-Agent string exceed 256 characters?
+    headers["User-Agent"] = f"{headers['User-Agent']} PT/no-op/{pt_version} PTEnv/{EXEC_ENV}"
+
+
+# creates the `add_feature_string` function with given feature parameter
+def _create_feature_function(feature):
+    def add_pt_feature(request, **kwargs):
+        headers = request.headers
+        # Actually, only one handler can be registered, registering a new one will replace the prev handler
+        # We don't need to replace/detect previous registered user-agent
+        headers["User-Agent"] = f"{headers['User-Agent']} PT/{feature}/{pt_version} PTEnv/{EXEC_ENV}"
+
+    # return created function
+    return add_pt_feature
+
+
+# add feature user-agent to given sdk session
+def register_feature_to_session(session, feature):
+    try:
+        session.events.register(target_sdk_event, _create_feature_function(feature))
+    except AttributeError as e:
+        logger.debug(f"session passed in doesn't have a event system:{e}")
+
+
+# add feature user-agent to given sdk client
+def register_feature_to_client(client, feature):
+    try:
+        client.meta.events.register(target_sdk_event, _create_feature_function(feature))
+    except AttributeError as e:
+        logger.debug(f"session passed in doesn't have a event system:{e}")
+
+
+# add feature user-agent to given sdk session.resource
+def register_feature_to_resource(resource, feature):
+    try:
+        resource.meta.client.meta.events.register(target_sdk_event, _create_feature_function(feature))
+    except AttributeError as e:
+        logger.debug(f"resource passed in doesn't have a event system:{e}")
+
+
+# register add_pt_version for all AWS SDK in runtime
+def inject_user_agent():
+    if handlers:
+        # register add_user_agent to BUILTIN_HANDLERS so every aws sdk session will have this event registered
+        handlers.BUILTIN_HANDLERS.append((target_sdk_event, _add_pt_version))

aws_lambda_powertools/utilities/idempotency/persistence/dynamodb.py

@@ -10,7 +10,7 @@
 from botocore.config import Config
 from botocore.exceptions import ClientError
 
-from aws_lambda_powertools.shared import constants
+from aws_lambda_powertools.shared import constants, user_agent
 from aws_lambda_powertools.utilities.idempotency import BasePersistenceLayer
 from aws_lambda_powertools.utilities.idempotency.exceptions import (
     IdempotencyItemAlreadyExistsError,
@@ -94,6 +94,8 @@ def __init__(
         else:
             self.client = boto3_client
 
+        user_agent.register_feature_to_client(self.client, "idempotency")
+
         if sort_key_attr == key_attr:
             raise ValueError(f"key_attr [{key_attr}] and sort_key_attr [{sort_key_attr}] cannot be the same!")
 

aws_lambda_powertools/utilities/parameters/base.py

@@ -25,7 +25,7 @@
 import boto3
 from botocore.config import Config
 
-from aws_lambda_powertools.shared import constants
+from aws_lambda_powertools.shared import constants, user_agent
 from aws_lambda_powertools.shared.functions import resolve_max_age
 from aws_lambda_powertools.utilities.parameters.types import TransformOptions
 
@@ -254,11 +254,14 @@ def _build_boto3_client(
             Instance of a boto3 client for Parameters feature (e.g., ssm, appconfig, secretsmanager, etc.)
         """
         if client is not None:
+            user_agent.register_feature_to_client(client, "parameters")
             return client
 
         session = session or boto3.Session()
         config = config or Config()
-        return session.client(service_name=service_name, config=config)
+        client_to_return = session.client(service_name=service_name, config=config)
+        user_agent.register_feature_to_client(client_to_return, "parameters")
+        return client_to_return
 
     # maintenance: change DynamoDBServiceResource type to ParameterResourceClients when we expand
     @staticmethod
@@ -288,11 +291,14 @@ def _build_boto3_resource_client(
             Instance of a boto3 resource client for Parameters feature (e.g., dynamodb, etc.)
         """
         if client is not None:
+            user_agent.register_feature_to_resource(client, "parameters")
             return client
 
         session = session or boto3.Session()
         config = config or Config()
-        return session.resource(service_name=service_name, config=config, endpoint_url=endpoint_url)
+        client_to_return = session.resource(service_name=service_name, config=config, endpoint_url=endpoint_url)
+        user_agent.register_feature_to_resource(client_to_return, "parameters")
+        return client_to_return
 
 
 def get_transform_method(value: str, transform: TransformOptions = None) -> Callable[..., Any]:

aws_lambda_powertools/utilities/streaming/_s3_seekable_io.py

@@ -15,6 +15,7 @@
 
 import boto3
 
+from aws_lambda_powertools.shared import user_agent
 from aws_lambda_powertools.utilities.streaming.compat import PowertoolsStreamingBody
 
 if TYPE_CHECKING:
@@ -67,6 +68,7 @@ def __init__(
         self._sdk_options = sdk_options
         self._sdk_options["Bucket"] = bucket
         self._sdk_options["Key"] = key
+        self._user_agent_added = False
         if version_id is not None:
             self._sdk_options["VersionId"] = version_id
 
@@ -77,6 +79,9 @@ def s3_client(self) -> "Client":
         """
         if self._s3_client is None:
             self._s3_client = boto3.client("s3")
+        if not self._user_agent_added:
+            user_agent.register_feature_to_client(self._s3_client, "streaming")
+            self._user_agent_added = True
         return self._s3_client
 
     @property