Merge branch 'master' into ivana/span-first-before-send-span

Author: sentrivana

.github/workflows/ai-integration-test.yml

@@ -34,7 +34,7 @@ jobs:
           token: ${{ secrets.GITHUB_TOKEN }}
 
       - name: Run Python SDK Tests
-        uses: getsentry/testing-ai-sdk-integrations@1dd9ee2a2d821f41473fc90809a758e8394c689c
+        uses: getsentry/testing-ai-sdk-integrations@18229d1da4bb8120ce923bdf7aac00fcf6f38311
         env:
           GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
         with:

sentry_sdk/_batcher.py

@@ -153,9 +153,10 @@ def _add_to_envelope(self, envelope: "Envelope") -> None:
                 },
                 payload=PayloadRef(
                     json={
+                        "version": 2,
                         "items": [
                             self._to_transport_format(item) for item in self._buffer
-                        ]
+                        ],
                     }
                 ),
             )

sentry_sdk/_log_batcher.py

@@ -45,7 +45,12 @@ def _record_lost(self, item: "Log") -> None:
             headers={
                 "item_count": 1,
             },
-            payload=PayloadRef(json={"items": [self._to_transport_format(item)]}),
+            payload=PayloadRef(
+                json={
+                    "version": 2,
+                    "items": [self._to_transport_format(item)],
+                }
+            ),
         )
 
         self._record_lost_func(

sentry_sdk/_span_batcher.py

@@ -214,10 +214,11 @@ def _flush(self, only_pending: bool = False) -> None:
                             },
                             payload=PayloadRef(
                                 json={
+                                    "version": 2,
                                     "items": [
                                         self._to_transport_format(spans[j])
                                         for j in range(start, end)
-                                    ]
+                                    ],
                                 }
                             ),
                         )

sentry_sdk/integrations/asyncpg.py

@@ -16,7 +16,6 @@
 )
 from sentry_sdk.utils import (
     capture_internal_exceptions,
-    ensure_integration_enabled,
     parse_version,
 )
 
@@ -52,8 +51,12 @@ def setup_once() -> None:
         asyncpg.Connection._executemany = _wrap_connection_method(
             asyncpg.Connection._executemany, executemany=True
         )
-        asyncpg.Connection.cursor = _wrap_cursor_creation(asyncpg.Connection.cursor)
         asyncpg.Connection.prepare = _wrap_connection_method(asyncpg.Connection.prepare)
+
+        BaseCursor._bind_exec = _wrap_cursor_method(BaseCursor._bind_exec)
+        BaseCursor._bind = _wrap_cursor_method(BaseCursor._bind)
+        BaseCursor._exec = _wrap_cursor_method(BaseCursor._exec)
+
         asyncpg.connect_utils._connect_addr = _wrap_connect_addr(
             asyncpg.connect_utils._connect_addr
         )
@@ -150,20 +153,26 @@ async def _inner(*args: "Any", **kwargs: "Any") -> "T":
     return _inner
 
 
-def _wrap_cursor_creation(f: "Callable[..., T]") -> "Callable[..., T]":
-    @ensure_integration_enabled(AsyncPGIntegration, f)
-    def _inner(*args: "Any", **kwargs: "Any") -> "T":  # noqa: N807
-        query = args[1]
-        params_list = args[2] if len(args) > 2 else None
+def _wrap_cursor_method(
+    f: "Callable[..., Awaitable[T]]",
+) -> "Callable[..., Awaitable[T]]":
+    async def _inner(*args: "Any", **kwargs: "Any") -> "T":
+        if sentry_sdk.get_client().get_integration(AsyncPGIntegration) is None:
+            return await f(*args, **kwargs)
 
-        with _record(
-            None,
-            query,
-            params_list,
+        cursor = args[0]
+        query = _normalize_query(cursor._query)
+        with record_sql_queries_supporting_streaming(
+            cursor=cursor,
+            query=query,
+            params_list=None,
+            paramstyle=None,
             executemany=False,
+            record_cursor_repr=True,
+            span_origin=AsyncPGIntegration.origin,
         ) as span:
-            _set_db_data(span, args[0])
-            res = f(*args, **kwargs)
+            _set_db_data(span, cursor._connection)
+            res = await f(*args, **kwargs)
 
         return res
 

sentry_sdk/integrations/pymongo.py

@@ -2,10 +2,12 @@
 import json
 
 import sentry_sdk
-from sentry_sdk.consts import SPANSTATUS, SPANDATA, OP
+from sentry_sdk.consts import OP, SPANDATA, SPANSTATUS
 from sentry_sdk.integrations import DidNotEnable, Integration
 from sentry_sdk.scope import should_send_default_pii
+from sentry_sdk.traces import SpanStatus, StreamedSpan
 from sentry_sdk.tracing import Span
+from sentry_sdk.tracing_utils import has_span_streaming_enabled
 from sentry_sdk.utils import capture_internal_exceptions
 
 try:
@@ -87,12 +89,11 @@ def _strip_pii(command: "Dict[str, Any]") -> "Dict[str, Any]":
 def _get_db_data(event: "Any") -> "Dict[str, Any]":
     data = {}
 
-    data[SPANDATA.DB_SYSTEM] = "mongodb"
-    data[SPANDATA.DB_DRIVER_NAME] = "pymongo"
+    client = sentry_sdk.get_client()
+    is_span_streaming_enabled = has_span_streaming_enabled(client.options)
 
+    data[SPANDATA.DB_DRIVER_NAME] = "pymongo"
     db_name = event.database_name
-    if db_name is not None:
-        data[SPANDATA.DB_NAME] = db_name
 
     server_address = event.connection_id[0]
     if server_address is not None:
@@ -102,12 +103,23 @@ def _get_db_data(event: "Any") -> "Dict[str, Any]":
     if server_port is not None:
         data[SPANDATA.SERVER_PORT] = server_port
 
+    if is_span_streaming_enabled:
+        data["db.system.name"] = "mongodb"
+
+        if db_name is not None:
+            data["db.namespace"] = db_name
+    else:
+        data[SPANDATA.DB_SYSTEM] = "mongodb"
+
+        if db_name is not None:
+            data[SPANDATA.DB_NAME] = db_name
+
     return data
 
 
 class CommandTracer(monitoring.CommandListener):
     def __init__(self) -> None:
-        self._ongoing_operations: "Dict[int, Span]" = {}
+        self._ongoing_operations: "Dict[int, Union[Span, StreamedSpan]]" = {}
 
     def _operation_key(
         self,
@@ -116,7 +128,8 @@ def _operation_key(
         return event.request_id
 
     def started(self, event: "CommandStartedEvent") -> None:
-        if sentry_sdk.get_client().get_integration(PyMongoIntegration) is None:
+        client = sentry_sdk.get_client()
+        if client.get_integration(PyMongoIntegration) is None:
             return
 
         with capture_internal_exceptions():
@@ -126,56 +139,88 @@ def started(self, event: "CommandStartedEvent") -> None:
             command.pop("$clusterTime", None)
             command.pop("$signature", None)
 
-            tags = {
-                "db.name": event.database_name,
-                SPANDATA.DB_SYSTEM: "mongodb",
-                SPANDATA.DB_DRIVER_NAME: "pymongo",
-                SPANDATA.DB_OPERATION: event.command_name,
-                SPANDATA.DB_MONGODB_COLLECTION: command.get(event.command_name),
-            }
-
-            try:
-                tags["net.peer.name"] = event.connection_id[0]
-                tags["net.peer.port"] = str(event.connection_id[1])
-            except TypeError:
-                pass
+            db_data = _get_db_data(event)
 
-            data: "Dict[str, Any]" = {"operation_ids": {}}
-            data["operation_ids"]["operation"] = event.operation_id
-            data["operation_ids"]["request"] = event.request_id
-
-            data.update(_get_db_data(event))
-
-            try:
-                lsid = command.pop("lsid")["id"]
-                data["operation_ids"]["session"] = str(lsid)
-            except KeyError:
-                pass
+            collection_name = command.get(event.command_name)
+            operation_name = event.command_name
+            db_name = event.database_name
 
+            lsid = command.pop("lsid", None)
             if not should_send_default_pii():
                 command = _strip_pii(command)
 
             query = json.dumps(command, default=str)
-            span = sentry_sdk.start_span(
-                op=OP.DB,
-                name=query,
-                origin=PyMongoIntegration.origin,
-            )
 
-            for tag, value in tags.items():
-                # set the tag for backwards-compatibility.
-                # TODO: remove the set_tag call in the next major release!
-                span.set_tag(tag, value)
+            if has_span_streaming_enabled(client.options):
+                span_first_data = {
+                    "db.operation.name": operation_name,
+                    "db.collection.name": collection_name,
+                    "sentry.op": OP.DB,
+                    "sentry.origin": PyMongoIntegration.origin,
+                    **db_data,
+                }
+
+                span = sentry_sdk.traces.start_span(
+                    name=query, attributes=span_first_data
+                )
+
+                with capture_internal_exceptions():
+                    sentry_sdk.add_breadcrumb(
+                        message=query,
+                        category="query",
+                        type=OP.DB,
+                        data=span_first_data,
+                    )
+
+            else:
+                tags = {
+                    "db.name": db_name,
+                    SPANDATA.DB_SYSTEM: "mongodb",
+                    SPANDATA.DB_DRIVER_NAME: "pymongo",
+                    SPANDATA.DB_OPERATION: operation_name,
+                    # The below is a deprecated field, but leaving for legacy reasons.
+                    # The v2 spans will use `db.collection.name` instead.
+                    SPANDATA.DB_MONGODB_COLLECTION: collection_name,
+                }
+
+                try:
+                    tags["net.peer.name"] = event.connection_id[0]
+                    tags["net.peer.port"] = str(event.connection_id[1])
+                except TypeError:
+                    pass
+
+                data: "Dict[str, Any]" = {"operation_ids": {}}
+                data["operation_ids"]["operation"] = event.operation_id
+                data["operation_ids"]["request"] = event.request_id
+
+                data.update(db_data)
+
+                try:
+                    if lsid:
+                        lsid_id = lsid["id"]
+                        data["operation_ids"]["session"] = str(lsid_id)
+                except KeyError:
+                    pass
+
+                span = sentry_sdk.start_span(
+                    op=OP.DB,
+                    name=query,
+                    origin=PyMongoIntegration.origin,
+                )
 
-                span.set_data(tag, value)
+                for tag, value in tags.items():
+                    # set the tag for backwards-compatibility.
+                    # TODO: remove the set_tag call in the next major release!
+                    span.set_tag(tag, value)
+                    span.set_data(tag, value)
 
-            for key, value in data.items():
-                span.set_data(key, value)
+                for key, value in data.items():
+                    span.set_data(key, value)
 
-            with capture_internal_exceptions():
-                sentry_sdk.add_breadcrumb(
-                    message=query, category="query", type=OP.DB, data=tags
-                )
+                with capture_internal_exceptions():
+                    sentry_sdk.add_breadcrumb(
+                        message=query, category="query", type=OP.DB, data=tags
+                    )
 
             self._ongoing_operations[self._operation_key(event)] = span.__enter__()
 
@@ -185,7 +230,11 @@ def failed(self, event: "CommandFailedEvent") -> None:
 
         try:
             span = self._ongoing_operations.pop(self._operation_key(event))
-            span.set_status(SPANSTATUS.INTERNAL_ERROR)
+            # Ignoring NoOpStreamedSpan as it will always have a status of "ok"
+            if type(span) is StreamedSpan:
+                span.status = SpanStatus.ERROR
+            elif type(span) is Span:
+                span.set_status(SPANSTATUS.INTERNAL_ERROR)
             span.__exit__(None, None, None)
         except KeyError:
             return
@@ -196,7 +245,8 @@ def succeeded(self, event: "CommandSucceededEvent") -> None:
 
         try:
             span = self._ongoing_operations.pop(self._operation_key(event))
-            span.set_status(SPANSTATUS.OK)
+            if type(span) is Span:
+                span.set_status(SPANSTATUS.OK)
             span.__exit__(None, None, None)
         except KeyError:
             pass

sentry_sdk/scrubber.py

@@ -137,7 +137,10 @@ def scrub_extra(self, event: "Event") -> None:
     def scrub_user(self, event: "Event") -> None:
         with capture_internal_exceptions():
             if "user" in event:
-                self.scrub_dict(event["user"])
+                user = event["user"]
+                if "ip_address" in self.denylist and isinstance(user, dict):
+                    user.pop("ip_address", None)
+                self.scrub_dict(user)
 
     def scrub_breadcrumbs(self, event: "Event") -> None:
         with capture_internal_exceptions():