[ISSUE #32871] extract trace message creation

airbyte-cdk/python/airbyte_cdk/exception_handler.py

@@ -5,16 +5,25 @@
 import logging
 import sys
 
+from types import TracebackType
+from typing import Any, Optional
+
 from airbyte_cdk.utils.traced_exception import AirbyteTracedException
 
 
+def assemble_uncaught_exception(exception_type: type[BaseException], exception_value: BaseException) -> AirbyteTracedException:
+    if issubclass(exception_type, AirbyteTracedException):
+        return exception_value  # type: ignore  # validated as part of the previous line
+    return AirbyteTracedException.from_exception(exception_value)
+
+
 def init_uncaught_exception_handler(logger: logging.Logger) -> None:
     """
     Handles uncaught exceptions by emitting an AirbyteTraceMessage and making sure they are not
     printed to the console without having secrets removed.
     """
 
-    def hook_fn(exception_type, exception_value, traceback_):
+    def hook_fn(exception_type: type[BaseException], exception_value: BaseException, traceback_: Optional[TracebackType]) -> Any:
         # For developer ergonomics, we want to see the stack trace in the logs when we do a ctrl-c
         if issubclass(exception_type, KeyboardInterrupt):
             sys.__excepthook__(exception_type, exception_value, traceback_)
@@ -23,11 +32,7 @@ def hook_fn(exception_type, exception_value, traceback_):
         logger.fatal(exception_value, exc_info=exception_value)
 
         # emit an AirbyteTraceMessage for any exception that gets to this spot
-        traced_exc = (
-            exception_value
-            if issubclass(exception_type, AirbyteTracedException)
-            else AirbyteTracedException.from_exception(exception_value)
-        )
+        traced_exc = assemble_uncaught_exception(exception_type, exception_value)
 
         traced_exc.emit_message()
 

airbyte-cdk/python/unit_tests/test_exception_handler.py

@@ -8,7 +8,21 @@
 import sys
 
 import pytest
+from airbyte_cdk.exception_handler import assemble_uncaught_exception
 from airbyte_cdk.models import AirbyteErrorTraceMessage, AirbyteLogMessage, AirbyteMessage, AirbyteTraceMessage
+from airbyte_cdk.utils.traced_exception import AirbyteTracedException
+
+
+def test_given_exception_is_traced_exception_when_assemble_uncaught_exception_then_return_same_exception():
+    exception = AirbyteTracedException()
+    assembled_exception = assemble_uncaught_exception(type(exception), exception)
+    assert exception == assembled_exception
+
+
+def test_given_exception_not_traced_exception_when_assemble_uncaught_exception_then_return_traced_exception():
+    exception = ValueError("any error")
+    assembled_exception = assemble_uncaught_exception(type(exception), exception)
+    assert isinstance(assembled_exception, AirbyteTracedException)
 
 
 def test_uncaught_exception_handler():