WIP: asyncio-based CDK

airbyte-cdk/python/airbyte_cdk/entrypoint.py

@@ -207,10 +207,14 @@ def _emit_queued_messages(self, source: Source) -> Iterable[AirbyteMessage]:
         return
 
 
-def launch(source: Source, args: List[str]) -> None:
+def get_source_iter(source: Source, args: List[str]) -> Iterable[str]:
     source_entrypoint = AirbyteEntrypoint(source)
     parsed_args = source_entrypoint.parse_args(args)
-    for message in source_entrypoint.run(parsed_args):
+    return source_entrypoint.run(parsed_args)
+
+
+def launch(source: Source, args: List[str]) -> None:
+    for message in get_source_iter(source, args):
         print(message)
 
 

airbyte-cdk/python/airbyte_cdk/sources/abstract_source.py

@@ -4,7 +4,17 @@
 
 import logging
 from abc import ABC, abstractmethod
-from typing import Any, Dict, Iterable, Iterator, List, Mapping, MutableMapping, Optional, Tuple, Union
+from typing import (
+    Any,
+    Iterable,
+    Iterator,
+    List,
+    Mapping,
+    MutableMapping,
+    Optional,
+    Tuple,
+    Union,
+)
 
 from airbyte_cdk.models import (
     AirbyteCatalog,
@@ -28,7 +38,9 @@
 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()
@@ -41,7 +53,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.
@@ -62,22 +76,26 @@ def streams(self, config: Mapping[str, Any]) -> List[Stream]:
         """
 
     # Stream name to instance map for applying output object transformation
-    _stream_to_instance_map: Dict[str, Stream] = {}
+    _stream_to_instance_map: Mapping[str, Stream] = {}
     _slice_logger: SliceLogger = DebugSliceLogger()
 
     @property
     def name(self) -> str:
         """Source name"""
         return self.__class__.__name__
 
-    def discover(self, logger: logging.Logger, config: Mapping[str, Any]) -> AirbyteCatalog:
+    def discover(
+        self, logger: logging.Logger, config: Mapping[str, Any]
+    ) -> AirbyteCatalog:
         """Implements the Discover operation from the Airbyte Specification.
         See https://docs.airbyte.com/understanding-airbyte/airbyte-protocol/#discover.
         """
         streams = [stream.as_airbyte_stream() for stream in self.streams(config=config)]
         return AirbyteCatalog(streams=streams)
 
-    def check(self, logger: logging.Logger, config: Mapping[str, Any]) -> AirbyteConnectionStatus:
+    def check(
+        self, logger: logging.Logger, config: Mapping[str, Any]
+    ) -> AirbyteConnectionStatus:
         """Implements the Check Connection operation from the Airbyte Specification.
         See https://docs.airbyte.com/understanding-airbyte/airbyte-protocol/#check.
         """
@@ -91,15 +109,19 @@ def read(
         logger: logging.Logger,
         config: Mapping[str, Any],
         catalog: ConfiguredAirbyteCatalog,
-        state: Optional[Union[List[AirbyteStateMessage], MutableMapping[str, Any]]] = None,
+        state: Optional[
+            Union[List[AirbyteStateMessage], MutableMapping[str, Any]]
+        ] = None,
     ) -> Iterator[AirbyteMessage]:
         """Implements the Read operation from the Airbyte Specification. See https://docs.airbyte.com/understanding-airbyte/airbyte-protocol/."""
         logger.info(f"Starting syncing {self.name}")
         config, internal_config = split_config(config)
         # TODO assert all streams exist in the connector
         # get the streams once in case the connector needs to make any queries to generate them
         stream_instances = {s.name: s for s in self.streams(config)}
-        state_manager = ConnectorStateManager(stream_instance_map=stream_instances, state=state)
+        state_manager = ConnectorStateManager(
+            stream_instance_map=stream_instances, state=state
+        )
         self._stream_to_instance_map = stream_instances
 
         stream_name_to_exception: MutableMapping[str, AirbyteTracedException] = {}
@@ -117,43 +139,69 @@ def read(
 
                 try:
                     timer.start_event(f"Syncing stream {configured_stream.stream.name}")
-                    stream_is_available, reason = stream_instance.check_availability(logger, self)
+                    stream_is_available, reason = stream_instance.check_availability(
+                        logger, self
+                    )
                     if not stream_is_available:
-                        logger.warning(f"Skipped syncing stream '{stream_instance.name}' because it was unavailable. {reason}")
+                        logger.warning(
+                            f"Skipped syncing stream '{stream_instance.name}' because it was unavailable. {reason}"
+                        )
                         continue
-                    logger.info(f"Marking stream {configured_stream.stream.name} as STARTED")
-                    yield stream_status_as_airbyte_message(configured_stream.stream, AirbyteStreamStatus.STARTED)
+                    logger.info(
+                        f"Marking stream {configured_stream.stream.name} as STARTED"
+                    )
+                    yield stream_status_as_airbyte_message(
+                        configured_stream.stream, AirbyteStreamStatus.STARTED
+                    )
                     yield from self._read_stream(
                         logger=logger,
                         stream_instance=stream_instance,
                         configured_stream=configured_stream,
                         state_manager=state_manager,
                         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)
+                    logger.info(
+                        f"Marking stream {configured_stream.stream.name} as STOPPED"
+                    )
+                    yield stream_status_as_airbyte_message(
+                        configured_stream.stream, AirbyteStreamStatus.COMPLETE
+                    )
                 except AirbyteTracedException as e:
-                    yield stream_status_as_airbyte_message(configured_stream.stream, AirbyteStreamStatus.INCOMPLETE)
+                    yield stream_status_as_airbyte_message(
+                        configured_stream.stream, AirbyteStreamStatus.INCOMPLETE
+                    )
                     if self.continue_sync_on_stream_failure:
                         stream_name_to_exception[stream_instance.name] = e
                     else:
                         raise e
                 except Exception as e:
                     yield from self._emit_queued_messages()
-                    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)
+                    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
+                    )
                     display_message = stream_instance.get_error_display_message(e)
                     if display_message:
-                        raise AirbyteTracedException.from_exception(e, message=display_message) from e
+                        raise AirbyteTracedException.from_exception(
+                            e, message=display_message
+                        ) from e
                     raise e
                 finally:
                     timer.finish_event()
                     logger.info(f"Finished syncing {configured_stream.stream.name}")
                     logger.info(timer.report())
 
         if self.continue_sync_on_stream_failure and len(stream_name_to_exception) > 0:
-            raise AirbyteTracedException(message=self._generate_failed_streams_error_message(stream_name_to_exception))
+            raise AirbyteTracedException(
+                message=self._generate_failed_streams_error_message(
+                    stream_name_to_exception
+                )
+            )
         logger.info(f"Finished syncing {self.name}")
 
     @property
@@ -173,7 +221,9 @@ def _read_stream(
         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}",
@@ -185,7 +235,10 @@ def _read_stream(
         )
         stream_instance.log_stream_sync_configuration()
 
-        use_incremental = configured_stream.sync_mode == SyncMode.incremental and stream_instance.supports_incremental
+        use_incremental = (
+            configured_stream.sync_mode == SyncMode.incremental
+            and stream_instance.supports_incremental
+        )
         if use_incremental:
             record_iterator = self._read_incremental(
                 logger,
@@ -195,7 +248,9 @@ def _read_stream(
                 internal_config,
             )
         else:
-            record_iterator = self._read_full_refresh(logger, stream_instance, configured_stream, internal_config)
+            record_iterator = self._read_full_refresh(
+                logger, stream_instance, configured_stream, internal_config
+            )
 
         record_counter = 0
         stream_name = configured_stream.stream.name
@@ -206,7 +261,9 @@ def _read_stream(
                 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
 
@@ -230,7 +287,9 @@ def _read_incremental(
         :return:
         """
         stream_name = configured_stream.stream.name
-        stream_state = state_manager.get_stream_state(stream_name, stream_instance.namespace)
+        stream_state = state_manager.get_stream_state(
+            stream_name, stream_instance.namespace
+        )
 
         if stream_state and "state" in dir(stream_instance):
             stream_instance.state = stream_state  # type: ignore # we check that state in the dir(stream_instance)
@@ -260,22 +319,31 @@ def _read_full_refresh(
         internal_config: InternalConfig,
     ) -> Iterator[AirbyteMessage]:
         total_records_counter = 0
-        for record_data_or_message in stream_instance.read_full_refresh(configured_stream.cursor_field, logger, self._slice_logger):
+        for record_data_or_message in stream_instance.read_full_refresh(
+            configured_stream.cursor_field, logger, self._slice_logger
+        ):
             message = self._get_message(record_data_or_message, stream_instance)
             yield message
             if message.type == MessageType.RECORD:
                 total_records_counter += 1
                 if internal_config.is_limit_reached(total_records_counter):
                     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
         """
         if isinstance(record_data_or_message, AirbyteMessage):
             return record_data_or_message
         else:
-            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]:
@@ -293,6 +361,13 @@ def continue_sync_on_stream_failure(self) -> bool:
         return False
 
     @staticmethod
-    def _generate_failed_streams_error_message(stream_failures: Mapping[str, AirbyteTracedException]) -> str:
-        failures = ", ".join([f"{stream}: {exception.__repr__()}" for stream, exception in stream_failures.items()])
+    def _generate_failed_streams_error_message(
+        stream_failures: Mapping[str, AirbyteTracedException]
+    ) -> str:
+        failures = ", ".join(
+            [
+                f"{stream}: {exception.__repr__()}"
+                for stream, exception in stream_failures.items()
+            ]
+        )
         return f"During the sync, the following streams did not sync successfully: {failures}"