Add org_id support (#5166)

### Description

implement part 1 of
https://develop.sentry.dev/sdk/telemetry/traces/#stricttracecontinuation

* Parse `org_id` in `Dsn` constructor
* Expose`parsed_dsn` property on `Client` and use it in some places to
avoid parsing `Dsn` again
* New explicit init option `org_id` that overrides the `org_id` from the
`parsed_dsn` if provided
* Add `org_id` onto populated DSC as head SDK

#### Issues

* part of: #5066 
* part of: PY-1963

Author: sl0thentr0py

sentry_sdk/client.py

@@ -70,6 +70,7 @@
     from sentry_sdk.transport import Transport, Item
     from sentry_sdk._log_batcher import LogBatcher
     from sentry_sdk._metrics_batcher import MetricsBatcher
+    from sentry_sdk.utils import Dsn
 
     I = TypeVar("I", bound=Integration)  # noqa: E741
 
@@ -201,6 +202,11 @@ def dsn(self):
         # type: () -> Optional[str]
         return None
 
+    @property
+    def parsed_dsn(self):
+        # type: () -> Optional[Dsn]
+        return None
+
     def should_send_default_pii(self):
         # type: () -> bool
         return False
@@ -512,6 +518,12 @@ def dsn(self):
         """Returns the configured DSN as string."""
         return self.options["dsn"]
 
+    @property
+    def parsed_dsn(self):
+        # type: () -> Optional[Dsn]
+        """Returns the configured parsed DSN object."""
+        return self.transport.parsed_dsn if self.transport else None
+
     def _prepare_event(
         self,
         event,  # type: Event

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/integrations/opentelemetry/span_processor.py

@@ -23,7 +23,6 @@
 )
 from sentry_sdk.scope import add_global_event_processor
 from sentry_sdk.tracing import Transaction, Span as SentrySpan
-from sentry_sdk.utils import Dsn
 
 from urllib3.util import parse_url as urlparse
 
@@ -113,12 +112,7 @@ def on_start(self, otel_span, parent_context=None):
         # type: (OTelSpan, Optional[context_api.Context]) -> None
         client = get_client()
 
-        if not client.dsn:
-            return
-
-        try:
-            _ = Dsn(client.dsn)
-        except Exception:
+        if not client.parsed_dsn:
             return
 
         if client.options["instrumenter"] != INSTRUMENTER.OTEL:
@@ -233,10 +227,8 @@ def _is_sentry_span(self, otel_span):
             otel_span_url = otel_span.attributes.get(SpanAttributes.HTTP_URL)
         otel_span_url = cast("Optional[str]", otel_span_url)
 
-        dsn_url = None
-        client = get_client()
-        if client.dsn:
-            dsn_url = Dsn(client.dsn).netloc
+        parsed_dsn = get_client().parsed_dsn
+        dsn_url = parsed_dsn.netloc if parsed_dsn else None
 
         if otel_span_url and dsn_url and dsn_url in otel_span_url:
             return True

sentry_sdk/tracing_utils.py

@@ -664,8 +664,10 @@ def from_options(cls, scope):
         if options.get("release"):
             sentry_items["release"] = options["release"]
 
-        if options.get("dsn"):
-            sentry_items["public_key"] = Dsn(options["dsn"]).public_key
+        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"])
@@ -696,8 +698,10 @@ def populate_from_transaction(cls, transaction):
         if options.get("release"):
             sentry_items["release"] = options["release"]
 
-        if options.get("dsn"):
-            sentry_items["public_key"] = Dsn(options["dsn"]).public_key
+        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/transport.py

@@ -70,7 +70,7 @@ def __init__(self, options=None):
         # type: (Self, Optional[Dict[str, Any]]) -> None
         self.options = options
         if options and options["dsn"] is not None and options["dsn"]:
-            self.parsed_dsn = Dsn(options["dsn"])
+            self.parsed_dsn = Dsn(options["dsn"], options.get("org_id"))
         else:
             self.parsed_dsn = None
 

sentry_sdk/utils.py

@@ -294,8 +294,10 @@ class BadDsn(ValueError):
 class Dsn:
     """Represents a DSN."""
 
-    def __init__(self, value):
-        # type: (Union[Dsn, str]) -> None
+    ORG_ID_REGEX = re.compile(r"^o(\d+)\.")
+
+    def __init__(self, value, org_id=None):
+        # type: (Union[Dsn, str], Optional[str]) -> None
         if isinstance(value, Dsn):
             self.__dict__ = dict(value.__dict__)
             return
@@ -310,6 +312,12 @@ def __init__(self, value):
 
         self.host = parts.hostname
 
+        if org_id is not None:
+            self.org_id = org_id  # type: Optional[str]
+        else:
+            org_id_match = Dsn.ORG_ID_REGEX.match(self.host)
+            self.org_id = org_id_match.group(1) if org_id_match else None
+
         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,18 @@ 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",
     [