airbyte-cdk 6.5.3rc2__py3-none-any.whl → 6.6.0__py3-none-any.whl

airbyte_cdk/exception_handler.py

@@ -11,7 +11,9 @@ from airbyte_cdk.utils.airbyte_secrets_utils import filter_secrets
 from airbyte_cdk.utils.traced_exception import AirbyteTracedException
 
 
-def assemble_uncaught_exception(exception_type: type[BaseException], exception_value: BaseException) -> AirbyteTracedException:
+def assemble_uncaught_exception(
+    exception_type: type[BaseException], exception_value: BaseException
+) -> AirbyteTracedException:
     if issubclass(exception_type, AirbyteTracedException):
         return exception_value  # type: ignore  # validated as part of the previous line
     return AirbyteTracedException.from_exception(exception_value)
@@ -23,7 +25,11 @@ def init_uncaught_exception_handler(logger: logging.Logger) -> None:
     printed to the console without having secrets removed.
     """
 
-    def hook_fn(exception_type: type[BaseException], exception_value: BaseException, traceback_: Optional[TracebackType]) -> Any:
+    def hook_fn(
+        exception_type: type[BaseException],
+        exception_value: BaseException,
+        traceback_: Optional[TracebackType],
+    ) -> Any:
         # For developer ergonomics, we want to see the stack trace in the logs when we do a ctrl-c
         if issubclass(exception_type, KeyboardInterrupt):
             sys.__excepthook__(exception_type, exception_value, traceback_)
@@ -41,6 +47,10 @@ def init_uncaught_exception_handler(logger: logging.Logger) -> None:
 
 def generate_failed_streams_error_message(stream_failures: Mapping[str, List[Exception]]) -> str:
     failures = "\n".join(
-        [f"{stream}: {filter_secrets(exception.__repr__())}" for stream, exceptions in stream_failures.items() for exception in exceptions]
+        [
+            f"{stream}: {filter_secrets(exception.__repr__())}"
+            for stream, exceptions in stream_failures.items()
+            for exception in exceptions
+        ]
     )
     return f"During the sync, the following streams did not sync successfully: {failures}"

airbyte_cdk/logger.py

@@ -7,7 +7,13 @@ import logging
 import logging.config
 from typing import Any, Callable, Mapping, Optional, Tuple
 
-from airbyte_cdk.models import AirbyteLogMessage, AirbyteMessage, AirbyteMessageSerializer, Level, Type
+from airbyte_cdk.models import (
+    AirbyteLogMessage,
+    AirbyteMessage,
+    AirbyteMessageSerializer,
+    Level,
+    Type,
+)
 from airbyte_cdk.utils.airbyte_secrets_utils import filter_secrets
 from orjson import orjson
 
@@ -68,7 +74,9 @@ class AirbyteLogFormatter(logging.Formatter):
         else:
             message = super().format(record)
             message = filter_secrets(message)
-            log_message = AirbyteMessage(type=Type.LOG, log=AirbyteLogMessage(level=airbyte_level, message=message))
+            log_message = AirbyteMessage(
+                type=Type.LOG, log=AirbyteLogMessage(level=airbyte_level, message=message)
+            )
             return orjson.dumps(AirbyteMessageSerializer.dump(log_message)).decode()  # type: ignore[no-any-return] # orjson.dumps(message).decode() always returns string
 
     @staticmethod

airbyte_cdk/models/airbyte_protocol.py

@@ -6,9 +6,11 @@ from dataclasses import InitVar, dataclass
 from typing import Annotated, Any, Dict, List, Mapping, Optional, Union
 
 from airbyte_cdk.models.file_transfer_record_message import AirbyteFileTransferRecordMessage
-from airbyte_protocol_dataclasses.models import *
+from airbyte_protocol_dataclasses.models import *  # noqa: F403  # Allow '*'
 from serpyco_rs.metadata import Alias
 
+# ruff: noqa: F405  # ignore fuzzy import issues with 'import *'
+
 
 @dataclass
 class AirbyteStateBlob:
@@ -42,7 +44,11 @@ class AirbyteStateBlob:
             setattr(self, key, value)
 
     def __eq__(self, other: object) -> bool:
-        return False if not isinstance(other, AirbyteStateBlob) else bool(self.__dict__ == other.__dict__)
+        return (
+            False
+            if not isinstance(other, AirbyteStateBlob)
+            else bool(self.__dict__ == other.__dict__)
+        )
 
 
 # The following dataclasses have been redeclared to include the new version of AirbyteStateBlob
@@ -62,9 +68,9 @@ class AirbyteGlobalState:
 class AirbyteStateMessage:
     type: Optional[AirbyteStateType] = None  # type: ignore [name-defined]
     stream: Optional[AirbyteStreamState] = None
-    global_: Annotated[
-        AirbyteGlobalState | None, Alias("global")
-    ] = None  # "global" is a reserved keyword in python ⇒ Alias is used for (de-)serialization
+    global_: Annotated[AirbyteGlobalState | None, Alias("global")] = (
+        None  # "global" is a reserved keyword in python ⇒ Alias is used for (de-)serialization
+    )
     data: Optional[Dict[str, Any]] = None
     sourceStats: Optional[AirbyteStateStats] = None  # type: ignore [name-defined]
     destinationStats: Optional[AirbyteStateStats] = None  # type: ignore [name-defined]

airbyte_cdk/models/airbyte_protocol_serializers.py

@@ -30,9 +30,15 @@ def custom_type_resolver(t: type) -> CustomType[AirbyteStateBlob, Dict[str, Any]
     return AirbyteStateBlobType() if t is AirbyteStateBlob else None
 
 
-AirbyteStreamStateSerializer = Serializer(AirbyteStreamState, omit_none=True, custom_type_resolver=custom_type_resolver)
-AirbyteStateMessageSerializer = Serializer(AirbyteStateMessage, omit_none=True, custom_type_resolver=custom_type_resolver)
-AirbyteMessageSerializer = Serializer(AirbyteMessage, omit_none=True, custom_type_resolver=custom_type_resolver)
+AirbyteStreamStateSerializer = Serializer(
+    AirbyteStreamState, omit_none=True, custom_type_resolver=custom_type_resolver
+)
+AirbyteStateMessageSerializer = Serializer(
+    AirbyteStateMessage, omit_none=True, custom_type_resolver=custom_type_resolver
+)
+AirbyteMessageSerializer = Serializer(
+    AirbyteMessage, omit_none=True, custom_type_resolver=custom_type_resolver
+)
 ConfiguredAirbyteCatalogSerializer = Serializer(ConfiguredAirbyteCatalog, omit_none=True)
 ConfiguredAirbyteStreamSerializer = Serializer(ConfiguredAirbyteStream, omit_none=True)
 ConnectorSpecificationSerializer = Serializer(ConnectorSpecification, omit_none=True)

airbyte_cdk/models/well_known_types.py

@@ -2,4 +2,4 @@
 # Copyright (c) 2023 Airbyte, Inc., all rights reserved.
 #
 
-from airbyte_protocol_dataclasses.models.well_known_types import *
+from airbyte_protocol_dataclasses.models.well_known_types import *  # noqa: F403  # Allow '*'

airbyte_cdk/sources/abstract_source.py

@@ -4,7 +4,18 @@
 
 import logging
 from abc import ABC, abstractmethod
-from typing import Any, Dict, Iterable, Iterator, List, Mapping, MutableMapping, Optional, Tuple, Union
+from typing import (
+    Any,
+    Dict,
+    Iterable,
+    Iterator,
+    List,
+    Mapping,
+    MutableMapping,
+    Optional,
+    Tuple,
+    Union,
+)
 
 from airbyte_cdk.exception_handler import generate_failed_streams_error_message
 from airbyte_cdk.models import (
@@ -30,7 +41,9 @@ from airbyte_cdk.sources.utils.record_helper import stream_data_to_airbyte_messa
 from airbyte_cdk.sources.utils.schema_helpers import InternalConfig, split_config
 from airbyte_cdk.sources.utils.slice_logger import DebugSliceLogger, SliceLogger
 from airbyte_cdk.utils.event_timing import create_timer
-from airbyte_cdk.utils.stream_status_utils import as_airbyte_message as stream_status_as_airbyte_message
+from airbyte_cdk.utils.stream_status_utils import (
+    as_airbyte_message as stream_status_as_airbyte_message,
+)
 from airbyte_cdk.utils.traced_exception import AirbyteTracedException
 
 _default_message_repository = InMemoryMessageRepository()
@@ -43,7 +56,9 @@ class AbstractSource(Source, ABC):
     """
 
     @abstractmethod
-    def check_connection(self, logger: logging.Logger, config: Mapping[str, Any]) -> Tuple[bool, Optional[Any]]:
+    def check_connection(
+        self, logger: logging.Logger, config: Mapping[str, Any]
+    ) -> Tuple[bool, Optional[Any]]:
         """
         :param logger: source logger
         :param config: The user-provided configuration as specified by the source's spec.
@@ -109,7 +124,9 @@ class AbstractSource(Source, ABC):
                     # Used direct reference to `stream_instance` instead of `is_stream_exist` to avoid mypy type checking errors
                     if not stream_instance:
                         if not self.raise_exception_on_missing_stream:
-                            yield stream_status_as_airbyte_message(configured_stream.stream, AirbyteStreamStatus.INCOMPLETE)
+                            yield stream_status_as_airbyte_message(
+                                configured_stream.stream, AirbyteStreamStatus.INCOMPLETE
+                            )
                             continue
 
                         error_message = (
@@ -129,7 +146,9 @@ class AbstractSource(Source, ABC):
 
                     timer.start_event(f"Syncing stream {configured_stream.stream.name}")
                     logger.info(f"Marking stream {configured_stream.stream.name} as STARTED")
-                    yield stream_status_as_airbyte_message(configured_stream.stream, AirbyteStreamStatus.STARTED)
+                    yield stream_status_as_airbyte_message(
+                        configured_stream.stream, AirbyteStreamStatus.STARTED
+                    )
                     yield from self._read_stream(
                         logger=logger,
                         stream_instance=stream_instance,
@@ -138,13 +157,19 @@ class AbstractSource(Source, ABC):
                         internal_config=internal_config,
                     )
                     logger.info(f"Marking stream {configured_stream.stream.name} as STOPPED")
-                    yield stream_status_as_airbyte_message(configured_stream.stream, AirbyteStreamStatus.COMPLETE)
+                    yield stream_status_as_airbyte_message(
+                        configured_stream.stream, AirbyteStreamStatus.COMPLETE
+                    )
 
                 except Exception as e:
                     yield from self._emit_queued_messages()
-                    logger.exception(f"Encountered an exception while reading stream {configured_stream.stream.name}")
+                    logger.exception(
+                        f"Encountered an exception while reading stream {configured_stream.stream.name}"
+                    )
                     logger.info(f"Marking stream {configured_stream.stream.name} as STOPPED")
-                    yield stream_status_as_airbyte_message(configured_stream.stream, AirbyteStreamStatus.INCOMPLETE)
+                    yield stream_status_as_airbyte_message(
+                        configured_stream.stream, AirbyteStreamStatus.INCOMPLETE
+                    )
 
                     stream_descriptor = StreamDescriptor(name=configured_stream.stream.name)
 
@@ -152,10 +177,14 @@ class AbstractSource(Source, ABC):
                         traced_exception = e
                         info_message = f"Stopping sync on error from stream {configured_stream.stream.name} because {self.name} does not support continuing syncs on error."
                     else:
-                        traced_exception = self._serialize_exception(stream_descriptor, e, stream_instance=stream_instance)
+                        traced_exception = self._serialize_exception(
+                            stream_descriptor, e, stream_instance=stream_instance
+                        )
                         info_message = f"{self.name} does not support continuing syncs on error from stream {configured_stream.stream.name}"
 
-                    yield traced_exception.as_sanitized_airbyte_message(stream_descriptor=stream_descriptor)
+                    yield traced_exception.as_sanitized_airbyte_message(
+                        stream_descriptor=stream_descriptor
+                    )
                     stream_name_to_exception[stream_instance.name] = traced_exception  # type: ignore # use configured_stream if stream_instance is None
                     if self.stop_sync_on_stream_failure:
                         logger.info(info_message)
@@ -169,12 +198,16 @@ class AbstractSource(Source, ABC):
                         logger.info(timer.report())
 
         if len(stream_name_to_exception) > 0:
-            error_message = generate_failed_streams_error_message({key: [value] for key, value in stream_name_to_exception.items()})  # type: ignore  # for some reason, mypy can't figure out the types for key and value
+            error_message = generate_failed_streams_error_message(
+                {key: [value] for key, value in stream_name_to_exception.items()}
+            )  # type: ignore  # for some reason, mypy can't figure out the types for key and value
             logger.info(error_message)
             # We still raise at least one exception when a stream raises an exception because the platform currently relies
             # on a non-zero exit code to determine if a sync attempt has failed. We also raise the exception as a config_error
             # type because this combined error isn't actionable, but rather the previously emitted individual errors.
-            raise AirbyteTracedException(message=error_message, failure_type=FailureType.config_error)
+            raise AirbyteTracedException(
+                message=error_message, failure_type=FailureType.config_error
+            )
         logger.info(f"Finished syncing {self.name}")
 
     @staticmethod
@@ -183,7 +216,9 @@ class AbstractSource(Source, ABC):
     ) -> AirbyteTracedException:
         display_message = stream_instance.get_error_display_message(e) if stream_instance else None
         if display_message:
-            return AirbyteTracedException.from_exception(e, message=display_message, stream_descriptor=stream_descriptor)
+            return AirbyteTracedException.from_exception(
+                e, message=display_message, stream_descriptor=stream_descriptor
+            )
         return AirbyteTracedException.from_exception(e, stream_descriptor=stream_descriptor)
 
     @property
@@ -199,7 +234,9 @@ class AbstractSource(Source, ABC):
         internal_config: InternalConfig,
     ) -> Iterator[AirbyteMessage]:
         if internal_config.page_size and isinstance(stream_instance, HttpStream):
-            logger.info(f"Setting page size for {stream_instance.name} to {internal_config.page_size}")
+            logger.info(
+                f"Setting page size for {stream_instance.name} to {internal_config.page_size}"
+            )
             stream_instance.page_size = internal_config.page_size
         logger.debug(
             f"Syncing configured stream: {configured_stream.stream.name}",
@@ -243,7 +280,9 @@ class AbstractSource(Source, ABC):
                 if record_counter == 1:
                     logger.info(f"Marking stream {stream_name} as RUNNING")
                     # If we just read the first record of the stream, emit the transition to the RUNNING state
-                    yield stream_status_as_airbyte_message(configured_stream.stream, AirbyteStreamStatus.RUNNING)
+                    yield stream_status_as_airbyte_message(
+                        configured_stream.stream, AirbyteStreamStatus.RUNNING
+                    )
             yield from self._emit_queued_messages()
             yield record
 
@@ -254,7 +293,9 @@ class AbstractSource(Source, ABC):
             yield from self.message_repository.consume_queue()
         return
 
-    def _get_message(self, record_data_or_message: Union[StreamData, AirbyteMessage], stream: Stream) -> AirbyteMessage:
+    def _get_message(
+        self, record_data_or_message: Union[StreamData, AirbyteMessage], stream: Stream
+    ) -> AirbyteMessage:
         """
         Converts the input to an AirbyteMessage if it is a StreamData. Returns the input as is if it is already an AirbyteMessage
         """
@@ -262,7 +303,12 @@ class AbstractSource(Source, ABC):
             case AirbyteMessage():
                 return record_data_or_message
             case _:
-                return stream_data_to_airbyte_message(stream.name, record_data_or_message, stream.transformer, stream.get_json_schema())
+                return stream_data_to_airbyte_message(
+                    stream.name,
+                    record_data_or_message,
+                    stream.transformer,
+                    stream.get_json_schema(),
+                )
 
     @property
     def message_repository(self) -> Union[None, MessageRepository]:

airbyte_cdk/sources/concurrent_source/concurrent_read_processor.py

@@ -7,7 +7,9 @@ from typing import Dict, Iterable, List, Optional, Set
 from airbyte_cdk.exception_handler import generate_failed_streams_error_message
 from airbyte_cdk.models import AirbyteMessage, AirbyteStreamStatus, FailureType, StreamDescriptor
 from airbyte_cdk.models import Type as MessageType
-from airbyte_cdk.sources.concurrent_source.partition_generation_completed_sentinel import PartitionGenerationCompletedSentinel
+from airbyte_cdk.sources.concurrent_source.partition_generation_completed_sentinel import (
+    PartitionGenerationCompletedSentinel,
+)
 from airbyte_cdk.sources.concurrent_source.stream_thread_exception import StreamThreadException
 from airbyte_cdk.sources.concurrent_source.thread_pool_manager import ThreadPoolManager
 from airbyte_cdk.sources.message import MessageRepository
@@ -20,7 +22,9 @@ from airbyte_cdk.sources.streams.concurrent.partitions.types import PartitionCom
 from airbyte_cdk.sources.utils.record_helper import stream_data_to_airbyte_message
 from airbyte_cdk.sources.utils.slice_logger import SliceLogger
 from airbyte_cdk.utils import AirbyteTracedException
-from airbyte_cdk.utils.stream_status_utils import as_airbyte_message as stream_status_as_airbyte_message
+from airbyte_cdk.utils.stream_status_utils import (
+    as_airbyte_message as stream_status_as_airbyte_message,
+)
 
 
 class ConcurrentReadProcessor:
@@ -61,7 +65,9 @@ class ConcurrentReadProcessor:
         self._streams_done: Set[str] = set()
         self._exceptions_per_stream_name: dict[str, List[Exception]] = {}
 
-    def on_partition_generation_completed(self, sentinel: PartitionGenerationCompletedSentinel) -> Iterable[AirbyteMessage]:
+    def on_partition_generation_completed(
+        self, sentinel: PartitionGenerationCompletedSentinel
+    ) -> Iterable[AirbyteMessage]:
         """
         This method is called when a partition generation is completed.
         1. Remove the stream from the list of streams currently generating partitions
@@ -72,7 +78,10 @@ class ConcurrentReadProcessor:
         self._streams_currently_generating_partitions.remove(sentinel.stream.name)
         # It is possible for the stream to already be done if no partitions were generated
         # If the partition generation process was completed and there are no partitions left to process, the stream is done
-        if self._is_stream_done(stream_name) or len(self._streams_to_running_partitions[stream_name]) == 0:
+        if (
+            self._is_stream_done(stream_name)
+            or len(self._streams_to_running_partitions[stream_name]) == 0
+        ):
             yield from self._on_stream_is_done(stream_name)
         if self._stream_instances_to_start_partition_generation:
             yield self.start_next_partition_generator()  # type:ignore # None may be yielded
@@ -87,10 +96,14 @@ class ConcurrentReadProcessor:
         stream_name = partition.stream_name()
         self._streams_to_running_partitions[stream_name].add(partition)
         if self._slice_logger.should_log_slice_message(self._logger):
-            self._message_repository.emit_message(self._slice_logger.create_slice_log_message(partition.to_slice()))
+            self._message_repository.emit_message(
+                self._slice_logger.create_slice_log_message(partition.to_slice())
+            )
         self._thread_pool_manager.submit(self._partition_reader.process_partition, partition)
 
-    def on_partition_complete_sentinel(self, sentinel: PartitionCompleteSentinel) -> Iterable[AirbyteMessage]:
+    def on_partition_complete_sentinel(
+        self, sentinel: PartitionCompleteSentinel
+    ) -> Iterable[AirbyteMessage]:
         """
         This method is called when a partition is completed.
         1. Close the partition
@@ -112,7 +125,10 @@ class ConcurrentReadProcessor:
             if partition in partitions_running:
                 partitions_running.remove(partition)
                 # If all partitions were generated and this was the last one, the stream is done
-                if partition.stream_name() not in self._streams_currently_generating_partitions and len(partitions_running) == 0:
+                if (
+                    partition.stream_name() not in self._streams_currently_generating_partitions
+                    and len(partitions_running) == 0
+                ):
                     yield from self._on_stream_is_done(partition.stream_name())
             yield from self._message_repository.consume_queue()
 
@@ -139,7 +155,9 @@ class ConcurrentReadProcessor:
         if message.type == MessageType.RECORD:
             if self._record_counter[stream.name] == 0:
                 self._logger.info(f"Marking stream {stream.name} as RUNNING")
-                yield stream_status_as_airbyte_message(stream.as_airbyte_stream(), AirbyteStreamStatus.RUNNING)
+                yield stream_status_as_airbyte_message(
+                    stream.as_airbyte_stream(), AirbyteStreamStatus.RUNNING
+                )
             self._record_counter[stream.name] += 1
             stream.cursor.observe(record)
         yield message
@@ -152,13 +170,17 @@ class ConcurrentReadProcessor:
         2. Raise the exception
         """
         self._flag_exception(exception.stream_name, exception.exception)
-        self._logger.exception(f"Exception while syncing stream {exception.stream_name}", exc_info=exception.exception)
+        self._logger.exception(
+            f"Exception while syncing stream {exception.stream_name}", exc_info=exception.exception
+        )
 
         stream_descriptor = StreamDescriptor(name=exception.stream_name)
         if isinstance(exception.exception, AirbyteTracedException):
             yield exception.exception.as_airbyte_message(stream_descriptor=stream_descriptor)
         else:
-            yield AirbyteTracedException.from_exception(exception, stream_descriptor=stream_descriptor).as_airbyte_message()
+            yield AirbyteTracedException.from_exception(
+                exception, stream_descriptor=stream_descriptor
+            ).as_airbyte_message()
 
     def _flag_exception(self, stream_name: str, exception: Exception) -> None:
         self._exceptions_per_stream_name.setdefault(stream_name, []).append(exception)
@@ -192,7 +214,12 @@ class ConcurrentReadProcessor:
         2. There are no more streams to read from
         3. All partitions for all streams are closed
         """
-        is_done = all([self._is_stream_done(stream_name) for stream_name in self._stream_name_to_instance.keys()])
+        is_done = all(
+            [
+                self._is_stream_done(stream_name)
+                for stream_name in self._stream_name_to_instance.keys()
+            ]
+        )
         if is_done and self._exceptions_per_stream_name:
             error_message = generate_failed_streams_error_message(self._exceptions_per_stream_name)
             self._logger.info(error_message)
@@ -200,7 +227,9 @@ class ConcurrentReadProcessor:
             # on a non-zero exit code to determine if a sync attempt has failed. We also raise the exception as a config_error
             # type because this combined error isn't actionable, but rather the previously emitted individual errors.
             raise AirbyteTracedException(
-                message=error_message, internal_message="Concurrent read failure", failure_type=FailureType.config_error
+                message=error_message,
+                internal_message="Concurrent read failure",
+                failure_type=FailureType.config_error,
             )
         return is_done
 
@@ -208,7 +237,9 @@ class ConcurrentReadProcessor:
         return stream_name in self._streams_done
 
     def _on_stream_is_done(self, stream_name: str) -> Iterable[AirbyteMessage]:
-        self._logger.info(f"Read {self._record_counter[stream_name]} records from {stream_name} stream")
+        self._logger.info(
+            f"Read {self._record_counter[stream_name]} records from {stream_name} stream"
+        )
         self._logger.info(f"Marking stream {stream_name} as STOPPED")
         stream = self._stream_name_to_instance[stream_name]
         stream.cursor.ensure_at_least_one_state_emitted()
@@ -216,6 +247,8 @@ class ConcurrentReadProcessor:
         self._logger.info(f"Finished syncing {stream.name}")
         self._streams_done.add(stream_name)
         stream_status = (
-            AirbyteStreamStatus.INCOMPLETE if self._exceptions_per_stream_name.get(stream_name, []) else AirbyteStreamStatus.COMPLETE
+            AirbyteStreamStatus.INCOMPLETE
+            if self._exceptions_per_stream_name.get(stream_name, [])
+            else AirbyteStreamStatus.COMPLETE
         )
         yield stream_status_as_airbyte_message(stream.as_airbyte_stream(), stream_status)

airbyte_cdk/sources/concurrent_source/concurrent_source.py

@@ -8,7 +8,9 @@ from typing import Iterable, Iterator, List
 
 from airbyte_cdk.models import AirbyteMessage
 from airbyte_cdk.sources.concurrent_source.concurrent_read_processor import ConcurrentReadProcessor
-from airbyte_cdk.sources.concurrent_source.partition_generation_completed_sentinel import PartitionGenerationCompletedSentinel
+from airbyte_cdk.sources.concurrent_source.partition_generation_completed_sentinel import (
+    PartitionGenerationCompletedSentinel,
+)
 from airbyte_cdk.sources.concurrent_source.stream_thread_exception import StreamThreadException
 from airbyte_cdk.sources.concurrent_source.thread_pool_manager import ThreadPoolManager
 from airbyte_cdk.sources.message import InMemoryMessageRepository, MessageRepository
@@ -17,7 +19,10 @@ from airbyte_cdk.sources.streams.concurrent.partition_enqueuer import PartitionE
 from airbyte_cdk.sources.streams.concurrent.partition_reader import PartitionReader
 from airbyte_cdk.sources.streams.concurrent.partitions.partition import Partition
 from airbyte_cdk.sources.streams.concurrent.partitions.record import Record
-from airbyte_cdk.sources.streams.concurrent.partitions.types import PartitionCompleteSentinel, QueueItem
+from airbyte_cdk.sources.streams.concurrent.partitions.types import (
+    PartitionCompleteSentinel,
+    QueueItem,
+)
 from airbyte_cdk.sources.utils.slice_logger import DebugSliceLogger, SliceLogger
 
 
@@ -41,14 +46,25 @@ class ConcurrentSource:
         timeout_seconds: int = DEFAULT_TIMEOUT_SECONDS,
     ) -> "ConcurrentSource":
         is_single_threaded = initial_number_of_partitions_to_generate == 1 and num_workers == 1
-        too_many_generator = not is_single_threaded and initial_number_of_partitions_to_generate >= num_workers
-        assert not too_many_generator, "It is required to have more workers than threads generating partitions"
+        too_many_generator = (
+            not is_single_threaded and initial_number_of_partitions_to_generate >= num_workers
+        )
+        assert (
+            not too_many_generator
+        ), "It is required to have more workers than threads generating partitions"
         threadpool = ThreadPoolManager(
-            concurrent.futures.ThreadPoolExecutor(max_workers=num_workers, thread_name_prefix="workerpool"),
+            concurrent.futures.ThreadPoolExecutor(
+                max_workers=num_workers, thread_name_prefix="workerpool"
+            ),
             logger,
         )
         return ConcurrentSource(
-            threadpool, logger, slice_logger, message_repository, initial_number_of_partitions_to_generate, timeout_seconds
+            threadpool,
+            logger,
+            slice_logger,
+            message_repository,
+            initial_number_of_partitions_to_generate,
+            timeout_seconds,
         )
 
     def __init__(
@@ -107,7 +123,9 @@ class ConcurrentSource:
         self._threadpool.check_for_errors_and_shutdown()
         self._logger.info("Finished syncing")
 
-    def _submit_initial_partition_generators(self, concurrent_stream_processor: ConcurrentReadProcessor) -> Iterable[AirbyteMessage]:
+    def _submit_initial_partition_generators(
+        self, concurrent_stream_processor: ConcurrentReadProcessor
+    ) -> Iterable[AirbyteMessage]:
         for _ in range(self._initial_number_partitions_to_generate):
             status_message = concurrent_stream_processor.start_next_partition_generator()
             if status_message:

airbyte_cdk/sources/concurrent_source/concurrent_source_adapter.py

@@ -15,8 +15,17 @@ from airbyte_cdk.sources.streams import Stream
 from airbyte_cdk.sources.streams.concurrent.abstract_stream import AbstractStream
 from airbyte_cdk.sources.streams.concurrent.abstract_stream_facade import AbstractStreamFacade
 from airbyte_cdk.sources.streams.concurrent.adapters import StreamFacade
-from airbyte_cdk.sources.streams.concurrent.cursor import ConcurrentCursor, Cursor, CursorField, CursorValueType, FinalStateCursor, GapType
-from airbyte_cdk.sources.streams.concurrent.state_converters.abstract_stream_state_converter import AbstractStreamStateConverter
+from airbyte_cdk.sources.streams.concurrent.cursor import (
+    ConcurrentCursor,
+    Cursor,
+    CursorField,
+    CursorValueType,
+    FinalStateCursor,
+    GapType,
+)
+from airbyte_cdk.sources.streams.concurrent.state_converters.abstract_stream_state_converter import (
+    AbstractStreamStateConverter,
+)
 
 DEFAULT_LOOKBACK_SECONDS = 0
 
@@ -43,14 +52,20 @@ class ConcurrentSourceAdapter(AbstractSource, ABC):
         abstract_streams = self._select_abstract_streams(config, catalog)
         concurrent_stream_names = {stream.name for stream in abstract_streams}
         configured_catalog_for_regular_streams = ConfiguredAirbyteCatalog(
-            streams=[stream for stream in catalog.streams if stream.stream.name not in concurrent_stream_names]
+            streams=[
+                stream
+                for stream in catalog.streams
+                if stream.stream.name not in concurrent_stream_names
+            ]
         )
         if abstract_streams:
             yield from self._concurrent_source.read(abstract_streams)
         if configured_catalog_for_regular_streams.streams:
             yield from super().read(logger, config, configured_catalog_for_regular_streams, state)
 
-    def _select_abstract_streams(self, config: Mapping[str, Any], configured_catalog: ConfiguredAirbyteCatalog) -> List[AbstractStream]:
+    def _select_abstract_streams(
+        self, config: Mapping[str, Any], configured_catalog: ConfiguredAirbyteCatalog
+    ) -> List[AbstractStream]:
         """
         Selects streams that can be processed concurrently and returns their abstract representations.
         """
@@ -67,7 +82,11 @@ class ConcurrentSourceAdapter(AbstractSource, ABC):
         return abstract_streams
 
     def convert_to_concurrent_stream(
-        self, logger: logging.Logger, stream: Stream, state_manager: ConnectorStateManager, cursor: Optional[Cursor] = None
+        self,
+        logger: logging.Logger,
+        stream: Stream,
+        state_manager: ConnectorStateManager,
+        cursor: Optional[Cursor] = None,
     ) -> Stream:
         """
         Prepares a stream for concurrent processing by initializing or assigning a cursor,
@@ -106,7 +125,9 @@ class ConcurrentSourceAdapter(AbstractSource, ABC):
 
         if cursor_field_name:
             if not isinstance(cursor_field_name, str):
-                raise ValueError(f"Cursor field type must be a string, but received {type(cursor_field_name).__name__}.")
+                raise ValueError(
+                    f"Cursor field type must be a string, but received {type(cursor_field_name).__name__}."
+                )
 
             return ConcurrentCursor(
                 stream.name,

airbyte_cdk/sources/concurrent_source/thread_pool_manager.py

@@ -37,7 +37,9 @@ class ThreadPoolManager:
     def prune_to_validate_has_reached_futures_limit(self) -> bool:
         self._prune_futures(self._futures)
         if len(self._futures) > self._logging_threshold:
-            self._logger.warning(f"ThreadPoolManager: The list of futures is getting bigger than expected ({len(self._futures)})")
+            self._logger.warning(
+                f"ThreadPoolManager: The list of futures is getting bigger than expected ({len(self._futures)})"
+            )
         return len(self._futures) >= self._max_concurrent_tasks
 
     def submit(self, function: Callable[..., Any], *args: Any) -> None:
@@ -92,14 +94,18 @@ class ThreadPoolManager:
             )
             self._stop_and_raise_exception(self._most_recently_seen_exception)
 
-        exceptions_from_futures = [f for f in [future.exception() for future in self._futures] if f is not None]
+        exceptions_from_futures = [
+            f for f in [future.exception() for future in self._futures] if f is not None
+        ]
         if exceptions_from_futures:
             exception = RuntimeError(f"Failed reading with errors: {exceptions_from_futures}")
             self._stop_and_raise_exception(exception)
         else:
             futures_not_done = [f for f in self._futures if not f.done()]
             if futures_not_done:
-                exception = RuntimeError(f"Failed reading with futures not done: {futures_not_done}")
+                exception = RuntimeError(
+                    f"Failed reading with futures not done: {futures_not_done}"
+                )
                 self._stop_and_raise_exception(exception)
             else:
                 self._shutdown()