feat(attachments): Add basic support for attachments (#856)

Author: mitsuhiko

sentry_sdk/attachments.py

@@ -0,0 +1,55 @@
+import os
+import mimetypes
+
+from sentry_sdk._types import MYPY
+from sentry_sdk.envelope import Item, PayloadRef
+
+if MYPY:
+    from typing import Optional, Union, Callable
+
+
+class Attachment(object):
+    def __init__(
+        self,
+        bytes=None,  # type: Union[None, bytes, Callable[[], bytes]]
+        filename=None,  # type: Optional[str]
+        path=None,  # type: Optional[str]
+        content_type=None,  # type: Optional[str]
+        add_to_transactions=False,  # type: bool
+    ):
+        # type: (...) -> None
+        if bytes is None and path is None:
+            raise TypeError("path or raw bytes required for attachment")
+        if filename is None and path is not None:
+            filename = os.path.basename(path)
+        if filename is None:
+            raise TypeError("filename is required for attachment")
+        if content_type is None:
+            content_type = mimetypes.guess_type(filename)[0]
+        self.bytes = bytes
+        self.filename = filename
+        self.path = path
+        self.content_type = content_type
+        self.add_to_transactions = add_to_transactions
+
+    def to_envelope_item(self):
+        # type: () -> Item
+        """Returns an envelope item for this attachment."""
+        payload = None  # type: Union[None, PayloadRef, bytes]
+        if self.bytes is not None:
+            if callable(self.bytes):
+                payload = self.bytes()
+            else:
+                payload = self.bytes
+        else:
+            payload = PayloadRef(path=self.path)
+        return Item(
+            payload=payload,
+            type="attachment",
+            content_type=self.content_type,
+            filename=self.filename,
+        )
+
+    def __repr__(self):
+        # type: () -> str
+        return "<Attachment %r>" % (self.filename,)

sentry_sdk/client.py

@@ -23,7 +23,7 @@
 from sentry_sdk.integrations import setup_integrations
 from sentry_sdk.utils import ContextVar
 from sentry_sdk.sessions import SessionFlusher
-from sentry_sdk.envelope import Envelope, Item, PayloadRef
+from sentry_sdk.envelope import Envelope
 
 from sentry_sdk._types import MYPY
 
@@ -146,16 +146,14 @@ def dsn(self):
     def _prepare_event(
         self,
         event,  # type: Event
-        hint,  # type: Optional[Hint]
+        hint,  # type: Hint
         scope,  # type: Optional[Scope]
     ):
         # type: (...) -> Optional[Event]
 
         if event.get("timestamp") is None:
             event["timestamp"] = datetime.utcnow()
 
-        hint = dict(hint or ())  # type: Hint
-
         if scope is not None:
             event_ = scope.apply_to_event(event, hint)
             if event_ is None:
@@ -322,10 +320,13 @@ def capture_event(
         if hint is None:
             hint = {}
         event_id = event.get("event_id")
+        hint = dict(hint or ())  # type: Hint
+
         if event_id is None:
             event["event_id"] = event_id = uuid.uuid4().hex
         if not self._should_capture(event, hint, scope):
             return None
+
         event_opt = self._prepare_event(event, hint, scope)
         if event_opt is None:
             return None
@@ -336,19 +337,27 @@ def capture_event(
         if session:
             self._update_session_from_event(session, event)
 
-        if event_opt.get("type") == "transaction":
-            # Transactions should go to the /envelope/ endpoint.
-            self.transport.capture_envelope(
-                Envelope(
-                    headers={
-                        "event_id": event_opt["event_id"],
-                        "sent_at": format_timestamp(datetime.utcnow()),
-                    },
-                    items=[
-                        Item(payload=PayloadRef(json=event_opt), type="transaction"),
-                    ],
-                )
+        attachments = hint.get("attachments")
+        is_transaction = event_opt.get("type") == "transaction"
+
+        if is_transaction or attachments:
+            # Transactions or events with attachments should go to the
+            # /envelope/ endpoint.
+            envelope = Envelope(
+                headers={
+                    "event_id": event_opt["event_id"],
+                    "sent_at": format_timestamp(datetime.utcnow()),
+                }
             )
+
+            if is_transaction:
+                envelope.add_transaction(event_opt)
+            else:
+                envelope.add_event(event_opt)
+
+            for attachment in attachments or ():
+                envelope.add_item(attachment.to_envelope_item())
+            self.transport.capture_envelope(envelope)
         else:
             # All other events go to the /store/ endpoint.
             self.transport.capture_event(event_opt)

sentry_sdk/envelope.py

@@ -1,17 +1,14 @@
 import io
 import json
-import shutil
 import mimetypes
 
 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
+from sentry_sdk.utils import json_dumps, capture_internal_exceptions
 
 if MYPY:
     from typing import Any
-    from typing import Tuple
     from typing import Optional
     from typing import Union
     from typing import Dict
@@ -24,7 +21,7 @@
 class Envelope(object):
     def __init__(
         self,
-        headers=None,  # type: Optional[Dict[str, str]]
+        headers=None,  # type: Optional[Dict[str, Any]]
         items=None,  # type: Optional[List[Item]]
     ):
         # type: (...) -> None
@@ -52,7 +49,7 @@ def add_event(
         self.add_item(Item(payload=PayloadRef(json=event), type="event"))
 
     def add_transaction(
-        self, transaction  # type: Transaction
+        self, transaction  # type: Event
     ):
         # type: (...) -> None
         self.add_item(Item(payload=PayloadRef(json=transaction), type="transaction"))
@@ -148,34 +145,15 @@ def get_bytes(self):
         # type: (...) -> bytes
         if self.bytes is None:
             if self.path is not None:
-                with open(self.path, "rb") as f:
-                    self.bytes = f.read()
+                with capture_internal_exceptions():
+                    with open(self.path, "rb") as f:
+                        self.bytes = f.read()
             elif self.json is not None:
                 self.bytes = json_dumps(self.json)
             else:
                 self.bytes = b""
         return self.bytes
 
-    def _prepare_serialize(self):
-        # type: (...) -> Tuple[Any, Any]
-        if self.path is not None and self.bytes is None:
-            f = open(self.path, "rb")
-            f.seek(0, 2)
-            length = f.tell()
-            f.seek(0, 0)
-
-            def writer(out):
-                # type: (Any) -> None
-                try:
-                    shutil.copyfileobj(f, out)
-                finally:
-                    f.close()
-
-            return length, writer
-
-        bytes = self.get_bytes()
-        return len(bytes), lambda f: f.write(bytes)
-
     @property
     def inferred_content_type(self):
         # type: (...) -> str
@@ -199,7 +177,7 @@ class Item(object):
     def __init__(
         self,
         payload,  # type: Union[bytes, text_type, PayloadRef]
-        headers=None,  # type: Optional[Dict[str, str]]
+        headers=None,  # type: Optional[Dict[str, Any]]
         type=None,  # type: Optional[str]
         content_type=None,  # type: Optional[str]
         filename=None,  # type: Optional[str]
@@ -279,11 +257,11 @@ def serialize_into(
     ):
         # type: (...) -> None
         headers = dict(self.headers)
-        length, writer = self.payload._prepare_serialize()
-        headers["length"] = length
+        bytes = self.get_bytes()
+        headers["length"] = len(bytes)
         f.write(json_dumps(headers))
         f.write(b"\n")
-        writer(f)
+        f.write(bytes)
         f.write(b"\n")
 
     def serialize(self):

sentry_sdk/scope.py

@@ -6,6 +6,7 @@
 from sentry_sdk._types import MYPY
 from sentry_sdk.utils import logger, capture_internal_exceptions
 from sentry_sdk.tracing import Transaction
+from sentry_sdk.attachments import Attachment
 
 if MYPY:
     from typing import Any
@@ -90,6 +91,7 @@ class Scope(object):
         "_should_capture",
         "_span",
         "_session",
+        "_attachments",
         "_force_auto_session_tracking",
     )
 
@@ -112,6 +114,7 @@ def clear(self):
         self._tags = {}  # type: Dict[str, Any]
         self._contexts = {}  # type: Dict[str, Dict[str, Any]]
         self._extras = {}  # type: Dict[str, Any]
+        self._attachments = []  # type: List[Attachment]
 
         self.clear_breadcrumbs()
         self._should_capture = True
@@ -251,6 +254,26 @@ def clear_breadcrumbs(self):
         """Clears breadcrumb buffer."""
         self._breadcrumbs = deque()  # type: Deque[Breadcrumb]
 
+    def add_attachment(
+        self,
+        bytes=None,  # type: Optional[bytes]
+        filename=None,  # type: Optional[str]
+        path=None,  # type: Optional[str]
+        content_type=None,  # type: Optional[str]
+        add_to_transactions=False,  # type: bool
+    ):
+        # type: (...) -> None
+        """Adds an attachment to future events sent."""
+        self._attachments.append(
+            Attachment(
+                bytes=bytes,
+                path=path,
+                filename=filename,
+                content_type=content_type,
+                add_to_transactions=add_to_transactions,
+            )
+        )
+
     def add_event_processor(
         self, func  # type: EventProcessor
     ):
@@ -310,10 +333,21 @@ def _drop(event, cause, ty):
             logger.info("%s (%s) dropped event (%s)", ty, cause, event)
             return None
 
+        is_transaction = event.get("type") == "transaction"
+
+        # put all attachments into the hint. This lets callbacks play around
+        # with attachments. We also later pull this out of the hint when we
+        # create the envelope.
+        attachments_to_send = hint.get("attachments") or []
+        for attachment in self._attachments:
+            if not is_transaction or attachment.add_to_transactions:
+                attachments_to_send.append(attachment)
+        hint["attachments"] = attachments_to_send
+
         if self._level is not None:
             event["level"] = self._level
 
-        if event.get("type") != "transaction":
+        if not is_transaction:
             event.setdefault("breadcrumbs", {}).setdefault("values", []).extend(
                 self._breadcrumbs
             )
@@ -379,6 +413,8 @@ def update_from_scope(self, scope):
             self._breadcrumbs.extend(scope._breadcrumbs)
         if scope._span:
             self._span = scope._span
+        if scope._attachments:
+            self._attachments.extend(scope._attachments)
 
     def update_from_kwargs(
         self,
@@ -425,6 +461,7 @@ def __copy__(self):
         rv._span = self._span
         rv._session = self._session
         rv._force_auto_session_tracking = self._force_auto_session_tracking
+        rv._attachments = list(self._attachments)
 
         return rv
 

tests/conftest.py

@@ -143,8 +143,10 @@ def check_string_keys(map):
     def check_envelope(envelope):
         with capture_internal_exceptions():
             # Assert error events are sent without envelope to server, for compat.
-            assert not any(item.data_category == "error" for item in envelope.items)
-            assert not any(item.get_event() is not None for item in envelope.items)
+            # This does not apply if any item in the envelope is an attachment.
+            if not any(x.type == "attachment" for x in envelope.items):
+                assert not any(item.data_category == "error" for item in envelope.items)
+                assert not any(item.get_event() is not None for item in envelope.items)
 
     def inner(client):
         monkeypatch.setattr(

tests/test_basics.py

@@ -1,3 +1,4 @@
+import os
 import logging
 
 import pytest
@@ -222,6 +223,39 @@ def test_breadcrumbs(sentry_init, capture_events):
     assert len(event["breadcrumbs"]["values"]) == 0
 
 
+def test_attachments(sentry_init, capture_envelopes):
+    sentry_init()
+    envelopes = capture_envelopes()
+
+    this_file = os.path.abspath(__file__.rstrip("c"))
+
+    with configure_scope() as scope:
+        scope.add_attachment(bytes=b"Hello World!", filename="message.txt")
+        scope.add_attachment(path=this_file)
+
+    capture_exception(ValueError())
+
+    (envelope,) = envelopes
+
+    assert len(envelope.items) == 3
+    assert envelope.get_event()["exception"] is not None
+
+    attachments = [x for x in envelope.items if x.type == "attachment"]
+    (message, pyfile) = attachments
+
+    assert message.headers["filename"] == "message.txt"
+    assert message.headers["type"] == "attachment"
+    assert message.headers["content_type"] == "text/plain"
+    assert message.payload.bytes == message.payload.get_bytes() == b"Hello World!"
+
+    assert pyfile.headers["filename"] == os.path.basename(this_file)
+    assert pyfile.headers["type"] == "attachment"
+    assert pyfile.headers["content_type"].startswith("text/")
+    assert pyfile.payload.bytes is None
+    with open(this_file, "rb") as f:
+        assert pyfile.payload.get_bytes() == f.read()
+
+
 def test_integration_scoping(sentry_init, capture_events):
     logger = logging.getLogger("test_basics")