fix: do not break existing transactions (e.g. aiohttp webhook integration)

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()