diff --git a/aws_xray_sdk/ext/pynamodb/patch.py b/aws_xray_sdk/ext/pynamodb/patch.py index 74b78259..1478808a 100644 --- a/aws_xray_sdk/ext/pynamodb/patch.py +++ b/aws_xray_sdk/ext/pynamodb/patch.py @@ -1,4 +1,3 @@ -import botocore.vendored.requests.sessions import json import wrapt import pynamodb @@ -9,20 +8,30 @@ PYNAMODB4 = int(pynamodb.__version__.split('.')[0]) >= 4 +if PYNAMODB4: + import botocore.httpsession +else: + import botocore.vendored.requests.sessions + def patch(): """Patch PynamoDB so it generates subsegements when calling DynamoDB.""" - if hasattr(botocore.vendored.requests.sessions, '_xray_enabled'): - return - setattr(botocore.vendored.requests.sessions, '_xray_enabled', True) - if PYNAMODB4: + if hasattr(botocore.httpsession, '_xray_enabled'): + return + setattr(botocore.httpsession, '_xray_enabled', True) + module = 'botocore.httpsession' name = 'URLLib3Session.send' else: + if hasattr(botocore.vendored.requests.sessions, '_xray_enabled'): + return + setattr(botocore.vendored.requests.sessions, '_xray_enabled', True) + module = 'botocore.vendored.requests.sessions' name = 'Session.send' + wrapt.wrap_function_wrapper( module, name, _xray_traced_pynamodb, )