Merge branch 'master' into webb/stdlib/code-source-tests

Author: alexander-alderman-webb

sentry_sdk/integrations/_wsgi_common.py

@@ -31,6 +31,7 @@
     "HTTP_SET_COOKIE",
     "HTTP_COOKIE",
     "HTTP_AUTHORIZATION",
+    "HTTP_PROXY_AUTHORIZATION",
     "HTTP_X_API_KEY",
     "HTTP_X_FORWARDED_FOR",
     "HTTP_X_REAL_IP",

sentry_sdk/integrations/celery/__init__.py

@@ -14,10 +14,12 @@
 )
 from sentry_sdk.integrations.celery.utils import _now_seconds_since_epoch
 from sentry_sdk.integrations.logging import ignore_logger
+from sentry_sdk.scope import should_send_default_pii
 from sentry_sdk.traces import StreamedSpan
 from sentry_sdk.tracing import BAGGAGE_HEADER_NAME, Span, TransactionSource
 from sentry_sdk.tracing_utils import Baggage, has_span_streaming_enabled
 from sentry_sdk.utils import (
+    SENSITIVE_DATA_SUBSTITUTE,
     capture_internal_exceptions,
     event_from_exception,
     reraise,
@@ -143,8 +145,12 @@ def event_processor(event: "Event", hint: "Hint") -> "Optional[Event]":
             extra = event.setdefault("extra", {})
             extra["celery-job"] = {
                 "task_name": task.name,
-                "args": args,
-                "kwargs": kwargs,
+                "args": (
+                    args if should_send_default_pii() else SENSITIVE_DATA_SUBSTITUTE
+                ),
+                "kwargs": (
+                    kwargs if should_send_default_pii() else SENSITIVE_DATA_SUBSTITUTE
+                ),
             }
 
         if "exc_info" in hint:

sentry_sdk/integrations/fastapi.py

@@ -7,6 +7,7 @@
 from sentry_sdk.scope import should_send_default_pii
 from sentry_sdk.traces import NoOpStreamedSpan, StreamedSpan
 from sentry_sdk.tracing import SOURCE_FOR_STYLE, TransactionSource
+from sentry_sdk.tracing_utils import has_span_streaming_enabled
 from sentry_sdk.utils import transaction_from_function
 
 from typing import TYPE_CHECKING
@@ -19,6 +20,7 @@
     from sentry_sdk.integrations.starlette import (
         StarletteIntegration,
         StarletteRequestExtractor,
+        _set_request_body_data_on_streaming_segment,
     )
 except DidNotEnable:
     raise DidNotEnable("Starlette is not installed")
@@ -109,7 +111,8 @@ def _sentry_call(*args: "Any", **kwargs: "Any") -> "Any":
         old_app = old_get_request_handler(*args, **kwargs)
 
         async def _sentry_app(*args: "Any", **kwargs: "Any") -> "Any":
-            integration = sentry_sdk.get_client().get_integration(FastApiIntegration)
+            client = sentry_sdk.get_client()
+            integration = client.get_integration(FastApiIntegration)
             if integration is None:
                 return await old_app(*args, **kwargs)
 
@@ -144,6 +147,9 @@ def event_processor(event: "Event", hint: "Dict[str, Any]") -> "Event":
                 _make_request_event_processor(request, integration)
             )
 
+            if has_span_streaming_enabled(client.options):
+                _set_request_body_data_on_streaming_segment(info)
+
             return await old_app(*args, **kwargs)
 
         return _sentry_app

sentry_sdk/integrations/graphene.py

@@ -143,7 +143,8 @@ def graphql_span(
 
     _graphql_span = sentry_sdk.start_span(op=op, name=operation_name)
 
-    _graphql_span.set_data("graphql.document", source)
+    if should_send_default_pii():
+        _graphql_span.set_data("graphql.document", source)
     _graphql_span.set_data("graphql.operation.name", operation_name)
     _graphql_span.set_data("graphql.operation.type", operation_type)
 

sentry_sdk/integrations/grpc/aio/server.py

@@ -25,7 +25,7 @@ def __init__(
         self: "ServerInterceptor",
         find_name: "Callable[[ServicerContext], str] | None" = None,
     ) -> None:
-        self._find_method_name = find_name or self._find_name
+        self._custom_find_name = find_name
 
         super().__init__()
 
@@ -34,17 +34,21 @@ async def intercept_service(
         continuation: "Callable[[HandlerCallDetails], Awaitable[RpcMethodHandler]]",
         handler_call_details: "HandlerCallDetails",
     ) -> "Optional[Awaitable[RpcMethodHandler]]":
-        self._handler_call_details = handler_call_details
         handler = await continuation(handler_call_details)
         if handler is None:
             return None
 
+        method_name = handler_call_details.method
+        custom_find_name = self._custom_find_name
+
         if not handler.request_streaming and not handler.response_streaming:
             handler_factory = grpc.unary_unary_rpc_method_handler
 
             async def wrapped(request: "Any", context: "ServicerContext") -> "Any":
                 with sentry_sdk.isolation_scope():
-                    name = self._find_method_name(context)
+                    name = (
+                        custom_find_name(context) if custom_find_name else method_name
+                    )
                     if not name:
                         return await handler(request, context)
 
@@ -96,6 +100,3 @@ async def wrapped(request: "Any", context: "ServicerContext") -> "Any":  # type:
             request_deserializer=handler.request_deserializer,
             response_serializer=handler.response_serializer,
         )
-
-    def _find_name(self, context: "ServicerContext") -> str:
-        return self._handler_call_details.method

sentry_sdk/integrations/rq.py

@@ -5,8 +5,10 @@
 from sentry_sdk.api import continue_trace
 from sentry_sdk.integrations import _check_minimum_version, DidNotEnable, Integration
 from sentry_sdk.integrations.logging import ignore_logger
+from sentry_sdk.scope import should_send_default_pii
 from sentry_sdk.tracing import TransactionSource
 from sentry_sdk.utils import (
+    SENSITIVE_DATA_SUBSTITUTE,
     capture_internal_exceptions,
     ensure_integration_enabled,
     event_from_exception,
@@ -140,8 +142,16 @@ def event_processor(event: "Event", hint: "dict[str, Any]") -> "Event":
                 rq_job = {
                     "job_id": job.id,
                     "func": job.func_name,
-                    "args": job.args,
-                    "kwargs": job.kwargs,
+                    "args": (
+                        job.args
+                        if should_send_default_pii()
+                        else SENSITIVE_DATA_SUBSTITUTE
+                    ),
+                    "kwargs": (
+                        job.kwargs
+                        if should_send_default_pii()
+                        else SENSITIVE_DATA_SUBSTITUTE
+                    ),
                     "description": job.description,
                 }
 

sentry_sdk/integrations/sqlalchemy.py

@@ -1,11 +1,16 @@
 from sentry_sdk.consts import SPANSTATUS, SPANDATA
 from sentry_sdk.integrations import _check_minimum_version, Integration, DidNotEnable
-from sentry_sdk.tracing_utils import add_query_source, record_sql_queries
+from sentry_sdk.tracing_utils import (
+    add_query_source,
+    record_sql_queries_supporting_streaming,
+)
 from sentry_sdk.utils import (
     capture_internal_exceptions,
     ensure_integration_enabled,
     parse_version,
 )
+from sentry_sdk.traces import StreamedSpan, SpanStatus
+from sentry_sdk.tracing import Span
 
 try:
     from sqlalchemy.engine import Engine  # type: ignore
@@ -20,8 +25,7 @@
     from typing import Any
     from typing import ContextManager
     from typing import Optional
-
-    from sentry_sdk.tracing import Span
+    from typing import Union
 
 
 class SqlalchemyIntegration(Integration):
@@ -48,7 +52,7 @@ def _before_cursor_execute(
     executemany: bool,
     *args: "Any",
 ) -> None:
-    ctx_mgr = record_sql_queries(
+    ctx_mgr = record_sql_queries_supporting_streaming(
         cursor,
         statement,
         parameters,
@@ -78,12 +82,19 @@ def _after_cursor_execute(
         context, "_sentry_sql_span_manager", None
     )
 
+    # Record query source immediately before span is finished: accurate end timestamp and before the span is flushed.
+    span: "Optional[Union[Span, StreamedSpan]]" = getattr(
+        context, "_sentry_sql_span", None
+    )
+    if isinstance(span, StreamedSpan):
+        with capture_internal_exceptions():
+            add_query_source(span)
+
     if ctx_mgr is not None:
         context._sentry_sql_span_manager = None
         ctx_mgr.__exit__(None, None, None)
 
-    span: "Optional[Span]" = getattr(context, "_sentry_sql_span", None)
-    if span is not None:
+    if isinstance(span, Span):
         with capture_internal_exceptions():
             add_query_source(span)
 
@@ -96,7 +107,10 @@ def _handle_error(context: "Any", *args: "Any") -> None:
     span: "Optional[Span]" = getattr(execution_context, "_sentry_sql_span", None)
 
     if span is not None:
-        span.set_status(SPANSTATUS.INTERNAL_ERROR)
+        if isinstance(span, StreamedSpan):
+            span.status = SpanStatus.ERROR
+        else:
+            span.set_status(SPANSTATUS.INTERNAL_ERROR)
 
     # _after_cursor_execute does not get called for crashing SQL stmts. Judging
     # from SQLAlchemy codebase it does seem like any error coming into this
@@ -132,29 +146,43 @@ def _get_db_system(name: str) -> "Optional[str]":
     return None
 
 
-def _set_db_data(span: "Span", conn: "Any") -> None:
+def _set_db_data(span: "Union[Span, StreamedSpan]", conn: "Any") -> None:
     db_system = _get_db_system(conn.engine.name)
-    if db_system is not None:
-        span.set_data(SPANDATA.DB_SYSTEM, db_system)
+
+    if isinstance(span, StreamedSpan):
+        if db_system is not None:
+            span.set_attribute(SPANDATA.DB_SYSTEM_NAME, db_system)
+    else:
+        if db_system is not None:
+            span.set_data(SPANDATA.DB_SYSTEM, db_system)
+
+    if isinstance(span, StreamedSpan):
+        set_on_span = span.set_attribute
+    else:
+        set_on_span = span.set_data
 
     try:
         driver = conn.dialect.driver
         if driver:
-            span.set_data(SPANDATA.DB_DRIVER_NAME, driver)
+            set_on_span(SPANDATA.DB_DRIVER_NAME, driver)
     except Exception:
         pass
 
     if conn.engine.url is None:
         return
 
     db_name = conn.engine.url.database
-    if db_name is not None:
-        span.set_data(SPANDATA.DB_NAME, db_name)
+    if isinstance(span, StreamedSpan):
+        if db_name is not None:
+            span.set_attribute(SPANDATA.DB_NAMESPACE, db_name)
+    else:
+        if db_name is not None:
+            span.set_data(SPANDATA.DB_NAME, db_name)
 
     server_address = conn.engine.url.host
     if server_address is not None:
-        span.set_data(SPANDATA.SERVER_ADDRESS, server_address)
+        set_on_span(SPANDATA.SERVER_ADDRESS, server_address)
 
     server_port = conn.engine.url.port
     if server_port is not None:
-        span.set_data(SPANDATA.SERVER_PORT, server_port)
+        set_on_span(SPANDATA.SERVER_PORT, server_port)