Add org_id support

Author: sl0thentr0py

sentry_sdk/client.py

@@ -382,6 +382,9 @@ def _record_lost_event(
             _client_init_debug.set(self.options["debug"])
             self.transport = make_transport(self.options)
 
+            if self.options["org_id"] and self.transport and self.transport.parsed_dsn:
+                self.transport.parsed_dsn.org_id = self.options["org_id"]
+
             self.monitor = None
             if self.transport:
                 if self.options["enable_backpressure_handling"]:

sentry_sdk/consts.py

@@ -1023,6 +1023,7 @@ def __init__(
         trace_ignore_status_codes=frozenset(),  # type: AbstractSet[int]
         enable_metrics=True,  # type: bool
         before_send_metric=None,  # type: Optional[Callable[[Metric, Hint], Optional[Metric]]]
+        org_id=None,  # type: Optional[str]
     ):
         # type: (...) -> None
         """Initialize the Sentry SDK with the given parameters. All parameters described here can be used in a call to `sentry_sdk.init()`.
@@ -1426,6 +1427,10 @@ def __init__(
             If `trace_ignore_status_codes` is not provided, requests with any status code
             may be traced.
 
+        :param org_id: An optional organization ID. The SDK will try to extract if from the DSN in most cases
+            but you can provide it explicitly for self-hosted and Relay setups. This value is used for
+            trace propagation and for features like `strict_trace_continuation`.
+
         :param _experiments:
         """
         pass

sentry_sdk/tracing_utils.py

@@ -666,6 +666,8 @@ def from_options(cls, scope):
 
         if client.parsed_dsn:
             sentry_items["public_key"] = client.parsed_dsn.public_key
+            if client.parsed_dsn.org_id:
+                sentry_items["org_id"] = client.parsed_dsn.org_id
 
         if options.get("traces_sample_rate"):
             sentry_items["sample_rate"] = str(options["traces_sample_rate"])
@@ -698,6 +700,8 @@ def populate_from_transaction(cls, transaction):
 
         if client.parsed_dsn:
             sentry_items["public_key"] = client.parsed_dsn.public_key
+            if client.parsed_dsn.org_id:
+                sentry_items["org_id"] = client.parsed_dsn.org_id
 
         if (
             transaction.name

sentry_sdk/utils.py

@@ -294,6 +294,8 @@ class BadDsn(ValueError):
 class Dsn:
     """Represents a DSN."""
 
+    ORG_ID_REGEX = re.compile(r"^o(\d+)\.")
+
     def __init__(self, value):
         # type: (Union[Dsn, str]) -> None
         if isinstance(value, Dsn):
@@ -310,6 +312,9 @@ def __init__(self, value):
 
         self.host = parts.hostname
 
+        org_id_match = Dsn.ORG_ID_REGEX.match(self.host)
+        self.org_id = org_id_match.group(1) if org_id_match else None  # type: Optional[str]
+
         if parts.port is None:
             self.port = self.scheme == "https" and 443 or 80  # type: int
         else:

tests/integrations/opentelemetry/test_span_processor.py

@@ -11,6 +11,7 @@
     SentrySpanProcessor,
     link_trace_context_to_error_event,
 )
+from sentry_sdk.utils import Dsn
 from sentry_sdk.tracing import Span, Transaction
 from sentry_sdk.tracing_utils import extract_sentrytrace_data
 
@@ -23,7 +24,7 @@ def test_is_sentry_span():
 
     client = MagicMock()
     client.options = {"instrumenter": "otel"}
-    client.dsn = "https://[email protected]/123456"
+    client.parsed_dsn = Dsn("https://[email protected]/123456")
     sentry_sdk.get_global_scope().set_client(client)
 
     assert not span_processor._is_sentry_span(otel_span)

tests/test_dsc.py

@@ -13,7 +13,19 @@
 import pytest
 
 import sentry_sdk
-import sentry_sdk.client
+from sentry_sdk.transport import Transport
+from sentry_sdk.envelope import Envelope
+
+
+class TransportWithOptions(Transport):
+    """conftest.TestTransport does not pass in the options so we need this here"""
+
+    def __init__(self, options=None):
+        Transport.__init__(self, options)
+
+    def capture_envelope(self, _: Envelope) -> None:
+        """No-op capture_envelope for tests"""
+        pass
 
 
 def test_dsc_head_of_trace(sentry_init, capture_envelopes):
@@ -22,10 +34,11 @@ def test_dsc_head_of_trace(sentry_init, capture_envelopes):
     and sends a transaction event to Sentry.
     """
     sentry_init(
-        dsn="https://mysecret@bla.ingest.sentry.io/12312012",
+        dsn="https://mysecret@o1234.ingest.sentry.io/12312012",
         release="[email protected]",
         environment="canary",
         traces_sample_rate=1.0,
+        transport=TransportWithOptions,
     )
     envelopes = capture_envelopes()
 
@@ -45,6 +58,10 @@ def test_dsc_head_of_trace(sentry_init, capture_envelopes):
     assert type(envelope_trace_header["public_key"]) == str
     assert envelope_trace_header["public_key"] == "mysecret"
 
+    assert "org_id" in envelope_trace_header
+    assert type(envelope_trace_header["org_id"]) == str
+    assert envelope_trace_header["org_id"] == "1234"
+
     assert "sample_rate" in envelope_trace_header
     assert type(envelope_trace_header["sample_rate"]) == str
     assert envelope_trace_header["sample_rate"] == "1.0"
@@ -66,16 +83,46 @@ def test_dsc_head_of_trace(sentry_init, capture_envelopes):
     assert envelope_trace_header["transaction"] == "foo"
 
 
+def test_dsc_head_of_trace_uses_custom_org_id(sentry_init, capture_envelopes):
+    """
+    Our service is the head of the trace (it starts a new trace)
+    and sends a transaction event to Sentry.
+    """
+    sentry_init(
+        dsn="https://[email protected]/12312012",
+        org_id="9999",
+        release="[email protected]",
+        environment="canary",
+        traces_sample_rate=1.0,
+        transport=TransportWithOptions,
+    )
+    envelopes = capture_envelopes()
+
+    # We start a new transaction
+    with sentry_sdk.start_transaction(name="foo"):
+        pass
+
+    assert len(envelopes) == 1
+
+    transaction_envelope = envelopes[0]
+    envelope_trace_header = transaction_envelope.headers["trace"]
+
+    assert "org_id" in envelope_trace_header
+    assert type(envelope_trace_header["org_id"]) == str
+    assert envelope_trace_header["org_id"] == "9999"
+
+
 def test_dsc_continuation_of_trace(sentry_init, capture_envelopes):
     """
     Another service calls our service and passes tracing information to us.
     Our service is continuing the trace and sends a transaction event to Sentry.
     """
     sentry_init(
-        dsn="https://mysecret@bla.ingest.sentry.io/12312012",
+        dsn="https://mysecret@o1234.ingest.sentry.io/12312012",
         release="[email protected]",
         environment="canary",
         traces_sample_rate=1.0,
+        transport=TransportWithOptions,
     )
     envelopes = capture_envelopes()
 
@@ -149,10 +196,11 @@ def my_traces_sampler(sampling_context):
         return 0.25
 
     sentry_init(
-        dsn="https://mysecret@bla.ingest.sentry.io/12312012",
+        dsn="https://mysecret@o1234.ingest.sentry.io/12312012",
         release="[email protected]",
         environment="canary",
         traces_sampler=my_traces_sampler,
+        transport=TransportWithOptions,
     )
     envelopes = capture_envelopes()
 
@@ -219,9 +267,10 @@ def test_dsc_issue(sentry_init, capture_envelopes):
     Our service is a standalone service that does not have tracing enabled. Just uses Sentry for error reporting.
     """
     sentry_init(
-        dsn="https://mysecret@bla.ingest.sentry.io/12312012",
+        dsn="https://mysecret@o1234.ingest.sentry.io/12312012",
         release="[email protected]",
         environment="canary",
+        transport=TransportWithOptions,
     )
     envelopes = capture_envelopes()
 
@@ -244,6 +293,10 @@ def test_dsc_issue(sentry_init, capture_envelopes):
     assert type(envelope_trace_header["public_key"]) == str
     assert envelope_trace_header["public_key"] == "mysecret"
 
+    assert "org_id" in envelope_trace_header
+    assert type(envelope_trace_header["org_id"]) == str
+    assert envelope_trace_header["org_id"] == "1234"
+
     assert "sample_rate" not in envelope_trace_header
 
     assert "sampled" not in envelope_trace_header
@@ -265,10 +318,11 @@ def test_dsc_issue_with_tracing(sentry_init, capture_envelopes):
     Envelopes containing errors also have the same DSC than the transaction envelopes.
     """
     sentry_init(
-        dsn="https://mysecret@bla.ingest.sentry.io/12312012",
+        dsn="https://mysecret@o1234.ingest.sentry.io/12312012",
         release="[email protected]",
         environment="canary",
         traces_sample_rate=1.0,
+        transport=TransportWithOptions,
     )
     envelopes = capture_envelopes()
 
@@ -294,6 +348,10 @@ def test_dsc_issue_with_tracing(sentry_init, capture_envelopes):
     assert type(envelope_trace_header["public_key"]) == str
     assert envelope_trace_header["public_key"] == "mysecret"
 
+    assert "org_id" in envelope_trace_header
+    assert type(envelope_trace_header["org_id"]) == str
+    assert envelope_trace_header["org_id"] == "1234"
+
     assert "sample_rate" in envelope_trace_header
     assert envelope_trace_header["sample_rate"] == "1.0"
     assert type(envelope_trace_header["sample_rate"]) == str
@@ -332,10 +390,11 @@ def test_dsc_issue_twp(sentry_init, capture_envelopes, traces_sample_rate):
     (This test would be service B in this scenario)
     """
     sentry_init(
-        dsn="https://mysecret@bla.ingest.sentry.io/12312012",
+        dsn="https://mysecret@o1234.ingest.sentry.io/12312012",
         release="[email protected]",
         environment="canary",
         traces_sample_rate=traces_sample_rate,
+        transport=TransportWithOptions,
     )
     envelopes = capture_envelopes()
 

tests/utils/test_general.py

@@ -119,6 +119,12 @@ def test_parse_dsn_paths(given, expected_envelope):
     assert auth.get_api_url(EndpointType.ENVELOPE) == expected_envelope
 
 
+@pytest.mark.parametrize("given,expected", [("https://[email protected]/123", None), ("https://[email protected]/123", "1234")])
+def test_parse_dsn_org_id(given, expected):
+    dsn = Dsn(given)
+    assert dsn.org_id == expected
+
+
 @pytest.mark.parametrize(
     "dsn",
     [