airbyte-cdk 0.31.0__tar.gz → 0.32.0__tar.gz

{airbyte-cdk-0.31.0 → airbyte-cdk-0.32.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: airbyte-cdk
-Version: 0.31.0
+Version: 0.32.0
 Summary: A framework for writing Airbyte Connectors.
 Home-page: https://github.com/airbytehq/airbyte
 Author: Airbyte

{airbyte-cdk-0.31.0 → airbyte-cdk-0.32.0}/airbyte_cdk/connector_builder/message_grouper.py

@@ -6,14 +6,30 @@ import json
 import logging
 from copy import deepcopy
 from json import JSONDecodeError
-from typing import Any, Iterable, Iterator, Mapping, Optional, Union
+from typing import Any, Iterable, Iterator, List, Mapping, Optional, Union
 from urllib.parse import parse_qs, urlparse
 
-from airbyte_cdk.connector_builder.models import HttpRequest, HttpResponse, StreamRead, StreamReadPages, StreamReadSlices
-from airbyte_cdk.models import AirbyteLogMessage, AirbyteMessage, Type
+from airbyte_cdk.connector_builder.models import (
+    HttpRequest,
+    HttpResponse,
+    LogMessage,
+    StreamRead,
+    StreamReadPages,
+    StreamReadSlices,
+    StreamReadSlicesInner,
+)
 from airbyte_cdk.sources.declarative.declarative_source import DeclarativeSource
+from airbyte_cdk.utils import AirbyteTracedException
 from airbyte_cdk.utils.schema_inferrer import SchemaInferrer
-from airbyte_protocol.models.airbyte_protocol import ConfiguredAirbyteCatalog
+from airbyte_protocol.models.airbyte_protocol import (
+    AirbyteLogMessage,
+    AirbyteMessage,
+    AirbyteTraceMessage,
+    ConfiguredAirbyteCatalog,
+    Level,
+    TraceType,
+)
+from airbyte_protocol.models.airbyte_protocol import Type as MessageType
 
 
 class MessageGrouper:
@@ -42,14 +58,18 @@ class MessageGrouper:
 
         slices = []
         log_messages = []
-        state = {}  # No support for incremental sync
         for message_group in self._get_message_groups(
-            source.read(self.logger, config, configured_catalog, state),
-            schema_inferrer,
-            record_limit,
+                self._read_stream(source, config, configured_catalog),
+                schema_inferrer,
+                record_limit,
         ):
             if isinstance(message_group, AirbyteLogMessage):
-                log_messages.append({"message": message_group.message})
+                log_messages.append(LogMessage(**{"message": message_group.message, "level": message_group.level.value}))
+            elif isinstance(message_group, AirbyteTraceMessage):
+                if message_group.type == TraceType.ERROR:
+                    error_message = f"{message_group.error.message} - {message_group.error.stack_trace}"
+                    log_messages.append(LogMessage(**{"message": error_message, "level": "ERROR"}))
+
             else:
                 slices.append(message_group)
 
@@ -63,8 +83,8 @@ class MessageGrouper:
         )
 
     def _get_message_groups(
-        self, messages: Iterator[AirbyteMessage], schema_inferrer: SchemaInferrer, limit: int
-    ) -> Iterable[Union[StreamReadPages, AirbyteLogMessage]]:
+            self, messages: Iterator[AirbyteMessage], schema_inferrer: SchemaInferrer, limit: int
+    ) -> Iterable[Union[StreamReadPages, AirbyteLogMessage, AirbyteTraceMessage]]:
         """
         Message groups are partitioned according to when request log messages are received. Subsequent response log messages
         and record messages belong to the prior request log message and when we encounter another request, append the latest
@@ -86,44 +106,56 @@ class MessageGrouper:
         current_slice_pages = []
         current_page_request: Optional[HttpRequest] = None
         current_page_response: Optional[HttpResponse] = None
+        had_error = False
 
         while records_count < limit and (message := next(messages, None)):
             if self._need_to_close_page(at_least_one_page_in_group, message):
-                self._close_page(current_page_request, current_page_response, current_slice_pages, current_page_records)
+                self._close_page(current_page_request, current_page_response, current_slice_pages, current_page_records, True)
                 current_page_request = None
                 current_page_response = None
 
-            if at_least_one_page_in_group and message.type == Type.LOG and message.log.message.startswith("slice:"):
+            if at_least_one_page_in_group and message.type == MessageType.LOG and message.log.message.startswith("slice:"):
                 yield StreamReadSlices(pages=current_slice_pages)
                 current_slice_pages = []
                 at_least_one_page_in_group = False
-            elif message.type == Type.LOG and message.log.message.startswith("request:"):
+            elif message.type == MessageType.LOG and message.log.message.startswith("request:"):
                 if not at_least_one_page_in_group:
                     at_least_one_page_in_group = True
                 current_page_request = self._create_request_from_log_message(message.log)
-            elif message.type == Type.LOG and message.log.message.startswith("response:"):
+            elif message.type == MessageType.LOG and message.log.message.startswith("response:"):
                 current_page_response = self._create_response_from_log_message(message.log)
-            elif message.type == Type.LOG:
+            elif message.type == MessageType.LOG:
+                if message.log.level == Level.ERROR:
+                    had_error = True
                 yield message.log
-            elif message.type == Type.RECORD:
+            elif message.type == MessageType.TRACE:
+                if message.trace.type == TraceType.ERROR:
+                    had_error = True
+                    yield message.trace
+            elif message.type == MessageType.RECORD:
                 current_page_records.append(message.record.data)
                 records_count += 1
                 schema_inferrer.accumulate(message.record)
         else:
-            self._close_page(current_page_request, current_page_response, current_slice_pages, current_page_records)
+            self._close_page(current_page_request, current_page_response, current_slice_pages, current_page_records, validate_page_complete=not had_error)
             yield StreamReadSlices(pages=current_slice_pages)
 
     @staticmethod
-    def _need_to_close_page(at_least_one_page_in_group, message):
+    def _need_to_close_page(at_least_one_page_in_group: bool, message: AirbyteMessage) -> bool:
         return (
             at_least_one_page_in_group
-            and message.type == Type.LOG
+            and message.type == MessageType.LOG
             and (message.log.message.startswith("request:") or message.log.message.startswith("slice:"))
         )
 
     @staticmethod
-    def _close_page(current_page_request, current_page_response, current_slice_pages, current_page_records):
-        if not current_page_request or not current_page_response:
+    def _close_page(current_page_request, current_page_response, current_slice_pages, current_page_records, validate_page_complete: bool):
+        """
+        Close a page when parsing message groups
+        @param validate_page_complete: in some cases, we expect the CDK to not return a response. As of today, this will only happen before
+        an uncaught exception and therefore, the assumption is that `validate_page_complete=True` only on the last page that is being closed
+        """
+        if validate_page_complete and (not current_page_request or not current_page_response):
             raise ValueError("Every message grouping should have at least one request and response")
 
         current_slice_pages.append(
@@ -131,6 +163,15 @@ class MessageGrouper:
         )
         current_page_records.clear()
 
+    def _read_stream(self, source: DeclarativeSource, config: Mapping[str, Any], configured_catalog: ConfiguredAirbyteCatalog) -> Iterator[AirbyteMessage]:
+        # the generator can raise an exception
+        # iterate over the generated messages. if next raise an exception, catch it and yield it as an AirbyteLogMessage
+        try:
+            yield from source.read(logger=self.logger, config=config, catalog=configured_catalog, state={})
+        except Exception as e:
+            error_message = f"{e.args[0] if len(e.args) > 0 else str(e)}"
+            yield AirbyteTracedException.from_exception(e, message=error_message).as_airbyte_message()
+
     def _create_request_from_log_message(self, log_message: AirbyteLogMessage) -> Optional[HttpRequest]:
         # TODO: As a temporary stopgap, the CDK emits request data as a log message string. Ideally this should come in the
         # form of a custom message object defined in the Airbyte protocol, but this unblocks us in the immediate while the
@@ -165,7 +206,7 @@ class MessageGrouper:
             self.logger.warning(f"Failed to parse log message into response object with error: {error}")
             return None
 
-    def _has_reached_limit(self, slices):
+    def _has_reached_limit(self, slices: List[StreamReadSlicesInner]):
         if len(slices) >= self._max_slices:
             return True
 

{airbyte-cdk-0.31.0 → airbyte-cdk-0.32.0}/airbyte_cdk/connector_builder/models.py

@@ -50,9 +50,15 @@ class StreamReadSlicesInner:
     state: Optional[Dict[str, Any]]
 
 
+@dataclass
+class LogMessage:
+    message: str
+    level: str
+
+
 @dataclass
 class StreamRead(object):
-    logs: List[object]
+    logs: List[LogMessage]
     slices: List[StreamReadSlicesInner]
     test_read_limit_reached: bool
     inferred_schema: Optional[Dict[str, Any]]

{airbyte-cdk-0.31.0 → airbyte-cdk-0.32.0}/airbyte_cdk/sources/abstract_source.py

@@ -79,13 +79,9 @@ class AbstractSource(Source, ABC):
         """Implements the Check Connection operation from the Airbyte Specification.
         See https://docs.airbyte.com/understanding-airbyte/airbyte-protocol/#check.
         """
-        try:
-            check_succeeded, error = self.check_connection(logger, config)
-            if not check_succeeded:
-                return AirbyteConnectionStatus(status=Status.FAILED, message=repr(error))
-        except Exception as e:
-            return AirbyteConnectionStatus(status=Status.FAILED, message=repr(e))
-
+        check_succeeded, error = self.check_connection(logger, config)
+        if not check_succeeded:
+            return AirbyteConnectionStatus(status=Status.FAILED, message=repr(error))
         return AirbyteConnectionStatus(status=Status.SUCCEEDED)
 
     def read(

{airbyte-cdk-0.31.0 → airbyte-cdk-0.32.0}/airbyte_cdk.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: airbyte-cdk
-Version: 0.31.0
+Version: 0.32.0
 Summary: A framework for writing Airbyte Connectors.
 Home-page: https://github.com/airbytehq/airbyte
 Author: Airbyte

{airbyte-cdk-0.31.0 → airbyte-cdk-0.32.0}/setup.py

@@ -17,7 +17,7 @@ setup(
     name="airbyte-cdk",
     # The version of the airbyte-cdk package is used at runtime to validate manifests. That validation must be
     # updated if our semver format changes such as using release candidate versions.
-    version="0.31.0",
+    version="0.32.0",
     description="A framework for writing Airbyte Connectors.",
     long_description=README,
     long_description_content_type="text/markdown",

{airbyte-cdk-0.31.0 → airbyte-cdk-0.32.0}/unit_tests/connector_builder/test_connector_builder_handler.py

@@ -3,6 +3,7 @@
 #
 
 import copy
+import dataclasses
 import json
 from unittest import mock
 from unittest.mock import patch
@@ -11,7 +12,7 @@ import pytest
 from airbyte_cdk import connector_builder
 from airbyte_cdk.connector_builder.connector_builder_handler import list_streams, resolve_manifest
 from airbyte_cdk.connector_builder.main import handle_connector_builder_request, handle_request, read_stream
-from airbyte_cdk.connector_builder.models import StreamRead, StreamReadSlicesInner, StreamReadSlicesInnerPagesInner
+from airbyte_cdk.connector_builder.models import LogMessage, StreamRead, StreamReadSlicesInner, StreamReadSlicesInnerPagesInner
 from airbyte_cdk.models import AirbyteMessage, AirbyteRecordMessage, ConfiguredAirbyteCatalog
 from airbyte_cdk.models import Type as MessageType
 from airbyte_cdk.sources.declarative.declarative_stream import DeclarativeStream
@@ -323,14 +324,31 @@ def test_read():
         assert output_record == expected_airbyte_message
 
 
-def test_read_returns_error_response():
+@patch("traceback.TracebackException.from_exception")
+def test_read_returns_error_response(mock_from_exception):
     class MockManifestDeclarativeSource:
         def read(self, logger, config, catalog, state):
-            raise ValueError
+            raise ValueError("error_message")
+
+    stack_trace = "a stack trace"
+    mock_from_exception.return_value = stack_trace
 
     source = MockManifestDeclarativeSource()
     response = read_stream(source, TEST_READ_CONFIG, ConfiguredAirbyteCatalog.parse_obj(CONFIGURED_CATALOG))
-    assert "Error reading" in response.trace.error.message
+
+    expected_stream_read = StreamRead(logs=[LogMessage("error_message - a stack trace", "ERROR")],
+                                      slices=[StreamReadSlicesInner(
+                                          pages=[StreamReadSlicesInnerPagesInner(records=[], request=None, response=None)],
+                                          slice_descriptor=None, state=None)],
+                                      test_read_limit_reached=False,
+                                      inferred_schema=None)
+
+    expected_message = AirbyteMessage(
+        type=MessageType.RECORD,
+        record=AirbyteRecordMessage(stream=_stream_name, data=dataclasses.asdict(expected_stream_read), emitted_at=1),
+    )
+    response.record.emitted_at = 1
+    assert response == expected_message
 
 
 @pytest.mark.parametrize(

{airbyte-cdk-0.31.0 → airbyte-cdk-0.32.0}/unit_tests/connector_builder/test_message_grouper.py

@@ -8,7 +8,7 @@ from unittest.mock import MagicMock
 
 import pytest
 from airbyte_cdk.connector_builder.message_grouper import MessageGrouper
-from airbyte_cdk.connector_builder.models import HttpRequest, HttpResponse, StreamRead, StreamReadPages
+from airbyte_cdk.connector_builder.models import HttpRequest, HttpResponse, LogMessage, StreamRead, StreamReadPages
 from airbyte_cdk.models import AirbyteLogMessage, AirbyteMessage, AirbyteRecordMessage, Level
 from airbyte_cdk.models import Type as MessageType
 from unit_tests.connector_builder.utils import create_configured_catalog
@@ -17,7 +17,7 @@ MAX_PAGES_PER_SLICE = 4
 MAX_SLICES = 3
 
 MANIFEST = {
-    "version": "0.1.0",
+    "version": "0.30.0",
     "type": "DeclarativeSource",
     "definitions": {
         "selector": {"extractor": {"field_path": ["items"], "type": "DpathExtractor"}, "type": "RecordSelector"},
@@ -167,9 +167,9 @@ def test_get_grouped_messages_with_logs():
         ),
     ]
     expected_logs = [
-        {"message": "log message before the request"},
-        {"message": "log message during the page"},
-        {"message": "log message after the response"},
+        LogMessage(**{"message": "log message before the request", "level": "INFO"}),
+        LogMessage(**{"message": "log message during the page", "level": "INFO"}),
+        LogMessage(**{"message": "log message after the response", "level": "INFO"}),
     ]
 
     mock_source = make_mock_source(
@@ -513,6 +513,21 @@ def test_get_grouped_messages_given_maximum_number_of_pages_then_test_read_limit
     assert stream_read.test_read_limit_reached
 
 
+def test_read_stream_returns_error_if_stream_does_not_exist():
+    mock_source = MagicMock()
+    mock_source.read.side_effect = ValueError("error")
+
+    full_config = {**CONFIG, **{"__injected_declarative_manifest": MANIFEST}}
+
+    message_grouper = MessageGrouper(MAX_PAGES_PER_SLICE, MAX_SLICES)
+    actual_response = message_grouper.get_message_groups(source=mock_source, config=full_config,
+                                                         configured_catalog=create_configured_catalog("not_in_manifest"))
+
+    assert 1 == len(actual_response.logs)
+    assert "Traceback" in actual_response.logs[0].message
+    assert "ERROR" in actual_response.logs[0].level
+
+
 def make_mock_source(return_value: Iterator) -> MagicMock:
     mock_source = MagicMock()
     mock_source.read.return_value = return_value

{airbyte-cdk-0.31.0 → airbyte-cdk-0.32.0}/unit_tests/sources/test_abstract_source.py

@@ -104,10 +104,11 @@ def test_failed_check():
     assert expected == MockSource(check_lambda=lambda: (False, "womp womp")).check(logger, {})
 
 
-def test_raising_check():
+def test_raising_check(mocker):
     """Tests that if a source raises an unexpected exception the appropriate connectionStatus failure message is returned."""
-    expected = AirbyteConnectionStatus(status=Status.FAILED, message="Exception('this should fail')")
-    assert expected == MockSource(check_lambda=lambda: exec('raise Exception("this should fail")')).check(logger, {})
+    check_lambda = mocker.Mock(side_effect=BaseException("this should fail"))
+    with pytest.raises(BaseException):
+        MockSource(check_lambda=check_lambda).check(logger, {})
 
 
 class MockStream(Stream):
@@ -334,10 +335,7 @@ def test_valid_full_refresh_read_with_slices(mocker):
 
 @pytest.mark.parametrize(
     "slices",
-    [
-        [{"1": "1"}, {"2": "2"}],
-        [{"date": datetime.date(year=2023, month=1, day=1)}, {"date": datetime.date(year=2023, month=1, day=1)}]
-    ]
+    [[{"1": "1"}, {"2": "2"}], [{"date": datetime.date(year=2023, month=1, day=1)}, {"date": datetime.date(year=2023, month=1, day=1)}]],
 )
 def test_read_full_refresh_with_slices_sends_slice_messages(mocker, slices):
     """Given the logger is debug and a full refresh, AirbyteMessages are sent for slices"""
@@ -369,7 +367,7 @@ def test_read_incremental_with_slices_sends_slice_messages(mocker):
     debug_logger.setLevel(logging.DEBUG)
     slices = [{"1": "1"}, {"2": "2"}]
     stream = MockStream(
-        [({"sync_mode": SyncMode.incremental, "stream_slice": s, 'stream_state': {}}, [s]) for s in slices],
+        [({"sync_mode": SyncMode.incremental, "stream_slice": s, "stream_state": {}}, [s]) for s in slices],
         name="s1",
     )