|
| 1 | +import json |
| 2 | +import logging |
| 3 | +import traceback |
| 4 | +from typing import Any, Optional, Type |
| 5 | + |
| 6 | +from temporalio import activity, workflow |
| 7 | +from temporalio.worker import ( |
| 8 | + ActivityInboundInterceptor, |
| 9 | + ExecuteActivityInput, |
| 10 | + ExecuteWorkflowInput, |
| 11 | + Interceptor, |
| 12 | + WorkflowInboundInterceptor, |
| 13 | + WorkflowInterceptorClassInput, |
| 14 | +) |
| 15 | + |
| 16 | +logger = logging.getLogger(__name__) |
| 17 | + |
| 18 | + |
| 19 | +class _MultilineLoggingActivityInboundInterceptor(ActivityInboundInterceptor): |
| 20 | + async def execute_activity(self, input: ExecuteActivityInput) -> Any: |
| 21 | + try: |
| 22 | + return await super().execute_activity(input) |
| 23 | + except Exception as e: |
| 24 | + exception_data = { |
| 25 | + "message": str(e), |
| 26 | + "type": type(e).__name__, |
| 27 | + "traceback": traceback.format_exc().replace("\n", " | ") |
| 28 | + } |
| 29 | + |
| 30 | + logger.error(f"Activity exception: {json.dumps(exception_data)}") |
| 31 | + |
| 32 | + raise e |
| 33 | + |
| 34 | + |
| 35 | +class _MultilineLoggingWorkflowInterceptor(WorkflowInboundInterceptor): |
| 36 | + async def execute_workflow(self, input: ExecuteWorkflowInput) -> Any: |
| 37 | + try: |
| 38 | + return await super().execute_workflow(input) |
| 39 | + except Exception as e: |
| 40 | + exception_data = { |
| 41 | + "message": str(e), |
| 42 | + "type": type(e).__name__, |
| 43 | + "traceback": traceback.format_exc().replace("\n", " | ") |
| 44 | + } |
| 45 | + |
| 46 | + if not workflow.unsafe.is_replaying(): |
| 47 | + with workflow.unsafe.sandbox_unrestricted(): |
| 48 | + logger.error(f"Workflow exception: {json.dumps(exception_data)}") |
| 49 | + |
| 50 | + raise e |
| 51 | + |
| 52 | + |
| 53 | +class MultilineLoggingInterceptor(Interceptor): |
| 54 | + """Temporal Interceptor that formats multiline exception logs as single-line JSON""" |
| 55 | + |
| 56 | + def intercept_activity( |
| 57 | + self, next: ActivityInboundInterceptor |
| 58 | + ) -> ActivityInboundInterceptor: |
| 59 | + return _MultilineLoggingActivityInboundInterceptor(super().intercept_activity(next)) |
| 60 | + |
| 61 | + def workflow_interceptor_class( |
| 62 | + self, input: WorkflowInterceptorClassInput |
| 63 | + ) -> Optional[Type[WorkflowInboundInterceptor]]: |
| 64 | + return _MultilineLoggingWorkflowInterceptor |
0 commit comments