Merge pull request #181 from conhealth/fix-pynamodb-with-botocore-1.13

Fix patching for PynamoDB 4.x with botocore 1.13

Author: chanchiem

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,
     )