From 1d6e7952eccb70d0696a567b46e8ffbdd2e94018 Mon Sep 17 00:00:00 2001
From: Milan Boers <milan.boers@datadoghq.com>
Date: Tue, 8 Apr 2025 17:42:21 +0200
Subject: [PATCH] add MultiplexedTransport

---
 sentry_sdk/transport.py | 76 +++++++++++++++++++++++++++++++++++++++++
 tests/test_transport.py | 42 +++++++++++++++++++++++
 2 files changed, 118 insertions(+)

diff --git a/sentry_sdk/transport.py b/sentry_sdk/transport.py
index efc955ca7b..cb3181a896 100644
--- a/sentry_sdk/transport.py
+++ b/sentry_sdk/transport.py
@@ -878,6 +878,82 @@ def capture_envelope(self, envelope: Envelope) -> None:
             self.capture_event(event)
 
 
+def make_multiplexed_transport(dsns):
+    # type: (List[str]) -> Type[Transport]
+    class MultiplexedTransport(Transport):
+        def __init__(self, options):
+            # type: (Self, Optional[Dict[str, Any]]) -> None
+            super().__init__(options)
+            self.transports = list(
+                filter(
+                    None,
+                    [
+                        make_transport(
+                            {**(options or {}), "dsn": dsn, "transport": None}
+                        )
+                        for dsn in dsns
+                    ],
+                )
+            )
+
+        @staticmethod
+        def _override_dsn(envelope, dsn):
+            # type: (Envelope, Optional[Dsn]) -> Envelope
+            if (
+                dsn is None
+                or "trace" not in envelope.headers
+                or "public_key" not in envelope.headers["trace"]
+            ):
+                return envelope
+            return Envelope(
+                {
+                    **envelope.headers,
+                    "trace": {
+                        **envelope.headers["trace"],
+                        "public_key": dsn.public_key,
+                    },
+                },
+                envelope.items,
+            )
+
+        def capture_envelope(self, envelope):
+            # type: (Self, Envelope) -> None
+            for transport in self.transports:
+                transport.capture_envelope(
+                    self._override_dsn(envelope, transport.parsed_dsn)
+                )
+
+        def flush(self, timeout, callback=None):
+            # type: (Self, float, Optional[Any]) -> None
+            for transport in self.transports:
+                transport.flush(timeout, callback)
+
+        def kill(self):
+            # type: (Self) -> None
+            for transport in self.transports:
+                transport.kill()
+
+        def record_lost_event(
+            self,
+            reason,  # type: str
+            data_category=None,  # type: Optional[EventDataCategory]
+            item=None,  # type: Optional[Item]
+            *,
+            quantity=1,  # type: int
+        ):
+            # type: (...) -> None
+            for transport in self.transports:
+                transport.record_lost_event(
+                    reason, data_category, item, quantity=quantity
+                )
+
+        def is_healthy(self):
+            # type: (Self) -> bool
+            return all(transport.is_healthy() for transport in self.transports)
+
+    return MultiplexedTransport
+
+
 def make_transport(options):
     # type: (Dict[str, Any]) -> Optional[Transport]
     ref_transport = options["transport"]
diff --git a/tests/test_transport.py b/tests/test_transport.py
index d24bea0491..6b6874e5f8 100644
--- a/tests/test_transport.py
+++ b/tests/test_transport.py
@@ -34,6 +34,7 @@
     KEEP_ALIVE_SOCKET_OPTIONS,
     _parse_rate_limits,
     HttpTransport,
+    make_multiplexed_transport,
 )
 from sentry_sdk.integrations.logging import LoggingIntegration, ignore_logger
 
@@ -204,6 +205,47 @@ def test_transport_works(
     assert any("Sending envelope" in record.msg for record in caplog.records) == debug
 
 
+@pytest.mark.forked
+def test_multiplexed_transport_works(capturing_server, request, make_client):
+    server_address = capturing_server.url[len("http://") :]
+    dsn1 = "http://publickeyone@{}/123".format(server_address)
+    dsn2 = "http://publickeytwo@{}/456".format(server_address)
+
+    client = make_client(transport=make_multiplexed_transport([dsn1, dsn2]))
+    sentry_sdk.get_global_scope().set_client(client)
+    request.addfinalizer(lambda: sentry_sdk.get_global_scope().set_client(None))
+
+    capture_message("some message")
+    client.close()
+
+    assert len(capturing_server.captured) == 2
+    assert {capturing_server.captured[0].path, capturing_server.captured[1].path} == {
+        "/api/123/envelope/",
+        "/api/456/envelope/",
+    }
+    assert {
+        capturing_server.captured[0].envelope.headers["trace"]["public_key"],
+        capturing_server.captured[1].envelope.headers["trace"]["public_key"],
+    } == {"publickeyone", "publickeytwo"}
+    assert (
+        capturing_server.captured[0].envelope.headers["event_id"]
+        == capturing_server.captured[1].envelope.headers["event_id"]
+    )
+    assert (
+        capturing_server.captured[0].envelope.headers["sent_at"]
+        == capturing_server.captured[1].envelope.headers["sent_at"]
+    )
+    # Key not sent to the wrong server
+    assert (
+        "publickeyone" in str(capturing_server.captured[0].envelope.serialize())
+    ) == ("publickeytwo" in str(capturing_server.captured[1].envelope.serialize()))
+    # Python SDK doesn't add a dsn header. Asserting in case it is added in the future.
+    assert {
+        capturing_server.captured[0].envelope.headers.get("dsn"),
+        capturing_server.captured[1].envelope.headers.get("dsn"),
+    } == {None, None}
+
+
 @pytest.mark.parametrize(
     "num_pools,expected_num_pools",
     (