From d19974138cbcf76d29ca0fd3fcd6d1deb2597dc8 Mon Sep 17 00:00:00 2001 From: Matthieu MN <10926130+gazorby@users.noreply.github.com> Date: Sat, 11 Mar 2023 00:47:59 +0100 Subject: [PATCH] feat(sentry): override traces sampler (#297) * feat(sentry): override traces sampler * chore: fix mypy errors. --------- Co-authored-by: Peter Schutt --- src/starlite_saqlalchemy/init_plugin.py | 6 +++-- src/starlite_saqlalchemy/sentry.py | 5 ++-- .../test_init_plugin_for_no_extras.py | 2 +- .../test_init_plugin_for_sentry.py | 26 +++++++++++++++++++ .../repository/test_sqlalchemy.py | 2 +- tests/unit/test_init_plugin.py | 3 +-- 6 files changed, 36 insertions(+), 8 deletions(-) diff --git a/src/starlite_saqlalchemy/init_plugin.py b/src/starlite_saqlalchemy/init_plugin.py index 616d8306..1c4ef8a7 100644 --- a/src/starlite_saqlalchemy/init_plugin.py +++ b/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. diff --git a/src/starlite_saqlalchemy/sentry.py b/src/starlite_saqlalchemy/sentry.py index 4941d68b..a204503f 100644 --- a/src/starlite_saqlalchemy/sentry.py +++ b/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, ) diff --git a/tests/unit/require_no_extras/test_init_plugin_for_no_extras.py b/tests/unit/require_no_extras/test_init_plugin_for_no_extras.py index 22cbf6bd..3b9f7c48 100644 --- a/tests/unit/require_no_extras/test_init_plugin_for_no_extras.py +++ b/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] diff --git a/tests/unit/require_sentry/test_init_plugin_for_sentry.py b/tests/unit/require_sentry/test_init_plugin_for_sentry.py index f1841705..aac32eb0 100644 --- a/tests/unit/require_sentry/test_init_plugin_for_sentry.py +++ b/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 diff --git a/tests/unit/require_sqlalchemy/repository/test_sqlalchemy.py b/tests/unit/require_sqlalchemy/repository/test_sqlalchemy.py index b52c0b5f..88c05adc 100644 --- a/tests/unit/require_sqlalchemy/repository/test_sqlalchemy.py +++ b/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): diff --git a/tests/unit/test_init_plugin.py b/tests/unit/test_init_plugin.py index ad6321b0..92966273 100644 --- a/tests/unit/test_init_plugin.py +++ b/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,