Skip to content

Commit 6fe881b

Browse files
authored
Emmett.butler/dd trace api wrapping (#12261)
## Checklist - [ ] PR author has checked that all the criteria below are met - The PR description includes an overview of the change - The PR description articulates the motivation for the change - The change includes tests OR the PR description describes a testing strategy - The PR description notes risks associated with the change, if any - Newly-added code is easy to change - The change follows the [library release note guidelines](https://ddtrace.readthedocs.io/en/stable/releasenotes.html) - The change includes or references documentation updates if necessary - Backport labels are set (if [applicable](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting)) ## Reviewer Checklist - [ ] Reviewer has checked that all the criteria below are met - Title is accurate - All changes are related to the pull request's stated goal - Avoids breaking [API](https://ddtrace.readthedocs.io/en/stable/versioning.html#interfaces) changes - Testing strategy adequately addresses listed risks - Newly-added code is easy to change - Release note makes sense to a user of the library - If necessary, author has acknowledged and discussed the performance implications of this PR as reported in the benchmarks PR comment - Backport labels are set in a manner that is consistent with the [release branch maintenance policy](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting)
1 parent 94a7494 commit 6fe881b

File tree

1 file changed

+73
-15
lines changed
  • ddtrace/contrib/internal/dd_trace_api

1 file changed

+73
-15
lines changed

ddtrace/contrib/internal/dd_trace_api/patch.py

Lines changed: 73 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -1,20 +1,62 @@
1-
from sys import addaudithook
1+
import inspect
22
from typing import Any
33
from typing import Dict
44
from typing import List
55
from typing import Optional
66
from typing import Tuple
7+
from typing import TypeVar
78
import weakref
89

910
import dd_trace_api
11+
from wrapt.importer import when_imported
1012

1113
import ddtrace
14+
from ddtrace.internal.logger import get_logger
15+
from ddtrace.internal.wrapping.context import WrappingContext
1216

1317

1418
_DD_HOOK_NAME = "dd.hook"
1519
_TRACER_KEY = "Tracer"
1620
_STUB_TO_REAL = weakref.WeakKeyDictionary()
1721
_STUB_TO_REAL[dd_trace_api.tracer] = ddtrace.tracer
22+
log = get_logger(__name__)
23+
T = TypeVar("T")
24+
25+
26+
class DDTraceAPIWrappingContextBase(WrappingContext):
27+
def _handle_return(self) -> None:
28+
_call_on_real_instance(
29+
self.get_local("self"),
30+
self.__frame__.f_code.co_name,
31+
self.get_local("retval"),
32+
**{
33+
param: self.get_local(param)
34+
for param in inspect.signature(self.__wrapped__).parameters.keys()
35+
if param != "self"
36+
},
37+
)
38+
39+
def _handle_enter(self) -> None:
40+
pass
41+
42+
def __enter__(self) -> "DDTraceAPIWrappingContextBase":
43+
super().__enter__()
44+
45+
try:
46+
self._handle_enter()
47+
except Exception: # noqa: E722
48+
log.debug("Error handling dd_trace_api instrumentation enter", exc_info=True)
49+
50+
return self
51+
52+
def __return__(self, value: T) -> T:
53+
"""Always return the original value no matter what our instrumentation does"""
54+
try:
55+
self._handle_return()
56+
except Exception: # noqa: E722
57+
log.debug("Error handling instrumentation return", exc_info=True)
58+
59+
return value
1860

1961

2062
def _proxy_span_arguments(args: List, kwargs: Dict) -> Tuple[List, Dict]:
@@ -41,31 +83,47 @@ def _call_on_real_instance(
4183
_STUB_TO_REAL[retval_from_api] = retval_from_impl
4284

4385

44-
def _hook(name, hook_args):
45-
"""Called in response to `sys.audit` events"""
46-
if name != _DD_HOOK_NAME or not dd_trace_api.__datadog_patch:
47-
return
48-
args = hook_args[0][0]
49-
api_return_value, stub_self, method_name = args[0:3]
50-
_call_on_real_instance(stub_self, method_name, api_return_value, *args[3:], **hook_args[0][1])
51-
52-
5386
def get_version() -> str:
5487
return getattr(dd_trace_api, "__version__", "")
5588

5689

5790
def patch(tracer=None):
5891
if getattr(dd_trace_api, "__datadog_patch", False):
5992
return
60-
dd_trace_api.__datadog_patch = True
6193
_STUB_TO_REAL[dd_trace_api.tracer] = tracer
62-
if not getattr(dd_trace_api, "__dd_has_audit_hook", False):
63-
addaudithook(_hook)
64-
dd_trace_api.__dd_has_audit_hook = True
94+
95+
@when_imported("dd_trace_api")
96+
def _(m):
97+
DDTraceAPIWrappingContextBase(m.Tracer.start_span).wrap()
98+
DDTraceAPIWrappingContextBase(m.Tracer.trace).wrap()
99+
DDTraceAPIWrappingContextBase(m.Tracer.current_span).wrap()
100+
DDTraceAPIWrappingContextBase(m.Tracer.current_root_span).wrap()
101+
DDTraceAPIWrappingContextBase(m.Span.finish).wrap()
102+
DDTraceAPIWrappingContextBase(m.Span.set_exc_info).wrap()
103+
DDTraceAPIWrappingContextBase(m.Span.finish_with_ancestors).wrap()
104+
DDTraceAPIWrappingContextBase(m.Span.set_tags).wrap()
105+
DDTraceAPIWrappingContextBase(m.Span.set_traceback).wrap()
106+
DDTraceAPIWrappingContextBase(m.Span.__enter__).wrap()
107+
DDTraceAPIWrappingContextBase(m.Span.__exit__).wrap()
108+
109+
dd_trace_api.__datadog_patch = True
65110

66111

67112
def unpatch():
68113
if not getattr(dd_trace_api, "__datadog_patch", False):
69114
return
70115
dd_trace_api.__datadog_patch = False
71-
# NB sys.addaudithook's cannot be removed
116+
117+
DDTraceAPIWrappingContextBase.extract(dd_trace_api.Tracer.start_span).unwrap()
118+
DDTraceAPIWrappingContextBase.extract(dd_trace_api.Tracer.trace).unwrap()
119+
DDTraceAPIWrappingContextBase.extract(dd_trace_api.Tracer.current_span).unwrap()
120+
DDTraceAPIWrappingContextBase.extract(dd_trace_api.Tracer.current_root_span).unwrap()
121+
DDTraceAPIWrappingContextBase.extract(dd_trace_api.Span.finish).unwrap()
122+
DDTraceAPIWrappingContextBase.extract(dd_trace_api.Span.set_exc_info).unwrap()
123+
DDTraceAPIWrappingContextBase.extract(dd_trace_api.Span.finish_with_ancestors).unwrap()
124+
DDTraceAPIWrappingContextBase.extract(dd_trace_api.Span.set_tags).unwrap()
125+
DDTraceAPIWrappingContextBase.extract(dd_trace_api.Span.set_traceback).unwrap()
126+
DDTraceAPIWrappingContextBase.extract(dd_trace_api.Span.__enter__).unwrap()
127+
DDTraceAPIWrappingContextBase.extract(dd_trace_api.Span.__exit__).unwrap()
128+
129+
dd_trace_api.__datadog_patch = False

0 commit comments

Comments
 (0)