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