some pr feedback to rework slice limiting, adding comments, and a few other tweaks

Author: brianjlai

airbyte_cdk/sources/declarative/concurrent_declarative_source.py

@@ -128,6 +128,8 @@ def __init__(
             disable_cache=True if limits else False,
         )
 
+        self._limits = limits
+
         super().__init__(
             source_config=source_config,
             config=config,
@@ -324,6 +326,9 @@ def _group_streams(
                                 self.message_repository,
                             ),
                             stream_slicer=declarative_stream.retriever.stream_slicer,
+                            slice_limit=self._limits.max_slices
+                            if self._limits
+                            else None,  # technically not needed because create_declarative_stream() -> create_simple_retriever() will apply the decorator. But for consistency and depending how we build create_default_stream, this may be needed later
                         )
                     else:
                         if (
@@ -355,6 +360,7 @@ def _group_streams(
                                 self.message_repository,
                             ),
                             stream_slicer=cursor,
+                            slice_limit=self._limits.max_slices if self._limits else None,
                         )
 
                     concurrent_streams.append(
@@ -386,6 +392,9 @@ def _group_streams(
                             self.message_repository,
                         ),
                         declarative_stream.retriever.stream_slicer,
+                        slice_limit=self._limits.max_slices
+                        if self._limits
+                        else None,  # technically not needed because create_declarative_stream() -> create_simple_retriever() will apply the decorator. But for consistency and depending how we build create_default_stream, this may be needed later
                     )
 
                     final_state_cursor = FinalStateCursor(
@@ -447,6 +456,7 @@ def _group_streams(
                             self.message_repository,
                         ),
                         perpartition_cursor,
+                        slice_limit=self._limits.max_slices if self._limits else None,
                     )
 
                     concurrent_streams.append(

airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py

@@ -1453,7 +1453,7 @@ def create_concurrent_cursor_from_datetime_based_cursor(
                         f"Invalid clamping target {evaluated_target}, expected DAY, WEEK, MONTH"
                     )
 
-        concurrent_cursor = ConcurrentCursor(
+        return ConcurrentCursor(
             stream_name=stream_name,
             stream_namespace=stream_namespace,
             stream_state=stream_state,
@@ -1470,17 +1470,6 @@ def create_concurrent_cursor_from_datetime_based_cursor(
             clamping_strategy=clamping_strategy,
         )
 
-        if self._should_limit_slices_fetched():
-            return cast(  # type: ignore  # For a test_read, this will return a StreamSlicer that wraps a cursor and should still work. Changing the signature creates even more type problems
-                StreamSlicer,
-                StreamSlicerTestReadDecorator(
-                    wrapped_slicer=concurrent_cursor,
-                    maximum_number_of_slices=self._limit_slices_fetched or 5,
-                ),
-            )
-        else:
-            return concurrent_cursor
-
     def create_concurrent_cursor_from_incrementing_count_cursor(
         self,
         model_type: Type[BaseModel],
@@ -1534,7 +1523,7 @@ def create_concurrent_cursor_from_incrementing_count_cursor(
             is_sequential_state=True,  # ConcurrentPerPartitionCursor only works with sequential state
         )
 
-        concurrent_cursor = ConcurrentCursor(
+        return ConcurrentCursor(
             stream_name=stream_name,
             stream_namespace=stream_namespace,
             stream_state=stream_state,
@@ -1547,17 +1536,6 @@ def create_concurrent_cursor_from_incrementing_count_cursor(
             end_provider=connector_state_converter.get_end_provider(),  # type: ignore  # Having issues w/ inspection for GapType and CursorValueType as shown in existing tests. Confirmed functionality is working in practice
         )
 
-        if self._should_limit_slices_fetched():
-            return cast(  # type: ignore  # For a test_read, this will return a StreamSlicer that wraps a cursor and should still work. Changing the signature creates even more type problems
-                StreamSlicer,
-                StreamSlicerTestReadDecorator(
-                    wrapped_slicer=concurrent_cursor,
-                    maximum_number_of_slices=self._limit_slices_fetched or MAX_SLICES,
-                ),
-            )
-        else:
-            return concurrent_cursor
-
     def _assemble_weekday(self, weekday: str) -> Weekday:
         match weekday:
             case "MONDAY":

airbyte_cdk/sources/declarative/stream_slicers/declarative_partition_generator.py

@@ -1,8 +1,11 @@
-# Copyright (c) 2024 Airbyte, Inc., all rights reserved.
+# Copyright (c) 2025 Airbyte, Inc., all rights reserved.
 
-from typing import Any, Iterable, Mapping, Optional
+from typing import Any, Iterable, Mapping, Optional, cast
 
 from airbyte_cdk.sources.declarative.retrievers import Retriever
+from airbyte_cdk.sources.declarative.stream_slicers.stream_slicer_test_read_decorator import (
+    StreamSlicerTestReadDecorator,
+)
 from airbyte_cdk.sources.message import MessageRepository
 from airbyte_cdk.sources.streams.concurrent.partitions.partition import Partition
 from airbyte_cdk.sources.streams.concurrent.partitions.partition_generator import PartitionGenerator
@@ -83,10 +86,23 @@ def __hash__(self) -> int:
 
 class StreamSlicerPartitionGenerator(PartitionGenerator):
     def __init__(
-        self, partition_factory: DeclarativePartitionFactory, stream_slicer: StreamSlicer
+        self,
+        partition_factory: DeclarativePartitionFactory,
+        stream_slicer: StreamSlicer,
+        slice_limit: Optional[int] = None,
     ) -> None:
         self._partition_factory = partition_factory
-        self._stream_slicer = stream_slicer
+
+        if slice_limit:
+            self._stream_slicer = cast(
+                StreamSlicer,
+                StreamSlicerTestReadDecorator(
+                    wrapped_slicer=stream_slicer,
+                    maximum_number_of_slices=slice_limit,
+                ),
+            )
+        else:
+            self._stream_slicer = stream_slicer
 
     def generate(self) -> Iterable[Partition]:
         for stream_slice in self._stream_slicer.stream_slices():

airbyte_cdk/sources/streams/concurrent/partition_reader.py

@@ -15,6 +15,10 @@
 from airbyte_cdk.sources.utils.slice_logger import SliceLogger
 
 
+# Since moving all the connector builder workflow to the concurrent CDK which required correct ordering
+# of grouping log messages onto the main write thread using the ConcurrentMessageRepository, this
+# separate flow and class that was used to log slices onto this partition's message_repository
+# should just be replaced by emitting messages directly onto the repository instead of an intermediary.
 class PartitionLogger:
     """
     Helper class that provides a mechanism for passing a log message onto the current

airbyte_cdk/sources/utils/slice_logger.py

@@ -11,6 +11,10 @@
 from airbyte_cdk.models import Type as MessageType
 
 
+# Once everything runs on the concurrent CDK and we've cleaned up the legacy flows, we should try to remove
+# this class and write messages directly to the message_repository instead of through the logger because for
+# cases like the connector builder where ordering of messages is important, using the logger can cause
+# messages to be grouped out of order. Alas work for a different day.
 class SliceLogger(ABC):
     """
     SliceLogger is an interface that allows us to log slices of data in a uniform way.

unit_tests/sources/declarative/schema/test_dynamic_schema_loader.py

@@ -4,23 +4,15 @@
 
 import json
 from copy import deepcopy
-from queue import Queue
 from unittest.mock import MagicMock
 
 import pytest
-from airbyte_protocol_dataclasses.models import Level
 
 from airbyte_cdk.sources.declarative.concurrent_declarative_source import (
     ConcurrentDeclarativeSource,
     TestLimits,
 )
-from airbyte_cdk.sources.declarative.parsers.model_to_component_factory import (
-    ModelToComponentFactory,
-)
 from airbyte_cdk.sources.declarative.schema import DynamicSchemaLoader, SchemaTypeIdentifier
-from airbyte_cdk.sources.message.concurrent_repository import ConcurrentMessageRepository
-from airbyte_cdk.sources.message.repository import InMemoryMessageRepository
-from airbyte_cdk.sources.streams.concurrent.partitions.types import QueueItem
 from airbyte_cdk.test.mock_http import HttpMocker, HttpRequest, HttpResponse
 
 _CONFIG = {
@@ -360,9 +352,6 @@ def test_dynamic_schema_loader_with_type_conditions():
         },
     }
 
-    # queue: Queue[QueueItem] = Queue(maxsize=10_000)
-    # message_repository = InMemoryMessageRepository(Level.INFO)
-
     source = ConcurrentDeclarativeSource(
         source_config=_MANIFEST_WITH_TYPE_CONDITIONS,
         config=_CONFIG,