[airbyte-cdk] refactor error handling in abstract source (#40571)

Author: lazebnyi

airbyte-cdk/python/airbyte_cdk/sources/abstract_source.py

@@ -139,35 +139,26 @@ def read(
                     )
                     logger.info(f"Marking stream {configured_stream.stream.name} as STOPPED")
                     yield stream_status_as_airbyte_message(configured_stream.stream, AirbyteStreamStatus.COMPLETE)
-                except AirbyteTracedException as e:
-                    yield from self._emit_queued_messages()
-                    logger.exception(f"Encountered an exception while reading stream {configured_stream.stream.name}")
-                    logger.info(f"Marking stream {configured_stream.stream.name} as STOPPED")
-                    yield stream_status_as_airbyte_message(configured_stream.stream, AirbyteStreamStatus.INCOMPLETE)
-                    yield e.as_sanitized_airbyte_message(stream_descriptor=StreamDescriptor(name=configured_stream.stream.name))
-                    stream_name_to_exception[stream_instance.name] = e  # type: ignore # use configured_stream if stream_instance is None
-                    if self.stop_sync_on_stream_failure:
-                        logger.info(
-                            f"Stopping sync on error from stream {configured_stream.stream.name} because {self.name} does not support continuing syncs on error."
-                        )
-                        break
+
                 except Exception as e:
                     yield from self._emit_queued_messages()
                     logger.exception(f"Encountered an exception while reading stream {configured_stream.stream.name}")
                     logger.info(f"Marking stream {configured_stream.stream.name} as STOPPED")
                     yield stream_status_as_airbyte_message(configured_stream.stream, AirbyteStreamStatus.INCOMPLETE)
-                    display_message = stream_instance.get_error_display_message(e)  # type: ignore[union-attr]
+
                     stream_descriptor = StreamDescriptor(name=configured_stream.stream.name)
-                    if display_message:
-                        traced_exception = AirbyteTracedException.from_exception(
-                            e, message=display_message, stream_descriptor=stream_descriptor
-                        )
+
+                    if isinstance(e, AirbyteTracedException):
+                        traced_exception = e
+                        info_message = f"Stopping sync on error from stream {configured_stream.stream.name} because {self.name} does not support continuing syncs on error."
                     else:
-                        traced_exception = AirbyteTracedException.from_exception(e, stream_descriptor=stream_descriptor)
-                    yield traced_exception.as_sanitized_airbyte_message()
-                    stream_name_to_exception[stream_instance.name] = traced_exception  # type: ignore
+                        traced_exception = self._serialize_exception(stream_descriptor, e, stream_instance=stream_instance)
+                        info_message = f"{self.name} does not support continuing syncs on error from stream {configured_stream.stream.name}"
+
+                    yield traced_exception.as_sanitized_airbyte_message(stream_descriptor=stream_descriptor)
+                    stream_name_to_exception[stream_instance.name] = traced_exception  # type: ignore # use configured_stream if stream_instance is None
                     if self.stop_sync_on_stream_failure:
-                        logger.info(f"{self.name} does not support continuing syncs on error from stream {configured_stream.stream.name}")
+                        logger.info(info_message)
                         break
                 finally:
                     # Finish read event only if the stream instance exists;
@@ -186,6 +177,15 @@ def read(
             raise AirbyteTracedException(message=error_message, failure_type=FailureType.config_error)
         logger.info(f"Finished syncing {self.name}")
 
+    @staticmethod
+    def _serialize_exception(
+        stream_descriptor: StreamDescriptor, e: Exception, stream_instance: Optional[Stream] = None
+    ) -> AirbyteTracedException:
+        display_message = stream_instance.get_error_display_message(e) if stream_instance else None
+        if display_message:
+            return AirbyteTracedException.from_exception(e, message=display_message, stream_descriptor=stream_descriptor)
+        return AirbyteTracedException.from_exception(e, stream_descriptor=stream_descriptor)
+
     @property
     def raise_exception_on_missing_stream(self) -> bool:
         return True