move limiting to global variable on DeclarativePartition

Author: brianjlai

airbyte_cdk/sources/declarative/concurrent_declarative_source.py

@@ -126,7 +126,6 @@ def __init__(
             max_concurrent_async_job_count=source_config.get("max_concurrent_async_job_count"),
             limit_pages_fetched_per_slice=limits.max_pages_per_slice if limits else None,
             limit_slices_fetched=limits.max_slices if limits else None,
-            limit_max_records=limits.max_records if limits else None,
             disable_retries=True if limits else False,
             disable_cache=True if limits else False,
         )
@@ -325,10 +324,13 @@ def _group_streams(
 
                         partition_generator = StreamSlicerPartitionGenerator(
                             partition_factory=DeclarativePartitionFactory(
-                                declarative_stream.name,
-                                declarative_stream.get_json_schema(),
-                                retriever,
-                                self.message_repository,
+                                stream_name=declarative_stream.name,
+                                json_schema=declarative_stream.get_json_schema(),
+                                retriever=retriever,
+                                message_repository=self.message_repository,
+                                max_records_limit=self._limits.max_records
+                                if self._limits
+                                else None,
                             ),
                             stream_slicer=declarative_stream.retriever.stream_slicer,
                             slice_limit=self._limits.max_slices
@@ -359,10 +361,13 @@ def _group_streams(
                             )
                         partition_generator = StreamSlicerPartitionGenerator(
                             partition_factory=DeclarativePartitionFactory(
-                                declarative_stream.name,
-                                declarative_stream.get_json_schema(),
-                                retriever,
-                                self.message_repository,
+                                stream_name=declarative_stream.name,
+                                json_schema=declarative_stream.get_json_schema(),
+                                retriever=retriever,
+                                message_repository=self.message_repository,
+                                max_records_limit=self._limits.max_records
+                                if self._limits
+                                else None,
                             ),
                             stream_slicer=cursor,
                             slice_limit=self._limits.max_slices if self._limits else None,
@@ -391,10 +396,11 @@ def _group_streams(
                 ) and hasattr(declarative_stream.retriever, "stream_slicer"):
                     partition_generator = StreamSlicerPartitionGenerator(
                         DeclarativePartitionFactory(
-                            declarative_stream.name,
-                            declarative_stream.get_json_schema(),
-                            declarative_stream.retriever,
-                            self.message_repository,
+                            stream_name=declarative_stream.name,
+                            json_schema=declarative_stream.get_json_schema(),
+                            retriever=declarative_stream.retriever,
+                            message_repository=self.message_repository,
+                            max_records_limit=self._limits.max_records if self._limits else None,
                         ),
                         declarative_stream.retriever.stream_slicer,
                         slice_limit=self._limits.max_slices
@@ -455,10 +461,11 @@ def _group_streams(
 
                     partition_generator = StreamSlicerPartitionGenerator(
                         DeclarativePartitionFactory(
-                            declarative_stream.name,
-                            declarative_stream.get_json_schema(),
-                            retriever,
-                            self.message_repository,
+                            stream_name=declarative_stream.name,
+                            json_schema=declarative_stream.get_json_schema(),
+                            retriever=retriever,
+                            message_repository=self.message_repository,
+                            max_records_limit=self._limits.max_records if self._limits else None,
                         ),
                         perpartition_cursor,
                         slice_limit=self._limits.max_slices if self._limits else None,

airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py

@@ -634,7 +634,6 @@ def __init__(
         self,
         limit_pages_fetched_per_slice: Optional[int] = None,
         limit_slices_fetched: Optional[int] = None,
-        limit_max_records: Optional[int] = None,
         emit_connector_builder_messages: bool = False,
         disable_retries: bool = False,
         disable_cache: bool = False,
@@ -646,7 +645,6 @@ def __init__(
         self._init_mappings()
         self._limit_pages_fetched_per_slice = limit_pages_fetched_per_slice
         self._limit_slices_fetched = limit_slices_fetched
-        self._limit_max_records = limit_max_records
         self._emit_connector_builder_messages = emit_connector_builder_messages
         self._disable_retries = disable_retries
         self._disable_cache = disable_cache
@@ -3400,7 +3398,6 @@ def _get_url() -> str:
             ignore_stream_slicer_parameters_on_paginated_requests=ignore_stream_slicer_parameters_on_paginated_requests,
             additional_query_properties=query_properties,
             log_formatter=self._get_log_formatter(log_formatter, name),
-            max_records=self._limit_max_records,
             parameters=model.parameters or {},
         )
 

airbyte_cdk/sources/declarative/retrievers/simple_retriever.py

@@ -92,7 +92,6 @@ class SimpleRetriever(Retriever):
     ignore_stream_slicer_parameters_on_paginated_requests: bool = False
     additional_query_properties: Optional[QueryProperties] = None
     log_formatter: Optional[Callable[[requests.Response], Any]] = None
-    max_records: Optional[int] = None
 
     def __post_init__(self, parameters: Mapping[str, Any]) -> None:
         self._paginator = self.paginator or NoPagination(parameters=parameters)
@@ -102,7 +101,6 @@ def __post_init__(self, parameters: Mapping[str, Any]) -> None:
             if isinstance(self._name, str)
             else self._name
         )
-        self._total_records_read = 0
 
     @property  # type: ignore
     def name(self) -> str:
@@ -503,12 +501,6 @@ def read_records(
         :param stream_slice: The stream slice to read data for
         :return: The records read from the API source
         """
-
-        # For Connector Builder test read operations, if the max number of records has already been
-        # reached, we just return without attempted to extract any more records
-        if self.max_records and self._total_records_read >= self.max_records:
-            return
-
         _slice = stream_slice or StreamSlice(partition={}, cursor_slice={})  # None-check
 
         most_recent_record_from_slice = None
@@ -537,13 +529,6 @@ def read_records(
 
                 yield stream_data
 
-                # For Connector Builder test read operations, if the max number of records is reached, we
-                # exit the process early without emitting more records or attempting to extract more
-                if self.max_records:
-                    self._total_records_read += 1
-                    if self._total_records_read >= self.max_records:
-                        break
-
             if self.cursor:
                 self.cursor.close_slice(_slice)
         return

airbyte_cdk/sources/declarative/stream_slicers/declarative_partition_generator.py

@@ -13,6 +13,11 @@
 from airbyte_cdk.sources.types import Record, StreamSlice
 from airbyte_cdk.utils.slice_hasher import SliceHasher
 
+# For Connector Builder test read operations, we track the total number of records
+# read for the stream at the global level so that we can stop reading early if we
+# exceed the record limit
+total_record_counter = 0
+
 
 class DeclarativePartitionFactory:
     def __init__(
@@ -21,6 +26,7 @@ def __init__(
         json_schema: Mapping[str, Any],
         retriever: Retriever,
         message_repository: MessageRepository,
+        max_records_limit: Optional[int] = None,
     ) -> None:
         """
         The DeclarativePartitionFactory takes a retriever_factory and not a retriever directly. The reason is that our components are not
@@ -31,14 +37,16 @@ def __init__(
         self._json_schema = json_schema
         self._retriever = retriever
         self._message_repository = message_repository
+        self._max_records_limit = max_records_limit
 
     def create(self, stream_slice: StreamSlice) -> Partition:
         return DeclarativePartition(
-            self._stream_name,
-            self._json_schema,
-            self._retriever,
-            self._message_repository,
-            stream_slice,
+            stream_name=self._stream_name,
+            json_schema=self._json_schema,
+            retriever=self._retriever,
+            message_repository=self._message_repository,
+            max_records_limit=self._max_records_limit,
+            stream_slice=stream_slice,
         )
 
 
@@ -49,17 +57,24 @@ def __init__(
         json_schema: Mapping[str, Any],
         retriever: Retriever,
         message_repository: MessageRepository,
+        max_records_limit: Optional[int],
         stream_slice: StreamSlice,
     ):
         self._stream_name = stream_name
         self._json_schema = json_schema
         self._retriever = retriever
         self._message_repository = message_repository
+        self._max_records_limit = max_records_limit
         self._stream_slice = stream_slice
         self._hash = SliceHasher.hash(self._stream_name, self._stream_slice)
 
     def read(self) -> Iterable[Record]:
         for stream_data in self._retriever.read_records(self._json_schema, self._stream_slice):
+            if self._max_records_limit:
+                global total_record_counter
+                if total_record_counter >= self._max_records_limit:
+                    break
+
             if isinstance(stream_data, Mapping):
                 record = (
                     stream_data
@@ -74,6 +89,9 @@ def read(self) -> Iterable[Record]:
             else:
                 self._message_repository.emit_message(stream_data)
 
+            if self._max_records_limit:
+                total_record_counter += 1
+
     def to_slice(self) -> Optional[Mapping[str, Any]]:
         return self._stream_slice
 
@@ -90,6 +108,7 @@ def __init__(
         partition_factory: DeclarativePartitionFactory,
         stream_slicer: StreamSlicer,
         slice_limit: Optional[int] = None,
+        max_records_limit: Optional[int] = None,
     ) -> None:
         self._partition_factory = partition_factory
 

unit_tests/sources/declarative/retrievers/test_simple_retriever.py

@@ -1566,122 +1566,3 @@ def test_simple_retriever_still_emit_records_if_no_merge_key():
 
     assert len(actual_records) == 10
     assert actual_records == expected_records
-
-
-def test_simple_retriever_max_records_reached():
-    expected_records = [
-        Record(data={"id": 1, "name": "Max"}, stream_name="stream_name"),
-        Record(data={"id": 1, "name": "Oscar"}, stream_name="stream_name"),
-        Record(data={"id": 1, "name": "Charles"}, stream_name="stream_name"),
-        Record(data={"id": 1, "name": "Alex"}, stream_name="stream_name"),
-        Record(data={"id": 1, "name": "Yuki"}, stream_name="stream_name"),
-    ]
-
-    mock_records = expected_records + [
-        Record(data={"id": 1, "name": "Lewis"}, stream_name="stream_name"),
-        Record(data={"id": 1, "name": "Lando"}, stream_name="stream_name"),
-    ]
-
-    record_selector = MagicMock()
-    record_selector.select_records.return_value = []
-
-    retriever = SimpleRetriever(
-        name="stream_name",
-        primary_key=primary_key,
-        requester=MagicMock(),
-        paginator=Mock(),
-        record_selector=record_selector,
-        max_records=5,
-        parameters={},
-        config={},
-    )
-
-    stream_slice = StreamSlice(cursor_slice={}, partition={"repository": "airbyte"})
-
-    with patch.object(
-        SimpleRetriever,
-        "_read_pages",
-        return_value=iter(mock_records),
-        # side_effect=retriever_read_pages,
-    ):
-        actual_records = list(retriever.read_records(stream_slice=stream_slice, records_schema={}))
-
-        assert len(actual_records) == 5
-        assert actual_records == expected_records
-
-
-def test_simple_retriever_max_records_already_reached_on_previous_read():
-    mock_records = [
-        Record(data={"id": 1, "name": "Max"}, stream_name="stream_name"),
-        Record(data={"id": 1, "name": "Oscar"}, stream_name="stream_name"),
-        Record(data={"id": 1, "name": "Charles"}, stream_name="stream_name"),
-        Record(data={"id": 1, "name": "Alex"}, stream_name="stream_name"),
-        Record(data={"id": 1, "name": "Yuki"}, stream_name="stream_name"),
-    ]
-
-    record_selector = MagicMock()
-    record_selector.select_records.return_value = []
-
-    retriever = SimpleRetriever(
-        name="stream_name",
-        primary_key=primary_key,
-        requester=MagicMock(),
-        paginator=Mock(),
-        record_selector=record_selector,
-        max_records=5,
-        parameters={},
-        config={},
-    )
-    retriever._total_records_read = 5
-
-    stream_slice = StreamSlice(cursor_slice={}, partition={"repository": "airbyte"})
-
-    with patch.object(
-        SimpleRetriever,
-        "_read_pages",
-        return_value=iter(mock_records),
-        # side_effect=retriever_read_pages,
-    ):
-        actual_records = list(retriever.read_records(stream_slice=stream_slice, records_schema={}))
-
-        assert len(actual_records) == 0
-
-
-def test_simple_retriever_read_some_records():
-    expected_records = [
-        Record(data={"id": 1, "name": "Max"}, stream_name="stream_name"),
-        Record(data={"id": 1, "name": "Oscar"}, stream_name="stream_name"),
-        Record(data={"id": 1, "name": "Charles"}, stream_name="stream_name"),
-    ]
-
-    mock_records = expected_records + [
-        Record(data={"id": 1, "name": "Alex"}, stream_name="stream_name"),
-        Record(data={"id": 1, "name": "Yuki"}, stream_name="stream_name"),
-    ]
-
-    record_selector = MagicMock()
-    record_selector.select_records.return_value = []
-
-    retriever = SimpleRetriever(
-        name="stream_name",
-        primary_key=primary_key,
-        requester=MagicMock(),
-        paginator=Mock(),
-        record_selector=record_selector,
-        max_records=5,
-        parameters={},
-        config={},
-    )
-    retriever._total_records_read = 2
-
-    stream_slice = StreamSlice(cursor_slice={}, partition={"repository": "airbyte"})
-
-    with patch.object(
-        SimpleRetriever,
-        "_read_pages",
-        return_value=iter(mock_records),
-    ):
-        actual_records = list(retriever.read_records(stream_slice=stream_slice, records_schema={}))
-
-        assert len(actual_records) == 3
-        assert actual_records == expected_records