add max records per-partition

Author: brianjlai

airbyte_cdk/sources/concurrent_source/concurrent_source.py

@@ -45,6 +45,7 @@ def create(
         message_repository: MessageRepository,
         queue: Optional[Queue[QueueItem]] = None,
         timeout_seconds: int = DEFAULT_TIMEOUT_SECONDS,
+        max_records_per_partition: Optional[int] = None,
     ) -> "ConcurrentSource":
         is_single_threaded = initial_number_of_partitions_to_generate == 1 and num_workers == 1
         too_many_generator = (
@@ -67,6 +68,7 @@ def create(
             message_repository=message_repository,
             initial_number_partitions_to_generate=initial_number_of_partitions_to_generate,
             timeout_seconds=timeout_seconds,
+            max_records_per_partition=max_records_per_partition,
         )
 
     def __init__(
@@ -78,6 +80,7 @@ def __init__(
         message_repository: MessageRepository = InMemoryMessageRepository(),
         initial_number_partitions_to_generate: int = 1,
         timeout_seconds: int = DEFAULT_TIMEOUT_SECONDS,
+        max_records_per_partition: Optional[int] = None,
     ) -> None:
         """
         :param threadpool: The threadpool to submit tasks to
@@ -93,6 +96,7 @@ def __init__(
         self._message_repository = message_repository
         self._initial_number_partitions_to_generate = initial_number_partitions_to_generate
         self._timeout_seconds = timeout_seconds
+        self._max_records_per_partition = max_records_per_partition
 
         # We set a maxsize to for the main thread to process record items when the queue size grows. This assumes that there are less
         # threads generating partitions that than are max number of workers. If it weren't the case, we could have threads only generating
@@ -115,6 +119,7 @@ def read(
             PartitionReader(
                 self._queue,
                 PartitionLogger(self._slice_logger, self._logger, self._message_repository),
+                self._max_records_per_partition,
             ),
         )
 

airbyte_cdk/sources/declarative/concurrent_declarative_source.py

@@ -170,6 +170,7 @@ def __init__(
             slice_logger=self._slice_logger,
             queue=queue,
             message_repository=self.message_repository,
+            max_records_per_partition=limits.max_records if limits else None,
         )
 
     # TODO: Remove this. This property is necessary to safely migrate Stripe during the transition state.

airbyte_cdk/sources/streams/concurrent/partition_reader.py

@@ -50,13 +50,17 @@ class PartitionReader:
     _IS_SUCCESSFUL = True
 
     def __init__(
-        self, queue: Queue[QueueItem], partition_logger: Optional[PartitionLogger] = None
+        self,
+        queue: Queue[QueueItem],
+        partition_logger: Optional[PartitionLogger] = None,
+        max_records_per_partition: Optional[int] = None,
     ) -> None:
         """
         :param queue: The queue to put the records in.
         """
         self._queue = queue
         self._partition_logger = partition_logger
+        self._max_records_per_partition = max_records_per_partition
 
     def process_partition(self, partition: Partition, cursor: Cursor) -> None:
         """
@@ -74,9 +78,18 @@ def process_partition(self, partition: Partition, cursor: Cursor) -> None:
             if self._partition_logger:
                 self._partition_logger.log(partition)
 
+            record_count = 0
             for record in partition.read():
                 self._queue.put(record)
                 cursor.observe(record)
+                record_count += 1
+                if (
+                    self._max_records_per_partition
+                    and record_count >= self._max_records_per_partition
+                ):
+                    # We stop processing a partition after exceeding the max_records for Connector
+                    # Builder test reads. The record limit only applies to an individual partition
+                    break
             cursor.close_partition(partition)
             self._queue.put(PartitionCompleteSentinel(partition, self._IS_SUCCESSFUL))
         except Exception as e: