Merge branch 'maxi297/bland_stream_instantiated_as_defaultstream' into maxi297/incremental_without_partition_router_as_defaultstream

Author: maxime.c

airbyte_cdk/connector_builder/test_reader/reader.py

@@ -120,7 +120,11 @@ def run_test_read(
         deprecation_warnings: List[LogMessage] = source.deprecation_warnings()
 
         schema_inferrer = SchemaInferrer(
-            self._pk_to_nested_and_composite_field(stream.primary_key) if stream else None,
+            self._pk_to_nested_and_composite_field(
+                stream.primary_key if hasattr(stream, "primary_key") else stream._primary_key  # type: ignore  # We are accessing the private property here as the primary key is not exposed. We should either expose it or use `as_airbyte_stream` to retrieve it as this is the "official" way where it is exposed in the Airbyte protocol
+            )
+            if stream
+            else None,
             self._cursor_field_to_nested_and_composite_field(stream.cursor_field)
             if stream
             else None,

airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py

@@ -2055,13 +2055,9 @@ def create_declarative_stream(
             schema_loader = DefaultSchemaLoader(config=config, parameters=options)
 
         if (
-            (
-                isinstance(combined_slicers, PartitionRouter)
-                or isinstance(concurrent_cursor, ConcurrentCursor)
-            )
-            and not is_parent
-            and not self._emit_connector_builder_messages
-        ):
+            isinstance(combined_slicers, PartitionRouter)
+            or isinstance(concurrent_cursor, ConcurrentCursor)
+        ) and not is_parent:
             # We are starting to migrate streams to instantiate directly the DefaultStream instead of instantiating the
             # DeclarativeStream and assembling the DefaultStream from that. The plan is the following:
             # * Streams without partition router nor cursors and streams with only partition router. This is the `isinstance(combined_slicers, PartitionRouter)` condition as the first kind with have a SinglePartitionRouter
@@ -2097,7 +2093,13 @@ def create_declarative_stream(
                     retriever,
                     self._message_repository,
                 ),
-                stream_slicer,
+                stream_slicer=cast(
+                    StreamSlicer,
+                    StreamSlicerTestReadDecorator(
+                        wrapped_slicer=stream_slicer,
+                        maximum_number_of_slices=self._limit_slices_fetched or 5,
+                    ),
+                ),
             )
 
             return DefaultStream(

unit_tests/connector_builder/test_connector_builder_handler.py

@@ -7,7 +7,7 @@
 import json
 import logging
 import os
-from typing import List, Literal
+from typing import List, Literal, Union
 from unittest import mock
 from unittest.mock import MagicMock, patch
 
@@ -17,7 +17,6 @@
 
 from airbyte_cdk import connector_builder
 from airbyte_cdk.connector_builder.connector_builder_handler import (
-    TestLimits,
     create_source,
     get_limits,
     resolve_manifest,
@@ -60,6 +59,7 @@
 from airbyte_cdk.sources.declarative.declarative_stream import DeclarativeStream
 from airbyte_cdk.sources.declarative.retrievers.simple_retriever import SimpleRetriever
 from airbyte_cdk.sources.declarative.stream_slicers import StreamSlicerTestReadDecorator
+from airbyte_cdk.sources.streams.concurrent.default_stream import DefaultStream
 from airbyte_cdk.test.mock_http import HttpMocker, HttpRequest, HttpResponse
 from airbyte_cdk.utils.airbyte_secrets_utils import filter_secrets, update_secrets
 from unit_tests.connector_builder.utils import create_configured_catalog
@@ -440,6 +440,14 @@
 }
 
 
+def get_retriever(stream: Union[DeclarativeStream, DefaultStream]):
+    return (
+        stream.retriever
+        if isinstance(stream, DeclarativeStream)
+        else stream._stream_partition_generator._partition_factory._retriever
+    )
+
+
 @pytest.fixture
 def valid_resolve_manifest_config_file(tmp_path):
     config_file = tmp_path / "config.json"
@@ -1130,8 +1138,9 @@ def test_read_source(mock_http_stream):
 
     streams = source.streams(config)
     for s in streams:
-        assert isinstance(s.retriever, SimpleRetriever)
-        assert isinstance(s.retriever.stream_slicer, StreamSlicerTestReadDecorator)
+        retriever = get_retriever(s)
+        assert isinstance(retriever, SimpleRetriever)
+        assert isinstance(retriever.stream_slicer, StreamSlicerTestReadDecorator)
 
 
 @patch.object(
@@ -1177,8 +1186,9 @@ def test_read_source_single_page_single_slice(mock_http_stream):
 
     streams = source.streams(config)
     for s in streams:
-        assert isinstance(s.retriever, SimpleRetriever)
-        assert isinstance(s.retriever.stream_slicer, StreamSlicerTestReadDecorator)
+        retriever = get_retriever(s)
+        assert isinstance(retriever, SimpleRetriever)
+        assert isinstance(retriever.stream_slicer, StreamSlicerTestReadDecorator)
 
 
 @pytest.mark.parametrize(