diff --git a/aiogram/contrib/middlewares/sentry.py b/aiogram/contrib/middlewares/sentry.py index 0e09d4f8..c6629b1e 100644 --- a/aiogram/contrib/middlewares/sentry.py +++ b/aiogram/contrib/middlewares/sentry.py @@ -1,6 +1,6 @@ -from sentry_sdk import (start_span, start_transaction, Hub, +from sentry_sdk import (start_transaction, Hub, set_user, set_tag, set_context, ) -from sentry_sdk.tracing import Transaction, Span +from sentry_sdk.tracing import Span from aiogram.dispatcher.middlewares import BaseMiddleware from aiogram.types import Chat, User @@ -24,12 +24,13 @@ class SentryMiddleware(BaseMiddleware): self._bot_username = me.username set_tag("bot.username", self._bot_username) - start_transaction(name="process_update") + self._start_span("process_update") async def on_post_process_update(self, *_, **__): - self._finish_transaction() + self._finish_span() async def on_pre_process_message(self, message, *_, **__): + self._save_base_context() self._start_span("process_message") set_context("message", message.to_python()) @@ -37,6 +38,7 @@ class SentryMiddleware(BaseMiddleware): self._finish_span() async def on_pre_process_edited_message(self, edited_message, *_, **__): + self._save_base_context() self._start_span("process_edited_message") set_context("message", edited_message.to_python()) @@ -44,6 +46,7 @@ class SentryMiddleware(BaseMiddleware): self._finish_span() async def on_pre_process_channel_post(self, channel_post, *_, **__): + self._save_base_context() self._start_span("process_channel_post") set_context("message", channel_post.to_python()) @@ -52,6 +55,7 @@ class SentryMiddleware(BaseMiddleware): async def on_pre_process_edited_channel_post(self, edited_channel_post, *_, **__): + self._save_base_context() self._start_span("process_edited_channel_post") set_context("message", edited_channel_post.to_python()) @@ -59,6 +63,7 @@ class SentryMiddleware(BaseMiddleware): self._finish_span() async def on_pre_process_inline_query(self, inline_query, *_, **__): + self._save_base_context() self._start_span("process_inline_query") set_context("message", inline_query.to_python()) @@ -67,6 +72,7 @@ class SentryMiddleware(BaseMiddleware): async def on_pre_process_chosen_inline_result(self, chosen_inline_result, *_, **__): + self._save_base_context() self._start_span("process_chosen_inline_result") set_context("inline_result", chosen_inline_result.to_python()) @@ -74,6 +80,7 @@ class SentryMiddleware(BaseMiddleware): self._finish_span() async def on_pre_process_callback_query(self, callback_query, *_, **__): + self._save_base_context() self._start_span("process_callback_query") set_context("callback_query", callback_query.to_python()) @@ -81,6 +88,7 @@ class SentryMiddleware(BaseMiddleware): self._finish_span() async def on_pre_process_shipping_query(self, shipping_query, *_, **__): + self._save_base_context() self._start_span("process_shipping_query") set_context("shipping_query", shipping_query.to_python()) @@ -88,6 +96,7 @@ class SentryMiddleware(BaseMiddleware): self._finish_span() async def on_pre_process_pre_checkout_query(self, pre_checkout_query, *_, **__): + self._save_base_context() self._start_span("process_pre_checkout_query") set_context("pre_checkout_query", pre_checkout_query.to_python()) @@ -95,6 +104,7 @@ class SentryMiddleware(BaseMiddleware): self._finish_span() async def on_pre_process_poll(self, poll, *_, **__): + self._save_base_context() self._start_span("process_poll") set_context("poll", poll.to_python()) @@ -102,6 +112,7 @@ class SentryMiddleware(BaseMiddleware): self._finish_span() async def on_pre_process_poll_answer(self, poll_answer, *_, **__): + self._save_base_context() self._start_span("process_poll_answer") set_context("poll_answer", poll_answer.to_python()) @@ -109,13 +120,21 @@ class SentryMiddleware(BaseMiddleware): self._finish_span() @staticmethod - def _finish_transaction(): - transaction = Hub.current.scope.transaction - if isinstance(transaction, Transaction): - return transaction.finish() + def _start_span(name): + """ + Trying to get current span, then: + - if no span in progress, create new transaction; + - if span exists - create new task span as child of current span; + """ + span = Hub.current.scope.span + if span is None: + start_transaction(name=name) + else: + span.start_child(op=name) @staticmethod - def _start_span(op): + def _save_base_context(): + """ Saving contexts if User and Chat. """ user = User.get_current() if isinstance(user, User): user_data = {"id": user.id} @@ -128,10 +147,9 @@ class SentryMiddleware(BaseMiddleware): if isinstance(chat, Chat): set_context("chat", chat.to_python()) - start_span(op=op) - @staticmethod def _finish_span(): + """ Finish current span. """ span = Hub.current.scope.span if isinstance(span, Span): return span.finish()