nameko_sentry.py

import os
import sys
from abc import ABC
from urllib.parse import urlsplit
from weakref import WeakKeyDictionary

import sentry_sdk
from nameko.extensions import DependencyProvider
from nameko.web.handlers import HttpRequestHandler
from sentry_sdk import Hub
from sentry_sdk.utils import event_from_exception
from werkzeug import Request
from werkzeug.exceptions import ClientDisconnected

from common.exceptions import AlreadyRanError, RetryableBadDataError
from common.logger import get_logger

logger = get_logger(__name__)


def before_send_filter(event, hint):
    if "exc_info" in hint:
        exc_type, exc_value, tb = hint["exc_info"]
        if isinstance(exc_value, (RetryableBadDataError, AlreadyRanError)):
            return None

    # also ignore retryable error log messages
    if "log_record" in hint:
        if "Retryable error:" in hint["log_record"].message:
            return None

    return event


app_env = os.environ.get("APP_ENV", "development").lower()
app_version = os.environ.get("APP_VERSION", None)
sentry_dsn = os.environ.get("SENTRY_DSN", None)

if sentry_dsn is None or app_env == "test":
    logger.info("Skipped sentry init; no DSN configured")
else:
    try:
        sentry_sdk.init(
            dsn=sentry_dsn,
            release=app_version,
            environment=app_env,
            traces_sample_rate=1.0,
            before_send=before_send_filter,
        )
    except sentry_sdk.utils.BadDsn as err:
        logger.exception(
            "Error initializing Sentry integration", extra=dict(error=str(err))
        )
        sys.exit(1)


class SentryDependencyProvider(DependencyProvider):
    def __init__(self):
        self.hubs = WeakKeyDictionary()
        self.transactions = WeakKeyDictionary()
        self.main_hub = None

    def setup(self):
        """
        When service is starting, init sentry for all our workers.
        """
        self.main_hub = Hub.current

    def worker_setup(self, worker_ctx):
        """
        When worker is about to start work, setup hub and capture a few things
        """
        worker_hub = self.get_worker_hub(worker_ctx)
        if not worker_hub:
            return  # sentry not setup

        worker_hub.add_breadcrumb(
            category="worker_lifecycle", message="worker_setup", level="debug"
        )
        self.http_context(worker_ctx)  # extract http context, if any

        # self.transactions[worker_ctx] = worker_hub.start_transaction(
        #    op="worker_lifecycle", name="do_work"
        # )
        # self.transactions[worker_ctx].__enter__()

    def worker_result(self, worker_ctx, result=None, exc_info=None):
        """
        When worker completed work, if the worker resulted in an error, capture it in sentry.
        """
        worker_hub = self.get_worker_hub(worker_ctx)
        if not worker_hub:
            return  # sentry not setup

        worker_hub.add_breadcrumb(
            category="worker_lifecycle", message="worker_result", level="debug"
        )

        if exc_info is None:
            return  # nothing to do

        self.tags_context(worker_ctx, exc_info)
        self.extra_context(worker_ctx, exc_info)

        # Capture the error in sentry.
        with worker_hub:
            client = worker_hub.client

            event, hint = event_from_exception(
                exc_info,
                client_options=client.options,
                mechanism={"type": "threading", "handled": False},
            )
            res = worker_hub.capture_event(event, hint=hint)
            logger.debug(f"capture_event result: {res}")

    def worker_teardown(self, worker_ctx):
        """On teardown"""
        # self.transactions[worker_ctx].__exit__(None, None, None)
        # del self.transactions[worker_ctx]
        if worker_ctx in self.hubs:
            del self.hubs[worker_ctx]

    def get_dependency(self, worker_ctx):
        return self.get_worker_hub(worker_ctx)

    def get_worker_hub(self, worker_ctx) -> Hub:
        """
        For this worker, get a Hub pointing to the main hub
        """
        if self.main_hub and not self.hubs.get(worker_ctx):
            self.hubs[worker_ctx] = Hub(self.main_hub)
        worker_hub = self.hubs.get(worker_ctx)
        return worker_hub if worker_hub else None

    def http_context(self, worker_ctx):
        """Attempt to extract HTTP context if an HTTP entrypoint was used."""
        http = {}
        if isinstance(worker_ctx.entrypoint, HttpRequestHandler):
            try:
                request: Request = worker_ctx.args[0]
                try:
                    if request.mimetype == "application/json":
                        data = request.data
                    else:
                        data = request.form
                except ClientDisconnected:
                    data = {}

                urlparts = urlsplit(request.url)
                http.update(
                    {
                        "url": "{}://{}{}".format(
                            urlparts.scheme, urlparts.netloc, urlparts.path
                        ),
                        "query_string": urlparts.query,
                        "method": request.method,
                        "data": data,
                        "headers": dict(request.headers),
                        "env": dict(request.environ),
                    }
                )
            except:
                pass  # probably not a compatible entrypoint

        self.get_worker_hub(worker_ctx).scope.set_context("http", http)

    def tags_context(self, worker_ctx, exc_info):
        """Merge any tags to include in the sentry payload"""

        worker_hub = self.get_worker_hub(worker_ctx)
        if not worker_hub:
            return  # sentry not setup

        tags = {
            "call_id": worker_ctx.call_id,
            "parent_call_id": worker_ctx.immediate_parent_call_id,
            "service_name": worker_ctx.container.service_name,
            "method_name": worker_ctx.entrypoint.method_name,
        }
        # for key in worker_ctx.context_data:
        #     for matcher in self.tag_type_context_keys:
        #         if re.search(matcher, key):
        #             tags[key] = worker_ctx.context_data[key]
        #             break

        for tag, value in tags.items():
            worker_hub.scope.set_tag(tag, value)

    def extra_context(self, worker_ctx, exc_info):
        """Merge any extra context to include in the sentry payload.

        Includes all available worker context data.
        """

        worker_hub = self.get_worker_hub(worker_ctx)
        if not worker_hub:
            return  # sentry not setup

        extra = {}
        extra.update(worker_ctx.context_data)

        worker_hub.scope.set_context("extra", extra)


class SentryMixin(ABC):
    sentry = SentryDependencyProvider()