airbyte-cdk 0.31.0__py3-none-any.whl → 0.31.1__py3-none-any.whl

airbyte_cdk/connector_builder/message_grouper.py

@@ -9,11 +9,19 @@ from json import JSONDecodeError
 from typing import Any, Iterable, Iterator, 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
 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 +50,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 +75,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 +98,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, message) -> 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 +155,15 @@ class MessageGrouper:
         )
         current_page_records.clear()
 
+    def _read_stream(self, source, config, configured_catalog) -> 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

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.dist-info → airbyte_cdk-0.31.1.dist-info}/METADATA

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

{airbyte_cdk-0.31.0.dist-info → airbyte_cdk-0.31.1.dist-info}/RECORD

@@ -8,8 +8,8 @@ airbyte_cdk/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 airbyte_cdk/connector_builder/__init__.py,sha256=4Hw-PX1-VgESLF16cDdvuYCzGJtHntThLF4qIiULWeo,61
 airbyte_cdk/connector_builder/connector_builder_handler.py,sha256=9iG2hmaJyBtLDtdwtQp8M19pG2HQjR7Yy6jnBDggXMk,4160
 airbyte_cdk/connector_builder/main.py,sha256=IOijgSQ4A9KhqJplHSVPTrxH_cyGWW_9uNmSIwBo0l0,3021
-airbyte_cdk/connector_builder/message_grouper.py,sha256=ZEc-NJlh9fTwyxYUc55pL8QoTzhLpXSyRFH5cmsliiE,9046
-airbyte_cdk/connector_builder/models.py,sha256=mf9cAZaxFW7uFp0gyHiJpuGlAuGuq2zFvs76VMYqUkA,1765
+airbyte_cdk/connector_builder/message_grouper.py,sha256=-0AYh21eZeNJyyFPvSB84NNKMZznSX7bf4Ygi3deWEU,10972
+airbyte_cdk/connector_builder/models.py,sha256=y0PJ-LwJk3e1RzRmMfjQSBP9ENx_a0wBcWNCjlW72Ks,1832
 airbyte_cdk/destinations/__init__.py,sha256=0Uxmz3iBAyZJdk_bqUVt2pb0UwRTpFjTnFE6fQFbWKY,126
 airbyte_cdk/destinations/destination.py,sha256=_tIMnKcRQbtIsjVvNOVjfbIxgCNLuBXQwQj8MyVm3BI,5420
 airbyte_cdk/models/__init__.py,sha256=LPQcYdDPwrCXiBPe_jexO4UAcbovIb1V9tHB6I7Un30,633
@@ -160,8 +160,8 @@ airbyte_cdk/utils/traced_exception.py,sha256=9G2sG9eYkvn6Aa7rMuUW_KIRszRaTc_xdnT
 source_declarative_manifest/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 source_declarative_manifest/main.py,sha256=HXzuRsRyhHwPrGU-hc4S7RrgoOoHImqkdfbmO2geBeE,1027
 unit_tests/connector_builder/__init__.py,sha256=4Hw-PX1-VgESLF16cDdvuYCzGJtHntThLF4qIiULWeo,61
-unit_tests/connector_builder/test_connector_builder_handler.py,sha256=EWF0hsr-mr4Q-hq_W4oYj7s4rQMvnhC4nIbkhNK9unU,18823
-unit_tests/connector_builder/test_message_grouper.py,sha256=hK6-CnEpxmPWpteJpt8k6L0t_FYQKdRv_uBRjZRzA5s,22181
+unit_tests/connector_builder/test_connector_builder_handler.py,sha256=fwOlQPdc9quPUoZYXg_GVXqdkw5TLWs0CuBjYEMAcfM,19745
+unit_tests/connector_builder/test_message_grouper.py,sha256=1_DXyuqHmr8B7b1t8PSpmYWaHvYYAB2mMwpp_YAbYwc,22959
 unit_tests/connector_builder/utils.py,sha256=AAggdGWP-mNuWOZUHLAVIbjTeIcdPo-3pbMm5zdYpS0,796
 unit_tests/destinations/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 unit_tests/destinations/test_destination.py,sha256=koG_j812KMkcIxoUH6XlAL3zsephZJmlHvyzJXm0dCs,10269
@@ -257,8 +257,8 @@ unit_tests/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,
 unit_tests/utils/test_schema_inferrer.py,sha256=ckl17GlNOZInqgxni7Z2A0bg_p6JDy0GVFAG8ph67pw,3288
 unit_tests/utils/test_secret_utils.py,sha256=XKe0f1RHYii8iwE6ATmBr5JGDI1pzzrnZUGdUSMJQP4,4886
 unit_tests/utils/test_traced_exception.py,sha256=bDFP5zMBizFenz6V2WvEZTRCKGB5ijh3DBezjbfoYIs,4198
-airbyte_cdk-0.31.0.dist-info/LICENSE.txt,sha256=Wfe61S4BaGPj404v8lrAbvhjYR68SHlkzeYrg3_bbuM,1051
-airbyte_cdk-0.31.0.dist-info/METADATA,sha256=Tjx6hFCTGl2NKZqynEK9dGwowRFkLGFueqEc_v_qU2s,8902
-airbyte_cdk-0.31.0.dist-info/WHEEL,sha256=pkctZYzUS4AYVn6dJ-7367OJZivF2e8RA9b_ZBjif18,92
-airbyte_cdk-0.31.0.dist-info/top_level.txt,sha256=edvsDKTnE6sD2wfCUaeTfKf5gQIL6CPVMwVL2sWZzqo,51
-airbyte_cdk-0.31.0.dist-info/RECORD,,
+airbyte_cdk-0.31.1.dist-info/LICENSE.txt,sha256=Wfe61S4BaGPj404v8lrAbvhjYR68SHlkzeYrg3_bbuM,1051
+airbyte_cdk-0.31.1.dist-info/METADATA,sha256=nSfqOPf-MueD0MHg9bTQU8PWrhP43lmxd6JK1IDNj-M,8902
+airbyte_cdk-0.31.1.dist-info/WHEEL,sha256=pkctZYzUS4AYVn6dJ-7367OJZivF2e8RA9b_ZBjif18,92
+airbyte_cdk-0.31.1.dist-info/top_level.txt,sha256=edvsDKTnE6sD2wfCUaeTfKf5gQIL6CPVMwVL2sWZzqo,51
+airbyte_cdk-0.31.1.dist-info/RECORD,,

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(

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