fix connector builder tests and format

Author: maxime.c

airbyte_cdk/sources/declarative/incremental/concurrent_partition_cursor.py

@@ -11,7 +11,13 @@
 from datetime import timedelta
 from typing import Any, Callable, Iterable, List, Mapping, MutableMapping, Optional
 
-from airbyte_cdk.models import AirbyteStateMessage, AirbyteStateBlob, AirbyteStreamState, AirbyteStateType, StreamDescriptor
+from airbyte_cdk.models import (
+    AirbyteStateMessage,
+    AirbyteStateBlob,
+    AirbyteStreamState,
+    AirbyteStateType,
+    StreamDescriptor,
+)
 from airbyte_cdk.sources.connector_state_manager import ConnectorStateManager
 from airbyte_cdk.sources.declarative.incremental.global_substream_cursor import (
     Timer,
@@ -548,21 +554,33 @@ def limit_reached(self) -> bool:
         return self._number_of_partitions > self.SWITCH_TO_GLOBAL_LIMIT
 
     @staticmethod
-    def get_parent_state(stream_state: Optional[StreamState], parent_stream_name: str) -> Optional[AirbyteStateMessage]:
-        return AirbyteStateMessage(
-            type=AirbyteStateType.STREAM,
-            stream=AirbyteStreamState(
-                stream_descriptor=StreamDescriptor(parent_stream_name, None),
-                stream_state=AirbyteStateBlob(stream_state["parent_state"][parent_stream_name])
+    def get_parent_state(
+        stream_state: Optional[StreamState], parent_stream_name: str
+    ) -> Optional[AirbyteStateMessage]:
+        return (
+            AirbyteStateMessage(
+                type=AirbyteStateType.STREAM,
+                stream=AirbyteStreamState(
+                    stream_descriptor=StreamDescriptor(parent_stream_name, None),
+                    stream_state=AirbyteStateBlob(stream_state["parent_state"][parent_stream_name]),
+                ),
             )
-        ) if stream_state and "parent_state" in stream_state else None
+            if stream_state and "parent_state" in stream_state
+            else None
+        )
 
     @staticmethod
-    def get_global_state(stream_state: Optional[StreamState], parent_stream_name: str) -> Optional[AirbyteStateMessage]:
-        return AirbyteStateMessage(
-            type=AirbyteStateType.STREAM,
-            stream=AirbyteStreamState(
-                stream_descriptor=StreamDescriptor(parent_stream_name, None),
-                stream_state=AirbyteStateBlob(stream_state["state"])
+    def get_global_state(
+        stream_state: Optional[StreamState], parent_stream_name: str
+    ) -> Optional[AirbyteStateMessage]:
+        return (
+            AirbyteStateMessage(
+                type=AirbyteStateType.STREAM,
+                stream=AirbyteStreamState(
+                    stream_descriptor=StreamDescriptor(parent_stream_name, None),
+                    stream_state=AirbyteStateBlob(stream_state["state"]),
+                ),
             )
-        ) if stream_state and "state" in stream_state else None
+            if stream_state and "state" in stream_state
+            else None
+        )

airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py

@@ -33,7 +33,15 @@
 from airbyte_cdk.connector_builder.models import (
     LogMessage as ConnectorBuilderLogMessage,
 )
-from airbyte_cdk.models import FailureType, Level, AirbyteStateMessage, AirbyteStreamState, AirbyteStateBlob, AirbyteStateType, StreamDescriptor
+from airbyte_cdk.models import (
+    FailureType,
+    Level,
+    AirbyteStateMessage,
+    AirbyteStreamState,
+    AirbyteStateBlob,
+    AirbyteStateType,
+    StreamDescriptor,
+)
 from airbyte_cdk.sources.connector_state_manager import ConnectorStateManager
 from airbyte_cdk.sources.declarative.async_job.job_orchestrator import AsyncJobOrchestrator
 from airbyte_cdk.sources.declarative.async_job.job_tracker import JobTracker
@@ -500,8 +508,9 @@
     InterpolatedRequestOptionsProvider,
     RequestOptionsProvider,
 )
-from airbyte_cdk.sources.declarative.requesters.request_options.per_partition_request_option_provider import \
-    PerPartitionRequestOptionsProvider
+from airbyte_cdk.sources.declarative.requesters.request_options.per_partition_request_option_provider import (
+    PerPartitionRequestOptionsProvider,
+)
 from airbyte_cdk.sources.declarative.requesters.request_path import RequestPath
 from airbyte_cdk.sources.declarative.requesters.requester import HttpMethod
 from airbyte_cdk.sources.declarative.resolvers import (
@@ -1282,7 +1291,9 @@ def create_concurrent_cursor_from_datetime_based_cursor(
 
         # TODO validate and explain why we need to do this...
         component_definition["$parameters"] = component_definition.get("parameters", {})
-        parameters = component_definition.get("parameters", component_definition.get("$parameters", {}))
+        parameters = component_definition.get(
+            "parameters", component_definition.get("$parameters", {})
+        )
         datetime_based_cursor_model = model_type.parse_obj(component_definition)
 
         if not isinstance(datetime_based_cursor_model, DatetimeBasedCursorModel):
@@ -1596,7 +1607,9 @@ def create_concurrent_cursor_from_perpartition_cursor(
 
         interpolated_cursor_field = InterpolatedString.create(
             datetime_based_cursor_model.cursor_field,
-            parameters=component_definition.get("parameters", component_definition.get("$parameters", {})),  # FIXME validate and explain why we need to do this
+            parameters=component_definition.get(
+                "parameters", component_definition.get("$parameters", {})
+            ),  # FIXME validate and explain why we need to do this
         )
         cursor_field = CursorField(interpolated_cursor_field.eval(config=config))
 
@@ -1973,13 +1986,17 @@ def create_declarative_stream(
             request_options_provider = (
                 datetime_request_options_provider
                 if not isinstance(concurrent_cursor, ConcurrentPerPartitionCursor)
-                else PerPartitionRequestOptionsProvider(partition_router, datetime_request_options_provider)
+                else PerPartitionRequestOptionsProvider(
+                    partition_router, datetime_request_options_provider
+                )
             )
         elif model.incremental_sync and isinstance(
             model.incremental_sync, IncrementingCountCursorModel
         ):
             if isinstance(concurrent_cursor, ConcurrentPerPartitionCursor):
-                raise ValueError("PerPartition does not support per partition states because switching to global state is time based")
+                raise ValueError(
+                    "PerPartition does not support per partition states because switching to global state is time based"
+                )
 
             cursor_model: IncrementingCountCursorModel = model.incremental_sync  # type: ignore
 
@@ -2019,7 +2036,9 @@ def create_declarative_stream(
             )
 
         stream_slicer: ConcurrentStreamSlicer = (
-            partition_router if isinstance(concurrent_cursor, FinalStateCursor) else concurrent_cursor
+            partition_router
+            if isinstance(concurrent_cursor, FinalStateCursor)
+            else concurrent_cursor
         )
         retriever = self._create_component_from_model(
             model=model.retriever,
@@ -2088,8 +2107,7 @@ def create_declarative_stream(
             logger=logging.getLogger(f"airbyte.{stream_name}"),
             # FIXME this is a breaking change compared to the old implementation which used the source name instead
             cursor=concurrent_cursor,
-            supports_file_transfer=hasattr(model, "file_uploader")
-                                   and bool(model.file_uploader),
+            supports_file_transfer=hasattr(model, "file_uploader") and bool(model.file_uploader),
         )
 
     def _is_stop_condition_on_cursor(self, model: DeclarativeStreamModel) -> bool:
@@ -3768,14 +3786,20 @@ def _create_message_repository_substream_wrapper(
         self, model: ParentStreamConfigModel, config: Config, **kwargs: Any
     ) -> Any:
         # getting the parent state
-        child_state = self._connector_state_manager.get_stream_state(kwargs["stream_name"], None)  # FIXME adding `stream_name` as a parameter means it will be a breaking change. I assume this is mostly called internally so I don't think we need to bother that much about this but still raising the flag
+        child_state = self._connector_state_manager.get_stream_state(
+            kwargs["stream_name"], None
+        )  # FIXME adding `stream_name` as a parameter means it will be a breaking change. I assume this is mostly called internally so I don't think we need to bother that much about this but still raising the flag
         if model.incremental_dependency and child_state:
             parent_stream_name = model.stream.name or ""
-            parent_state = ConcurrentPerPartitionCursor.get_parent_state(child_state, parent_stream_name)
+            parent_state = ConcurrentPerPartitionCursor.get_parent_state(
+                child_state, parent_stream_name
+            )
 
             if not parent_state:
                 # there are two migration cases: state value from child stream or from global state
-                parent_state = ConcurrentPerPartitionCursor.get_global_state(child_state, parent_stream_name)
+                parent_state = ConcurrentPerPartitionCursor.get_global_state(
+                    child_state, parent_stream_name
+                )
 
                 if not parent_state and not isinstance(parent_state, dict):
                     cursor_field = InterpolatedString.create(
@@ -3787,8 +3811,12 @@ def _create_message_repository_substream_wrapper(
                         parent_state = AirbyteStateMessage(
                             type=AirbyteStateType.STREAM,
                             stream=AirbyteStreamState(
-                                stream_descriptor=StreamDescriptor(name=parent_stream_name, namespace=None),
-                                stream_state=AirbyteStateBlob({cursor_field: list(cursor_values)[0]}),
+                                stream_descriptor=StreamDescriptor(
+                                    name=parent_stream_name, namespace=None
+                                ),
+                                stream_state=AirbyteStateBlob(
+                                    {cursor_field: list(cursor_values)[0]}
+                                ),
                             ),
                         )
             connector_state_manager = ConnectorStateManager([parent_state] if parent_state else [])
@@ -3804,7 +3832,10 @@ def _create_message_repository_substream_wrapper(
             disable_cache=self._disable_cache,
             message_repository=StateFilteringMessageRepository(
                 LogAppenderMessageRepositoryDecorator(
-                    {"airbyte_cdk": {"stream": {"is_substream": True}}, "http": {"is_auxiliary": True}},
+                    {
+                        "airbyte_cdk": {"stream": {"is_substream": True}},
+                        "http": {"is_auxiliary": True},
+                    },
                     self._message_repository,
                     self._evaluate_log_level(self._emit_connector_builder_messages),
                 ),
@@ -4127,7 +4158,9 @@ def create_grouping_partition_router(
         self, model: GroupingPartitionRouterModel, config: Config, **kwargs: Any
     ) -> GroupingPartitionRouter:
         underlying_router = self._create_component_from_model(
-            model=model.underlying_partition_router, config=config, **kwargs,
+            model=model.underlying_partition_router,
+            config=config,
+            **kwargs,
         )
         if model.group_size < 1:
             raise ValueError(f"Group size must be greater than 0, got {model.group_size}")

airbyte_cdk/sources/declarative/partition_routers/substream_partition_router.py

@@ -7,7 +7,17 @@
 import json
 import logging
 from dataclasses import InitVar, dataclass
-from typing import TYPE_CHECKING, Any, Iterable, List, Mapping, MutableMapping, Optional, Union, TypeVar
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Iterable,
+    List,
+    Mapping,
+    MutableMapping,
+    Optional,
+    Union,
+    TypeVar,
+)
 
 import dpath
 import requests
@@ -27,7 +37,6 @@
 
 
 def iterate_with_last_flag(generator: Iterable[Partition]) -> Iterable[tuple[Partition, bool]]:
-
     iterator = iter(generator)
 
     try:
@@ -191,8 +200,12 @@ def stream_slices(self) -> Iterable[StreamSlice]:
                         for field_path in parent_stream_config.extra_fields
                     ]
 
-                for partition, is_last_slice in iterate_with_last_flag(parent_stream.generate_partitions()):
-                    for parent_record, is_last_record_in_slice in iterate_with_last_flag(partition.read()):
+                for partition, is_last_slice in iterate_with_last_flag(
+                    parent_stream.generate_partitions()
+                ):
+                    for parent_record, is_last_record_in_slice in iterate_with_last_flag(
+                        partition.read()
+                    ):
                         parent_stream.cursor.observe(parent_record)
                         parent_partition = (
                             parent_record.associated_slice.partition
@@ -211,7 +224,9 @@ def stream_slices(self) -> Iterable[StreamSlice]:
                             continue
 
                         # Add extra fields
-                        extracted_extra_fields = self._extract_extra_fields(record_data, extra_fields)
+                        extracted_extra_fields = self._extract_extra_fields(
+                            record_data, extra_fields
+                        )
 
                         if parent_stream_config.lazy_read_pointer:
                             extracted_extra_fields = {
@@ -421,7 +436,9 @@ def get_stream_state(self) -> Optional[Mapping[str, StreamState]]:
         parent_state = {}
         for parent_config in self.parent_stream_configs:
             if parent_config.incremental_dependency:
-                parent_state[parent_config.stream.name] = copy.deepcopy(parent_config.stream.cursor.state)
+                parent_state[parent_config.stream.name] = copy.deepcopy(
+                    parent_config.stream.cursor.state
+                )
         return parent_state
 
     @property

unit_tests/connector_builder/test_connector_builder_handler.py

@@ -1140,7 +1140,9 @@ def test_read_source(mock_http_stream):
     for s in streams:
         retriever = get_retriever(s)
         assert isinstance(retriever, SimpleRetriever)
-        assert isinstance(retriever.stream_slicer, StreamSlicerTestReadDecorator)
+        assert isinstance(
+            s._stream_partition_generator._stream_slicer, StreamSlicerTestReadDecorator
+        )
 
 
 @patch.object(
@@ -1188,7 +1190,9 @@ def test_read_source_single_page_single_slice(mock_http_stream):
     for s in streams:
         retriever = get_retriever(s)
         assert isinstance(retriever, SimpleRetriever)
-        assert isinstance(retriever.stream_slicer, StreamSlicerTestReadDecorator)
+        assert isinstance(
+            s._stream_partition_generator._stream_slicer, StreamSlicerTestReadDecorator
+        )
 
 
 @pytest.mark.parametrize(

unit_tests/sources/declarative/incremental/test_concurrent_perpartitioncursor.py

@@ -1185,8 +1185,8 @@ def run_incremental_parent_state_test(
                 ),
                 # FIXME this is an interesting case. The previous solution would not update the parent state until `ensure_at_least_one_state_emitted` but the concurrent cursor does just before which is probably fine too
                 (
-                        f"https://api.example.com/community/posts?per_page=100&start_time={POST_1_UPDATED_AT}",
-                        {"posts": [{"id": 1, "updated_at": POST_1_UPDATED_AT}]},
+                    f"https://api.example.com/community/posts?per_page=100&start_time={POST_1_UPDATED_AT}",
+                    {"posts": [{"id": 1, "updated_at": POST_1_UPDATED_AT}]},
                 ),
                 # Fetch the first page of comments for post 1
                 (
@@ -1483,8 +1483,8 @@ def run_incremental_parent_state_test(
                 ),
                 # FIXME this is an interesting case. The previous solution would not update the parent state until `ensure_at_least_one_state_emitted` but the concurrent cursor does just before which is probably fine too
                 (
-                        f"https://api.example.com/community/posts?per_page=100&start_time={POST_1_UPDATED_AT}",
-                        {"posts": [{"id": 1, "updated_at": POST_1_UPDATED_AT}]},
+                    f"https://api.example.com/community/posts?per_page=100&start_time={POST_1_UPDATED_AT}",
+                    {"posts": [{"id": 1, "updated_at": POST_1_UPDATED_AT}]},
                 ),
                 # Fetch the first page of comments for post 1
                 (
@@ -1629,8 +1629,8 @@ def run_incremental_parent_state_test(
                 ),
                 # FIXME this is an interesting case. The previous solution would not update the parent state until `ensure_at_least_one_state_emitted` but the concurrent cursor does just before which is probably fine too
                 (
-                        f"https://api.example.com/community/posts?per_page=100&start_time={POST_1_UPDATED_AT}",
-                        {"posts": [{"id": 1, "updated_at": POST_1_UPDATED_AT}]},
+                    f"https://api.example.com/community/posts?per_page=100&start_time={POST_1_UPDATED_AT}",
+                    {"posts": [{"id": 1, "updated_at": POST_1_UPDATED_AT}]},
                 ),
                 # Fetch the first page of comments for post 1
                 (
@@ -2130,7 +2130,9 @@ def test_incremental_parent_state_migration(
                         "states": [
                             {
                                 "partition": {"id": 1, "parent_slice": {}},
-                                "cursor": {"updated_at": START_DATE},  # FIXME this happens because the concurrent framework gets the start date as the max between the state value and the start value. In this case, the start value is higher
+                                "cursor": {
+                                    "updated_at": START_DATE
+                                },  # FIXME this happens because the concurrent framework gets the start date as the max between the state value and the start value. In this case, the start value is higher
                             }
                         ],
                         "lookback_window": 0,  # FIXME the concurrent framework sets the lookback window to 0 as opposed to the declarative framework which would set not define it