remove SimpleRetrieverTestReadDecorator and use StreamSlicerTestReadDecorator for scenarios wehere we want to limit the number of requests that are made to the backend source.

Author: aldogonzalez8

airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py

@@ -510,7 +510,6 @@
     AsyncRetriever,
     LazySimpleRetriever,
     SimpleRetriever,
-    SimpleRetrieverTestReadDecorator,
 )
 from airbyte_cdk.sources.declarative.retrievers.file_uploader import (
     ConnectorBuilderFileUploader,
@@ -530,7 +529,10 @@
 )
 from airbyte_cdk.sources.declarative.schema.composite_schema_loader import CompositeSchemaLoader
 from airbyte_cdk.sources.declarative.spec import ConfigMigration, Spec
-from airbyte_cdk.sources.declarative.stream_slicers import StreamSlicer
+from airbyte_cdk.sources.declarative.stream_slicers import (
+    StreamSlicer,
+    StreamSlicerTestReadDecorator,
+)
 from airbyte_cdk.sources.declarative.transformations import (
     AddFields,
     RecordTransformation,
@@ -3241,6 +3243,11 @@ def _get_url() -> str:
             request_options_provider = DefaultRequestOptionsProvider(parameters={})
 
         stream_slicer = stream_slicer or SinglePartitionRouter(parameters={})
+        if self._should_limit_slices_fetched():
+            stream_slicer = StreamSlicerTestReadDecorator(
+                wrapped_slicer=stream_slicer,
+                maximum_number_of_slices=self._limit_slices_fetched or 5,
+            )
 
         cursor_used_for_stop_condition = cursor if stop_condition_on_cursor else None
         paginator = (
@@ -3299,10 +3306,7 @@ def _get_url() -> str:
                 parameters=model.parameters or {},
             )
 
-        test_read_enabled = self._limit_slices_fetched or self._emit_connector_builder_messages
-        maximum_number_of_slices = self._limit_slices_fetched or 5 if test_read_enabled else 0
-        retriever_log_formatter = log_formatter if test_read_enabled else None
-        return SimpleRetrieverTestReadDecorator(
+        return SimpleRetriever(
             name=name,
             paginator=paginator,
             primary_key=primary_key,
@@ -3312,14 +3316,37 @@ def _get_url() -> str:
             request_option_provider=request_options_provider,
             cursor=cursor,
             config=config,
-            maximum_number_of_slices=maximum_number_of_slices,
-            emit_connector_builder_messages=self._emit_connector_builder_messages,
             ignore_stream_slicer_parameters_on_paginated_requests=ignore_stream_slicer_parameters_on_paginated_requests,
             additional_query_properties=query_properties,
-            log_formatter=retriever_log_formatter,
+            log_formatter=self._get_log_formatter(log_formatter, name),
             parameters=model.parameters or {},
         )
 
+    def _get_log_formatter(
+        self, log_formatter: Callable[[Response], Any] | None, name: str
+    ) -> Callable[[Response], Any] | None:
+        if self._should_limit_slices_fetched():
+            return (
+                (
+                    lambda response: format_http_message(
+                        response,
+                        f"Stream '{name}' request",
+                        f"Request performed in order to extract records for stream '{name}'",
+                        name,
+                    )
+                )
+                if not log_formatter
+                else log_formatter
+            )
+        return None
+
+    def _should_limit_slices_fetched(self) -> bool:
+        """
+        Returns True if the number of slices fetched should be limited, False otherwise.
+        This is used to limit the number of slices fetched during tests.
+        """
+        return bool(self._limit_slices_fetched or self._emit_connector_builder_messages)
+
     @staticmethod
     def _query_properties_in_request_parameters(
         requester: Union[HttpRequesterModel, CustomRequesterModel],
@@ -3410,7 +3437,7 @@ def create_async_retriever(
         transformations: List[RecordTransformation],
         **kwargs: Any,
     ) -> AsyncRetriever:
-        def _get_download_retriever() -> SimpleRetrieverTestReadDecorator | SimpleRetriever:
+        def _get_download_retriever() -> SimpleRetriever:
             record_selector = RecordSelector(
                 extractor=download_extractor,
                 name=name,
@@ -3430,19 +3457,15 @@ def _get_download_retriever() -> SimpleRetrieverTestReadDecorator | SimpleRetrie
                 if model.download_paginator
                 else NoPagination(parameters={})
             )
-            test_read_enabled = self._limit_slices_fetched or self._emit_connector_builder_messages
-            maximum_number_of_slices = self._limit_slices_fetched or 5 if test_read_enabled else 0
 
-            return SimpleRetrieverTestReadDecorator(
+            return SimpleRetriever(
                 requester=download_requester,
                 record_selector=record_selector,
                 primary_key=None,
                 name=job_download_components_name,
                 paginator=paginator,
                 config=config,
                 parameters={},
-                maximum_number_of_slices=maximum_number_of_slices,
-                emit_connector_builder_messages=self._emit_connector_builder_messages,
             )
 
         def _get_job_timeout() -> datetime.timedelta:
@@ -3479,7 +3502,14 @@ def _get_job_timeout() -> datetime.timedelta:
             transformations=transformations,
             client_side_incremental_sync=client_side_incremental_sync,
         )
+
         stream_slicer = stream_slicer or SinglePartitionRouter(parameters={})
+        if self._should_limit_slices_fetched():
+            stream_slicer = StreamSlicerTestReadDecorator(
+                wrapped_slicer=stream_slicer,
+                maximum_number_of_slices=self._limit_slices_fetched or 5,
+            )
+
         creation_requester = self._create_component_from_model(
             model=model.creation_requester,
             decoder=decoder,

airbyte_cdk/sources/declarative/retrievers/__init__.py

@@ -7,13 +7,11 @@
 from airbyte_cdk.sources.declarative.retrievers.simple_retriever import (
     LazySimpleRetriever,
     SimpleRetriever,
-    SimpleRetrieverTestReadDecorator,
 )
 
 __all__ = [
     "Retriever",
     "SimpleRetriever",
-    "SimpleRetrieverTestReadDecorator",
     "AsyncRetriever",
     "LazySimpleRetriever",
 ]

airbyte_cdk/sources/declarative/retrievers/simple_retriever.py

@@ -645,54 +645,6 @@ def _deep_merge(
             target[key] = value
 
 
-@dataclass
-class SimpleRetrieverTestReadDecorator(SimpleRetriever):
-    """
-    In some cases, we want to limit the number of requests that are made to the backend source. This class allows for limiting the number of
-    slices that are queried throughout a read command.
-
-    maximum_number_of_slices must be provided when test read is enabled.
-    """
-
-    maximum_number_of_slices: int = 0
-    emit_connector_builder_messages: bool = False
-
-    def __post_init__(self, parameters: Mapping[str, Any]) -> None:
-        super().__post_init__(parameters)
-        if self.test_read_enabled():
-            self.log_formatter = (
-                (
-                    lambda response: format_http_message(
-                        response,
-                        f"Stream '{self.name}' request",
-                        f"Request performed in order to extract records for stream '{self.name}'",
-                        self.name,
-                    )
-                )
-                if not self.log_formatter
-                else self.log_formatter
-            )
-
-            if self.maximum_number_of_slices and self.maximum_number_of_slices < 1:
-                raise ValueError(
-                    f"The maximum number of slices on a test read needs to be strictly positive. Got {self.maximum_number_of_slices}"
-                )
-
-    def test_read_enabled(self) -> bool:
-        """
-        Indicates whether the retriever is in test read mode.
-        This is used to limit the number of slices processed during a test read.
-        """
-        return bool(self.maximum_number_of_slices or self.emit_connector_builder_messages)
-
-    # stream_slices is defined with arguments on http stream and fixing this has a long tail of dependencies. Will be resolved by the decoupling of http stream and simple retriever
-    def stream_slices(self) -> Iterable[Optional[StreamSlice]]:  # type: ignore
-        if not self.test_read_enabled():
-            return super().stream_slices()
-        else:
-            return islice(super().stream_slices(), self.maximum_number_of_slices)
-
-
 @deprecated(
     "This class is experimental. Use at your own risk.",
     category=ExperimentalClassWarning,

airbyte_cdk/sources/declarative/stream_slicers/__init__.py

@@ -3,5 +3,8 @@
 #
 
 from airbyte_cdk.sources.declarative.stream_slicers.stream_slicer import StreamSlicer
+from airbyte_cdk.sources.declarative.stream_slicers.stream_slicer_test_read_decorator import (
+    StreamSlicerTestReadDecorator,
+)
 
-__all__ = ["StreamSlicer"]
+__all__ = ["StreamSlicer", "StreamSlicerTestReadDecorator"]

airbyte_cdk/sources/declarative/stream_slicers/stream_slicer_test_read_decorator.py

@@ -0,0 +1,81 @@
+#
+# Copyright (c) 2025 Airbyte, Inc., all rights reserved.
+#
+
+from dataclasses import dataclass
+from itertools import islice
+from typing import Any, Iterable, Mapping, Optional, Union
+
+from airbyte_cdk.sources.types import StreamSlice, StreamState
+
+from .stream_slicer import StreamSlicer
+
+
+@dataclass
+class StreamSlicerTestReadDecorator(StreamSlicer):
+    """
+    In some cases, we want to limit the number of requests that are made to the backend source. This class allows for limiting the number of
+    slices that are queried throughout a read command.
+    """
+
+    wrapped_slicer: StreamSlicer
+    maximum_number_of_slices: int = 5
+
+    def stream_slices(self) -> Iterable[StreamSlice]:
+        return islice(self.wrapped_slicer.stream_slices(), self.maximum_number_of_slices)
+
+    def get_request_params(
+        self,
+        *,
+        stream_state: Optional[StreamState] = None,
+        stream_slice: Optional[StreamSlice] = None,
+        next_page_token: Optional[Mapping[str, Any]] = None,
+    ) -> Mapping[str, Any]:
+        return self.wrapped_slicer.get_request_params(
+            stream_state=stream_state,
+            stream_slice=stream_slice,
+            next_page_token=next_page_token,
+        )
+
+    def get_request_headers(
+        self,
+        *,
+        stream_state: Optional[StreamState] = None,
+        stream_slice: Optional[StreamSlice] = None,
+        next_page_token: Optional[Mapping[str, Any]] = None,
+    ) -> Mapping[str, Any]:
+        return self.wrapped_slicer.get_request_headers(
+            stream_state=stream_state,
+            stream_slice=stream_slice,
+            next_page_token=next_page_token,
+        )
+
+    def get_request_body_data(
+        self,
+        *,
+        stream_state: Optional[StreamState] = None,
+        stream_slice: Optional[StreamSlice] = None,
+        next_page_token: Optional[Mapping[str, Any]] = None,
+    ) -> Union[Mapping[str, Any], str]:
+        return self.wrapped_slicer.get_request_body_data(
+            stream_state=stream_state,
+            stream_slice=stream_slice,
+            next_page_token=next_page_token,
+        )
+
+    def get_request_body_json(
+        self,
+        *,
+        stream_state: Optional[StreamState] = None,
+        stream_slice: Optional[StreamSlice] = None,
+        next_page_token: Optional[Mapping[str, Any]] = None,
+    ) -> Mapping[str, Any]:
+        return self.wrapped_slicer.get_request_body_json(
+            stream_state=stream_state,
+            stream_slice=stream_slice,
+            next_page_token=next_page_token,
+        )
+
+    def __getattr__(self, name: str) -> Any:
+        # Delegate everything else to the wrapped object
+        return getattr(self.wrapped_slicer, name)

unit_tests/connector_builder/test_connector_builder_handler.py

@@ -58,7 +58,6 @@
 from airbyte_cdk.models import Type as MessageType
 from airbyte_cdk.sources.declarative.declarative_stream import DeclarativeStream
 from airbyte_cdk.sources.declarative.manifest_declarative_source import ManifestDeclarativeSource
-from airbyte_cdk.sources.declarative.retrievers import SimpleRetrieverTestReadDecorator
 from airbyte_cdk.sources.declarative.retrievers.simple_retriever import SimpleRetriever
 from airbyte_cdk.test.mock_http import HttpMocker, HttpRequest, HttpResponse
 from airbyte_cdk.utils.airbyte_secrets_utils import filter_secrets, update_secrets
@@ -1112,7 +1111,7 @@ def test_read_source(mock_http_stream):
 
     streams = source.streams(config)
     for s in streams:
-        assert isinstance(s.retriever, SimpleRetrieverTestReadDecorator)
+        assert isinstance(s.retriever, SimpleRetriever)
 
 
 @patch.object(
@@ -1158,7 +1157,7 @@ def test_read_source_single_page_single_slice(mock_http_stream):
 
     streams = source.streams(config)
     for s in streams:
-        assert isinstance(s.retriever, SimpleRetrieverTestReadDecorator)
+        assert isinstance(s.retriever, SimpleRetriever)
 
 
 @pytest.mark.parametrize(

unit_tests/sources/declarative/parsers/test_model_to_component_factory.py

@@ -151,15 +151,12 @@
 )
 from airbyte_cdk.sources.declarative.requesters.request_path import RequestPath
 from airbyte_cdk.sources.declarative.requesters.requester import HttpMethod
-from airbyte_cdk.sources.declarative.retrievers import (
-    AsyncRetriever,
-    SimpleRetriever,
-    SimpleRetrieverTestReadDecorator,
-)
+from airbyte_cdk.sources.declarative.retrievers import AsyncRetriever, SimpleRetriever
 from airbyte_cdk.sources.declarative.schema import InlineSchemaLoader, JsonFileSchemaLoader
 from airbyte_cdk.sources.declarative.schema.composite_schema_loader import CompositeSchemaLoader
 from airbyte_cdk.sources.declarative.schema.schema_loader import SchemaLoader
 from airbyte_cdk.sources.declarative.spec import Spec
+from airbyte_cdk.sources.declarative.stream_slicers import StreamSlicerTestReadDecorator
 from airbyte_cdk.sources.declarative.transformations import AddFields, RemoveFields
 from airbyte_cdk.sources.declarative.transformations.add_fields import AddedFieldDefinition
 from airbyte_cdk.sources.declarative.yaml_declarative_source import YamlDeclarativeSource
@@ -2715,6 +2712,13 @@ def test_simple_retriever_emit_log_messages():
             "path": "/v1/api",
         },
     }
+    request = requests.PreparedRequest()
+    request.headers = {"header": "value"}
+    request.url = "http://byrde.enterprises.com/casinos"
+
+    response = requests.Response()
+    response.request = request
+    response.status_code = 200
 
     connector_builder_factory = ModelToComponentFactory(emit_connector_builder_messages=True)
     retriever = connector_builder_factory.create_component(
@@ -2727,8 +2731,12 @@ def test_simple_retriever_emit_log_messages():
         transformations=[],
     )
 
-    assert isinstance(retriever, SimpleRetrieverTestReadDecorator)
+    assert isinstance(retriever, SimpleRetriever)
     assert connector_builder_factory._message_repository._log_level == Level.DEBUG
+    assert retriever.log_formatter is not None
+    assert retriever.log_formatter(response) == connector_builder_factory._get_log_formatter(
+        None, retriever.name
+    )(response)
 
 
 def test_create_page_increment():
@@ -3078,7 +3086,8 @@ def test_use_request_options_provider_for_datetime_based_cursor():
     assert retriever.name == "Test"
 
     assert isinstance(retriever.cursor, DatetimeBasedCursor)
-    assert isinstance(retriever.stream_slicer, DatetimeBasedCursor)
+    assert isinstance(retriever.stream_slicer, StreamSlicerTestReadDecorator)
+    assert isinstance(retriever.stream_slicer.wrapped_slicer, DatetimeBasedCursor)
 
     assert isinstance(retriever.request_option_provider, DatetimeBasedRequestOptionsProvider)
     assert (
@@ -3166,7 +3175,8 @@ def test_do_not_separate_request_options_provider_for_non_datetime_based_cursor(
     assert retriever.name == "Test"
 
     assert isinstance(retriever.cursor, PerPartitionCursor)
-    assert isinstance(retriever.stream_slicer, PerPartitionCursor)
+    assert isinstance(retriever.stream_slicer, StreamSlicerTestReadDecorator)
+    assert isinstance(retriever.stream_slicer.wrapped_slicer, PerPartitionCursor)
 
     assert isinstance(retriever.request_option_provider, PerPartitionCursor)
     assert isinstance(retriever.request_option_provider._cursor_factory, CursorFactory)
@@ -3207,7 +3217,8 @@ def test_use_default_request_options_provider():
     assert retriever.primary_key == "id"
     assert retriever.name == "Test"
 
-    assert isinstance(retriever.stream_slicer, SinglePartitionRouter)
+    assert isinstance(retriever.stream_slicer, StreamSlicerTestReadDecorator)
+    assert isinstance(retriever.stream_slicer.wrapped_slicer, SinglePartitionRouter)
     assert isinstance(retriever.request_option_provider, DefaultRequestOptionsProvider)