import importlib import inspect import logging import os import pkgutil import re import sys import wrapt from aws_xray_sdk import global_sdk_config from .utils.compat import PY2, is_classmethod, is_instance_method log = logging.getLogger(__name__) SUPPORTED_MODULES = ( 'aiobotocore', 'botocore', 'pynamodb', 'requests', 'sqlite3', 'mysql', 'httplib', 'pymongo', 'pymysql', 'psycopg2', 'pg8000', 'sqlalchemy_core', ) NO_DOUBLE_PATCH = ( 'aiobotocore', 'botocore', 'pynamodb', 'requests', 'sqlite3', 'mysql', 'pymongo', 'pymysql', 'psycopg2', 'pg8000', 'sqlalchemy_core', ) _PATCHED_MODULES = set() def patch_all(double_patch=False): if double_patch: patch(SUPPORTED_MODULES, raise_errors=False) else: patch(NO_DOUBLE_PATCH, raise_errors=False) 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): enabled = global_sdk_config.sdk_enabled() if not enabled: log.debug("Skipped patching modules %s because the SDK is currently disabled." % ', '.join(modules_to_patch)) return # Disable module patching if the SDK is disabled. modules = set() for module_to_patch in modules_to_patch: # boto3 depends on botocore and patching botocore is sufficient if module_to_patch == 'boto3': modules.add('botocore') # aioboto3 depends on aiobotocore and patching aiobotocore is sufficient elif module_to_patch == 'aioboto3': modules.add('aiobotocore') # pynamodb requires botocore to be patched as well elif module_to_patch == 'pynamodb': modules.add('botocore') modules.add(module_to_patch) else: modules.add(module_to_patch) 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 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: _patch(module_to_patch) except Exception: if raise_errors: raise log.debug('failed to patch module %s', module_to_patch) def _patch(module_to_patch): path = 'aws_xray_sdk.ext.%s' % module_to_patch if module_to_patch in _PATCHED_MODULES: log.debug('%s already patched', module_to_patch) return imported_module = importlib.import_module(path) imported_module.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)