airbyte_cdk.sources.streams.concurrent.partition_reader

 1# Copyright (c) 2025 Airbyte, Inc., all rights reserved.
 2
 3import logging
 4from queue import Queue
 5from typing import Optional
 6
 7from airbyte_cdk.sources.concurrent_source.stream_thread_exception import StreamThreadException
 8from airbyte_cdk.sources.message.repository import MessageRepository
 9from airbyte_cdk.sources.streams.concurrent.cursor import Cursor
10from airbyte_cdk.sources.streams.concurrent.partitions.partition import Partition
11from airbyte_cdk.sources.streams.concurrent.partitions.types import (
12    PartitionCompleteSentinel,
13    QueueItem,
14)
15from airbyte_cdk.sources.utils.slice_logger import SliceLogger
16
17
18# Since moving all the connector builder workflow to the concurrent CDK which required correct ordering
19# of grouping log messages onto the main write thread using the ConcurrentMessageRepository, this
20# separate flow and class that was used to log slices onto this partition's message_repository
21# should just be replaced by emitting messages directly onto the repository instead of an intermediary.
22class PartitionLogger:
23    """
24    Helper class that provides a mechanism for passing a log message onto the current
25    partitions message repository
26    """
27
28    def __init__(
29        self,
30        slice_logger: SliceLogger,
31        logger: logging.Logger,
32        message_repository: MessageRepository,
33    ):
34        self._slice_logger = slice_logger
35        self._logger = logger
36        self._message_repository = message_repository
37
38    def log(self, partition: Partition) -> None:
39        if self._slice_logger.should_log_slice_message(self._logger):
40            self._message_repository.emit_message(
41                self._slice_logger.create_slice_log_message(partition.to_slice())
42            )
43
44
45class PartitionReader:
46    """
47    Generates records from a partition and puts them in a queue.
48    """
49
50    _IS_SUCCESSFUL = True
51
52    def __init__(
53        self,
54        queue: Queue[QueueItem],
55        partition_logger: Optional[PartitionLogger] = None,
56    ) -> None:
57        """
58        :param queue: The queue to put the records in.
59        """
60        self._queue = queue
61        self._partition_logger = partition_logger
62
63    def process_partition(self, partition: Partition, cursor: Cursor) -> None:
64        """
65        Process a partition and put the records in the output queue.
66        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.
67
68        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
69        main thread will have no way to know that something when wrong and will wait until the timeout is reached
70
71        This method is meant to be called from a thread.
72        :param partition: The partition to read data from
73        :return: None
74        """
75        try:
76            if self._partition_logger:
77                self._partition_logger.log(partition)
78
79            for record in partition.read():
80                self._queue.put(record)
81                cursor.observe(record)
82            cursor.close_partition(partition)
83            self._queue.put(PartitionCompleteSentinel(partition, self._IS_SUCCESSFUL))
84        except Exception as e:
85            self._queue.put(StreamThreadException(e, partition.stream_name()))
86            self._queue.put(PartitionCompleteSentinel(partition, not self._IS_SUCCESSFUL))