ref: Instrument SQL, and change schema as discussed

untitaker · untitaker · commit 270ff43ce9b3 · 2019-05-29T11:00:16.000+02:00
diff --git a/sentry_sdk/hub.py b/sentry_sdk/hub.py
@@ -366,8 +366,8 @@ def span(self, span=None, **kwargs):
             span.finish()
             self.capture_trace(span)
 
-    def trace(self, *args, **kwargs):
-        return self.span(self.start_trace(*args, **kwargs))
+    def trace(self, **kwargs):
+        return self.span(self.start_trace(**kwargs))
 
     def start_span(self, **kwargs):
         _, scope = self._stack[-1]
@@ -376,26 +376,34 @@ def start_span(self, **kwargs):
             return span.new_span(**kwargs)
         return None
 
-    def start_trace(self, transaction, **kwargs):
+    def start_trace(self, **kwargs):
+        span = Span.start_trace(**kwargs)
+
         _, scope = self._stack[-1]
-        scope.span = span = Span.start_trace(transaction, **kwargs)
+        scope.span = span
         return span
 
     def capture_trace(self, span):
-        if span.transaction is None:
-            return None
-
-        client = self.client
-
-        if client is None:
+        if (
+            span.transaction is None
+            or span.timestamp is None
+            or self.client is None
+            or span.sampled is False
+        ):
             return None
 
-        sample_rate = client.options["traces_sample_rate"]
+        sample_rate = self.client.options["traces_sample_rate"]
         if sample_rate < 1.0 and random.random() >= sample_rate:
             return None
 
         return self.capture_event(
-            {"type": "none", "spans": [s.to_json() for s in span._finished_spans]}
+            {
+                "type": "transaction",
+                "contexts": {"trace": span.get_trace_context()},
+                "timestamp": span.timestamp,
+                "start_timestamp": span.start_timestamp,
+                "spans": [s.to_json() for s in span._finished_spans if s is not span],
+            }
         )
 
     @overload  # noqa
diff --git a/sentry_sdk/integrations/django/__init__.py b/sentry_sdk/integrations/django/__init__.py
@@ -1,6 +1,7 @@
 # -*- coding: utf-8 -*-
 from __future__ import absolute_import
 
+import contextlib
 import sys
 import weakref
 
@@ -318,10 +319,12 @@ def format_sql(sql, params):
     return sql, rv
 
 
+@contextlib.contextmanager
 def record_sql(sql, params, cursor=None):
     # type: (Any, Any, Any) -> None
     hub = Hub.current
     if hub.get_integration(DjangoIntegration) is None:
+        yield
         return
 
     real_sql = None
@@ -353,9 +356,33 @@ def record_sql(sql, params, cursor=None):
         except Exception:
             pass
 
+    span = None
+
     if real_sql:
         with capture_internal_exceptions():
             hub.add_breadcrumb(message=real_sql, category="query")
+            span = hub.start_span(op="sql.query", description=real_sql)
+
+    if span is None:
+        yield
+    else:
+        try:
+            yield
+        finally:
+            span.set_tag("status", sys.exc_info()[1] is None)
+            span.finish()
+
+
+@contextlib.contextmanager
+def record_many_sql(sql, param_list, cursor):
+    ctxs = [record_sql(sql, params, cursor).__enter__() for params in param_list]
+
+    try:
+        yield
+    finally:
+        einfo = sys.exc_info()
+        for ctx in ctxs:
+            ctx.__exit__(*einfo)
 
 
 def install_sql_hook():
@@ -373,21 +400,13 @@ def install_sql_hook():
         # This won't work on Django versions < 1.6
         return
 
-    def record_many_sql(sql, param_list, cursor):
-        for params in param_list:
-            record_sql(sql, params, cursor)
-
     def execute(self, sql, params=None):
-        try:
+        with record_sql(sql, params, self.cursor):
             return real_execute(self, sql, params)
-        finally:
-            record_sql(sql, params, self.cursor)
 
     def executemany(self, sql, param_list):
-        try:
+        with record_many_sql(sql, param_list, self.cursor):
             return real_executemany(self, sql, param_list)
-        finally:
-            record_many_sql(sql, param_list, self.cursor)
 
     CursorWrapper.execute = execute
     CursorWrapper.executemany = executemany
diff --git a/sentry_sdk/integrations/stdlib.py b/sentry_sdk/integrations/stdlib.py
@@ -28,9 +28,6 @@ def putrequest(self, method, url, *args, **kwargs):
         if hub.get_integration(StdlibIntegration) is None:
             return rv
 
-        self._sentrysdk_data_dict = data = {}
-        self._sentrysdk_span = hub.start_span()
-
         host = self.host
         port = self.port
         default_port = self.default_port
@@ -44,6 +41,11 @@ def putrequest(self, method, url, *args, **kwargs):
                 url,
             )
 
+        self._sentrysdk_data_dict = data = {}
+        self._sentrysdk_span = hub.start_span(
+            op="http", description="%s %s" % (real_url, method)
+        )
+
         for key, value in hub.iter_trace_propagation_headers():
             self.putheader(key, value)
 
@@ -66,6 +68,8 @@ def getresponse(self, *args, **kwargs):
         span = self._sentrysdk_span
         if span is not None:
             span.set_tag("status_code", rv.status)
+            for k, v in data.items():
+                span.set_data(k, v)
             span.finish()
 
         hub.add_breadcrumb(
diff --git a/sentry_sdk/scope.py b/sentry_sdk/scope.py
@@ -212,10 +212,9 @@ def _drop(event, cause, ty):
             event.setdefault("contexts", {}).update(self._contexts)
 
         if self._span is not None:
-            event.setdefault("contexts", {})["trace"] = {
-                "trace_id": self._span.trace_id,
-                "span_id": self._span.span_id,
-            }
+            contexts = event.setdefault("contexts", {})
+            if not contexts.get("trace"):
+                contexts["trace"] = self._span.get_trace_context()
 
         exc_info = hint.get("exc_info") if hint is not None else None
         if exc_info is not None:
diff --git a/sentry_sdk/tracing.py b/sentry_sdk/tracing.py
@@ -24,29 +24,41 @@ class Span(object):
     __slots__ = (
         "trace_id",
         "span_id",
-        "ref",
-        "ref_type",
+        "parent_span_id",
+        "same_process_as_parent",
         "sampled",
         "transaction",
+        "op",
+        "description",
+        "start_timestamp",
+        "timestamp",
         "_tags",
+        "_data",
         "_finished_spans",
-        "start",
-        "end",
     )
 
     def __init__(
-        self, trace_id, span_id, transaction=None, ref=None, ref_type=None, sampled=None
+        self,
+        trace_id,
+        span_id,
+        parent_span_id=None,
+        same_process_as_parent=True,
+        sampled=None,
+        transaction=None,
+        op=None,
+        description=None,
     ):
         self.trace_id = trace_id
         self.span_id = span_id
-        self.ref = ref
-        self.ref_type = ref_type
+        self.parent_span_id = parent_span_id
+        self.same_process_as_parent = same_process_as_parent
         self.sampled = sampled
         self.transaction = transaction
         self._tags = {}
+        self._data = {}
         self._finished_spans = []
-        self.start = datetime.now()
-        self.end = None
+        self.start_timestamp = datetime.now()
+        self.timestamp = None
 
     def __repr__(self):
         return "<%s(transaction=%r, trace_id=%r, span_id=%r, ref=%r)>" % (
@@ -58,24 +70,18 @@ def __repr__(self):
         )
 
     @classmethod
-    def start_trace(cls, transaction=None, sampled=None):
-        return cls(
-            transaction=transaction,
-            trace_id=uuid.uuid4().hex,
-            span_id=uuid.uuid4().hex[16:],
-            sampled=sampled,
-        )
+    def start_trace(cls, **kwargs):
+        return cls(trace_id=uuid.uuid4().hex, span_id=uuid.uuid4().hex[16:], **kwargs)
 
-    def new_span(self, ref_type="child"):
+    def new_span(self, **kwargs):
         if self.trace_id is None:
             return Span.start_trace()
 
         rv = Span(
             trace_id=self.trace_id,
             span_id=uuid.uuid4().hex[16:],
-            ref=self,
-            ref_type=ref_type,
-            sampled=self.sampled,
+            parent_span_id=self.span_id,
+            **kwargs
         )
         rv._finished_spans = self._finished_spans
         return rv
@@ -89,7 +95,7 @@ def continue_from_headers(cls, headers):
         parent = cls.from_traceparent(headers.get("sentry-trace"))
         if parent is None:
             return cls.start_trace()
-        return parent.new_span("follows_from")
+        return parent.new_span(same_process_as_parent=False)
 
     def iter_headers(self):
         yield "sentry-trace", self.to_traceparent()
@@ -110,31 +116,34 @@ def from_traceparent(cls, traceparent):
         if span_id is not None:
             span_id = int(span_id, 16)
         if sampled is not None:
-            sampled = sampled == "1"
+            sampled = sampled != "0"
 
         return cls(trace_id=trace_id, span_id=span_id, sampled=sampled)
 
     def to_traceparent(self):
-        return "%s-%s-%s" % (
-            self.trace_id,
-            self.span_id,
-            "1" if not self.sampled else "0",
-        )
+        return "%s-%s-%s" % (self.trace_id, self.span_id, "1" if self.sampled else "0")
 
     def set_tag(self, key, value):
         self._tags[key] = value
 
+    def set_data(self, key, value):
+        self._data[key] = value
+
     def finish(self):
-        self.end = datetime.now()
+        self.timestamp = datetime.now()
         self._finished_spans.append(self)
 
     def to_json(self):
         return {
             "trace_id": self.trace_id,
             "span_id": self.span_id,
-            "ref_span_id": self.ref and self.ref.span_id or None,
+            "parent_span_id": self.parent_span_id,
             "transaction": self.transaction,
             "tags": self._tags,
-            "start": self.start,
-            "end": self.end,
+            "data": self._data,
+            "start_timestamp": self.start_timestamp,
+            "timestamp": self.timestamp,
         }
+
+    def get_trace_context(self):
+        return {"trace_id": self.trace_id, "span_id": self.span_id}
diff --git a/tests/test_tracing.py b/tests/test_tracing.py
@@ -8,9 +8,10 @@ def test_basic(sentry_init, capture_events, sample_rate):
     sentry_init(traces_sample_rate=sample_rate)
     events = capture_events()
 
-    with Hub.current.trace("hi"):
-        with Hub.current.span():
-            1 / 0
+    with Hub.current.trace(transaction="hi"):
+        with pytest.raises(ZeroDivisionError):
+            with Hub.current.span():
+                1 / 0
 
         with Hub.current.span():
             pass
