feat(sentry): override traces sampler (#297)

* feat(sentry): override traces sampler

* chore: fix mypy errors.

---------

Co-authored-by: Peter Schutt <[email protected]>

src/starlite_saqlalchemy/init_plugin.py

@@ -104,7 +104,7 @@ class PluginConfig(BaseModel):
     Add configuration for gzip compression to
     [`AppConfig.compression_config`][starlite.config.app.AppConfig.compression_config].
     """
-    do_collection_dependencies = True
+    do_collection_dependencies: bool = True
     """Add collection route dependencies.
 
     Add the [`Provide`][starlite.datastructures.Provide]'s for collection route dependencies to
@@ -143,6 +143,8 @@ class PluginConfig(BaseModel):
     Configure the application to initialize Sentry on startup. Adds a handler to
     [`AppConfig.on_startup`][starlite.config.app.AppConfig.on_startup].
     """
+    sentry_traces_sampler: Callable[[dict], float] | None = None
+    """Override sentry traces sampler."""
     do_set_debug: bool = True
     """Configure Starlite debug mode.
 
@@ -379,7 +381,7 @@ def configure_sentry(self) -> None:
         if self.config.do_sentry:
             from starlite_saqlalchemy import sentry
 
-            sentry.configure()
+            sentry.configure(traces_sampler=self.config.sentry_traces_sampler)
 
     def configure_sqlalchemy_plugin(self, app_config: AppConfig) -> None:
         """Configure `SQLAlchemy` for the application.

src/starlite_saqlalchemy/sentry.py

@@ -13,6 +13,7 @@
     from collections.abc import Mapping
     from typing import Any
 
+    from sentry_sdk._types import TracesSampler
     from starlite.types import HTTPScope
 
 
@@ -32,7 +33,7 @@ def sentry_traces_sampler(sampling_context: Mapping[str, Any]) -> float:
     return settings.sentry.TRACES_SAMPLE_RATE
 
 
-def configure() -> None:
+def configure(traces_sampler: TracesSampler | None = None) -> None:
     """Configure sentry on app startup.
 
     See [SentrySettings][starlite_saqlalchemy.settings.SentrySettings].
@@ -42,5 +43,5 @@ def configure() -> None:
         environment=settings.app.ENVIRONMENT,
         release=settings.app.BUILD_NUMBER,
         integrations=[StarliteIntegration(), SqlalchemyIntegration()],
-        traces_sampler=sentry_traces_sampler,
+        traces_sampler=sentry_traces_sampler if traces_sampler is None else traces_sampler,
     )

tests/unit/require_no_extras/test_init_plugin_for_no_extras.py

@@ -79,4 +79,4 @@ def test_extra_dependencies_not_installed(enabled_config: str, error_pattern: st
         **{enabled_config: True},
     }
     with pytest.raises(ValidationError, match=error_pattern):
-        init_plugin.PluginConfig(**kwargs)
+        init_plugin.PluginConfig(**kwargs)  # type:ignore[arg-type]

tests/unit/require_sentry/test_init_plugin_for_sentry.py

@@ -6,6 +6,7 @@
 from unittest.mock import MagicMock
 
 import pytest
+import sentry_sdk
 from sentry_sdk.integrations.starlite import SentryStarliteASGIMiddleware
 from sentry_sdk.integrations.starlite import (
     exception_handler as sentry_after_exception_handler,
@@ -18,6 +19,13 @@
 
 if TYPE_CHECKING:
     from pytest import MonkeyPatch
+    from sentry_sdk._types import TracesSampler
+
+    from starlite_saqlalchemy.sentry import SamplingContext
+
+
+def _custom_traces_sampler(_: SamplingContext) -> float:
+    return 0.42
 
 
 @pytest.mark.parametrize(
@@ -57,3 +65,21 @@ def test_do_sentry() -> None:
     Starlite.__init__ = old_init  # type: ignore
     HTTPRoute.handle = old_route_handle  # type: ignore
     BaseRouteHandler.resolve_middleware = old_resolve_middleware  # type: ignore
+
+
+@pytest.mark.parametrize(
+    ("traces_sampler", "exp_traces_sampler"),
+    [(None, sentry.sentry_traces_sampler), (_custom_traces_sampler, _custom_traces_sampler)],
+)
+def test_sentry_init(
+    traces_sampler: TracesSampler | None,
+    exp_traces_sampler: TracesSampler | None,
+    monkeypatch: MonkeyPatch,
+) -> None:
+    sentry_init_mock = MagicMock()
+    monkeypatch.setattr(sentry_sdk, "init", sentry_init_mock)
+    init_plugin.ConfigureApp(
+        config=init_plugin.PluginConfig(do_sentry=True, sentry_traces_sampler=traces_sampler)
+    )
+    sentry_init_mock.assert_called_once()
+    assert sentry_init_mock.call_args.kwargs["traces_sampler"] is exp_traces_sampler

tests/unit/require_sqlalchemy/repository/test_sqlalchemy.py

@@ -239,7 +239,7 @@ def test_filter_collection_by_kwargs_raises_repository_exception_for_attribute_e
 ) -> None:
     """Test that we raise a repository exception if an attribute name is
     incorrect."""
-    mock_repo._select.filter_by = MagicMock(  # type:ignore[assignment]
+    mock_repo._select.filter_by = MagicMock(  # type:ignore[method-assign]
         side_effect=InvalidRequestError,
     )
     with pytest.raises(StarliteSaqlalchemyError):

tests/unit/test_init_plugin.py

@@ -23,8 +23,7 @@ def test_config_switches() -> None:
         do_after_exception=False,
         do_cache=False,
         do_compression=False,
-        # pyright reckons this parameter doesn't exist, I beg to differ
-        do_collection_dependencies=False,  # pyright:ignore
+        do_collection_dependencies=False,
         do_exception_handlers=False,
         do_health_check=False,
         do_logging=False,