airbyte_cdk.utils.traced_exception

  1#
  2# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
  3#
  4import time
  5import traceback
  6from typing import Any, Optional
  7
  8import orjson
  9
 10from airbyte_cdk.models import (
 11    AirbyteConnectionStatus,
 12    AirbyteErrorTraceMessage,
 13    AirbyteMessage,
 14    AirbyteMessageSerializer,
 15    AirbyteTraceMessage,
 16    FailureType,
 17    Status,
 18    StreamDescriptor,
 19    TraceType,
 20)
 21from airbyte_cdk.models import Type as MessageType
 22from airbyte_cdk.utils.airbyte_secrets_utils import filter_secrets
 23
 24
 25class AirbyteTracedException(Exception):
 26    """
 27    An exception that should be emitted as an AirbyteTraceMessage
 28    """
 29
 30    def __init__(
 31        self,
 32        internal_message: Optional[str] = None,
 33        message: Optional[str] = None,
 34        failure_type: FailureType = FailureType.system_error,
 35        exception: Optional[BaseException] = None,
 36        stream_descriptor: Optional[StreamDescriptor] = None,
 37    ):
 38        """
 39        :param internal_message: the internal error that caused the failure
 40        :param message: a user-friendly message that indicates the cause of the error
 41        :param failure_type: the type of error
 42        :param exception: the exception that caused the error, from which the stack trace should be retrieved
 43        :param stream_descriptor: describe the stream from which the exception comes from
 44        """
 45        self.internal_message = internal_message
 46        self.message = message
 47        self.failure_type = failure_type
 48        self._exception = exception
 49        self._stream_descriptor = stream_descriptor
 50        super().__init__(internal_message)
 51
 52    def __str__(self) -> str:
 53        """Return the user-facing message, falling back to internal_message."""
 54        if self.message is not None:
 55            return self.message
 56        if self.internal_message is not None:
 57            return self.internal_message
 58        return ""
 59
 60    def as_airbyte_message(
 61        self, stream_descriptor: Optional[StreamDescriptor] = None
 62    ) -> AirbyteMessage:
 63        """
 64        Builds an AirbyteTraceMessage from the exception
 65
 66        :param stream_descriptor is deprecated, please use the stream_description in `__init__ or `from_exception`. If many
 67          stream_descriptors are defined, the one from `as_airbyte_message` will be discarded.
 68        """
 69        now_millis = time.time_ns() // 1_000_000
 70
 71        trace_exc = self._exception or self
 72        stack_trace_str = "".join(traceback.TracebackException.from_exception(trace_exc).format())
 73
 74        trace_message = AirbyteTraceMessage(
 75            type=TraceType.ERROR,
 76            emitted_at=now_millis,
 77            error=AirbyteErrorTraceMessage(
 78                message=self.message
 79                or "Something went wrong in the connector. See the logs for more details.",
 80                internal_message=self.internal_message,
 81                failure_type=self.failure_type,
 82                stack_trace=stack_trace_str,
 83                stream_descriptor=self._stream_descriptor
 84                if self._stream_descriptor is not None
 85                else stream_descriptor,
 86            ),
 87        )
 88
 89        return AirbyteMessage(type=MessageType.TRACE, trace=trace_message)
 90
 91    def as_connection_status_message(self) -> Optional[AirbyteMessage]:
 92        """Build a failed CONNECTION_STATUS from a config error, with secrets masked.
 93
 94        The message can quote config values, as a config validation error does, so it goes
 95        through the same secret filter as logs and traces.
 96        """
 97        if self.failure_type == FailureType.config_error:
 98            return AirbyteMessage(
 99                type=MessageType.CONNECTION_STATUS,
100                connectionStatus=AirbyteConnectionStatus(
101                    status=Status.FAILED,
102                    message=filter_secrets(self.message) if self.message else self.message,
103                ),
104            )
105        return None
106
107    def emit_message(self) -> None:
108        """
109        Prints the exception as an AirbyteTraceMessage.
110        Note that this will be called automatically on uncaught exceptions when using the airbyte_cdk entrypoint.
111        """
112        message = orjson.dumps(AirbyteMessageSerializer.dump(self.as_airbyte_message())).decode()
113        filtered_message = filter_secrets(message)
114        print(filtered_message)
115
116    @classmethod
117    def from_exception(
118        cls,
119        exc: BaseException,
120        stream_descriptor: Optional[StreamDescriptor] = None,
121        *args: Any,
122        **kwargs: Any,
123    ) -> "AirbyteTracedException":
124        """
125        Helper to create an AirbyteTracedException from an existing exception
126        :param exc: the exception that caused the error
127        :param stream_descriptor: describe the stream from which the exception comes from
128        """
129        if isinstance(exc, AirbyteTracedException):
130            internal_message = exc.internal_message
131            # Preserve the original user-facing message if the caller didn't provide one
132            if "message" not in kwargs:
133                kwargs["message"] = exc.message
134        else:
135            internal_message = str(exc)
136        return cls(
137            internal_message=internal_message,
138            exception=exc,
139            stream_descriptor=stream_descriptor,
140            *args,
141            **kwargs,
142        )  # type: ignore  # ignoring because of args and kwargs
143
144    def as_sanitized_airbyte_message(
145        self, stream_descriptor: Optional[StreamDescriptor] = None
146    ) -> AirbyteMessage:
147        """
148        Builds an AirbyteTraceMessage from the exception and sanitizes any secrets from the message body
149
150        :param stream_descriptor is deprecated, please use the stream_description in `__init__ or `from_exception`. If many
151          stream_descriptors are defined, the one from `as_sanitized_airbyte_message` will be discarded.
152        """
153        error_message = self.as_airbyte_message(stream_descriptor=stream_descriptor)
154        if error_message.trace.error.message:  # type: ignore[union-attr] # AirbyteMessage with MessageType.TRACE has AirbyteTraceMessage
155            error_message.trace.error.message = filter_secrets(  # type: ignore[union-attr]
156                error_message.trace.error.message,  # type: ignore[union-attr]
157            )
158        if error_message.trace.error.internal_message:  # type: ignore[union-attr] # AirbyteMessage with MessageType.TRACE has AirbyteTraceMessage
159            error_message.trace.error.internal_message = filter_secrets(  # type: ignore[union-attr] # AirbyteMessage with MessageType.TRACE has AirbyteTraceMessage
160                error_message.trace.error.internal_message  # type: ignore[union-attr] # AirbyteMessage with MessageType.TRACE has AirbyteTraceMessage
161            )
162        if error_message.trace.error.stack_trace:  # type: ignore[union-attr] # AirbyteMessage with MessageType.TRACE has AirbyteTraceMessage
163            error_message.trace.error.stack_trace = filter_secrets(  # type: ignore[union-attr] # AirbyteMessage with MessageType.TRACE has AirbyteTraceMessage
164                error_message.trace.error.stack_trace  # type: ignore[union-attr] # AirbyteMessage with MessageType.TRACE has AirbyteTraceMessage
165            )
166        return error_message