sentry-sdk 0.7.5__py2.py3-none-any.whl → 2.46.0__py2.py3-none-any.whl

sentry_sdk/integrations/celery/beat.py

@@ -0,0 +1,293 @@
+import sentry_sdk
+from sentry_sdk.crons import capture_checkin, MonitorStatus
+from sentry_sdk.integrations import DidNotEnable
+from sentry_sdk.integrations.celery.utils import (
+    _get_humanized_interval,
+    _now_seconds_since_epoch,
+)
+from sentry_sdk.utils import (
+    logger,
+    match_regex_list,
+)
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+    from typing import Any, Optional, TypeVar, Union
+    from sentry_sdk._types import (
+        MonitorConfig,
+        MonitorConfigScheduleType,
+        MonitorConfigScheduleUnit,
+    )
+
+    F = TypeVar("F", bound=Callable[..., Any])
+
+
+try:
+    from celery import Task, Celery  # type: ignore
+    from celery.beat import Scheduler  # type: ignore
+    from celery.schedules import crontab, schedule  # type: ignore
+    from celery.signals import (  # type: ignore
+        task_failure,
+        task_success,
+        task_retry,
+    )
+except ImportError:
+    raise DidNotEnable("Celery not installed")
+
+try:
+    from redbeat.schedulers import RedBeatScheduler  # type: ignore
+except ImportError:
+    RedBeatScheduler = None
+
+
+def _get_headers(task):
+    # type: (Task) -> dict[str, Any]
+    headers = task.request.get("headers") or {}
+
+    # flatten nested headers
+    if "headers" in headers:
+        headers.update(headers["headers"])
+        del headers["headers"]
+
+    headers.update(task.request.get("properties") or {})
+
+    return headers
+
+
+def _get_monitor_config(celery_schedule, app, monitor_name):
+    # type: (Any, Celery, str) -> MonitorConfig
+    monitor_config = {}  # type: MonitorConfig
+    schedule_type = None  # type: Optional[MonitorConfigScheduleType]
+    schedule_value = None  # type: Optional[Union[str, int]]
+    schedule_unit = None  # type: Optional[MonitorConfigScheduleUnit]
+
+    if isinstance(celery_schedule, crontab):
+        schedule_type = "crontab"
+        schedule_value = (
+            "{0._orig_minute} "
+            "{0._orig_hour} "
+            "{0._orig_day_of_month} "
+            "{0._orig_month_of_year} "
+            "{0._orig_day_of_week}".format(celery_schedule)
+        )
+    elif isinstance(celery_schedule, schedule):
+        schedule_type = "interval"
+        (schedule_value, schedule_unit) = _get_humanized_interval(
+            celery_schedule.seconds
+        )
+
+        if schedule_unit == "second":
+            logger.warning(
+                "Intervals shorter than one minute are not supported by Sentry Crons. Monitor '%s' has an interval of %s seconds. Use the `exclude_beat_tasks` option in the celery integration to exclude it.",
+                monitor_name,
+                schedule_value,
+            )
+            return {}
+
+    else:
+        logger.warning(
+            "Celery schedule type '%s' not supported by Sentry Crons.",
+            type(celery_schedule),
+        )
+        return {}
+
+    monitor_config["schedule"] = {}
+    monitor_config["schedule"]["type"] = schedule_type
+    monitor_config["schedule"]["value"] = schedule_value
+
+    if schedule_unit is not None:
+        monitor_config["schedule"]["unit"] = schedule_unit
+
+    monitor_config["timezone"] = (
+        (
+            hasattr(celery_schedule, "tz")
+            and celery_schedule.tz is not None
+            and str(celery_schedule.tz)
+        )
+        or app.timezone
+        or "UTC"
+    )
+
+    return monitor_config
+
+
+def _apply_crons_data_to_schedule_entry(scheduler, schedule_entry, integration):
+    # type: (Any, Any, sentry_sdk.integrations.celery.CeleryIntegration) -> None
+    """
+    Add Sentry Crons information to the schedule_entry headers.
+    """
+    if not integration.monitor_beat_tasks:
+        return
+
+    monitor_name = schedule_entry.name
+
+    task_should_be_excluded = match_regex_list(
+        monitor_name, integration.exclude_beat_tasks
+    )
+    if task_should_be_excluded:
+        return
+
+    celery_schedule = schedule_entry.schedule
+    app = scheduler.app
+
+    monitor_config = _get_monitor_config(celery_schedule, app, monitor_name)
+
+    is_supported_schedule = bool(monitor_config)
+    if not is_supported_schedule:
+        return
+
+    headers = schedule_entry.options.pop("headers", {})
+    headers.update(
+        {
+            "sentry-monitor-slug": monitor_name,
+            "sentry-monitor-config": monitor_config,
+        }
+    )
+
+    check_in_id = capture_checkin(
+        monitor_slug=monitor_name,
+        monitor_config=monitor_config,
+        status=MonitorStatus.IN_PROGRESS,
+    )
+    headers.update({"sentry-monitor-check-in-id": check_in_id})
+
+    # Set the Sentry configuration in the options of the ScheduleEntry.
+    # Those will be picked up in `apply_async` and added to the headers.
+    schedule_entry.options["headers"] = headers
+
+
+def _wrap_beat_scheduler(original_function):
+    # type: (Callable[..., Any]) -> Callable[..., Any]
+    """
+    Makes sure that:
+    - a new Sentry trace is started for each task started by Celery Beat and
+      it is propagated to the task.
+    - the Sentry Crons information is set in the Celery Beat task's
+      headers so that is is monitored with Sentry Crons.
+
+    After the patched function is called,
+    Celery Beat will call apply_async to put the task in the queue.
+    """
+    # Patch only once
+    # Can't use __name__ here, because some of our tests mock original_apply_entry
+    already_patched = "sentry_patched_scheduler" in str(original_function)
+    if already_patched:
+        return original_function
+
+    from sentry_sdk.integrations.celery import CeleryIntegration
+
+    def sentry_patched_scheduler(*args, **kwargs):
+        # type: (*Any, **Any) -> None
+        integration = sentry_sdk.get_client().get_integration(CeleryIntegration)
+        if integration is None:
+            return original_function(*args, **kwargs)
+
+        # Tasks started by Celery Beat start a new Trace
+        scope = sentry_sdk.get_isolation_scope()
+        scope.set_new_propagation_context()
+        scope._name = "celery-beat"
+
+        scheduler, schedule_entry = args
+        _apply_crons_data_to_schedule_entry(scheduler, schedule_entry, integration)
+
+        return original_function(*args, **kwargs)
+
+    return sentry_patched_scheduler
+
+
+def _patch_beat_apply_entry():
+    # type: () -> None
+    Scheduler.apply_entry = _wrap_beat_scheduler(Scheduler.apply_entry)
+
+
+def _patch_redbeat_apply_async():
+    # type: () -> None
+    if RedBeatScheduler is None:
+        return
+
+    RedBeatScheduler.apply_async = _wrap_beat_scheduler(RedBeatScheduler.apply_async)
+
+
+def _setup_celery_beat_signals(monitor_beat_tasks):
+    # type: (bool) -> None
+    if monitor_beat_tasks:
+        task_success.connect(crons_task_success)
+        task_failure.connect(crons_task_failure)
+        task_retry.connect(crons_task_retry)
+
+
+def crons_task_success(sender, **kwargs):
+    # type: (Task, dict[Any, Any]) -> None
+    logger.debug("celery_task_success %s", sender)
+    headers = _get_headers(sender)
+
+    if "sentry-monitor-slug" not in headers:
+        return
+
+    monitor_config = headers.get("sentry-monitor-config", {})
+
+    start_timestamp_s = headers.get("sentry-monitor-start-timestamp-s")
+
+    capture_checkin(
+        monitor_slug=headers["sentry-monitor-slug"],
+        monitor_config=monitor_config,
+        check_in_id=headers["sentry-monitor-check-in-id"],
+        duration=(
+            _now_seconds_since_epoch() - float(start_timestamp_s)
+            if start_timestamp_s
+            else None
+        ),
+        status=MonitorStatus.OK,
+    )
+
+
+def crons_task_failure(sender, **kwargs):
+    # type: (Task, dict[Any, Any]) -> None
+    logger.debug("celery_task_failure %s", sender)
+    headers = _get_headers(sender)
+
+    if "sentry-monitor-slug" not in headers:
+        return
+
+    monitor_config = headers.get("sentry-monitor-config", {})
+
+    start_timestamp_s = headers.get("sentry-monitor-start-timestamp-s")
+
+    capture_checkin(
+        monitor_slug=headers["sentry-monitor-slug"],
+        monitor_config=monitor_config,
+        check_in_id=headers["sentry-monitor-check-in-id"],
+        duration=(
+            _now_seconds_since_epoch() - float(start_timestamp_s)
+            if start_timestamp_s
+            else None
+        ),
+        status=MonitorStatus.ERROR,
+    )
+
+
+def crons_task_retry(sender, **kwargs):
+    # type: (Task, dict[Any, Any]) -> None
+    logger.debug("celery_task_retry %s", sender)
+    headers = _get_headers(sender)
+
+    if "sentry-monitor-slug" not in headers:
+        return
+
+    monitor_config = headers.get("sentry-monitor-config", {})
+
+    start_timestamp_s = headers.get("sentry-monitor-start-timestamp-s")
+
+    capture_checkin(
+        monitor_slug=headers["sentry-monitor-slug"],
+        monitor_config=monitor_config,
+        check_in_id=headers["sentry-monitor-check-in-id"],
+        duration=(
+            _now_seconds_since_epoch() - float(start_timestamp_s)
+            if start_timestamp_s
+            else None
+        ),
+        status=MonitorStatus.ERROR,
+    )

sentry_sdk/integrations/celery/utils.py

@@ -0,0 +1,43 @@
+import time
+from typing import TYPE_CHECKING, cast
+
+if TYPE_CHECKING:
+    from typing import Any, Tuple
+    from sentry_sdk._types import MonitorConfigScheduleUnit
+
+
+def _now_seconds_since_epoch():
+    # type: () -> float
+    # We cannot use `time.perf_counter()` when dealing with the duration
+    # of a Celery task, because the start of a Celery task and
+    # the end are recorded in different processes.
+    # Start happens in the Celery Beat process,
+    # the end in a Celery Worker process.
+    return time.time()
+
+
+def _get_humanized_interval(seconds):
+    # type: (float) -> Tuple[int, MonitorConfigScheduleUnit]
+    TIME_UNITS = (  # noqa: N806
+        ("day", 60 * 60 * 24.0),
+        ("hour", 60 * 60.0),
+        ("minute", 60.0),
+    )
+
+    seconds = float(seconds)
+    for unit, divider in TIME_UNITS:
+        if seconds >= divider:
+            interval = int(seconds / divider)
+            return (interval, cast("MonitorConfigScheduleUnit", unit))
+
+    return (int(seconds), "second")
+
+
+class NoOpMgr:
+    def __enter__(self):
+        # type: () -> None
+        return None
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        # type: (Any, Any, Any) -> None
+        return None

sentry_sdk/integrations/chalice.py

@@ -0,0 +1,134 @@
+import sys
+from functools import wraps
+
+import sentry_sdk
+from sentry_sdk.integrations import Integration, DidNotEnable
+from sentry_sdk.integrations.aws_lambda import _make_request_event_processor
+from sentry_sdk.tracing import TransactionSource
+from sentry_sdk.utils import (
+    capture_internal_exceptions,
+    event_from_exception,
+    parse_version,
+    reraise,
+)
+
+try:
+    import chalice  # type: ignore
+    from chalice import __version__ as CHALICE_VERSION
+    from chalice import Chalice, ChaliceViewError
+    from chalice.app import EventSourceHandler as ChaliceEventSourceHandler  # type: ignore
+except ImportError:
+    raise DidNotEnable("Chalice is not installed")
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from typing import Any
+    from typing import Dict
+    from typing import TypeVar
+    from typing import Callable
+
+    F = TypeVar("F", bound=Callable[..., Any])
+
+
+class EventSourceHandler(ChaliceEventSourceHandler):  # type: ignore
+    def __call__(self, event, context):
+        # type: (Any, Any) -> Any
+        client = sentry_sdk.get_client()
+
+        with sentry_sdk.isolation_scope() as scope:
+            with capture_internal_exceptions():
+                configured_time = context.get_remaining_time_in_millis()
+                scope.add_event_processor(
+                    _make_request_event_processor(event, context, configured_time)
+                )
+            try:
+                return ChaliceEventSourceHandler.__call__(self, event, context)
+            except Exception:
+                exc_info = sys.exc_info()
+                event, hint = event_from_exception(
+                    exc_info,
+                    client_options=client.options,
+                    mechanism={"type": "chalice", "handled": False},
+                )
+                sentry_sdk.capture_event(event, hint=hint)
+                client.flush()
+                reraise(*exc_info)
+
+
+def _get_view_function_response(app, view_function, function_args):
+    # type: (Any, F, Any) -> F
+    @wraps(view_function)
+    def wrapped_view_function(**function_args):
+        # type: (**Any) -> Any
+        client = sentry_sdk.get_client()
+        with sentry_sdk.isolation_scope() as scope:
+            with capture_internal_exceptions():
+                configured_time = app.lambda_context.get_remaining_time_in_millis()
+                scope.set_transaction_name(
+                    app.lambda_context.function_name,
+                    source=TransactionSource.COMPONENT,
+                )
+
+                scope.add_event_processor(
+                    _make_request_event_processor(
+                        app.current_request.to_dict(),
+                        app.lambda_context,
+                        configured_time,
+                    )
+                )
+            try:
+                return view_function(**function_args)
+            except Exception as exc:
+                if isinstance(exc, ChaliceViewError):
+                    raise
+                exc_info = sys.exc_info()
+                event, hint = event_from_exception(
+                    exc_info,
+                    client_options=client.options,
+                    mechanism={"type": "chalice", "handled": False},
+                )
+                sentry_sdk.capture_event(event, hint=hint)
+                client.flush()
+                raise
+
+    return wrapped_view_function  # type: ignore
+
+
+class ChaliceIntegration(Integration):
+    identifier = "chalice"
+
+    @staticmethod
+    def setup_once():
+        # type: () -> None
+
+        version = parse_version(CHALICE_VERSION)
+
+        if version is None:
+            raise DidNotEnable("Unparsable Chalice version: {}".format(CHALICE_VERSION))
+
+        if version < (1, 20):
+            old_get_view_function_response = Chalice._get_view_function_response
+        else:
+            from chalice.app import RestAPIEventHandler
+
+            old_get_view_function_response = (
+                RestAPIEventHandler._get_view_function_response
+            )
+
+        def sentry_event_response(app, view_function, function_args):
+            # type: (Any, F, Dict[str, Any]) -> Any
+            wrapped_view_function = _get_view_function_response(
+                app, view_function, function_args
+            )
+
+            return old_get_view_function_response(
+                app, wrapped_view_function, function_args
+            )
+
+        if version < (1, 20):
+            Chalice._get_view_function_response = sentry_event_response
+        else:
+            RestAPIEventHandler._get_view_function_response = sentry_event_response
+        # for everything else (like events)
+        chalice.app.EventSourceHandler = EventSourceHandler

sentry_sdk/integrations/clickhouse_driver.py

@@ -0,0 +1,177 @@
+import sentry_sdk
+from sentry_sdk.consts import OP, SPANDATA
+from sentry_sdk.integrations import _check_minimum_version, Integration, DidNotEnable
+from sentry_sdk.tracing import Span
+from sentry_sdk.scope import should_send_default_pii
+from sentry_sdk.utils import capture_internal_exceptions, ensure_integration_enabled
+
+from typing import TYPE_CHECKING, TypeVar
+
+# Hack to get new Python features working in older versions
+# without introducing a hard dependency on `typing_extensions`
+# from: https://stackoverflow.com/a/71944042/300572
+if TYPE_CHECKING:
+    from collections.abc import Iterator
+    from typing import Any, ParamSpec, Callable
+else:
+    # Fake ParamSpec
+    class ParamSpec:
+        def __init__(self, _):
+            self.args = None
+            self.kwargs = None
+
+    # Callable[anything] will return None
+    class _Callable:
+        def __getitem__(self, _):
+            return None
+
+    # Make instances
+    Callable = _Callable()
+
+
+try:
+    import clickhouse_driver  # type: ignore[import-not-found]
+
+except ImportError:
+    raise DidNotEnable("clickhouse-driver not installed.")
+
+
+class ClickhouseDriverIntegration(Integration):
+    identifier = "clickhouse_driver"
+    origin = f"auto.db.{identifier}"
+
+    @staticmethod
+    def setup_once() -> None:
+        _check_minimum_version(ClickhouseDriverIntegration, clickhouse_driver.VERSION)
+
+        # Every query is done using the Connection's `send_query` function
+        clickhouse_driver.connection.Connection.send_query = _wrap_start(
+            clickhouse_driver.connection.Connection.send_query
+        )
+
+        # If the query contains parameters then the send_data function is used to send those parameters to clickhouse
+        _wrap_send_data()
+
+        # Every query ends either with the Client's `receive_end_of_query` (no result expected)
+        # or its `receive_result` (result expected)
+        clickhouse_driver.client.Client.receive_end_of_query = _wrap_end(
+            clickhouse_driver.client.Client.receive_end_of_query
+        )
+        if hasattr(clickhouse_driver.client.Client, "receive_end_of_insert_query"):
+            # In 0.2.7, insert queries are handled separately via `receive_end_of_insert_query`
+            clickhouse_driver.client.Client.receive_end_of_insert_query = _wrap_end(
+                clickhouse_driver.client.Client.receive_end_of_insert_query
+            )
+        clickhouse_driver.client.Client.receive_result = _wrap_end(
+            clickhouse_driver.client.Client.receive_result
+        )
+
+
+P = ParamSpec("P")
+T = TypeVar("T")
+
+
+def _wrap_start(f: Callable[P, T]) -> Callable[P, T]:
+    @ensure_integration_enabled(ClickhouseDriverIntegration, f)
+    def _inner(*args: P.args, **kwargs: P.kwargs) -> T:
+        connection = args[0]
+        query = args[1]
+        query_id = args[2] if len(args) > 2 else kwargs.get("query_id")
+        params = args[3] if len(args) > 3 else kwargs.get("params")
+
+        span = sentry_sdk.start_span(
+            op=OP.DB,
+            name=query,
+            origin=ClickhouseDriverIntegration.origin,
+        )
+
+        connection._sentry_span = span  # type: ignore[attr-defined]
+
+        _set_db_data(span, connection)
+
+        span.set_data("query", query)
+
+        if query_id:
+            span.set_data("db.query_id", query_id)
+
+        if params and should_send_default_pii():
+            span.set_data("db.params", params)
+
+        # run the original code
+        ret = f(*args, **kwargs)
+
+        return ret
+
+    return _inner
+
+
+def _wrap_end(f: Callable[P, T]) -> Callable[P, T]:
+    def _inner_end(*args: P.args, **kwargs: P.kwargs) -> T:
+        res = f(*args, **kwargs)
+        instance = args[0]
+        span = getattr(instance.connection, "_sentry_span", None)  # type: ignore[attr-defined]
+
+        if span is not None:
+            if res is not None and should_send_default_pii():
+                span.set_data("db.result", res)
+
+            with capture_internal_exceptions():
+                span.scope.add_breadcrumb(
+                    message=span._data.pop("query"), category="query", data=span._data
+                )
+
+            span.finish()
+
+        return res
+
+    return _inner_end
+
+
+def _wrap_send_data() -> None:
+    original_send_data = clickhouse_driver.client.Client.send_data
+
+    def _inner_send_data(  # type: ignore[no-untyped-def] # clickhouse-driver does not type send_data
+        self, sample_block, data, types_check=False, columnar=False, *args, **kwargs
+    ):
+        span = getattr(self.connection, "_sentry_span", None)
+
+        if span is not None:
+            _set_db_data(span, self.connection)
+
+            if should_send_default_pii():
+                db_params = span._data.get("db.params", [])
+
+                if isinstance(data, (list, tuple)):
+                    db_params.extend(data)
+
+                else:  # data is a generic iterator
+                    orig_data = data
+
+                    # Wrap the generator to add items to db.params as they are yielded.
+                    # This allows us to send the params to Sentry without needing to allocate
+                    # memory for the entire generator at once.
+                    def wrapped_generator() -> "Iterator[Any]":
+                        for item in orig_data:
+                            db_params.append(item)
+                            yield item
+
+                    # Replace the original iterator with the wrapped one.
+                    data = wrapped_generator()
+
+                span.set_data("db.params", db_params)
+
+        return original_send_data(
+            self, sample_block, data, types_check, columnar, *args, **kwargs
+        )
+
+    clickhouse_driver.client.Client.send_data = _inner_send_data
+
+
+def _set_db_data(
+    span: Span, connection: clickhouse_driver.connection.Connection
+) -> None:
+    span.set_data(SPANDATA.DB_SYSTEM, "clickhouse")
+    span.set_data(SPANDATA.SERVER_ADDRESS, connection.host)
+    span.set_data(SPANDATA.SERVER_PORT, connection.port)
+    span.set_data(SPANDATA.DB_NAME, connection.database)
+    span.set_data(SPANDATA.DB_USER, connection.user)