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

aws_lambda_powertools/__init__.py

@@ -4,11 +4,14 @@
 
 from pathlib import Path
 
-from .logging import Logger
-from .metrics import Metrics, single_metric
-from .package_logger import set_package_logger_handler
-from .tracing import Tracer
+from aws_lambda_powertools.logging import Logger
+from aws_lambda_powertools.metrics import Metrics, single_metric
+from aws_lambda_powertools.package_logger import set_package_logger_handler
+from aws_lambda_powertools.shared.user_agent import inject_user_agent
+from aws_lambda_powertools.shared.version import VERSION
+from aws_lambda_powertools.tracing import Tracer
 
+__version__ = VERSION
 __author__ = """Amazon Web Services"""
 __all__ = [
     "Logger",
@@ -20,3 +23,5 @@
 PACKAGE_PATH = Path(__file__).parent
 
 set_package_logger_handler()
+
+inject_user_agent()

aws_lambda_powertools/shared/user_agent.py

@@ -0,0 +1,149 @@
+import logging
+import os
+
+from aws_lambda_powertools.shared.version import VERSION
+
+powertools_version = VERSION
+inject_header = True
+
+try:
+    import botocore
+except ImportError:
+    # if botocore failed to import, user might be using custom runtime and we can't inject header
+    inject_header = False
+
+logger = logging.getLogger(__name__)
+
+EXEC_ENV: str = os.environ.get("AWS_EXECUTION_ENV", "NA")
+TARGET_SDK_EVENT: str = "request-created"
+FEATURE_PREFIX: str = "PT"
+DEFAULT_FEATURE: str = "no-op"
+HEADER_NO_OP: str = f"{FEATURE_PREFIX}/{DEFAULT_FEATURE}/{powertools_version} PTEnv/{EXEC_ENV}"
+
+
+def _initializer_botocore_session(session):
+    """
+    Pass in Botocore session and register _create_feature_function to append user-agent, for more details see
+    https://github.com/boto/botocore/pull/2682
+
+    session.register:
+    https://github.com/boto/botocore/blob/develop/botocore/session.py#L703
+
+    Parameters
+    ----------
+    session: Botocore.session.Session
+        A botocore session passed in to register user-agent function
+
+    """
+    try:
+        session.register(TARGET_SDK_EVENT, _create_feature_function(DEFAULT_FEATURE))
+    except Exception:
+        logger.debug("Can't add extra header User-Agent")
+
+
+def _create_feature_function(feature):
+    """
+    Create a add_powertools_feature function using the given feature paramter
+    add_powertools_feature will be returned and to be regiestered in boto3's event system
+    once registered, add_powertools_feature will append the given feature string to user-agent of AWS SDK's request
+
+    Parameters
+    ----------
+    feature: str
+
+    Returns:
+    -------
+    add_powertools_feature: Callable
+
+    """
+
+    def add_powertools_feature(request, **kwargs):
+        """
+        Signiture of this function is required at Boto3's event system. See:
+        https://boto3.amazonaws.com/v1/documentation/api/latest/guide/events.html
+
+        """
+        try:
+            headers = request.headers
+            header_user_agent: str = (
+                f"{headers['User-Agent']} {FEATURE_PREFIX}/{feature}/{powertools_version} PTEnv/{EXEC_ENV}"
+            )
+
+            # This function is exclusive to client and resources objects created in Powertools
+            # and must remove the no-op header, if present
+            if HEADER_NO_OP in headers["User-Agent"] and feature != DEFAULT_FEATURE:
+                # Remove HEADER_NO_OP + space
+                header_user_agent = header_user_agent.replace(f"{HEADER_NO_OP} ", "")
+
+            headers["User-Agent"] = f"{header_user_agent}"
+        except Exception:
+            logger.debug("Can't find User-Agent header")
+
+    return add_powertools_feature
+
+
+# Add feature user-agent to given sdk boto3.session
+def register_feature_to_session(session, feature):
+    """
+    Register the given feature string given session's event system
+    When this session makes an api request, the given feature will be appended to request's user-agent
+
+    Parameters
+    ----------
+    session: boto3.session.Session
+        session to register
+    feature: str
+        the feature string in Powertools, e.g.:streaming
+
+    """
+    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 boto3.client
+def register_feature_to_client(client, feature):
+    """
+    Register the given feature string given client's event system
+    When this client makes an api request, the given feature will be appended to request's user-agent
+
+    Parameters
+    ----------
+    client: boto3.session.Session.client
+        client to register
+    feature: str
+        the feature string in Powertools, e.g.:streaming
+
+    """
+    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 boto3.resource
+def register_feature_to_resource(resource, feature):
+    """
+    Register the given feature string given resource's event system
+    When this resource makes an api request, the given feature will be appended to request's user-agent
+
+    Parameters
+    ----------
+    resource: boto3.session.Session.resource
+        resource to register
+    feature: str
+        the feature string in Powertools, e.g.:streaming
+
+    """
+    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}")
+
+
+def inject_user_agent():
+    if inject_header:
+        # Customize botocore session to inject Powertools header
+        # See: https://github.com/boto/botocore/pull/2682
+        botocore.register_initializer(_initializer_botocore_session)

aws_lambda_powertools/shared/version.py

@@ -0,0 +1,16 @@
+"""
+    This file serves to create a constant that informs
+    the current version of the Powertools package and exposes it in the main module
+
+    Since Python 3.8 there the built-in importlib.metadata
+    When support for Python3.7 is dropped, we can remove the optional importlib_metadata dependency
+    See: https://docs.python.org/3/library/importlib.metadata.html
+"""
+import sys
+
+if sys.version_info >= (3, 8):
+    from importlib.metadata import version
+else:
+    from importlib_metadata import version
+
+VERSION = version("aws-lambda-powertools")

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

mypy.ini

@@ -63,4 +63,5 @@ ignore_missing_imports = True
 [mypy-ijson]
 ignore_missing_imports = True
 
-
+[mypy-importlib.metadata]
+ignore_missing_imports = True