From 96e7d758f8339dc36272e379319332dd4c5c8f15 Mon Sep 17 00:00:00 2001 From: Roger Zhang Date: Thu, 29 Jun 2023 13:55:39 -0700 Subject: [PATCH 1/3] support patching botocore session. Which is used in Datamasking's kms client --- aws_lambda_powertools/shared/user_agent.py | 25 ++++++++++++++++++++++ 1 file changed, 25 insertions(+) diff --git a/aws_lambda_powertools/shared/user_agent.py b/aws_lambda_powertools/shared/user_agent.py index 21818bca2b8..9ced7a67b91 100644 --- a/aws_lambda_powertools/shared/user_agent.py +++ b/aws_lambda_powertools/shared/user_agent.py @@ -108,6 +108,31 @@ def register_feature_to_session(session, feature): logger.debug(f"session passed in doesn't have a event system:{e}") +# Add feature user-agent to given sdk botocore.session.Session +def register_feature_to_botocore_session(botocore_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 + ---------- + botocore_session : botocore.session.Session + The botocore session to which the feature will be registered. + feature : str + The feature string to be appended to the User-Agent header, e.g., "data-masking" in Powertools. + + Raises + ------ + AttributeError + If the provided session does not have an event system. + + """ + try: + botocore_session.register(TARGET_SDK_EVENT, _create_feature_function(feature)) + except AttributeError as e: + logger.debug(f"botocore 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): """ From fadad137049ea63f3b29dba9641b9426b240cfc4 Mon Sep 17 00:00:00 2001 From: Roger Zhang Date: Thu, 6 Jul 2023 21:57:02 -0700 Subject: [PATCH 2/3] fix doc-string Signed-off-by: Roger Zhang --- aws_lambda_powertools/shared/user_agent.py | 20 +++++++++++++++++--- 1 file changed, 17 insertions(+), 3 deletions(-) diff --git a/aws_lambda_powertools/shared/user_agent.py b/aws_lambda_powertools/shared/user_agent.py index 9ced7a67b91..cadbe6d8eeb 100644 --- a/aws_lambda_powertools/shared/user_agent.py +++ b/aws_lambda_powertools/shared/user_agent.py @@ -111,8 +111,10 @@ def register_feature_to_session(session, feature): # Add feature user-agent to given sdk botocore.session.Session def register_feature_to_botocore_session(botocore_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 + Register the given feature string to the event system of the provided botocore session + + Please Notice this is for patching botocore session and is different from the previous one + which is for patching boto3 session Parameters ---------- @@ -125,7 +127,19 @@ def register_feature_to_botocore_session(botocore_session, feature): ------ AttributeError If the provided session does not have an event system. - + + Examples + -------- + **register data-masking user-agent to botocore session** + + >>> from aws_lambda_powertools.shared.user_agent import ( + >>> register_feature_to_botocore_session + >>> ) + >>> + >>> session = botocore.session.Session() + >>> register_feature_to_botocore_session(botocore_session=session, feature="data-masking") + >>> key_provider = StrictAwsKmsMasterKeyProvider(key_ids=self.keys, botocore_session=self.session) + """ try: botocore_session.register(TARGET_SDK_EVENT, _create_feature_function(feature)) From d92fe09efddbffd18d8cc403681bc76ab58c3b18 Mon Sep 17 00:00:00 2001 From: Roger Zhang Date: Thu, 6 Jul 2023 22:02:04 -0700 Subject: [PATCH 3/3] Update user_agent.py Signed-off-by: Roger Zhang --- aws_lambda_powertools/shared/user_agent.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/aws_lambda_powertools/shared/user_agent.py b/aws_lambda_powertools/shared/user_agent.py index cadbe6d8eeb..098be7a503a 100644 --- a/aws_lambda_powertools/shared/user_agent.py +++ b/aws_lambda_powertools/shared/user_agent.py @@ -113,8 +113,8 @@ def register_feature_to_botocore_session(botocore_session, feature): """ Register the given feature string to the event system of the provided botocore session - Please Notice this is for patching botocore session and is different from the previous one - which is for patching boto3 session + Please notice this function is for patching botocore session and is different from + previous one which is for patching boto3 session Parameters ---------- @@ -138,7 +138,7 @@ def register_feature_to_botocore_session(botocore_session, feature): >>> >>> session = botocore.session.Session() >>> register_feature_to_botocore_session(botocore_session=session, feature="data-masking") - >>> key_provider = StrictAwsKmsMasterKeyProvider(key_ids=self.keys, botocore_session=self.session) + >>> key_provider = StrictAwsKmsMasterKeyProvider(key_ids=self.keys, botocore_session=session) """ try: