add additional check to skip calling read_records() if max records is already reached

Author: brianjlai

airbyte_cdk/sources/declarative/stream_slicers/declarative_partition_generator.py

@@ -69,9 +69,12 @@ def __init__(
         self._hash = SliceHasher.hash(self._stream_name, self._stream_slice)
 
     def read(self) -> Iterable[Record]:
+        if self._max_records_limit:
+            global total_record_counter
+            if total_record_counter >= self._max_records_limit:
+                return
         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
 

unit_tests/sources/declarative/stream_slicers/test_declarative_partition_generator.py

@@ -142,6 +142,11 @@ def test_max_records_reached_on_previous_partition(self) -> None:
         second_partition_records = list(partition.read())
         assert len(second_partition_records) == 0
 
+        # The DeclarativePartition exits out of the read before attempting to read_records() if
+        # the max_records_limit has already been reached. So we only expect to see read_records()
+        # called for the first partition read and not the second
+        retriever.read_records.assert_called_once()
+
     @staticmethod
     def _mock_retriever(read_return_value: List[StreamData]) -> Mock:
         retriever = Mock(spec=Retriever)