feat(tracing): Add tracestate feature flag (#1182)

lobsterkatie · web-flow · commit 53de6aa1cc3e · 2021-09-07T15:07:35.000-07:00
This adds a feature flag, `"propagate_tracestate"`, to `_experiments`, which enables the tracestate header handling behavior.

In particular, if the flag is `False` or missing:

- `tracestate` headers won't be attached to outgoing HTTP requests, and
- `trace` data will not be added to envelope headers.
diff --git a/sentry_sdk/client.py b/sentry_sdk/client.py
@@ -22,7 +22,7 @@
 from sentry_sdk.utils import ContextVar
 from sentry_sdk.sessions import SessionFlusher
 from sentry_sdk.envelope import Envelope
-from sentry_sdk.tracing_utils import reinflate_tracestate
+from sentry_sdk.tracing_utils import has_tracestate_enabled, reinflate_tracestate
 
 from sentry_sdk._types import MYPY
 
@@ -349,7 +349,7 @@ def capture_event(
             tracestate_data = raw_tracestate and reinflate_tracestate(
                 raw_tracestate.replace("sentry=", "")
             )
-            if tracestate_data:
+            if tracestate_data and has_tracestate_enabled():
                 headers["trace"] = tracestate_data
 
             envelope = Envelope(headers=headers)
diff --git a/sentry_sdk/consts.py b/sentry_sdk/consts.py
@@ -32,6 +32,7 @@
             "max_spans": Optional[int],
             "record_sql_params": Optional[bool],
             "smart_transaction_trimming": Optional[bool],
+            "propagate_tracestate": Optional[bool],
         },
         total=False,
     )
diff --git a/sentry_sdk/tracing.py b/sentry_sdk/tracing.py
@@ -12,6 +12,7 @@
     compute_tracestate_entry,
     extract_sentrytrace_data,
     extract_tracestate_data,
+    has_tracestate_enabled,
     has_tracing_enabled,
     is_valid_sample_rate,
     maybe_create_breadcrumbs_from_span,
@@ -270,8 +271,10 @@ def iter_headers(self):
         """
         yield "sentry-trace", self.to_traceparent()
 
-        tracestate = self.to_tracestate()
+        tracestate = self.to_tracestate() if has_tracestate_enabled(self) else None
         # `tracestate` will only be `None` if there's no client or no DSN
+        # TODO (kmclb) the above will be true once the feature is no longer
+        # behind a flag
         if tracestate:
             yield "tracestate", tracestate
 
diff --git a/sentry_sdk/tracing_utils.py b/sentry_sdk/tracing_utils.py
@@ -396,3 +396,12 @@ def _format_sql(cursor, sql):
         real_sql = None
 
     return real_sql or to_string(sql)
+
+
+def has_tracestate_enabled(span=None):
+    # type: (Optional[Span]) -> bool
+
+    client = ((span and span.hub) or sentry_sdk.Hub.current).client
+    options = client and client.options
+
+    return bool(options and options["_experiments"].get("propagate_tracestate"))
diff --git a/tests/test_envelope.py b/tests/test_envelope.py
@@ -4,6 +4,13 @@
 from sentry_sdk.tracing_utils import compute_tracestate_value
 import sentry_sdk.client
 
+import pytest
+
+try:
+    from unittest import mock  # python 3.3 and above
+except ImportError:
+    import mock  # python < 3.3
+
 
 def generate_transaction_item():
     return {
@@ -25,6 +32,8 @@ def generate_transaction_item():
                         "environment": "dogpark",
                         "release": "off.leash.park",
                         "public_key": "dogsarebadatkeepingsecrets",
+                        "user": {"id": 12312013, "segment": "bigs"},
+                        "transaction": "/interactions/other-dogs/new-dog",
                     }
                 ),
             }
@@ -79,13 +88,23 @@ def test_add_and_get_session():
             assert item.payload.json == expected.to_json()
 
 
-def test_envelope_headers(sentry_init, capture_envelopes, monkeypatch):
+# TODO (kmclb) remove this parameterization once tracestate is a real feature
+@pytest.mark.parametrize("tracestate_enabled", [True, False])
+def test_envelope_headers(
+    sentry_init, capture_envelopes, monkeypatch, tracestate_enabled
+):
     monkeypatch.setattr(
         sentry_sdk.client,
         "format_timestamp",
         lambda x: "2012-11-21T12:31:12.415908Z",
     )
 
+    monkeypatch.setattr(
+        sentry_sdk.client,
+        "has_tracestate_enabled",
+        mock.Mock(return_value=tracestate_enabled),
+    )
+
     sentry_init(
         dsn="https://dogsarebadatkeepingsecrets@squirrelchasers.ingest.sentry.io/12312012",
     )
@@ -95,13 +114,21 @@ def test_envelope_headers(sentry_init, capture_envelopes, monkeypatch):
 
     assert len(envelopes) == 1
 
-    assert envelopes[0].headers == {
-        "event_id": "15210411201320122115110420122013",
-        "sent_at": "2012-11-21T12:31:12.415908Z",
-        "trace": {
-            "trace_id": "12312012123120121231201212312012",
-            "environment": "dogpark",
-            "release": "off.leash.park",
-            "public_key": "dogsarebadatkeepingsecrets",
-        },
-    }
+    if tracestate_enabled:
+        assert envelopes[0].headers == {
+            "event_id": "15210411201320122115110420122013",
+            "sent_at": "2012-11-21T12:31:12.415908Z",
+            "trace": {
+                "trace_id": "12312012123120121231201212312012",
+                "environment": "dogpark",
+                "release": "off.leash.park",
+                "public_key": "dogsarebadatkeepingsecrets",
+                "user": {"id": 12312013, "segment": "bigs"},
+                "transaction": "/interactions/other-dogs/new-dog",
+            },
+        }
+    else:
+        assert envelopes[0].headers == {
+            "event_id": "15210411201320122115110420122013",
+            "sent_at": "2012-11-21T12:31:12.415908Z",
+        }
diff --git a/tests/tracing/test_http_headers.py b/tests/tracing/test_http_headers.py
@@ -282,7 +282,9 @@ def test_tracestate_extraction(
     }
 
 
-def test_iter_headers(sentry_init, monkeypatch):
+# TODO (kmclb) remove this parameterization once tracestate is a real feature
+@pytest.mark.parametrize("tracestate_enabled", [True, False])
+def test_iter_headers(sentry_init, monkeypatch, tracestate_enabled):
     monkeypatch.setattr(
         Transaction,
         "to_traceparent",
@@ -293,6 +295,11 @@ def test_iter_headers(sentry_init, monkeypatch):
         "to_tracestate",
         mock.Mock(return_value="sentry=doGsaREgReaT,charlie=goofy"),
     )
+    monkeypatch.setattr(
+        sentry_sdk.tracing,
+        "has_tracestate_enabled",
+        mock.Mock(return_value=tracestate_enabled),
+    )
 
     transaction = Transaction(
         name="/interactions/other-dogs/new-dog",
@@ -303,7 +310,11 @@ def test_iter_headers(sentry_init, monkeypatch):
     assert (
         headers["sentry-trace"] == "12312012123120121231201212312012-0415201309082013-0"
     )
-    assert headers["tracestate"] == "sentry=doGsaREgReaT,charlie=goofy"
+    if tracestate_enabled:
+        assert "tracestate" in headers
+        assert headers["tracestate"] == "sentry=doGsaREgReaT,charlie=goofy"
+    else:
+        assert "tracestate" not in headers
 
 
 @pytest.mark.parametrize(
diff --git a/tests/tracing/test_misc.py b/tests/tracing/test_misc.py
@@ -2,6 +2,7 @@
 
 from sentry_sdk import Hub, start_span, start_transaction
 from sentry_sdk.tracing import Span, Transaction
+from sentry_sdk.tracing_utils import has_tracestate_enabled
 
 
 def test_span_trimming(sentry_init, capture_events):
@@ -149,3 +150,19 @@ def test_finds_non_orphan_span_on_scope(sentry_init):
     assert scope._span is not None
     assert isinstance(scope._span, Span)
     assert scope._span.op == "sniffing"
+
+
+# TODO (kmclb) remove this test once tracestate is a real feature
+@pytest.mark.parametrize("tracestate_enabled", [True, False, None])
+def test_has_tracestate_enabled(sentry_init, tracestate_enabled):
+    experiments = (
+        {"propagate_tracestate": tracestate_enabled}
+        if tracestate_enabled is not None
+        else {}
+    )
+    sentry_init(_experiments=experiments)
+
+    if tracestate_enabled is True:
+        assert has_tracestate_enabled() is True
+    else:
+        assert has_tracestate_enabled() is False

Original file line number	Diff line number	Diff line change
`@@ -32,6 +32,7 @@`
`32`	`32`	`"max_spans": Optional[int],`
`33`	`33`	`"record_sql_params": Optional[bool],`
`34`	`34`	`"smart_transaction_trimming": Optional[bool],`
	`35`	`+ "propagate_tracestate": Optional[bool],`
`35`	`36`	`},`
`36`	`37`	`total=False,`
`37`	`38`	`)`