Add type annotations and attach graphene query again

Author: alexander-alderman-webb

sentry_sdk/integrations/graphene.py

@@ -20,7 +20,7 @@
 
 if TYPE_CHECKING:
     from collections.abc import Generator
-    from typing import Any, Dict, Union
+    from typing import Any, Dict, Union, Callable
     from graphene.language.source import Source  # type: ignore
     from graphql.execution import ExecutionResult
     from graphql.type import GraphQLSchema
@@ -48,7 +48,7 @@ def _patch_graphql():
     def _sentry_patched_graphql_sync(schema, source, *args, **kwargs):
         # type: (GraphQLSchema, Union[str, Source], Any, Any) -> ExecutionResult
         scope = sentry_sdk.get_isolation_scope()
-        scope.add_event_processor(_event_processor)
+        scope.add_event_processor(_make_event_processor(source))
 
         with graphql_span(schema, source, kwargs):
             result = old_graphql_sync(schema, source, *args, **kwargs)
@@ -75,7 +75,7 @@ async def _sentry_patched_graphql_async(schema, source, *args, **kwargs):
             return await old_graphql_async(schema, source, *args, **kwargs)
 
         scope = sentry_sdk.get_isolation_scope()
-        scope.add_event_processor(_event_processor)
+        scope.add_event_processor(_make_event_processor(source))
 
         with graphql_span(schema, source, kwargs):
             result = await old_graphql_async(schema, source, *args, **kwargs)
@@ -99,16 +99,21 @@ async def _sentry_patched_graphql_async(schema, source, *args, **kwargs):
     graphene_schema.graphql = _sentry_patched_graphql_async
 
 
-def _event_processor(event, hint):
-    # type: (Event, Dict[str, Any]) -> Event
-    if should_send_default_pii():
-        request_info = event.setdefault("request", {})
-        request_info["api_target"] = "graphql"
+def _make_event_processor(source):
+    # type: (Any) -> Callable[[Event, dict[str, Any]], Event]
+    def _event_processor(event, hint):
+        # type: (Event, Dict[str, Any]) -> Event
+        if should_send_default_pii():
+            request_info = event.setdefault("request", {})
+            request_info["api_target"] = "graphql"
+            request_info["data"] = {"query": source}
 
-    elif event.get("request", {}).get("data"):
-        del event["request"]["data"]
+        elif event.get("request", {}).get("data"):
+            del event["request"]["data"]
 
-    return event
+        return event
+
+    return _event_processor
 
 
 @contextmanager

sentry_sdk/integrations/starlette.py

@@ -423,33 +423,38 @@ def _is_async_callable(obj):
 
 
 def _patch_request(request):
+    # type: (Request) -> None
     _original_body = request.body
     _original_json = request.json
     _original_form = request.form
 
     def restore_original_methods():
+        # type: () -> None
         request.body = _original_body
         request.json = _original_json
         request.form = _original_form
 
     async def sentry_body():
+        # type: () -> bytes
         request.scope.setdefault("state", {})["sentry_sdk.is_body_cached"] = True
         restore_original_methods()
         return await _original_body()
 
     async def sentry_json():
+        # type: () -> Any
         request.scope.setdefault("state", {})["sentry_sdk.is_body_cached"] = True
         restore_original_methods()
         return await _original_json()
 
     async def sentry_form():
+        # type: () -> Any
         request.scope.setdefault("state", {})["sentry_sdk.is_body_cached"] = True
         restore_original_methods()
         return await _original_form()
 
-    request.body = sentry_body
-    request.json = sentry_json
-    request.form = sentry_form
+    request.body = sentry_body  # type: ignore
+    request.json = sentry_json  # type: ignore
+    request.form = sentry_form  # type: ignore
 
 
 def patch_request_response():
@@ -673,6 +678,7 @@ async def extract_request_info(self):
             json = await self.json()
             if json:
                 request_info["data"] = json
+                print("in fastapi json", request_info)
                 return request_info
 
             # Add form as key/value pairs, if request has form data