def instrument(): httplib = six.moves.http_client def dynatrace_putrequest(wrapped, instance, args, kwargs): method, path = args[:2] scheme = "https" if "HTTPS" in instance.__class__.__name__ else "http" url = "{}://{}{}{}".format( scheme, instance.host, ":{}".format(instance.port) if str(instance.port) not in ["80", "443"] else "", path ) tracer = sdk.trace_outgoing_web_request(url, method) tracer.start() setattr(instance, "__dynatrace_tracer", tracer) ret = wrapped(*args, **kwargs) tag = tracer.outgoing_dynatrace_string_tag logger.debug("Tracing urllib, url: '{}', tag: '{}'".format(url, tag)) instance.putheader("x-dynatrace", tag) return ret def dynatrace_getresponse(wrapped, instance, args, kwargs): tracer = getattr(instance, "__dynatrace_tracer", None) response = wrapped(*args, **kwargs) # print(traceback.print_stack()) if tracer is not None: tracer.set_status_code(response.status) tracer.end() delattr(instance, "__dynatrace_tracer") return response setattr(httplib.HTTPConnection, "putrequest", wrapt.FunctionWrapper(httplib.HTTPConnection.putrequest, dynatrace_putrequest)) setattr( httplib.HTTPConnection, "getresponse", wrapt.FunctionWrapper(httplib.HTTPConnection.getresponse, dynatrace_getresponse) )
def traced_execute_async(func, instance, args, kwargs): cluster = getattr(instance, 'cluster', None) pin = Pin.get_from(cluster) if not pin or not pin.enabled(): return func(*args, **kwargs) query = kwargs.get("query") or args[0] span = _start_span_and_set_tags(pin, query, instance, cluster) try: result = func(*args, **kwargs) setattr(result, CURRENT_SPAN, span) setattr(result, PAGE_NUMBER, 1) setattr( result, '_set_final_result', wrapt.FunctionWrapper( result._set_final_result, traced_set_final_result ) ) setattr( result, '_set_final_exception', wrapt.FunctionWrapper( result._set_final_exception, traced_set_final_exception ) ) setattr( result, 'start_fetching_next_page', wrapt.FunctionWrapper( result.start_fetching_next_page, traced_start_fetching_next_page ) ) # Since we cannot be sure that the previous methods were overwritten # before the call ended, we add callbacks that will be run # synchronously if the call already returned and we remove them right # after. result.add_callbacks( _close_span_on_success, _close_span_on_error, callback_args=(result,), errback_args=(result,) ) result.clear_callbacks() return result except: with span: span.set_exc_info(*sys.exc_info()) raise
def patch(): """ patch the built-in urllib/httplib/httplib.client methods for tracing""" if getattr(httplib, '__datadog_patch', False): return setattr(httplib, '__datadog_patch', True) # Patch the desired methods setattr(httplib.HTTPConnection, '__init__', wrapt.FunctionWrapper(httplib.HTTPConnection.__init__, _wrap_init)) setattr(httplib.HTTPConnection, 'getresponse', wrapt.FunctionWrapper(httplib.HTTPConnection.getresponse, _wrap_getresponse)) setattr(httplib.HTTPConnection, 'putrequest', wrapt.FunctionWrapper(httplib.HTTPConnection.putrequest, _wrap_putrequest))
def test_re_bind_after_none(self): def function(): pass def wrapper(wrapped, instance, args, kwargs): return wrapped(*args, **kwargs) _wrapper = wrapt.FunctionWrapper(function, wrapper) self.assertTrue(isinstance(_wrapper, wrapt.FunctionWrapper)) instance = object() _bound_wrapper_1 = _wrapper.__get__(None, type(instance)) self.assertTrue(_bound_wrapper_1._self_parent is _wrapper) self.assertTrue(isinstance(_bound_wrapper_1, wrapt.BoundFunctionWrapper)) self.assertEqual(_bound_wrapper_1._self_instance, None) _bound_wrapper_2 = _bound_wrapper_1.__get__(instance, type(instance)) self.assertTrue(_bound_wrapper_2._self_parent is _wrapper) self.assertTrue(isinstance(_bound_wrapper_2, wrapt.BoundFunctionWrapper)) self.assertEqual(_bound_wrapper_2._self_instance, instance) self.assertTrue(_bound_wrapper_1 is not _bound_wrapper_2)
def traced_connect(func, instance, args, kwargs): session = func(*args, **kwargs) if not isinstance(session.execute, wrapt.FunctionWrapper): # FIXME[matt] this should probably be private. setattr(session, 'execute', wrapt.FunctionWrapper(session.execute, traced_execute)) return session
def patch_app(app, pin=None): """ patch_app will add tracing to a celery app """ pin = pin or Pin(service=SERVICE, app=APP, app_type=AppTypes.worker) patch_methods = [ ('task', _app_task), ] for method_name, wrapper in patch_methods: # Get the original method method = getattr(app, method_name, None) if method is None: continue # Do not patch if method is already patched if isinstance(method, wrapt.ObjectProxy): continue # Patch method setattr(app, method_name, wrapt.FunctionWrapper(method, wrapper)) # patch the Task class if available setattr(app, 'Task', patch_task(app.Task)) # Attach our pin to the app pin.onto(app) return app
def handle_resource_class(wrapped, instance, args, kwargs): for method in ['__call__', 'on_get', 'on_post', 'on_put', 'on_delete']: m = getattr(wrapped, method, None) if m: setattr(wrapped, method, wrapt.FunctionWrapper(m, handle_resource_func)) return wrapped(*args, **kwargs)
async def read(self): with CeilTimeout(self.__wrapped_read_timeout, loop=self._loop): resp_msg, stream_reader = await super().read() if hasattr(stream_reader, '_wait'): stream_reader._wait = wrapt.FunctionWrapper( stream_reader._wait, self._wrapped_wait) return resp_msg, stream_reader
def _app_task(pin, func, app, args, kwargs): task = func(*args, **kwargs) # `app.task` is a decorator which may return a function wrapper if isinstance(task, types.FunctionType): def wrapper(func, instance, args, kwargs): return patch_task(func(*args, **kwargs), pin=pin) return wrapt.FunctionWrapper(task, wrapper) return patch_task(task, pin=pin)
def test_update_module_modified_on_original(self): def function(): pass def wrapper(wrapped, instance, args, kwargs): return wrapped(*args, **kwargs) instance = wrapt.FunctionWrapper(function, wrapper) self.assertEqual(instance.__module__, __name__) instance.__module__ = 'override_module' self.assertEqual(function.__module__, 'override_module') self.assertEqual(instance.__module__, 'override_module')
def test_update_doc_modified_on_original(self): def function(): """documentation""" pass def wrapper(wrapped, instance, args, kwargs): return wrapped(*args, **kwargs) instance = wrapt.FunctionWrapper(function, wrapper) self.assertEqual(instance.__doc__, "documentation") instance.__doc__ = 'override_doc' self.assertEqual(function.__doc__, 'override_doc') self.assertEqual(instance.__doc__, 'override_doc')
def test_update_qualname_modified_on_original(self): def function(): pass def wrapper(wrapped, instance, args, kwargs): return wrapped(*args, **kwargs) instance = wrapt.FunctionWrapper(function, wrapper) if PY3: method = self.test_update_qualname_modified_on_original self.assertEqual(instance.__qualname__, (method.__qualname__ + '.<locals>.function')) instance.__qualname__ = 'override_qualname' self.assertEqual(function.__qualname__, 'override_qualname') self.assertEqual(instance.__qualname__, 'override_qualname')
def test_update_annotations_modified_on_original(self): def function(): pass def wrapper(wrapped, instance, args, kwargs): return wrapped(*args, **kwargs) instance = wrapt.FunctionWrapper(function, wrapper) if PY3: self.assertEqual(instance.__annotations__, {}) else: def run(*args): instance.__annotations__ self.assertRaises(AttributeError, run, ()) override_annotations = { 'override_annotations': '' } instance.__annotations__ = override_annotations self.assertEqual(function.__annotations__, override_annotations) self.assertEqual(instance.__annotations__, override_annotations)
def patch_task(task, pin=None): """ patch_task will add tracing to a celery task """ # The service set here is actually ignored, because it's not possible to # be certain whether this process is being used as a worker, a producer, # or both. So the service as recorded in traces is set based on the actual # work being done (ie. apply/apply_async vs run). pin = pin or Pin(service=WORKER_SERVICE, app=APP, app_type=AppTypes.worker) patch_methods = [ ('__init__', _task_init), ('run', _task_run), ('apply', _task_apply), ('apply_async', _task_apply_async), ] for method_name, wrapper in patch_methods: # Get original method method = getattr(task, method_name, None) if method is None: continue # Do not patch if method is already patched if isinstance(method, wrapt.ObjectProxy): continue # If the function as been applied as a decorator for v1 Celery tasks, then a different patching is needed if inspect.isclass(task) and issubclass(task, celery.task.Task): wrapped = wrapt.FunctionWrapper(method, wrapper) setattr(task, method_name, wrapped) continue # Patch method # DEV: Using `BoundFunctionWrapper` ensures our `task` wrapper parameter is properly set setattr(task, method_name, wrapt.BoundFunctionWrapper(method, task, wrapper)) # Attach our pin to the app pin.onto(task) return task
def _wrap(cls, method_name, wrapper): original = getattr(cls, method_name) wrapper = wrapt.FunctionWrapper(original, wrapper) wrapt.apply_patch(cls, method_name, wrapper)
def patch(): wrapped = wrapt.FunctionWrapper(_connect, traced_connect) setattr(sqlite3, 'connect', wrapped) setattr(sqlite3.dbapi2, 'connect', wrapped)
def patch(): """ patch will add tracing to the cassandra library. """ setattr(cassandra.cluster.Cluster, 'connect', wrapt.FunctionWrapper(_connect, traced_connect)) Pin(service=SERVICE, app=SERVICE, app_type="db").onto(cassandra.cluster.Cluster)
def _build(wrapped, wrapper, enabled=None, user_function=None): return wrapt.FunctionWrapper(wrapped=wrapped, wrapper=wrapper, enabled=enabled)
def run(*args): def _wrapper(wrapped, instance, args, kwargs): return wrapped(*args, **kwargs) wrapper = wrapt.FunctionWrapper(None, _wrapper) wrapper.__get__(list(), list)()