Improve sentry workflow & activity interceptors (temporalio#46)

nediamond · web-flow · commit 0800779ba4a0 · 2023-01-23T16:28:58.000-06:00
diff --git a/sentry/interceptor.py b/sentry/interceptor.py
@@ -1,7 +1,6 @@
 from dataclasses import asdict, is_dataclass
 from typing import Any, Optional, Type, Union
 
-import sentry_sdk
 from temporalio import activity, workflow
 from temporalio.worker import (
     ActivityInboundInterceptor,
@@ -12,66 +11,60 @@
     WorkflowInterceptorClassInput,
 )
 
+with workflow.unsafe.imports_passed_through():
+    from sentry_sdk import Hub, capture_exception, set_context, set_tag
 
-def _set_common_workflow_tags(
-    info: Union[workflow.Info, activity.Info], scope: sentry_sdk.Scope
-):
-    scope.set_tag("temporal.workflow.type", info.workflow_type)
-    scope.set_tag("temporal.workflow.id", info.workflow_id)
+
+def _set_common_workflow_tags(info: Union[workflow.Info, activity.Info]):
+    set_tag("temporal.workflow.type", info.workflow_type)
+    set_tag("temporal.workflow.id", info.workflow_id)
 
 
 class _SentryActivityInboundInterceptor(ActivityInboundInterceptor):
     async def execute_activity(self, input: ExecuteActivityInput) -> Any:
-        transaction_name = input.fn.__module__ + "." + input.fn.__qualname__
-        scope_ctx_manager = sentry_sdk.configure_scope()
-        with scope_ctx_manager as scope, sentry_sdk.start_transaction(
-            name=transaction_name
-        ):
-            scope.set_tag("temporal.execution_type", "activity")
+        # https://docs.sentry.io/platforms/python/troubleshooting/#addressing-concurrency-issues
+        with Hub(Hub.current):
+            set_tag("temporal.execution_type", "activity")
+            set_tag("module", input.fn.__module__ + "." + input.fn.__qualname__)
+
             activity_info = activity.info()
-            _set_common_workflow_tags(activity_info, scope)
-            scope.set_tag("temporal.activity.id", activity_info.activity_id)
-            scope.set_tag("temporal.activity.type", activity_info.activity_type)
-            scope.set_tag("temporal.activity.task_queue", activity_info.task_queue)
-            scope.set_tag(
-                "temporal.workflow.namespace", activity_info.workflow_namespace
-            )
-            scope.set_tag("temporal.workflow.run_id", activity_info.workflow_run_id)
+            _set_common_workflow_tags(activity_info)
+            set_tag("temporal.activity.id", activity_info.activity_id)
+            set_tag("temporal.activity.type", activity_info.activity_type)
+            set_tag("temporal.activity.task_queue", activity_info.task_queue)
+            set_tag("temporal.workflow.namespace", activity_info.workflow_namespace)
+            set_tag("temporal.workflow.run_id", activity_info.workflow_run_id)
             try:
                 return await super().execute_activity(input)
             except Exception as e:
                 if len(input.args) == 1 and is_dataclass(input.args[0]):
-                    scope.set_context("temporal.activity.input", asdict(input.args[0]))
-                scope.set_context("temporal.activity.info", activity.info().__dict__)
-                sentry_sdk.capture_exception(e)
+                    set_context("temporal.activity.input", asdict(input.args[0]))
+                set_context("temporal.activity.info", activity.info().__dict__)
+                capture_exception()
                 raise e
-            finally:
-                scope.clear()
 
 
 class _SentryWorkflowInterceptor(WorkflowInboundInterceptor):
     async def execute_workflow(self, input: ExecuteWorkflowInput) -> Any:
-        transaction_name = input.run_fn.__module__ + "." + input.run_fn.__qualname__
-        scope_ctx_manager = sentry_sdk.configure_scope()
-        with scope_ctx_manager as scope, sentry_sdk.start_transaction(
-            name=transaction_name
-        ):
-            scope.set_tag("temporal.execution_type", "workflow")
+        # https://docs.sentry.io/platforms/python/troubleshooting/#addressing-concurrency-issues
+        with Hub(Hub.current):
+            set_tag("temporal.execution_type", "workflow")
+            set_tag("module", input.run_fn.__module__ + "." + input.run_fn.__qualname__)
             workflow_info = workflow.info()
-            _set_common_workflow_tags(workflow_info, scope)
-            scope.set_tag("temporal.workflow.task_queue", workflow_info.task_queue)
-            scope.set_tag("temporal.workflow.namespace", workflow_info.namespace)
-            scope.set_tag("temporal.workflow.run_id", workflow_info.run_id)
+            _set_common_workflow_tags(workflow_info)
+            set_tag("temporal.workflow.task_queue", workflow_info.task_queue)
+            set_tag("temporal.workflow.namespace", workflow_info.namespace)
+            set_tag("temporal.workflow.run_id", workflow_info.run_id)
             try:
                 return await super().execute_workflow(input)
             except Exception as e:
                 if len(input.args) == 1 and is_dataclass(input.args[0]):
-                    scope.set_context("temporal.workflow.input", asdict(input.args[0]))
-                scope.set_context("temporal.workflow.info", workflow.info().__dict__)
-                sentry_sdk.capture_exception(e)
+                    set_context("temporal.workflow.input", asdict(input.args[0]))
+                set_context("temporal.workflow.info", workflow.info().__dict__)
+
+                if not workflow.unsafe.is_replaying():
+                    capture_exception()
                 raise e
-            finally:
-                scope.clear()
 
 
 class SentryInterceptor(Interceptor):
