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

Co-authored-by: Leandro Damascena <[email protected]>

Author: roger-zhangg

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,165 @@
+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 = os.environ.get("AWS_EXECUTION_ENV", "NA")
+TARGET_SDK_EVENT = "request-created"
+FEATURE_PREFIX = "PT"
+DEFAULT_FEATURE = "no-op"
+HEADER_NO_OP = f"{FEATURE_PREFIX}/{DEFAULT_FEATURE}/{powertools_version} PTEnv/{EXEC_ENV}"
+
+
+def _initializer_botocore_session(session):
+    """
+    This function is used to add an extra header for the User-Agent in the Botocore session,
+    as described in the pull request: https://github.com/boto/botocore/pull/2682
+
+    Parameters
+    ----------
+    session : botocore.session.Session
+        The Botocore session to which the user-agent function will be registered.
+
+    Raises
+    ------
+    Exception
+        If there is an issue while adding the extra header for the User-Agent.
+
+    """
+    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 and return the `add_powertools_feature` function.
+
+    The `add_powertools_feature` function is designed to be registered in boto3's event system.
+    When registered, it appends the given feature string to the User-Agent header of AWS SDK requests.
+
+    Parameters
+    ----------
+    feature : str
+        The feature string to be appended to the User-Agent header.
+
+    Returns
+    -------
+    add_powertools_feature : Callable
+        The `add_powertools_feature` function that modifies the User-Agent header.
+
+
+    """
+
+    def add_powertools_feature(request, **kwargs):
+        try:
+            headers = request.headers
+            header_user_agent = (
+                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 to the event system of the provided boto3 session
+    and append the feature to the User-Agent header of the request
+
+    Parameters
+    ----------
+    session : boto3.session.Session
+        The boto3 session to which the feature will be registered.
+    feature : str
+        The feature string to be appended to the User-Agent header, e.g., "streaming" in Powertools.
+
+    Raises
+    ------
+    AttributeError
+        If the provided session does not have an event system.
+
+    """
+    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 to the event system of the provided boto3 client
+    and append the feature to the User-Agent header of the request
+
+    Parameters
+    ----------
+    client : boto3.session.Session.client
+        The boto3 client to which the feature will be registered.
+    feature : str
+        The feature string to be appended to the User-Agent header, e.g., "streaming" in Powertools.
+
+    Raises
+    ------
+    AttributeError
+        If the provided client does not have an event system.
+
+    """
+    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 to the event system of the provided boto3 resource
+    and append the feature to the User-Agent header of the request
+
+    Parameters
+    ----------
+    resource : boto3.session.Session.resource
+        The boto3 resource to which the feature will be registered.
+    feature : str
+        The feature string to be appended to the User-Agent header, e.g., "streaming" in Powertools.
+
+    Raises
+    ------
+    AttributeError
+        If the provided resource does not have an event system.
+
+    """
+    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/data_classes/code_pipeline_job_event.py

@@ -6,6 +6,7 @@
 
 import boto3
 
+from aws_lambda_powertools.shared import user_agent
 from aws_lambda_powertools.utilities.data_classes.common import DictWrapper
 
 
@@ -203,12 +204,14 @@ def setup_s3_client(self):
         BaseClient
             An S3 client with the appropriate credentials
         """
-        return boto3.client(
+        s3 = boto3.client(
             "s3",
             aws_access_key_id=self.data.artifact_credentials.access_key_id,
             aws_secret_access_key=self.data.artifact_credentials.secret_access_key,
             aws_session_token=self.data.artifact_credentials.session_token,
         )
+        user_agent.register_feature_to_client(client=s3, feature="data_classes")
+        return s3
 
     def find_input_artifact(self, artifact_name: str) -> Optional[CodePipelineArtifact]:
         """Find an input artifact by artifact name

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(client=self.client, feature="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=client, feature="parameters")
             return client
 
         session = session or boto3.Session()
         config = config or Config()
-        return session.client(service_name=service_name, config=config)
+        client = session.client(service_name=service_name, config=config)
+        user_agent.register_feature_to_client(client=client, feature="parameters")
+        return client
 
     # 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(resource=client, feature="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 = session.resource(service_name=service_name, config=config, endpoint_url=endpoint_url)
+        user_agent.register_feature_to_resource(resource=client, feature="parameters")
+        return client
 
 
 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._has_user_agent = 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._has_user_agent:
+            user_agent.register_feature_to_client(client=self._s3_client, feature="streaming")
+            self._has_user_agent = 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

pyproject.toml

@@ -32,6 +32,7 @@ aws-xray-sdk = { version = "^2.8.0", optional = true }
 fastjsonschema = { version = "^2.14.5", optional = true }
 pydantic = { version = "^1.8.2", optional = true }
 boto3 = { version = "^1.20.32", optional = true }
+importlib-metadata = {version = "^6.6.0", python = "<3.8"}
 typing-extensions = "^4.6.2"
 
 [tool.poetry.dev-dependencies]
@@ -86,7 +87,6 @@ mkdocs-material = "^9.1.15"
 filelock = "^3.12.0"
 checksumdir = "^1.2.0"
 mypy-boto3-appconfigdata = "^1.26.70"
-importlib-metadata = "^6.6"
 ijson = "^3.2.0"
 typed-ast = { version = "^1.5.4", python = "< 3.8"}
 hvac = "^1.1.0"