feat(envelope): Add some useful envelope methods (#793)

Co-authored-by: Rodolfo Carvalho <[email protected]>
Co-authored-by: Mark Story <[email protected]>

Author: 3 people

sentry_sdk/envelope.py

@@ -6,6 +6,7 @@
 from sentry_sdk._compat import text_type
 from sentry_sdk._types import MYPY
 from sentry_sdk.sessions import Session
+from sentry_sdk.tracing import Transaction
 from sentry_sdk.utils import json_dumps
 
 if MYPY:
@@ -50,6 +51,12 @@ def add_event(
         # type: (...) -> None
         self.add_item(Item(payload=PayloadRef(json=event), type="event"))
 
+    def add_transaction(
+        self, transaction  # type: Transaction
+    ):
+        # type: (...) -> None
+        self.add_item(Item(payload=PayloadRef(json=transaction), type="transaction"))
+
     def add_session(
         self, session  # type: Union[Session, Any]
     ):
@@ -72,6 +79,14 @@ def get_event(self):
                 return event
         return None
 
+    def get_transaction_event(self):
+        # type: (...) -> Optional[Event]
+        for item in self.items:
+            event = item.get_transaction_event()
+            if event is not None:
+                return event
+        return None
+
     def __iter__(self):
         # type: (...) -> Iterator[Item]
         return iter(self.items)
@@ -220,6 +235,11 @@ def __repr__(self):
             self.data_category,
         )
 
+    @property
+    def type(self):
+        # type: (...) -> Optional[str]
+        return self.headers.get("type")
+
     @property
     def data_category(self):
         # type: (...) -> EventDataCategory
@@ -244,7 +264,13 @@ def get_event(self):
         """
         Returns an error event if there is one.
         """
-        if self.headers.get("type") == "event" and self.payload.json is not None:
+        if self.type == "event" and self.payload.json is not None:
+            return self.payload.json
+        return None
+
+    def get_transaction_event(self):
+        # type: (...) -> Optional[Event]
+        if self.type == "transaction" and self.payload.json is not None:
             return self.payload.json
         return None
 
@@ -277,7 +303,7 @@ def deserialize_from(
         headers = json.loads(line)
         length = headers["length"]
         payload = f.read(length)
-        if headers.get("type") == "event":
+        if headers.get("type") in ("event", "transaction"):
             rv = cls(headers=headers, payload=PayloadRef(json=json.loads(payload)))
         else:
             rv = cls(headers=headers, payload=payload)

tests/test_envelope.py

@@ -0,0 +1,66 @@
+from sentry_sdk.envelope import Envelope
+from sentry_sdk.sessions import Session
+
+
+def generate_transaction_item():
+    return {
+        "event_id": "d2132d31b39445f1938d7e21b6bf0ec4",
+        "type": "transaction",
+        "transaction": "/organizations/:orgId/performance/:eventSlug/",
+        "start_timestamp": 1597976392.6542819,
+        "timestamp": 1597976400.6189718,
+        "contexts": {
+            "trace": {
+                "trace_id": "4C79F60C11214EB38604F4AE0781BFB2",
+                "span_id": "FA90FDEAD5F74052",
+                "type": "trace",
+            }
+        },
+        "spans": [
+            {
+                "description": "<OrganizationContext>",
+                "op": "react.mount",
+                "parent_span_id": "8f5a2b8768cafb4e",
+                "span_id": "bd429c44b67a3eb4",
+                "start_timestamp": 1597976393.4619668,
+                "timestamp": 1597976393.4718769,
+                "trace_id": "ff62a8b040f340bda5d830223def1d81",
+            }
+        ],
+    }
+
+
+def test_basic_event():
+    envelope = Envelope()
+
+    expected = {"message": "Hello, World!"}
+    envelope.add_event(expected)
+
+    assert envelope.get_event() == {"message": "Hello, World!"}
+
+
+def test_transaction_event():
+    envelope = Envelope()
+
+    transaction_item = generate_transaction_item()
+    transaction_item.update({"event_id": "a" * 32})
+    envelope.add_transaction(transaction_item)
+
+    # typically it should not be possible to be able to add a second transaction;
+    # but we do it anyways
+    another_transaction_item = generate_transaction_item()
+    envelope.add_transaction(another_transaction_item)
+
+    # should only fetch the first inserted transaction event
+    assert envelope.get_transaction_event() == transaction_item
+
+
+def test_session():
+    envelope = Envelope()
+
+    expected = Session()
+    envelope.add_session(expected)
+
+    for item in envelope:
+        if item.type == "session":
+            assert item.payload.json == expected.to_json()