move print logic to AirbyteTracedException.emit_message()

airbytehq · May 6, 2022 · b00fc78 · b00fc78
1 parent 420c6d9
commit b00fc78
Show file tree

Hide file tree

Showing 3 changed files with 43 additions and 5 deletions.
diff --git a/airbyte-cdk/python/airbyte_cdk/exception_handler.py b/airbyte-cdk/python/airbyte_cdk/exception_handler.py
@@ -5,7 +5,6 @@
 import logging
 import sys
 
-from airbyte_cdk.utils.airbyte_secrets_utils import filter_secrets
 from airbyte_cdk.utils.traced_exception import AirbyteTracedException
 
 
@@ -29,8 +28,7 @@ def hook_fn(exception_type, exception_value, traceback_):
             if issubclass(exception_type, AirbyteTracedException)
             else AirbyteTracedException.from_exception(exception_value)
         )
-        message = traced_exc.as_airbyte_message()
-        message_json = message.json(exclude_unset=True)
-        print(filter_secrets(message_json))
+
+        traced_exc.emit_message()
 
     sys.excepthook = hook_fn
diff --git a/airbyte-cdk/python/airbyte_cdk/utils/traced_exception.py b/airbyte-cdk/python/airbyte_cdk/utils/traced_exception.py
@@ -7,6 +7,7 @@
 
 from airbyte_cdk.models import AirbyteErrorTraceMessage, AirbyteMessage, AirbyteTraceMessage, FailureType, TraceType
 from airbyte_cdk.models import Type as MessageType
+from airbyte_cdk.utils.airbyte_secrets_utils import filter_secrets
 
 
 class AirbyteTracedException(Exception):
@@ -55,6 +56,15 @@ def as_airbyte_message(self) -> AirbyteMessage:
 
         return AirbyteMessage(type=MessageType.TRACE, trace=trace_message)
 
+    def emit_message(self):
+        """
+        Prints the exception as an AirbyteTraceMessage.
+        Note that this will be called automatically on uncaught exceptions when using the airbyte_cdk entrypoint.
+        """
+        message = self.as_airbyte_message().json(exclude_unset=True)
+        filtered_message = filter_secrets(message)
+        print(filtered_message)
+
     @classmethod
     def from_exception(cls, exc: Exception, *args, **kwargs) -> "AirbyteTracedException":
         """

diff --git a/airbyte-cdk/python/unit_tests/utils/test_traced_exception.py b/airbyte-cdk/python/unit_tests/utils/test_traced_exception.py
@@ -2,8 +2,10 @@
 # Copyright (c) 2021 Airbyte, Inc., all rights reserved.
 #
 
+import json
+
 import pytest
-from airbyte_cdk.models.airbyte_protocol import AirbyteMessage, FailureType, TraceType
+from airbyte_cdk.models.airbyte_protocol import AirbyteErrorTraceMessage, AirbyteMessage, AirbyteTraceMessage, FailureType, TraceType
 from airbyte_cdk.models.airbyte_protocol import Type as MessageType
 from airbyte_cdk.utils.traced_exception import AirbyteTracedException
 
@@ -51,3 +53,31 @@ def test_existing_exception_as_airbyte_message(raised_exception):
     assert airbyte_message.trace.error.stack_trace.endswith(
         'raise RuntimeError("an error has occurred")\n' "RuntimeError: an error has occurred\n"
     )
+
+
+def test_emit_message(capsys):
+    traced_exc = AirbyteTracedException(
+        internal_message="internal message", message="user-friendly message", exception=RuntimeError("oh no")
+    )
+
+    expected_message = AirbyteMessage(
+        type="TRACE",
+        trace=AirbyteTraceMessage(
+            type="ERROR",
+            emitted_at=0.0,
+            error=AirbyteErrorTraceMessage(
+                failure_type="system_error",
+                message="user-friendly message",
+                internal_message="internal message",
+                stack_trace="RuntimeError: oh no\n",
+            ),
+        ),
+    )
+
+    traced_exc.emit_message()
+
+    stdout = capsys.readouterr().out
+    printed_message = AirbyteMessage.parse_obj(json.loads(stdout))
+    printed_message.trace.emitted_at = 0.0
+
+    assert printed_message == expected_message