43 lines
1.7 KiB
Python
43 lines
1.7 KiB
Python
#
|
|
# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
|
|
#
|
|
from queue import Queue
|
|
|
|
from airbyte_cdk.sources.concurrent_source.stream_thread_exception import StreamThreadException
|
|
from airbyte_cdk.sources.streams.concurrent.partitions.partition import Partition
|
|
from airbyte_cdk.sources.streams.concurrent.partitions.types import PartitionCompleteSentinel, QueueItem
|
|
|
|
|
|
class PartitionReader:
|
|
"""
|
|
Generates records from a partition and puts them in a queue.
|
|
"""
|
|
|
|
_IS_SUCCESSFUL = True
|
|
|
|
def __init__(self, queue: Queue[QueueItem]) -> None:
|
|
"""
|
|
:param queue: The queue to put the records in.
|
|
"""
|
|
self._queue = queue
|
|
|
|
def process_partition(self, partition: Partition) -> None:
|
|
"""
|
|
Process a partition and put the records in the output queue.
|
|
When all the partitions are added to the queue, a sentinel is added to the queue to indicate that all the partitions have been generated.
|
|
|
|
If an exception is encountered, the exception will be caught and put in the queue. This is very important because if we don't, the
|
|
main thread will have no way to know that something when wrong and will wait until the timeout is reached
|
|
|
|
This method is meant to be called from a thread.
|
|
:param partition: The partition to read data from
|
|
:return: None
|
|
"""
|
|
try:
|
|
for record in partition.read():
|
|
self._queue.put(record)
|
|
self._queue.put(PartitionCompleteSentinel(partition, self._IS_SUCCESSFUL))
|
|
except Exception as e:
|
|
self._queue.put(StreamThreadException(e, partition.stream_name()))
|
|
self._queue.put(PartitionCompleteSentinel(partition, not self._IS_SUCCESSFUL))
|