airbyte-cdk 0.35.4__tar.gz → 0.36.1__tar.gz

{airbyte-cdk-0.35.4 → airbyte-cdk-0.36.1}/PKG-INFO

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

{airbyte-cdk-0.35.4 → airbyte-cdk-0.36.1}/airbyte_cdk/connector_builder/connector_builder_handler.py

@@ -50,7 +50,9 @@ def create_source(config: Mapping[str, Any], limits: TestReadLimits) -> Manifest
         component_factory=ModelToComponentFactory(
             emit_connector_builder_messages=True,
             limit_pages_fetched_per_slice=limits.max_pages_per_slice,
-            limit_slices_fetched=limits.max_slices)
+            limit_slices_fetched=limits.max_slices,
+            disable_retries=True
+        )
     )
 
 

{airbyte-cdk-0.35.4 → airbyte-cdk-0.36.1}/airbyte_cdk/sources/abstract_source.py

@@ -13,6 +13,7 @@ from airbyte_cdk.models import (
     AirbyteLogMessage,
     AirbyteMessage,
     AirbyteStateMessage,
+    AirbyteStreamStatus,
     ConfiguredAirbyteCatalog,
     ConfiguredAirbyteStream,
     Level,
@@ -28,6 +29,7 @@ from airbyte_cdk.sources.streams.http.http import HttpStream
 from airbyte_cdk.sources.utils.record_helper import stream_data_to_airbyte_message
 from airbyte_cdk.sources.utils.schema_helpers import InternalConfig, split_config
 from airbyte_cdk.utils.event_timing import create_timer
+from airbyte_cdk.utils.stream_status_utils import as_airbyte_message as stream_status_as_airbyte_message
 from airbyte_cdk.utils.traced_exception import AirbyteTracedException
 
 
@@ -113,6 +115,8 @@ class AbstractSource(Source, ABC):
                     continue
                 try:
                     timer.start_event(f"Syncing stream {configured_stream.stream.name}")
+                    logger.info(f"Marking stream {configured_stream.stream.name} as STARTED")
+                    yield stream_status_as_airbyte_message(configured_stream, AirbyteStreamStatus.STARTED)
                     yield from self._read_stream(
                         logger=logger,
                         stream_instance=stream_instance,
@@ -120,10 +124,15 @@ class AbstractSource(Source, ABC):
                         state_manager=state_manager,
                         internal_config=internal_config,
                     )
+                    logger.info(f"Marking stream {configured_stream.stream.name} as STOPPED")
+                    yield stream_status_as_airbyte_message(configured_stream, AirbyteStreamStatus.COMPLETE)
                 except AirbyteTracedException as e:
+                    yield stream_status_as_airbyte_message(configured_stream, AirbyteStreamStatus.INCOMPLETE)
                     raise e
                 except Exception as e:
                     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, AirbyteStreamStatus.INCOMPLETE)
                     display_message = stream_instance.get_error_display_message(e)
                     if display_message:
                         raise AirbyteTracedException.from_exception(e, message=display_message) from e
@@ -185,6 +194,10 @@ class AbstractSource(Source, ABC):
         for record in record_iterator:
             if record.type == MessageType.RECORD:
                 record_counter += 1
+                if record_counter == 1:
+                    logger.info(f"Marking stream {stream_name} as RUNNING")
+                    # If we just read the first record of the stream, emit the transition to the RUNNING state
+                    yield stream_status_as_airbyte_message(configured_stream, AirbyteStreamStatus.RUNNING)
             yield record
 
         logger.info(f"Read {record_counter} records from {stream_name} stream")

{airbyte-cdk-0.35.4 → airbyte-cdk-0.36.1}/airbyte_cdk/sources/declarative/declarative_stream.py

@@ -5,14 +5,14 @@
 from dataclasses import InitVar, dataclass, field
 from typing import Any, Iterable, List, Mapping, MutableMapping, Optional, Union
 
-from airbyte_cdk.models import AirbyteLogMessage, AirbyteMessage, AirbyteTraceMessage, SyncMode
+from airbyte_cdk.models import AirbyteMessage, SyncMode
 from airbyte_cdk.sources.declarative.interpolation import InterpolatedString
 from airbyte_cdk.sources.declarative.retrievers.retriever import Retriever
 from airbyte_cdk.sources.declarative.schema import DefaultSchemaLoader
 from airbyte_cdk.sources.declarative.schema.schema_loader import SchemaLoader
 from airbyte_cdk.sources.declarative.transformations import RecordTransformation
 from airbyte_cdk.sources.declarative.types import Config, StreamSlice
-from airbyte_cdk.sources.streams.core import Stream
+from airbyte_cdk.sources.streams.core import Stream, StreamData
 
 
 @dataclass
@@ -102,17 +102,27 @@ class DeclarativeStream(Stream):
 
     def _apply_transformations(
         self,
-        message_or_record_data: Union[AirbyteMessage, AirbyteLogMessage, AirbyteTraceMessage, Mapping[str, Any]],
+        message_or_record_data: StreamData,
         config: Config,
         stream_slice: StreamSlice,
     ):
-        # If the input is an AirbyteRecord, transform the record's data
-        # If the input is another type of Airbyte Message, return it as is
+        # If the input is an AirbyteMessage with a record, transform the record's data
+        # If the input is another type of AirbyteMessage, return it as is
         # If the input is a dict, transform it
-        if isinstance(message_or_record_data, AirbyteLogMessage) or isinstance(message_or_record_data, AirbyteTraceMessage):
-            return message_or_record_data
+        if isinstance(message_or_record_data, AirbyteMessage):
+            if message_or_record_data.record:
+                record = message_or_record_data.record.data
+            else:
+                return message_or_record_data
+        elif isinstance(message_or_record_data, dict):
+            record = message_or_record_data
+        else:
+            # Raise an error because this is unexpected and indicative of a typing problem in the CDK
+            raise ValueError(
+                f"Unexpected record type. Expected {StreamData}. Got {type(message_or_record_data)}. This is probably due to a bug in the CDK."
+            )
         for transformation in self.transformations:
-            transformation.transform(message_or_record_data, config=config, stream_state=self.state, stream_slice=stream_slice)
+            transformation.transform(record, config=config, stream_state=self.state, stream_slice=stream_slice)
 
         return message_or_record_data
 

{airbyte-cdk-0.35.4 → airbyte-cdk-0.36.1}/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py

@@ -112,12 +112,17 @@ DEFAULT_BACKOFF_STRATEGY = ExponentialBackoffStrategy
 
 class ModelToComponentFactory:
     def __init__(
-        self, limit_pages_fetched_per_slice: int = None, limit_slices_fetched: int = None, emit_connector_builder_messages: bool = False
+        self,
+        limit_pages_fetched_per_slice: int = None,
+        limit_slices_fetched: int = None,
+        emit_connector_builder_messages: bool = False,
+        disable_retries=False,
     ):
         self._init_mappings()
         self._limit_pages_fetched_per_slice = limit_pages_fetched_per_slice
         self._limit_slices_fetched = limit_slices_fetched
         self._emit_connector_builder_messages = emit_connector_builder_messages
+        self._disable_retries = disable_retries
 
     def _init_mappings(self):
         self.PYDANTIC_MODEL_TO_CONSTRUCTOR: [Type[BaseModel], Callable] = {
@@ -779,6 +784,7 @@ class ModelToComponentFactory:
                 config=config,
                 maximum_number_of_slices=self._limit_slices_fetched,
                 parameters=model.parameters,
+                disable_retries=self._disable_retries,
             )
         return SimpleRetriever(
             name=name,
@@ -789,6 +795,7 @@ class ModelToComponentFactory:
             stream_slicer=stream_slicer or SinglePartitionRouter(parameters={}),
             config=config,
             parameters=model.parameters,
+            disable_retries=self._disable_retries,
         )
 
     @staticmethod

{airbyte-cdk-0.35.4 → airbyte-cdk-0.36.1}/airbyte_cdk/sources/declarative/retrievers/simple_retriever.py

@@ -50,6 +50,8 @@ class SimpleRetriever(Retriever, HttpStream):
         parameters (Mapping[str, Any]): Additional runtime parameters to be used for string interpolation
     """
 
+    _DEFAULT_MAX_RETRY = 5
+
     requester: Requester
     record_selector: HttpSelector
     config: Config
@@ -61,6 +63,7 @@ class SimpleRetriever(Retriever, HttpStream):
     paginator: Optional[Paginator] = None
     stream_slicer: Optional[StreamSlicer] = SinglePartitionRouter(parameters={})
     emit_connector_builder_messages: bool = False
+    disable_retries: bool = False
 
     def __post_init__(self, parameters: Mapping[str, Any]):
         self.paginator = self.paginator or NoPagination(parameters=parameters)
@@ -95,6 +98,14 @@ class SimpleRetriever(Retriever, HttpStream):
         # never raise on http_errors because this overrides the error handler logic...
         return False
 
+    @property
+    def max_retries(self) -> Union[int, None]:
+        if self.disable_retries:
+            return 0
+        if hasattr(self.requester.error_handler, "max_retries"):
+            return self.requester.error_handler.max_retries
+        return self._DEFAULT_MAX_RETRY
+
     def should_retry(self, response: requests.Response) -> bool:
         """
         Specifies conditions for backoff based on the response from the server.

{airbyte-cdk-0.35.4 → airbyte-cdk-0.36.1}/airbyte_cdk/sources/streams/core.py

@@ -11,7 +11,7 @@ from functools import lru_cache
 from typing import Any, Iterable, List, Mapping, MutableMapping, Optional, Tuple, Union
 
 import airbyte_cdk.sources.utils.casing as casing
-from airbyte_cdk.models import AirbyteLogMessage, AirbyteStream, AirbyteTraceMessage, SyncMode
+from airbyte_cdk.models import AirbyteMessage, AirbyteStream, SyncMode
 
 # list of all possible HTTP methods which can be used for sending of request bodies
 from airbyte_cdk.sources.utils.schema_helpers import ResourceSchemaLoader
@@ -24,10 +24,8 @@ if typing.TYPE_CHECKING:
 
 # A stream's read method can return one of the following types:
 # Mapping[str, Any]: The content of an AirbyteRecordMessage
-# AirbyteRecordMessage: An AirbyteRecordMessage
-# AirbyteLogMessage: A log message
-# AirbyteTraceMessage: A trace message
-StreamData = Union[Mapping[str, Any], AirbyteLogMessage, AirbyteTraceMessage]
+# AirbyteMessage: An AirbyteMessage. Could be of any type
+StreamData = Union[Mapping[str, Any], AirbyteMessage]
 
 
 def package_name_from_class(cls: object) -> str:

airbyte-cdk-0.36.1/airbyte_cdk/utils/stream_status_utils.py

@@ -0,0 +1,36 @@
+#
+# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
+#
+
+
+from datetime import datetime
+
+from airbyte_cdk.models import (
+    AirbyteMessage,
+    AirbyteStreamStatus,
+    AirbyteStreamStatusTraceMessage,
+    AirbyteTraceMessage,
+    ConfiguredAirbyteStream,
+    StreamDescriptor,
+    TraceType,
+)
+from airbyte_cdk.models import Type as MessageType
+
+
+def as_airbyte_message(stream: ConfiguredAirbyteStream, current_status: AirbyteStreamStatus) -> AirbyteMessage:
+    """
+    Builds an AirbyteStreamStatusTraceMessage for the provided stream
+    """
+
+    now_millis = datetime.now().timestamp() * 1000.0
+
+    trace_message = AirbyteTraceMessage(
+        type=TraceType.STREAM_STATUS,
+        emitted_at=now_millis,
+        stream_status=AirbyteStreamStatusTraceMessage(
+            stream_descriptor=StreamDescriptor(name=stream.stream.name, namespace=stream.stream.namespace),
+            status=current_status,
+        ),
+    )
+
+    return AirbyteMessage(type=MessageType.TRACE, trace=trace_message)

{airbyte-cdk-0.35.4 → airbyte-cdk-0.36.1}/airbyte_cdk.egg-info/PKG-INFO

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

{airbyte-cdk-0.35.4 → airbyte-cdk-0.36.1}/airbyte_cdk.egg-info/SOURCES.txt

@@ -165,6 +165,7 @@ airbyte_cdk/utils/__init__.py
 airbyte_cdk/utils/airbyte_secrets_utils.py
 airbyte_cdk/utils/event_timing.py
 airbyte_cdk/utils/schema_inferrer.py
+airbyte_cdk/utils/stream_status_utils.py
 airbyte_cdk/utils/traced_exception.py
 source_declarative_manifest/__init__.py
 source_declarative_manifest/main.py
@@ -265,4 +266,5 @@ unit_tests/sources/streams/http/requests_native_auth/test_requests_native_auth.p
 unit_tests/utils/__init__.py
 unit_tests/utils/test_schema_inferrer.py
 unit_tests/utils/test_secret_utils.py
+unit_tests/utils/test_stream_status_utils.py
 unit_tests/utils/test_traced_exception.py

{airbyte-cdk-0.35.4 → airbyte-cdk-0.36.1}/airbyte_cdk.egg-info/requires.txt

@@ -1,4 +1,4 @@
-airbyte-protocol-models==1.0.0
+airbyte-protocol-models==0.3.6
 backoff
 dpath~=2.0.1
 isodate~=0.6.1

{airbyte-cdk-0.35.4 → airbyte-cdk-0.36.1}/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.35.4",
+    version="0.36.1",
     description="A framework for writing Airbyte Connectors.",
     long_description=README,
     long_description_content_type="text/markdown",
@@ -46,7 +46,7 @@ setup(
     packages=find_packages(exclude=("unit_tests",)),
     package_data={"airbyte_cdk": ["py.typed", "sources/declarative/declarative_component_schema.yaml"]},
     install_requires=[
-        "airbyte-protocol-models==1.0.0",
+        "airbyte-protocol-models==0.3.6",
         "backoff",
         "dpath~=2.0.1",
         "isodate~=0.6.1",

{airbyte-cdk-0.35.4 → airbyte-cdk-0.36.1}/unit_tests/connector_builder/test_connector_builder_handler.py

@@ -566,6 +566,7 @@ def test_create_source():
     assert isinstance(source, ManifestDeclarativeSource)
     assert source._constructor._limit_pages_fetched_per_slice == limits.max_pages_per_slice
     assert source._constructor._limit_slices_fetched == limits.max_slices
+    assert source.streams(config={})[0].retriever.max_retries == 0
 
 
 def request_log_message(request: dict) -> AirbyteMessage:

{airbyte-cdk-0.35.4 → airbyte-cdk-0.36.1}/unit_tests/sources/declarative/parsers/test_model_to_component_factory.py

@@ -1360,3 +1360,29 @@ def test_simple_retriever_emit_log_messages():
     )
 
     assert isinstance(retriever, SimpleRetrieverTestReadDecorator)
+
+
+def test_ignore_retry():
+    requester_model = {
+        "type": "SimpleRetriever",
+        "record_selector": {
+            "type": "RecordSelector",
+            "extractor": {
+                "type": "DpathExtractor",
+                "field_path": [],
+            },
+        },
+        "requester": {"type": "HttpRequester", "name": "list", "url_base": "orange.com", "path": "/v1/api"},
+    }
+
+    connector_builder_factory = ModelToComponentFactory(disable_retries=True)
+    retriever = connector_builder_factory.create_component(
+        model_type=SimpleRetrieverModel,
+        component_definition=requester_model,
+        config={},
+        name="Test",
+        primary_key="id",
+        stream_slicer=None,
+    )
+
+    assert retriever.max_retries == 0

{airbyte-cdk-0.35.4 → airbyte-cdk-0.36.1}/unit_tests/sources/declarative/retrievers/test_simple_retriever.py

@@ -256,6 +256,48 @@ def test_parse_response(test_name, status_code, response_status, len_expected_re
         assert len(records) == len_expected_records
 
 
+def test_max_retries_given_error_handler_has_max_retries():
+    requester = MagicMock()
+    requester.error_handler = MagicMock()
+    requester.error_handler.max_retries = 10
+    retriever = SimpleRetriever(
+        name="stream_name",
+        primary_key=primary_key,
+        requester=requester,
+        record_selector=MagicMock(),
+        parameters={},
+        config={}
+    )
+    assert retriever.max_retries == 10
+
+
+def test_max_retries_given_error_handler_without_max_retries():
+    requester = MagicMock()
+    requester.error_handler = MagicMock(spec=[u'without_max_retries_attribute'])
+    retriever = SimpleRetriever(
+        name="stream_name",
+        primary_key=primary_key,
+        requester=requester,
+        record_selector=MagicMock(),
+        parameters={},
+        config={}
+    )
+    assert retriever.max_retries == 5
+
+
+def test_max_retries_given_disable_retries():
+    retriever = SimpleRetriever(
+        name="stream_name",
+        primary_key=primary_key,
+        requester=MagicMock(),
+        record_selector=MagicMock(),
+        disable_retries=True,
+        parameters={},
+        config={}
+    )
+    assert retriever.max_retries == 0
+
+
 @pytest.mark.parametrize(
     "test_name, response_action, retry_in, expected_backoff_time",
     [

{airbyte-cdk-0.35.4 → airbyte-cdk-0.36.1}/unit_tests/sources/declarative/test_declarative_stream.py

@@ -5,7 +5,16 @@
 from unittest import mock
 from unittest.mock import MagicMock, call
 
-from airbyte_cdk.models import AirbyteLogMessage, AirbyteTraceMessage, Level, SyncMode, TraceType
+from airbyte_cdk.models import (
+    AirbyteLogMessage,
+    AirbyteMessage,
+    AirbyteRecordMessage,
+    AirbyteTraceMessage,
+    Level,
+    SyncMode,
+    TraceType,
+    Type,
+)
 from airbyte_cdk.sources.declarative.declarative_stream import DeclarativeStream
 from airbyte_cdk.sources.declarative.transformations import AddFields, RecordTransformation
 from airbyte_cdk.sources.declarative.transformations.add_fields import AddedFieldDefinition
@@ -24,8 +33,8 @@ def test_declarative_stream():
     records = [
         {"pk": 1234, "field": "value"},
         {"pk": 4567, "field": "different_value"},
-        AirbyteLogMessage(level=Level.INFO, message="This is a log  message"),
-        AirbyteTraceMessage(type=TraceType.ERROR, emitted_at=12345),
+        AirbyteMessage(type=Type.LOG, log=AirbyteLogMessage(level=Level.INFO, message="This is a log  message")),
+        AirbyteMessage(type=Type.TRACE, trace=AirbyteTraceMessage(type=TraceType.ERROR, emitted_at=12345)),
     ]
     stream_slices = [
         {"date": "2021-01-01"},
@@ -84,15 +93,17 @@ def test_declarative_stream_with_add_fields_transform():
     retriever_records = [
         {"pk": 1234, "field": "value"},
         {"pk": 4567, "field": "different_value"},
-        AirbyteLogMessage(level=Level.INFO, message="This is a log  message"),
-        AirbyteTraceMessage(type=TraceType.ERROR, emitted_at=12345),
+        AirbyteMessage(type=Type.RECORD, record=AirbyteRecordMessage(data={"pk": 1357, "field": "a_value"}, emitted_at=12344, stream="stream")),
+        AirbyteMessage(type=Type.LOG, log=AirbyteLogMessage(level=Level.INFO, message="This is a log  message")),
+        AirbyteMessage(type=Type.TRACE, trace=AirbyteTraceMessage(type=TraceType.ERROR, emitted_at=12345)),
     ]
 
     expected_records = [
         {"pk": 1234, "field": "value", "added_key": "added_value"},
         {"pk": 4567, "field": "different_value", "added_key": "added_value"},
-        AirbyteLogMessage(level=Level.INFO, message="This is a log  message"),
-        AirbyteTraceMessage(type=TraceType.ERROR, emitted_at=12345),
+        AirbyteMessage(type=Type.RECORD, record=AirbyteRecordMessage(data={"pk": 1357, "field": "a_value", "added_key": "added_value"}, emitted_at=12344, stream="stream")),
+        AirbyteMessage(type=Type.LOG, log=AirbyteLogMessage(level=Level.INFO, message="This is a log  message")),
+        AirbyteMessage(type=Type.TRACE, trace=AirbyteTraceMessage(type=TraceType.ERROR, emitted_at=12345)),
     ]
     stream_slices = [
         {"date": "2021-01-01"},

{airbyte-cdk-0.35.4 → airbyte-cdk-0.36.1}/unit_tests/sources/declarative/test_manifest_declarative_source.py

@@ -1242,7 +1242,7 @@ def _create_page(response_body):
 def test_read_manifest_declarative_source(test_name, manifest, pages, expected_records, expected_calls):
     _stream_name = "Rates"
     with patch.object(HttpStream, "_fetch_next_page", side_effect=pages) as mock_http_stream:
-        output_data = [message.record.data for message in _run_read(manifest, _stream_name)]
+        output_data = [message.record.data for message in _run_read(manifest, _stream_name) if message.record]
         assert expected_records == output_data
         mock_http_stream.assert_has_calls(expected_calls)