Merge branch 'master' into antonpirker/django-caching-module

getsentry · May 23, 2024 · 941ca77 · 941ca77
2 parents 891b453 + ec23396
commit 941ca77
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 9 deletions.
diff --git a/sentry_sdk/integrations/clickhouse_driver.py b/sentry_sdk/integrations/clickhouse_driver.py
@@ -107,7 +107,7 @@ def _wrap_end(f: Callable[P, T]) -> Callable[P, T]:
     def _inner_end(*args: P.args, **kwargs: P.kwargs) -> T:
         res = f(*args, **kwargs)
         instance = args[0]
-        span = instance.connection._sentry_span  # type: ignore[attr-defined]
+        span = getattr(instance.connection, "_sentry_span", None)  # type: ignore[attr-defined]
 
         if span is not None:
             if res is not None and should_send_default_pii():
@@ -129,14 +129,15 @@ def _wrap_send_data(f: Callable[P, T]) -> Callable[P, T]:
     def _inner_send_data(*args: P.args, **kwargs: P.kwargs) -> T:
         instance = args[0]  # type: clickhouse_driver.client.Client
         data = args[2]
-        span = instance.connection._sentry_span
+        span = getattr(instance.connection, "_sentry_span", None)
 
-        _set_db_data(span, instance.connection)
+        if span is not None:
+            _set_db_data(span, instance.connection)
 
-        if should_send_default_pii():
-            db_params = span._data.get("db.params", [])
-            db_params.extend(data)
-            span.set_data("db.params", db_params)
+            if should_send_default_pii():
+                db_params = span._data.get("db.params", [])
+                db_params.extend(data)
+                span.set_data("db.params", db_params)
 
         return f(*args, **kwargs)
 

diff --git a/sentry_sdk/integrations/cohere.py b/sentry_sdk/integrations/cohere.py
@@ -22,7 +22,11 @@
 try:
     from cohere.client import Client
     from cohere.base_client import BaseCohere
-    from cohere import ChatStreamEndEvent, NonStreamedChatResponse
+    from cohere import (
+        ChatStreamEndEvent,
+        NonStreamedChatResponse,
+        StreamedChatResponse_StreamEnd,
+    )
 
     if TYPE_CHECKING:
         from cohere import StreamedChatResponse
@@ -181,7 +185,9 @@ def new_iterator():
 
                     with capture_internal_exceptions():
                         for x in old_iterator:
-                            if isinstance(x, ChatStreamEndEvent):
+                            if isinstance(x, ChatStreamEndEvent) or isinstance(
+                                x, StreamedChatResponse_StreamEnd
+                            ):
                                 collect_chat_response_fields(
                                     span,
                                     x.response,