airbyte_cdk.sources.concurrent_source.concurrent_read_processor

  1#
  2# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
  3#
  4import logging
  5import os
  6from typing import Dict, Iterable, List, Optional, Set
  7
  8from airbyte_cdk.exception_handler import generate_failed_streams_error_message
  9from airbyte_cdk.models import AirbyteMessage, AirbyteStreamStatus, FailureType, StreamDescriptor
 10from airbyte_cdk.models import Type as MessageType
 11from airbyte_cdk.sources.concurrent_source.partition_generation_completed_sentinel import (
 12    PartitionGenerationCompletedSentinel,
 13)
 14from airbyte_cdk.sources.concurrent_source.stream_thread_exception import StreamThreadException
 15from airbyte_cdk.sources.concurrent_source.thread_pool_manager import ThreadPoolManager
 16from airbyte_cdk.sources.declarative.partition_routers.cartesian_product_stream_slicer import (
 17    CartesianProductStreamSlicer,
 18)
 19from airbyte_cdk.sources.declarative.partition_routers.grouping_partition_router import (
 20    GroupingPartitionRouter,
 21)
 22from airbyte_cdk.sources.declarative.partition_routers.substream_partition_router import (
 23    SubstreamPartitionRouter,
 24)
 25from airbyte_cdk.sources.declarative.partition_routers.union_partition_router import (
 26    UnionPartitionRouter,
 27)
 28from airbyte_cdk.sources.message import MessageRepository
 29from airbyte_cdk.sources.streams.concurrent.abstract_stream import AbstractStream
 30from airbyte_cdk.sources.streams.concurrent.default_stream import DefaultStream
 31from airbyte_cdk.sources.streams.concurrent.partition_enqueuer import PartitionEnqueuer
 32from airbyte_cdk.sources.streams.concurrent.partition_reader import PartitionReader
 33from airbyte_cdk.sources.streams.concurrent.partitions.partition import Partition
 34from airbyte_cdk.sources.streams.concurrent.partitions.types import PartitionCompleteSentinel
 35from airbyte_cdk.sources.types import Record
 36from airbyte_cdk.sources.utils.record_helper import stream_data_to_airbyte_message
 37from airbyte_cdk.sources.utils.slice_logger import SliceLogger
 38from airbyte_cdk.utils import AirbyteTracedException
 39from airbyte_cdk.utils.stream_status_utils import (
 40    as_airbyte_message as stream_status_as_airbyte_message,
 41)
 42
 43
 44class ConcurrentReadProcessor:
 45    def __init__(
 46        self,
 47        stream_instances_to_read_from: List[AbstractStream],
 48        partition_enqueuer: PartitionEnqueuer,
 49        thread_pool_manager: ThreadPoolManager,
 50        logger: logging.Logger,
 51        slice_logger: SliceLogger,
 52        message_repository: MessageRepository,
 53        partition_reader: PartitionReader,
 54        max_concurrent_partition_generators: Optional[int] = None,
 55    ):
 56        """
 57        This class is responsible for handling items from a concurrent stream read process.
 58        :param stream_instances_to_read_from: List of streams to read from
 59        :param partition_enqueuer: PartitionEnqueuer instance
 60        :param thread_pool_manager: ThreadPoolManager instance
 61        :param logger: Logger instance
 62        :param slice_logger: SliceLogger instance
 63        :param message_repository: MessageRepository instance
 64        :param partition_reader: PartitionReader instance
 65        :param max_concurrent_partition_generators: Maximum number of partition generators allowed
 66            to run concurrently. None means no limit. When set, should be less than the number of
 67            workers in multi-worker mode so at least one worker slot is always available for
 68            partition reading, preventing thread pool starvation. In single-threaded mode
 69            (num_workers=1) the value may equal num_workers; ConcurrentSource.create() handles
 70            this distinction. ConcurrentSource.read() passes this value explicitly.
 71        """
 72        self._stream_name_to_instance = {s.name: s for s in stream_instances_to_read_from}
 73        self._record_counter = {}
 74        self._streams_to_running_partitions: Dict[str, Set[Partition]] = {}
 75        for stream in stream_instances_to_read_from:
 76            self._streams_to_running_partitions[stream.name] = set()
 77            self._record_counter[stream.name] = 0
 78        if (
 79            max_concurrent_partition_generators is not None
 80            and max_concurrent_partition_generators < 1
 81        ):
 82            raise ValueError(
 83                f"max_concurrent_partition_generators must be >= 1 or None, got {max_concurrent_partition_generators}"
 84            )
 85        self._thread_pool_manager = thread_pool_manager
 86        self._partition_enqueuer = partition_enqueuer
 87        self._max_concurrent_partition_generators = max_concurrent_partition_generators
 88        self._stream_instances_to_start_partition_generation = stream_instances_to_read_from
 89        self._streams_currently_generating_partitions: List[str] = []
 90        self._logger = logger
 91        self._slice_logger = slice_logger
 92        self._message_repository = message_repository
 93        self._partition_reader = partition_reader
 94        self._streams_done: Set[str] = set()
 95        self._exceptions_per_stream_name: dict[str, List[Exception]] = {}
 96
 97        # Track which streams (by name) are currently active
 98        # A stream is "active" if it's generating partitions or has partitions being read
 99        self._active_stream_names: Set[str] = set()
100
101        # Store blocking group names for streams that require blocking simultaneous reads
102        # Maps stream name -> group name (empty string means no blocking)
103        self._stream_block_simultaneous_read: Dict[str, str] = {
104            stream.name: stream.block_simultaneous_read for stream in stream_instances_to_read_from
105        }
106
107        # Track which groups are currently active
108        # Maps group name -> set of stream names in that group
109        self._active_groups: Dict[str, Set[str]] = {}
110
111        for stream in stream_instances_to_read_from:
112            if stream.block_simultaneous_read:
113                self._logger.info(
114                    f"Stream '{stream.name}' is in blocking group '{stream.block_simultaneous_read}'. "
115                    f"Will defer starting this stream if another stream in the same group or its parents are active."
116                )
117
118    def on_partition_generation_completed(
119        self, sentinel: PartitionGenerationCompletedSentinel
120    ) -> Iterable[AirbyteMessage]:
121        """
122        This method is called when a partition generation is completed.
123        1. Remove the stream from the list of streams currently generating partitions
124        2. Deactivate parent streams (they were only needed for partition generation)
125        3. If the stream is done, mark it as such and return a stream status message
126        4. If there are more streams to read from, start the next partition generator
127        """
128        stream_name = sentinel.stream.name
129        self._streams_currently_generating_partitions.remove(sentinel.stream.name)
130
131        # Deactivate all parent streams now that partition generation is complete
132        # Parents were only needed to generate slices, they can now be reused
133        parent_streams = self._collect_all_parent_stream_names(stream_name)
134        for parent_stream_name in parent_streams:
135            if parent_stream_name in self._active_stream_names:
136                self._logger.debug(f"Removing '{parent_stream_name}' from active streams")
137                self._active_stream_names.discard(parent_stream_name)
138
139                # Remove from active groups
140                parent_group = self._stream_block_simultaneous_read.get(parent_stream_name, "")
141                if parent_group:
142                    if parent_group in self._active_groups:
143                        self._active_groups[parent_group].discard(parent_stream_name)
144                        if not self._active_groups[parent_group]:
145                            del self._active_groups[parent_group]
146                    self._logger.info(
147                        f"Parent stream '{parent_stream_name}' (group '{parent_group}') deactivated after "
148                        f"partition generation completed for child '{stream_name}'. "
149                        f"Blocked streams in the queue will be retried on next start_next_partition_generator call."
150                    )
151
152        # It is possible for the stream to already be done if no partitions were generated
153        # If the partition generation process was completed and there are no partitions left to process, the stream is done
154        if (
155            self._is_stream_done(stream_name)
156            or len(self._streams_to_running_partitions[stream_name]) == 0
157        ):
158            yield from self._on_stream_is_done(stream_name)
159        if self._stream_instances_to_start_partition_generation:
160            status_message = self.start_next_partition_generator()
161            if status_message:
162                yield status_message
163
164    def on_partition(self, partition: Partition) -> None:
165        """
166        This method is called when a partition is generated.
167        1. Add the partition to the set of partitions for the stream
168        2. Log the slice if necessary
169        3. Submit the partition to the thread pool manager
170        """
171        stream_name = partition.stream_name()
172        self._streams_to_running_partitions[stream_name].add(partition)
173        cursor = self._stream_name_to_instance[stream_name].cursor
174        if self._slice_logger.should_log_slice_message(self._logger):
175            self._message_repository.emit_message(
176                self._slice_logger.create_slice_log_message(partition.to_slice())
177            )
178        self._thread_pool_manager.submit(
179            self._partition_reader.process_partition, partition, cursor
180        )
181
182    def on_partition_complete_sentinel(
183        self, sentinel: PartitionCompleteSentinel
184    ) -> Iterable[AirbyteMessage]:
185        """
186        This method is called when a partition is completed.
187        1. Close the partition
188        2. If the stream is done, mark it as such and return a stream status message
189        3. Emit messages that were added to the message repository
190        4. If there are more streams to read from, start the next partition generator
191        """
192        partition = sentinel.partition
193
194        partitions_running = self._streams_to_running_partitions[partition.stream_name()]
195        if partition in partitions_running:
196            partitions_running.remove(partition)
197            # If all partitions were generated and this was the last one, the stream is done
198            if (
199                partition.stream_name() not in self._streams_currently_generating_partitions
200                and len(partitions_running) == 0
201            ):
202                yield from self._on_stream_is_done(partition.stream_name())
203                # Try to start the next stream in the queue (may be a deferred stream)
204                if self._stream_instances_to_start_partition_generation:
205                    status_message = self.start_next_partition_generator()
206                    if status_message:
207                        yield status_message
208        yield from self._message_repository.consume_queue()
209
210    def on_record(self, record: Record) -> Iterable[AirbyteMessage]:
211        """
212        This method is called when a record is read from a partition.
213        1. Convert the record to an AirbyteMessage
214        2. If this is the first record for the stream, mark the stream as RUNNING
215        3. Increment the record counter for the stream
216        4. Ensures the cursor knows the record has been successfully emitted
217        5. Emit the message
218        6. Emit messages that were added to the message repository
219        """
220        # Do not pass a transformer or a schema
221        # AbstractStreams are expected to return data as they are expected.
222        # Any transformation on the data should be done before reaching this point
223        message = stream_data_to_airbyte_message(
224            stream_name=record.stream_name,
225            data_or_message=record.data,
226            file_reference=record.file_reference,
227        )
228        stream = self._stream_name_to_instance[record.stream_name]
229
230        if message.type == MessageType.RECORD:
231            if self._record_counter[stream.name] == 0:
232                self._logger.info(f"Marking stream {stream.name} as RUNNING")
233                yield stream_status_as_airbyte_message(
234                    stream.as_airbyte_stream(), AirbyteStreamStatus.RUNNING
235                )
236            self._record_counter[stream.name] += 1
237        yield message
238        yield from self._message_repository.consume_queue()
239
240    def on_exception(self, exception: StreamThreadException) -> Iterable[AirbyteMessage]:
241        """
242        This method is called when an exception is raised.
243        1. Stop all running streams
244        2. Raise the exception
245        """
246        self._flag_exception(exception.stream_name, exception.exception)
247        self._logger.exception(
248            f"Exception while syncing stream {exception.stream_name}", exc_info=exception.exception
249        )
250
251        stream_descriptor = StreamDescriptor(name=exception.stream_name)
252        if isinstance(exception.exception, AirbyteTracedException):
253            yield exception.exception.as_airbyte_message(stream_descriptor=stream_descriptor)
254        else:
255            yield AirbyteTracedException.from_exception(
256                exception.exception,
257                stream_descriptor=stream_descriptor,
258                message=f"An unexpected error occurred in stream {exception.stream_name}: {type(exception.exception).__name__}",
259            ).as_airbyte_message()
260
261    def _flag_exception(self, stream_name: str, exception: Exception) -> None:
262        self._exceptions_per_stream_name.setdefault(stream_name, []).append(exception)
263
264    def start_next_partition_generator(self) -> Optional[AirbyteMessage]:
265        """
266        Submits the next partition generator to the thread pool.
267
268        A stream will be deferred (moved to end of queue) if:
269        1. The stream itself has block_simultaneous_read=True AND is already active
270        2. Any parent stream has block_simultaneous_read=True AND is currently active
271
272        This prevents simultaneous reads of streams that shouldn't be accessed concurrently.
273
274        :return: A status message if a partition generator was started, otherwise None
275        """
276        if not self._stream_instances_to_start_partition_generation:
277            return None
278
279        # Enforce the concurrent generator cap so at least one worker slot is always available
280        # for partition reading. Recovery is guaranteed: on_partition_generation_completed
281        # decrements the count before calling here, so the guard always passes there.
282        if (
283            self._max_concurrent_partition_generators is not None
284            and len(self._streams_currently_generating_partitions)
285            >= self._max_concurrent_partition_generators
286        ):
287            self._logger.debug(
288                f"Concurrent partition generator cap ({self._max_concurrent_partition_generators}) reached "
289                f"({len(self._streams_currently_generating_partitions)} active). Deferring next generator start."
290            )
291            return None
292
293        # Remember initial queue size to avoid infinite loops if all streams are blocked
294        max_attempts = len(self._stream_instances_to_start_partition_generation)
295        attempts = 0
296
297        while self._stream_instances_to_start_partition_generation and attempts < max_attempts:
298            attempts += 1
299
300            # Pop the first stream from the queue
301            stream = self._stream_instances_to_start_partition_generation.pop(0)
302            stream_name = stream.name
303            stream_group = self._stream_block_simultaneous_read.get(stream_name, "")
304
305            # Check if this stream has a blocking group and is already active as parent stream
306            # (i.e. being read from during partition generation for another stream)
307            if stream_group and stream_name in self._active_stream_names:
308                # Add back to the END of the queue for retry later
309                self._stream_instances_to_start_partition_generation.append(stream)
310                self._logger.info(
311                    f"Deferring stream '{stream_name}' (group '{stream_group}') because it's already active. Trying next stream."
312                )
313                continue  # Try the next stream in the queue
314
315            # Check if this stream's group is already active (another stream in the same group is running)
316            if (
317                stream_group
318                and stream_group in self._active_groups
319                and self._active_groups[stream_group]
320            ):
321                # Add back to the END of the queue for retry later
322                self._stream_instances_to_start_partition_generation.append(stream)
323                active_streams_in_group = self._active_groups[stream_group]
324                self._logger.info(
325                    f"Deferring stream '{stream_name}' (group '{stream_group}') because other stream(s) "
326                    f"{active_streams_in_group} in the same group are active. Trying next stream."
327                )
328                continue  # Try the next stream in the queue
329
330            # Check if any parent streams have a blocking group and are currently active
331            parent_streams = self._collect_all_parent_stream_names(stream_name)
332            blocked_by_parents = [
333                p
334                for p in parent_streams
335                if self._stream_block_simultaneous_read.get(p, "")
336                and p in self._active_stream_names
337            ]
338
339            if blocked_by_parents:
340                # Add back to the END of the queue for retry later
341                self._stream_instances_to_start_partition_generation.append(stream)
342                parent_groups = {
343                    self._stream_block_simultaneous_read.get(p, "") for p in blocked_by_parents
344                }
345                self._logger.info(
346                    f"Deferring stream '{stream_name}' because parent stream(s) "
347                    f"{blocked_by_parents} (groups {parent_groups}) are active. Trying next stream."
348                )
349                continue  # Try the next stream in the queue
350
351            # No blocking - start this stream
352            # Mark stream as active before starting
353            self._active_stream_names.add(stream_name)
354            self._streams_currently_generating_partitions.append(stream_name)
355
356            # Track this stream in its group if it has one
357            if stream_group:
358                if stream_group not in self._active_groups:
359                    self._active_groups[stream_group] = set()
360                self._active_groups[stream_group].add(stream_name)
361                self._logger.debug(f"Added '{stream_name}' to active group '{stream_group}'")
362
363            # Also mark all parent streams as active (they will be read from during partition generation)
364            for parent_stream_name in parent_streams:
365                parent_group = self._stream_block_simultaneous_read.get(parent_stream_name, "")
366                if parent_group:
367                    self._active_stream_names.add(parent_stream_name)
368                    if parent_group not in self._active_groups:
369                        self._active_groups[parent_group] = set()
370                    self._active_groups[parent_group].add(parent_stream_name)
371                    self._logger.info(
372                        f"Marking parent stream '{parent_stream_name}' (group '{parent_group}') as active "
373                        f"(will be read during partition generation for '{stream_name}')"
374                    )
375
376            self._thread_pool_manager.submit(self._partition_enqueuer.generate_partitions, stream)
377            self._logger.info(f"Marking stream {stream_name} as STARTED")
378            self._logger.info(f"Syncing stream: {stream_name}")
379            return stream_status_as_airbyte_message(
380                stream.as_airbyte_stream(),
381                AirbyteStreamStatus.STARTED,
382            )
383
384        # All streams in the queue are currently blocked
385        return None
386
387    def is_done(self) -> bool:
388        """
389        This method is called to check if the sync is done.
390        The sync is done when:
391        1. There are no more streams generating partitions
392        2. There are no more streams to read from
393        3. All partitions for all streams are closed
394        """
395        is_done = all(
396            [
397                self._is_stream_done(stream_name)
398                for stream_name in self._stream_name_to_instance.keys()
399            ]
400        )
401        if is_done and self._stream_instances_to_start_partition_generation:
402            stuck_stream_names = [
403                s.name for s in self._stream_instances_to_start_partition_generation
404            ]
405            raise AirbyteTracedException(
406                message="Partition generation queue is not empty after all streams completed.",
407                internal_message=f"Streams {stuck_stream_names} remained in the partition generation queue after all streams were marked done.",
408                failure_type=FailureType.system_error,
409            )
410        if is_done and self._active_groups:
411            raise AirbyteTracedException(
412                message="Active stream groups are not empty after all streams completed.",
413                internal_message=f"Groups {dict(self._active_groups)} still active after all streams were marked done.",
414                failure_type=FailureType.system_error,
415            )
416        if is_done and self._exceptions_per_stream_name:
417            error_message = generate_failed_streams_error_message(self._exceptions_per_stream_name)
418            self._logger.info(error_message)
419            # We still raise at least one exception when a stream raises an exception because the platform currently relies
420            # on a non-zero exit code to determine if a sync attempt has failed. We also raise the exception as a config_error
421            # type because this combined error isn't actionable, but rather the previously emitted individual errors.
422            raise AirbyteTracedException(
423                message=error_message,
424                internal_message="Concurrent read failure",
425                failure_type=FailureType.config_error,
426            )
427        return is_done
428
429    def _is_stream_done(self, stream_name: str) -> bool:
430        return stream_name in self._streams_done
431
432    def _collect_all_parent_stream_names(self, stream_name: str) -> Set[str]:
433        """Recursively collect all parent stream names for a given stream.
434
435        For example, if we have: epics -> issues -> comments
436        Then for comments, this returns {issues, epics}.
437        """
438        parent_names: Set[str] = set()
439        stream = self._stream_name_to_instance.get(stream_name)
440
441        if not stream:
442            return parent_names
443
444        partition_router = (
445            stream.get_partition_router() if isinstance(stream, DefaultStream) else None
446        )
447        routers = [partition_router] if partition_router is not None else []
448        while routers:
449            router = routers.pop()
450            if isinstance(router, GroupingPartitionRouter):
451                routers.append(router.underlying_partition_router)
452            elif isinstance(router, UnionPartitionRouter):
453                routers.extend(router.partition_routers)
454            elif isinstance(router, CartesianProductStreamSlicer):
455                routers.extend(router.stream_slicers)
456            elif isinstance(router, SubstreamPartitionRouter):
457                for parent_config in router.parent_stream_configs:
458                    parent_name = parent_config.stream.name
459                    parent_names.add(parent_name)
460                    parent_names.update(self._collect_all_parent_stream_names(parent_name))
461
462        return parent_names
463
464    def _on_stream_is_done(self, stream_name: str) -> Iterable[AirbyteMessage]:
465        self._logger.info(
466            f"Read {self._record_counter[stream_name]} records from {stream_name} stream"
467        )
468        self._logger.info(f"Marking stream {stream_name} as STOPPED")
469        stream = self._stream_name_to_instance[stream_name]
470        stream.cursor.ensure_at_least_one_state_emitted()
471        yield from self._message_repository.consume_queue()
472        self._logger.info(f"Finished syncing {stream.name}")
473        self._streams_done.add(stream_name)
474        stream_status = (
475            AirbyteStreamStatus.INCOMPLETE
476            if self._exceptions_per_stream_name.get(stream_name, [])
477            else AirbyteStreamStatus.COMPLETE
478        )
479        yield stream_status_as_airbyte_message(stream.as_airbyte_stream(), stream_status)
480
481        # Remove only this stream from active set (NOT parents)
482        if stream_name in self._active_stream_names:
483            self._active_stream_names.discard(stream_name)
484
485            # Remove from active groups
486            stream_group = self._stream_block_simultaneous_read.get(stream_name, "")
487            if stream_group:
488                if stream_group in self._active_groups:
489                    self._active_groups[stream_group].discard(stream_name)
490                    if not self._active_groups[stream_group]:
491                        del self._active_groups[stream_group]
492                self._logger.info(
493                    f"Stream '{stream_name}' (group '{stream_group}') is no longer active. "
494                    f"Blocked streams in the queue will be retried on next start_next_partition_generator call."
495                )