Recursively patch any given module functions with capture (#113)

* Add recursive function and method patching with capture

* Patch with only import hooks

* Remove setup patching, add envvar to patch modules on Django app ready

* Prevent double patching and add startup patching segment

* Add module pattern ignore

* Add function attribute and check to detect whether already decorated

* Add tests

* Fix PY3 patching + importing and amend tests

* Add ignore tests

* Fix class methods, classmethods and staticmethods, and add tests

* Update docs

* Fix referenced inherited methods

* Avoid automatically patching classmethods

* Add settings variable to opt-in for import-time parent segment

* Patch pg8000 tests only once per module

Author: hasier

CHANGELOG.rst

@@ -2,6 +2,10 @@
 CHANGELOG
 =========
 
+unreleased
+==========
+* feature: Recursively patch any given module functions with capture
+
 2.2.0
 =====
 * feature: Added context managers on segment/subsegment capture. `PR97 <https://github.com/aws/aws-xray-sdk-python/pull/97>`_.
@@ -32,11 +36,11 @@ CHANGELOG
 * **Breaking**: The original sampling modules for local defined rules are moved from `models.sampling` to `models.sampling.local`.
 * **Breaking**: The default behavior of `patch_all` changed to selectively patches libraries to avoid double patching. You can use `patch_all(double_patch=True)` to force it to patch ALL supported libraries. See more details on `ISSUE63 <https://github.com/aws/aws-xray-sdk-python/issues/63>`_
 * **Breaking**: The latest `botocore` that has new X-Ray service API `GetSamplingRules` and `GetSamplingTargets` are required.
-* **Breaking**: Version 2.x doesn't support pynamodb and aiobotocore as it requires botocore >= 1.11.3 which isn’t currently supported by the pynamodb and aiobotocore libraries. Please continue to use version 1.x if you’re using pynamodb or aiobotocore until those haven been updated to use botocore > = 1.11.3. 
+* **Breaking**: Version 2.x doesn't support pynamodb and aiobotocore as it requires botocore >= 1.11.3 which isn’t currently supported by the pynamodb and aiobotocore libraries. Please continue to use version 1.x if you’re using pynamodb or aiobotocore until those haven been updated to use botocore > = 1.11.3.
 * feature: Environment variable `AWS_XRAY_DAEMON_ADDRESS` now takes an additional notation in `tcp:127.0.0.1:2000 udp:127.0.0.2:2001` to set TCP and UDP destination separately. By default it assumes a X-Ray daemon listening to both UDP and TCP traffic on `127.0.0.1:2000`.
 * feature: Added MongoDB python client support. `PR65 <https://github.com/aws/aws-xray-sdk-python/pull/65>`_.
-* bugfix: Support binding connection in sqlalchemy as well as engine. `PR78 <https://github.com/aws/aws-xray-sdk-python/pull/78>`_. 
-* bugfix: Flask middleware safe request teardown. `ISSUE75 <https://github.com/aws/aws-xray-sdk-python/issues/75>`_. 
+* bugfix: Support binding connection in sqlalchemy as well as engine. `PR78 <https://github.com/aws/aws-xray-sdk-python/pull/78>`_.
+* bugfix: Flask middleware safe request teardown. `ISSUE75 <https://github.com/aws/aws-xray-sdk-python/issues/75>`_.
 
 
 1.1.2
@@ -68,7 +72,7 @@ CHANGELOG
 * bugfix: Fixed an issue where arbitrary fields in trace header being dropped when calling downstream.
 * bugfix: Fixed a compatibility issue between botocore and httplib patcher. `ISSUE48 <https://github.com/aws/aws-xray-sdk-python/issues/48>`_.
 * bugfix: Fixed a typo in sqlalchemy decorators. `PR50 <https://github.com/aws/aws-xray-sdk-python/pull/50>`_.
-* Updated `README` with more usage examples. 
+* Updated `README` with more usage examples.
 
 0.97
 ====

README.md

@@ -260,7 +260,29 @@ libs_to_patch = ('boto3', 'mysql', 'requests')
 patch(libs_to_patch)
 ```
 
-### Add Django middleware
+#### Automatic module patching
+
+Full modules in the local codebase can be recursively patched by providing the module references
+to the patch function.
+```python
+from aws_xray_sdk.core import patch
+
+libs_to_patch = ('boto3', 'requests', 'local.module.ref', 'other_module')
+patch(libs_to_patch)
+```
+An `xray_recorder.capture()` decorator will be applied to all functions and class methods in the
+given module and all the modules inside them recursively. Some files/modules can be excluded by
+providing to the `patch` function a regex that matches them.
+```python
+from aws_xray_sdk.core import patch
+
+libs_to_patch = ('boto3', 'requests', 'local.module.ref', 'other_module')
+ignore = ('local.module.ref.some_file', 'other_module.some_module\.*')
+patch(libs_to_patch, ignore_module_patterns=ignore)
+```
+
+### Django
+#### Add Django middleware
 
 In django settings.py, use the following.
 
@@ -276,6 +298,27 @@ MIDDLEWARE = [
 ]
 ```
 
+#### Automatic patching
+The automatic module patching can also be configured through Django settings.
+```python
+XRAY_RECORDER = {
+    'PATCH_MODULES': [
+        'boto3',
+        'requests',
+        'local.module.ref',
+        'other_module',
+    ],
+    'IGNORE_MODULE_PATTERNS': [
+        'local.module.ref.some_file',
+        'other_module.some_module\.*',
+    ],
+    ...
+}
+```
+If `AUTO_PATCH_PARENT_SEGMENT_NAME` is also specified, then a segment parent will be created 
+with the supplied name, wrapping the automatic patching so that it captures any dangling
+subsegments created on the import patching.
+
 ### Add Flask middleware
 
 ```python

aws_xray_sdk/core/async_recorder.py

@@ -1,10 +1,8 @@
 import time
 
-import wrapt
-
 from aws_xray_sdk.core.recorder import AWSXRayRecorder
 from aws_xray_sdk.core.utils import stacktrace
-from aws_xray_sdk.core.models.subsegment import SubsegmentContextManager
+from aws_xray_sdk.core.models.subsegment import SubsegmentContextManager, is_already_recording, subsegment_decorator
 from aws_xray_sdk.core.models.segment import SegmentContextManager
 
 
@@ -17,8 +15,13 @@ async def __aexit__(self, exc_type, exc_val, exc_tb):
 
 class AsyncSubsegmentContextManager(SubsegmentContextManager):
 
-    @wrapt.decorator
+    @subsegment_decorator
     async def __call__(self, wrapped, instance, args, kwargs):
+        if is_already_recording(wrapped):
+            # The wrapped function is already decorated, the subsegment will be created later,
+            # just return the result
+            return await wrapped(*args, **kwargs)
+
         func_name = self.name
         if not func_name:
             func_name = wrapped.__name__

aws_xray_sdk/core/models/subsegment.py

@@ -7,6 +7,28 @@
 from ..exceptions.exceptions import SegmentNotFoundException
 
 
+# Attribute starts with _self_ to prevent wrapt proxying to underlying function
+SUBSEGMENT_RECORDING_ATTRIBUTE = '_self___SUBSEGMENT_RECORDING_ATTRIBUTE__'
+
+
+def set_as_recording(decorated_func, wrapped):
+    # If the wrapped function has the attribute, then it has already been patched
+    setattr(decorated_func, SUBSEGMENT_RECORDING_ATTRIBUTE, hasattr(wrapped, SUBSEGMENT_RECORDING_ATTRIBUTE))
+
+
+def is_already_recording(func):
+    # The function might have the attribute, but its value might still be false
+    # as it might be the first decorator
+    return getattr(func, SUBSEGMENT_RECORDING_ATTRIBUTE, False)
+
+
+@wrapt.decorator
+def subsegment_decorator(wrapped, instance, args, kwargs):
+    decorated_func = wrapt.decorator(wrapped)(*args, **kwargs)
+    set_as_recording(decorated_func, wrapped)
+    return decorated_func
+
+
 class SubsegmentContextManager:
     """
     Wrapper for segment and recorder to provide segment context manager.
@@ -18,8 +40,13 @@ def __init__(self, recorder, name=None, **subsegment_kwargs):
         self.recorder = recorder
         self.subsegment = None
 
-    @wrapt.decorator
+    @subsegment_decorator
     def __call__(self, wrapped, instance, args, kwargs):
+        if is_already_recording(wrapped):
+            # The wrapped function is already decorated, the subsegment will be created later,
+            # just return the result
+            return wrapped(*args, **kwargs)
+
         func_name = self.name
         if not func_name:
             func_name = wrapped.__name__

aws_xray_sdk/core/patcher.py

@@ -1,5 +1,13 @@
-import logging
 import importlib
+import inspect
+import logging
+import os
+import pkgutil
+import re
+import sys
+import wrapt
+
+from .utils.compat import PY2, is_classmethod, is_instance_method
 
 log = logging.getLogger(__name__)
 
@@ -36,7 +44,22 @@ def patch_all(double_patch=False):
         patch(NO_DOUBLE_PATCH, raise_errors=False)
 
 
-def patch(modules_to_patch, raise_errors=True):
+def _is_valid_import(module):
+    module = module.replace('.', '/')
+    if PY2:
+        return bool(pkgutil.get_loader(module))
+    else:
+        realpath = os.path.realpath(module)
+        is_module = os.path.isdir(realpath) and (
+            os.path.isfile('{}/__init__.py'.format(module)) or os.path.isfile('{}/__init__.pyc'.format(module))
+        )
+        is_file = not is_module and (
+                os.path.isfile('{}.py'.format(module)) or os.path.isfile('{}.pyc'.format(module))
+        )
+        return is_module or is_file
+
+
+def patch(modules_to_patch, raise_errors=True, ignore_module_patterns=None):
     modules = set()
     for module_to_patch in modules_to_patch:
         # boto3 depends on botocore and patching botocore is sufficient
@@ -51,14 +74,24 @@ def patch(modules_to_patch, raise_errors=True):
             modules.add(module_to_patch)
         else:
             modules.add(module_to_patch)
-    unsupported_modules = modules - set(SUPPORTED_MODULES)
+
+    unsupported_modules = set(module for module in modules if module not in SUPPORTED_MODULES)
+    native_modules = modules - unsupported_modules
+
+    external_modules = set(module for module in unsupported_modules if _is_valid_import(module))
+    unsupported_modules = unsupported_modules - external_modules
+
     if unsupported_modules:
         raise Exception('modules %s are currently not supported for patching'
                         % ', '.join(unsupported_modules))
 
-    for m in modules:
+    for m in native_modules:
         _patch_module(m, raise_errors)
 
+    ignore_module_patterns = [re.compile(pattern) for pattern in ignore_module_patterns or []]
+    for m in external_modules:
+        _external_module_patch(m, ignore_module_patterns)
+
 
 def _patch_module(module_to_patch, raise_errors=True):
     try:
@@ -82,3 +115,93 @@ def _patch(module_to_patch):
 
     _PATCHED_MODULES.add(module_to_patch)
     log.info('successfully patched module %s', module_to_patch)
+
+
+def _patch_func(parent, func_name, func, modifier=lambda x: x):
+    if func_name not in parent.__dict__:
+        # Ignore functions not directly defined in parent, i.e. exclude inherited ones
+        return
+
+    from aws_xray_sdk.core import xray_recorder
+
+    capture_name = func_name
+    if func_name.startswith('__') and func_name.endswith('__'):
+        capture_name = '{}.{}'.format(parent.__name__, capture_name)
+    setattr(parent, func_name, modifier(xray_recorder.capture(name=capture_name)(func)))
+
+
+def _patch_class(module, cls):
+    for member_name, member in inspect.getmembers(cls, inspect.isclass):
+        if member.__module__ == module.__name__:
+            # Only patch classes of the module, ignore imports
+            _patch_class(module, member)
+
+    for member_name, member in inspect.getmembers(cls, inspect.ismethod):
+        if member.__module__ == module.__name__:
+            # Only patch methods of the class defined in the module, ignore other modules
+            if is_classmethod(member):
+                # classmethods are internally generated through descriptors. The classmethod
+                # decorator must be the last applied, so we cannot apply another one on top
+                log.warning('Cannot automatically patch classmethod %s.%s, '
+                            'please apply decorator manually', cls.__name__, member_name)
+            else:
+                _patch_func(cls, member_name, member)
+
+    for member_name, member in inspect.getmembers(cls, inspect.isfunction):
+        if member.__module__ == module.__name__:
+            # Only patch static methods of the class defined in the module, ignore other modules
+            if is_instance_method(cls, member_name, member):
+                _patch_func(cls, member_name, member)
+            else:
+                _patch_func(cls, member_name, member, modifier=staticmethod)
+
+
+def _on_import(module):
+    for member_name, member in inspect.getmembers(module, inspect.isfunction):
+        if member.__module__ == module.__name__:
+            # Only patch functions of the module, ignore imports
+            _patch_func(module, member_name, member)
+
+    for member_name, member in inspect.getmembers(module, inspect.isclass):
+        if member.__module__ == module.__name__:
+            # Only patch classes of the module, ignore imports
+            _patch_class(module, member)
+
+
+def _external_module_patch(module, ignore_module_patterns):
+    if module.startswith('.'):
+        raise Exception('relative packages not supported for patching: {}'.format(module))
+
+    if module in _PATCHED_MODULES:
+        log.debug('%s already patched', module)
+    elif any(pattern.match(module) for pattern in ignore_module_patterns):
+        log.debug('%s ignored due to rules: %s', module, ignore_module_patterns)
+    else:
+        if module in sys.modules:
+            _on_import(sys.modules[module])
+        else:
+            wrapt.importer.when_imported(module)(_on_import)
+
+    for loader, submodule_name, is_module in pkgutil.iter_modules([module.replace('.', '/')]):
+        submodule = '.'.join([module, submodule_name])
+        if is_module:
+            _external_module_patch(submodule, ignore_module_patterns)
+        else:
+            if submodule in _PATCHED_MODULES:
+                log.debug('%s already patched', submodule)
+                continue
+            elif any(pattern.match(submodule) for pattern in ignore_module_patterns):
+                log.debug('%s ignored due to rules: %s', submodule, ignore_module_patterns)
+                continue
+
+            if submodule in sys.modules:
+                _on_import(sys.modules[submodule])
+            else:
+                wrapt.importer.when_imported(submodule)(_on_import)
+
+            _PATCHED_MODULES.add(submodule)
+            log.info('successfully patched module %s', submodule)
+
+    if module not in _PATCHED_MODULES:
+        _PATCHED_MODULES.add(module)
+        log.info('successfully patched module %s', module)

aws_xray_sdk/core/recorder.py

@@ -190,7 +190,7 @@ def in_subsegment(self, name=None, **subsegment_kwargs):
         Return a subsegment context manger.
 
         :param str name: the name of the subsegment
-        :param dict segment_kwargs: remaining arguments passed directly to `begin_subsegment`
+        :param dict subsegment_kwargs: remaining arguments passed directly to `begin_subsegment`
         """
         return SubsegmentContextManager(self, name=name, **subsegment_kwargs)
 

aws_xray_sdk/core/utils/compat.py

@@ -1,3 +1,4 @@
+import inspect
 import sys
 
 
@@ -10,3 +11,21 @@
 else:
     annotation_value_types = (int, float, bool, str)
     string_types = str
+
+
+def is_classmethod(func):
+    return getattr(func, '__self__', None) is not None
+
+
+def is_instance_method(parent_class, func_name, func):
+    try:
+        func_from_dict = parent_class.__dict__[func_name]
+    except KeyError:
+        for base in inspect.getmro(parent_class):
+            if func_name in base.__dict__:
+                func_from_dict = base.__dict__[func_name]
+                break
+        else:
+            return True
+
+    return not is_classmethod(func) and not isinstance(func_from_dict, staticmethod)

aws_xray_sdk/ext/django/apps.py

@@ -5,7 +5,7 @@
 from .conf import settings
 from .db import patch_db
 from .templates import patch_template
-from aws_xray_sdk.core import xray_recorder
+from aws_xray_sdk.core import patch, xray_recorder
 from aws_xray_sdk.core.exceptions.exceptions import SegmentNameMissingException
 
 
@@ -38,6 +38,13 @@ def ready(self):
             max_trace_back=settings.MAX_TRACE_BACK,
         )
 
+        if settings.PATCH_MODULES:
+            if settings.AUTO_PATCH_PARENT_SEGMENT_NAME is not None:
+                with xray_recorder.in_segment(settings.AUTO_PATCH_PARENT_SEGMENT_NAME):
+                    patch(settings.PATCH_MODULES, ignore_module_patterns=settings.IGNORE_MODULE_PATTERNS)
+            else:
+                patch(settings.PATCH_MODULES, ignore_module_patterns=settings.IGNORE_MODULE_PATTERNS)
+
         # if turned on subsegment will be generated on
         # built-in database and template rendering
         if settings.AUTO_INSTRUMENT:

aws_xray_sdk/ext/django/conf.py

@@ -14,6 +14,9 @@
     'DYNAMIC_NAMING': None,
     'STREAMING_THRESHOLD': None,
     'MAX_TRACE_BACK': None,
+    'PATCH_MODULES': [],
+    'AUTO_PATCH_PARENT_SEGMENT_NAME': None,
+    'IGNORE_MODULE_PATTERNS': [],
 }
 
 XRAY_NAMESPACE = 'XRAY_RECORDER'