diff --git a/aws_lambda_powertools/__init__.py b/aws_lambda_powertools/__init__.py index 574c9b257f1..14237bc7119 100644 --- a/aws_lambda_powertools/__init__.py +++ b/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() diff --git a/aws_lambda_powertools/shared/user_agent.py b/aws_lambda_powertools/shared/user_agent.py new file mode 100644 index 00000000000..62cdc16601d --- /dev/null +++ b/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) diff --git a/aws_lambda_powertools/shared/version.py b/aws_lambda_powertools/shared/version.py new file mode 100644 index 00000000000..0db71627089 --- /dev/null +++ b/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") diff --git a/aws_lambda_powertools/utilities/data_classes/code_pipeline_job_event.py b/aws_lambda_powertools/utilities/data_classes/code_pipeline_job_event.py index c502aacb090..434df509deb 100644 --- a/aws_lambda_powertools/utilities/data_classes/code_pipeline_job_event.py +++ b/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 diff --git a/aws_lambda_powertools/utilities/idempotency/persistence/dynamodb.py b/aws_lambda_powertools/utilities/idempotency/persistence/dynamodb.py index 654f8ca99d4..d69e42f9287 100644 --- a/aws_lambda_powertools/utilities/idempotency/persistence/dynamodb.py +++ b/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!") diff --git a/aws_lambda_powertools/utilities/parameters/base.py b/aws_lambda_powertools/utilities/parameters/base.py index 8ec1052ae37..4357b5d520e 100644 --- a/aws_lambda_powertools/utilities/parameters/base.py +++ b/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]: diff --git a/aws_lambda_powertools/utilities/streaming/_s3_seekable_io.py b/aws_lambda_powertools/utilities/streaming/_s3_seekable_io.py index 8e280f3f7d7..de9b77410a3 100644 --- a/aws_lambda_powertools/utilities/streaming/_s3_seekable_io.py +++ b/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 diff --git a/mypy.ini b/mypy.ini index 4af89217fdc..2b50293b561 100644 --- a/mypy.ini +++ b/mypy.ini @@ -63,4 +63,5 @@ ignore_missing_imports = True [mypy-ijson] ignore_missing_imports = True - +[mypy-importlib.metadata] +ignore_missing_imports = True diff --git a/pyproject.toml b/pyproject.toml index 81fc80567cd..07fd2c28f97 100644 --- a/pyproject.toml +++ b/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"