Feat/sentry tracing

fastapi_jsonrpc/__init__.py

@@ -32,6 +32,7 @@
 from starlette.routing import Match, request_response, compile_path
 import fastapi.params
 import aiojobs
+import warnings
 
 logger = logging.getLogger(__name__)
 
@@ -53,6 +54,11 @@ def _normalize_errors(errors: Sequence[Any]) -> List[Dict[str, Any]]:
 if hasattr(sentry_sdk, 'new_scope'):
     # sentry_sdk 2.x
     sentry_new_scope = sentry_sdk.new_scope
+
+    def get_sentry_integration():
+        return sentry_sdk.get_client().get_integration(
+            "FastApiJsonRPCIntegration"
+        )
 else:
     # sentry_sdk 1.x
     @contextmanager
@@ -62,6 +68,8 @@ def sentry_new_scope():
             with hub.configure_scope() as scope:
                 yield scope
 
+    get_sentry_integration = lambda : None
+
 
 class Params(fastapi.params.Body):
     def __init__(
@@ -642,6 +650,14 @@ async def _handle_exception(self, reraise=True):
 
     @contextmanager
     def _enter_sentry_scope(self):
+        if get_sentry_integration() is not None:
+            yield
+            return
+
+        warnings.warn(
+            "You are using implicit sentry integration. This feature might be removed in a future major release."
+            "Use explicit `FastApiJsonRPCIntegration` with sentry-sdk 2.* instead.",
+        )
         with sentry_new_scope() as scope:
             # Actually we can use set_transaction_name
             #             scope.set_transaction_name(

fastapi_jsonrpc/contrib/__init__.py

@@ -0,0 +1,7 @@
+from .sentry import TransactionNameGenerator, FastApiJsonRPCIntegration, jrpc_transaction_middleware
+
+__all__ = [
+    "FastApiJsonRPCIntegration",
+    "TransactionNameGenerator",
+    "jrpc_transaction_middleware",
+]

fastapi_jsonrpc/contrib/sentry/__init__.py

@@ -0,0 +1,8 @@
+from .jrpc import TransactionNameGenerator, jrpc_transaction_middleware
+from .integration import FastApiJsonRPCIntegration
+
+__all__ = [
+    "FastApiJsonRPCIntegration",
+    "TransactionNameGenerator",
+    "jrpc_transaction_middleware",
+]

fastapi_jsonrpc/contrib/sentry/http.py

@@ -0,0 +1,20 @@
+import uuid
+from functools import wraps
+from contextvars import ContextVar
+
+from starlette.requests import Request
+from sentry_sdk.integrations.asgi import _get_headers
+
+sentry_asgi_context: ContextVar[dict] = ContextVar("_sentry_asgi_context")
+
+
+def set_shared_sentry_context(cls):
+    original_handle_body = cls.handle_body
+
+    @wraps(original_handle_body)
+    async def _patched_handle_body(self, http_request: Request, *args, **kwargs):
+        headers = _get_headers(http_request.scope)
+        sentry_asgi_context.set({"sampled_sentry_trace_id": uuid.uuid4(), "asgi_headers": headers})
+        return await original_handle_body(self, http_request, *args, **kwargs)
+
+    cls.handle_body = _patched_handle_body

fastapi_jsonrpc/contrib/sentry/integration.py

@@ -0,0 +1,25 @@
+from typing import Optional
+from fastapi_jsonrpc import MethodRoute, EntrypointRoute
+from sentry_sdk.integrations import Integration
+
+from .http import set_shared_sentry_context
+from .jrpc import TransactionNameGenerator, default_transaction_name_generator, prepend_jrpc_transaction_middleware
+
+
+class FastApiJsonRPCIntegration(Integration):
+    identifier = "fastapi_jsonrpc"
+    _already_enabled: bool = False
+
+    def __init__(self, transaction_name_generator: Optional[TransactionNameGenerator] = None):
+        self.transaction_name_generator = transaction_name_generator or default_transaction_name_generator
+
+    @staticmethod
+    def setup_once():
+        if FastApiJsonRPCIntegration._already_enabled:
+            return
+
+        prepend_jrpc_transaction_middleware()
+        set_shared_sentry_context(MethodRoute)
+        set_shared_sentry_context(EntrypointRoute)
+
+        FastApiJsonRPCIntegration._already_enabled = True

fastapi_jsonrpc/contrib/sentry/jrpc.py

@@ -0,0 +1,132 @@
+from random import Random
+from typing import TYPE_CHECKING, Callable
+from contextlib import asynccontextmanager
+
+import sentry_sdk
+from fastapi_jsonrpc import BaseError, Entrypoint, JsonRpcContext
+from sentry_sdk.utils import event_from_exception, is_valid_sample_rate
+from sentry_sdk.consts import OP
+from sentry_sdk.tracing import SENTRY_TRACE_HEADER_NAME, Transaction
+from sentry_sdk.tracing_utils import normalize_incoming_data
+
+from .http import sentry_asgi_context
+
+if TYPE_CHECKING:
+    from .integration import FastApiJsonRPCIntegration
+
+_DEFAULT_TRANSACTION_NAME = "generic JRPC request"
+TransactionNameGenerator = Callable[[JsonRpcContext], str]
+
+
+@asynccontextmanager
+async def jrpc_transaction_middleware(ctx: JsonRpcContext):
+    """
+    Start new transaction for each JRPC request. Applies same sampling decision for every transaction in the batch.
+    """
+
+    current_asgi_context = sentry_asgi_context.get()
+    headers = current_asgi_context["asgi_headers"]
+    transaction_params = dict(
+        # this name is replaced by event processor
+        name=_DEFAULT_TRANSACTION_NAME,
+        op=OP.HTTP_SERVER,
+        source=sentry_sdk.tracing.TRANSACTION_SOURCE_CUSTOM,
+        origin="manual",
+    )
+    with sentry_sdk.isolation_scope() as jrpc_request_scope:
+        jrpc_request_scope.clear()
+
+        if SENTRY_TRACE_HEADER_NAME in headers:
+            # continue existing trace
+            # https://github.com/getsentry/sentry-python/blob/2.19.2/sentry_sdk/scope.py#L471
+            jrpc_request_scope.generate_propagation_context(headers)
+            transaction = JrpcTransaction.continue_from_headers(
+                normalize_incoming_data(headers),
+                **transaction_params,
+            )
+        else:
+            # no parent transaction, start a new trace
+            transaction = JrpcTransaction(
+                trace_id=current_asgi_context["sampled_sentry_trace_id"].hex,
+                **transaction_params,  # type: ignore
+            )
+
+        integration: FastApiJsonRPCIntegration | None = sentry_sdk.get_client().get_integration(  # type: ignore
+            "FastApiJsonRPCIntegration"
+        )
+        name_generator = integration.transaction_name_generator if integration else default_transaction_name_generator
+
+        with jrpc_request_scope.start_transaction(
+            transaction,
+            scope=jrpc_request_scope,
+        ):
+            jrpc_request_scope.add_event_processor(make_transaction_info_event_processor(ctx, name_generator))
+            try:
+                yield
+            except Exception as exc:
+                if isinstance(exc, BaseError):
+                    raise
+
+                # attaching event to current transaction
+                event, hint = event_from_exception(
+                    exc,
+                    client_options=sentry_sdk.get_client().options,
+                    mechanism={"type": "asgi", "handled": False},
+                )
+                sentry_sdk.capture_event(event, hint=hint)
+
+
+class JrpcTransaction(Transaction):
+    """
+    Overrides `_set_initial_sampling_decision` to apply same sampling decision for transactions with same `trace_id`.
+    """
+
+    def _set_initial_sampling_decision(self, sampling_context):
+        super()._set_initial_sampling_decision(sampling_context)
+        # https://github.com/getsentry/sentry-python/blob/2.19.2/sentry_sdk/tracing.py#L1125
+        if self.sampled or not is_valid_sample_rate(self.sample_rate, source="Tracing"):
+            return
+
+        if not self.sample_rate:
+            return
+
+        # https://github.com/getsentry/sentry-python/blob/2.19.2/sentry_sdk/tracing.py#L1158
+        self.sampled = Random(self.trace_id).random() < self.sample_rate  # noqa: S311
+
+
+def make_transaction_info_event_processor(ctx: JsonRpcContext, name_generator: TransactionNameGenerator) -> Callable:
+    def _event_processor(event, _):
+        event["transaction_info"]["source"] = sentry_sdk.tracing.TRANSACTION_SOURCE_CUSTOM
+        if ctx.method_route is not None:
+            event["transaction"] = name_generator(ctx)
+
+        return event
+
+    return _event_processor
+
+
+def default_transaction_name_generator(ctx: JsonRpcContext) -> str:
+    return f"JRPC:{ctx.method_route.name}"
+
+
+def prepend_jrpc_transaction_middleware():  # noqa: C901
+    # prepend the jrpc_sentry_transaction_middleware to the middlewares list.
+    # we cannot patch Entrypoint _init_ directly,  since objects can be created before invoking this integration
+
+    def _prepend_transaction_middleware(self: Entrypoint):
+        if not hasattr(self, "__patched_middlewares__"):
+            original_middlewares = self.__dict__.get("middlewares", [])
+            self.__patched_middlewares__ = original_middlewares
+
+        # middleware was passed manually
+        if any(middleware is jrpc_transaction_middleware for middleware in self.__patched_middlewares__):
+            return self.__patched_middlewares__
+
+        self.__patched_middlewares__ = [jrpc_transaction_middleware, *self.__patched_middlewares__]
+        return self.__patched_middlewares__
+
+    def _middleware_setter(self: Entrypoint, value):
+        self.__patched_middlewares__ = value
+        _prepend_transaction_middleware(self)
+
+    Entrypoint.middlewares = property(_prepend_transaction_middleware, _middleware_setter)

fastapi_jsonrpc/contrib/sentry/test_utils.py

@@ -0,0 +1,23 @@
+from collections import Counter
+
+
+def assert_jrpc_batch_sentry_items(envelops, expected_items):
+    items = [item.type for e in envelops for item in e.items]
+    actual_items = Counter(items)
+    assert all(item in actual_items.items() for item in expected_items.items()), actual_items.items()
+    transactions = get_captured_transactions(envelops)
+    # same trace_id across jrpc batch
+    trace_ids = set()
+    for transaction in transactions:
+        trace_ids.add(get_transaction_trace_id(transaction))
+
+    assert len(trace_ids) == 1, trace_ids
+    return actual_items
+
+
+def get_transaction_trace_id(transaction):
+    return transaction.payload.json["contexts"]["trace"]["trace_id"]
+
+
+def get_captured_transactions(envelops):
+    return [item for e in envelops for item in e.items if item.type == "transaction"]

tests/conftest.py

@@ -9,15 +9,11 @@
 from _pytest.python_api import RaisesContext
 from starlette.testclient import TestClient
 import fastapi_jsonrpc as jsonrpc
-
-
 # Workaround for osx systems
 # https://stackoverflow.com/questions/58597334/unittest-performance-issue-when-using-requests-mock-on-osx
 if platform.system() == 'Darwin':
     import socket
     socket.gethostbyname = lambda x: '127.0.0.1'
-
-
 pytest_plugins = 'pytester'
 
 
@@ -90,21 +86,24 @@ def app_client(app):
 
 @pytest.fixture
 def raw_request(app_client, ep_path):
-    def requester(body, path_postfix='', auth=None):
+    def requester(body, path_postfix='', auth=None, headers=None):
         resp = app_client.post(
             url=ep_path + path_postfix,
             content=body,
+            headers=headers,
             auth=auth,
         )
         return resp
+
     return requester
 
 
 @pytest.fixture
 def json_request(raw_request):
-    def requester(data, path_postfix=''):
-        resp = raw_request(json_dumps(data), path_postfix=path_postfix)
+    def requester(data, path_postfix='', headers=None):
+        resp = raw_request(json_dumps(data), path_postfix=path_postfix, headers=headers)
         return resp.json()
+
     return requester
 
 
@@ -126,6 +125,7 @@ def requester(method, params, request_id=0):
             'method': method,
             'params': params,
         }, path_postfix=path_postfix)
+
     return requester