more fixes for DefaultStream in Connector Builder

Author: maxime.c

airbyte_cdk/connector_builder/test_reader/reader.py

@@ -120,7 +120,7 @@ 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) 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

@@ -2072,7 +2072,6 @@ def create_declarative_stream(
         if (
             isinstance(combined_slicers, PartitionRouter)
             and not is_parent
-            and not self._emit_connector_builder_messages
         ):
             # 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:
@@ -2089,7 +2088,13 @@ def create_declarative_stream(
                     retriever,
                     self._message_repository,
                 ),
-                combined_slicers,
+                stream_slicer=cast(
+                    StreamSlicer,
+                    StreamSlicerTestReadDecorator(
+                        wrapped_slicer=combined_slicers,
+                        maximum_number_of_slices=self._limit_slices_fetched or 5,
+                    ),
+                ),
             )
             return DefaultStream(
                 partition_generator=partition_generator,

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(