Merge branch 'main' into maxi297/incremental_without_partition_router_as_defaultstream

Author: maxime.c

airbyte_cdk/connector_builder/connector_builder_handler.py

@@ -3,8 +3,8 @@
 #
 
 
-from dataclasses import asdict
-from typing import Any, Dict, List, Mapping, Optional
+from dataclasses import asdict, dataclass, field
+from typing import Any, ClassVar, Dict, List, Mapping
 
 from airbyte_cdk.connector_builder.test_reader import TestReader
 from airbyte_cdk.models import (
@@ -15,32 +15,45 @@
     Type,
 )
 from airbyte_cdk.models import Type as MessageType
-from airbyte_cdk.sources.declarative.concurrent_declarative_source import (
-    ConcurrentDeclarativeSource,
-    TestLimits,
-)
 from airbyte_cdk.sources.declarative.declarative_source import DeclarativeSource
 from airbyte_cdk.sources.declarative.manifest_declarative_source import ManifestDeclarativeSource
+from airbyte_cdk.sources.declarative.parsers.model_to_component_factory import (
+    ModelToComponentFactory,
+)
 from airbyte_cdk.utils.airbyte_secrets_utils import filter_secrets
 from airbyte_cdk.utils.datetime_helpers import ab_datetime_now
 from airbyte_cdk.utils.traced_exception import AirbyteTracedException
 
+DEFAULT_MAXIMUM_NUMBER_OF_PAGES_PER_SLICE = 5
+DEFAULT_MAXIMUM_NUMBER_OF_SLICES = 5
+DEFAULT_MAXIMUM_RECORDS = 100
+DEFAULT_MAXIMUM_STREAMS = 100
+
 MAX_PAGES_PER_SLICE_KEY = "max_pages_per_slice"
 MAX_SLICES_KEY = "max_slices"
 MAX_RECORDS_KEY = "max_records"
 MAX_STREAMS_KEY = "max_streams"
 
 
+@dataclass
+class TestLimits:
+    __test__: ClassVar[bool] = False  # Tell Pytest this is not a Pytest class, despite its name
+
+    max_records: int = field(default=DEFAULT_MAXIMUM_RECORDS)
+    max_pages_per_slice: int = field(default=DEFAULT_MAXIMUM_NUMBER_OF_PAGES_PER_SLICE)
+    max_slices: int = field(default=DEFAULT_MAXIMUM_NUMBER_OF_SLICES)
+    max_streams: int = field(default=DEFAULT_MAXIMUM_STREAMS)
+
+
 def get_limits(config: Mapping[str, Any]) -> TestLimits:
     command_config = config.get("__test_read_config", {})
-    return TestLimits(
-        max_records=command_config.get(MAX_RECORDS_KEY, TestLimits.DEFAULT_MAX_RECORDS),
-        max_pages_per_slice=command_config.get(
-            MAX_PAGES_PER_SLICE_KEY, TestLimits.DEFAULT_MAX_PAGES_PER_SLICE
-        ),
-        max_slices=command_config.get(MAX_SLICES_KEY, TestLimits.DEFAULT_MAX_SLICES),
-        max_streams=command_config.get(MAX_STREAMS_KEY, TestLimits.DEFAULT_MAX_STREAMS),
+    max_pages_per_slice = (
+        command_config.get(MAX_PAGES_PER_SLICE_KEY) or DEFAULT_MAXIMUM_NUMBER_OF_PAGES_PER_SLICE
     )
+    max_slices = command_config.get(MAX_SLICES_KEY) or DEFAULT_MAXIMUM_NUMBER_OF_SLICES
+    max_records = command_config.get(MAX_RECORDS_KEY) or DEFAULT_MAXIMUM_RECORDS
+    max_streams = command_config.get(MAX_STREAMS_KEY) or DEFAULT_MAXIMUM_STREAMS
+    return TestLimits(max_records, max_pages_per_slice, max_slices, max_streams)
 
 
 def should_migrate_manifest(config: Mapping[str, Any]) -> bool:
@@ -62,30 +75,21 @@ def should_normalize_manifest(config: Mapping[str, Any]) -> bool:
     return config.get("__should_normalize", False)
 
 
-def create_source(
-    config: Mapping[str, Any],
-    limits: TestLimits,
-    catalog: Optional[ConfiguredAirbyteCatalog],
-    state: Optional[List[AirbyteStateMessage]],
-) -> ConcurrentDeclarativeSource[Optional[List[AirbyteStateMessage]]]:
+def create_source(config: Mapping[str, Any], limits: TestLimits) -> ManifestDeclarativeSource:
     manifest = config["__injected_declarative_manifest"]
-
-    # We enforce a concurrency level of 1 so that the stream is processed on a single thread
-    # to retain ordering for the grouping of the builder message responses.
-    if "concurrency_level" in manifest:
-        manifest["concurrency_level"]["default_concurrency"] = 1
-    else:
-        manifest["concurrency_level"] = {"type": "ConcurrencyLevel", "default_concurrency": 1}
-
-    return ConcurrentDeclarativeSource(
-        catalog=catalog,
+    return ManifestDeclarativeSource(
         config=config,
-        state=state,
-        source_config=manifest,
         emit_connector_builder_messages=True,
+        source_config=manifest,
         migrate_manifest=should_migrate_manifest(config),
         normalize_manifest=should_normalize_manifest(config),
-        limits=limits,
+        component_factory=ModelToComponentFactory(
+            emit_connector_builder_messages=True,
+            limit_pages_fetched_per_slice=limits.max_pages_per_slice,
+            limit_slices_fetched=limits.max_slices,
+            disable_retries=True,
+            disable_cache=True,
+        ),
     )
 
 

airbyte_cdk/connector_builder/main.py

@@ -91,12 +91,12 @@ def handle_connector_builder_request(
 def handle_request(args: List[str]) -> str:
     command, config, catalog, state = get_config_and_catalog_from_args(args)
     limits = get_limits(config)
-    source = create_source(config=config, limits=limits, catalog=catalog, state=state)
-    return orjson.dumps(  # type: ignore[no-any-return] # Serializer.dump() always returns AirbyteMessage
+    source = create_source(config, limits)
+    return orjson.dumps(
         AirbyteMessageSerializer.dump(
             handle_connector_builder_request(source, command, config, catalog, state, limits)
         )
-    ).decode()
+    ).decode()  # type: ignore[no-any-return] # Serializer.dump() always returns AirbyteMessage
 
 
 if __name__ == "__main__":

airbyte_cdk/connector_builder/test_reader/helpers.py

@@ -5,7 +5,7 @@
 import json
 from copy import deepcopy
 from json import JSONDecodeError
-from typing import Any, Dict, List, Mapping, Optional, Union
+from typing import Any, Dict, List, Mapping, Optional
 
 from airbyte_cdk.connector_builder.models import (
     AuxiliaryRequest,
@@ -17,8 +17,6 @@
 from airbyte_cdk.models import (
     AirbyteLogMessage,
     AirbyteMessage,
-    AirbyteStateBlob,
-    AirbyteStateMessage,
     OrchestratorType,
     TraceType,
 )
@@ -468,7 +466,7 @@ def handle_current_slice(
     return StreamReadSlices(
         pages=current_slice_pages,
         slice_descriptor=current_slice_descriptor,
-        state=[convert_state_blob_to_mapping(latest_state_message)] if latest_state_message else [],
+        state=[latest_state_message] if latest_state_message else [],
         auxiliary_requests=auxiliary_requests if auxiliary_requests else [],
     )
 
@@ -720,23 +718,3 @@ def get_auxiliary_request_type(stream: dict, http: dict) -> str:  # type: ignore
     Determines the type of the auxiliary request based on the stream and HTTP properties.
     """
     return "PARENT_STREAM" if stream.get("is_substream", False) else str(http.get("type", None))
-
-
-def convert_state_blob_to_mapping(
-    state_message: Union[AirbyteStateMessage, Dict[str, Any]],
-) -> Dict[str, Any]:
-    """
-    The AirbyteStreamState stores state as an AirbyteStateBlob which deceivingly is not
-    a dictionary, but rather a list of kwargs fields. This in turn causes it to not be
-    properly turned into a dictionary when translating this back into response output
-    by the connector_builder_handler using asdict()
-    """
-
-    if isinstance(state_message, AirbyteStateMessage) and state_message.stream:
-        state_value = state_message.stream.stream_state
-        if isinstance(state_value, AirbyteStateBlob):
-            state_value_mapping = {k: v for k, v in state_value.__dict__.items()}
-            state_message.stream.stream_state = state_value_mapping  # type: ignore  # we intentionally set this as a Dict so that StreamReadSlices is translated properly in the resulting HTTP response
-        return state_message  # type: ignore  # See above, but when this is an AirbyteStateMessage we must convert AirbyteStateBlob to a Dict
-    else:
-        return state_message  # type: ignore  # This is guaranteed to be a Dict since we check isinstance AirbyteStateMessage above

airbyte_cdk/connector_builder/test_reader/message_grouper.py

@@ -95,7 +95,7 @@ def get_message_groups(
     latest_state_message: Optional[Dict[str, Any]] = None
     slice_auxiliary_requests: List[AuxiliaryRequest] = []
 
-    while message := next(messages, None):
+    while records_count < limit and (message := next(messages, None)):
         json_message = airbyte_message_to_json(message)
 
         if is_page_http_request_for_different_stream(json_message, stream_name):

airbyte_cdk/entrypoint.py

@@ -22,7 +22,7 @@
 
 from airbyte_cdk.connector import TConfig
 from airbyte_cdk.exception_handler import init_uncaught_exception_handler
-from airbyte_cdk.logger import PRINT_BUFFER, init_logger
+from airbyte_cdk.logger import PRINT_BUFFER, init_logger, is_platform_debug_log_enabled
 from airbyte_cdk.models import (
     AirbyteConnectionStatus,
     AirbyteMessage,
@@ -158,7 +158,7 @@ def run(self, parsed_args: argparse.Namespace) -> Iterable[str]:
         if not cmd:
             raise Exception("No command passed")
 
-        if hasattr(parsed_args, "debug") and parsed_args.debug:
+        if (hasattr(parsed_args, "debug") and parsed_args.debug) or is_platform_debug_log_enabled():
             self.logger.setLevel(logging.DEBUG)
             logger.setLevel(logging.DEBUG)
             self.logger.debug("Debug logs enabled")

airbyte_cdk/logger.py

@@ -1,10 +1,10 @@
 #
 # Copyright (c) 2023 Airbyte, Inc., all rights reserved.
 #
-
 import json
 import logging
 import logging.config
+import os
 from typing import Any, Callable, Mapping, Optional, Tuple
 
 import orjson
@@ -40,6 +40,10 @@
 }
 
 
+def is_platform_debug_log_enabled() -> bool:
+    return os.environ.get("LOG_LEVEL", "info").lower() == "debug"
+
+
 def init_logger(name: Optional[str] = None) -> logging.Logger:
     """Initial set up of logger"""
     logger = logging.getLogger(name)
@@ -73,8 +77,22 @@ def format(self, record: logging.LogRecord) -> str:
         airbyte_level = self.level_mapping.get(record.levelno, "INFO")
         if airbyte_level == Level.DEBUG:
             extras = self.extract_extra_args_from_record(record)
-            debug_dict = {"type": "DEBUG", "message": record.getMessage(), "data": extras}
-            return filter_secrets(json.dumps(debug_dict))
+            if is_platform_debug_log_enabled():
+                # We have a different behavior between debug logs enabled through `--debug` argument and debug logs
+                # enabled through environment variable. The reason is that for platform logs, we need to have these
+                # printed as AirbyteMessage which is not the case with the current previous implementation.
+                # Why not migrate both to AirbyteMessages then? AirbyteMessages do not support having structured logs.
+                # which means that the DX would be degraded compared to the current solution (devs will need to identify
+                # the `log.message` field and figure out where in this field is the response while the current solution
+                # have a specific field that is structured for extras.
+                message = f"{filter_secrets(record.getMessage())} ///\nExtra logs: {filter_secrets(json.dumps(extras))}"
+                log_message = AirbyteMessage(
+                    type=Type.LOG, log=AirbyteLogMessage(level=airbyte_level, message=message)
+                )
+                return orjson.dumps(AirbyteMessageSerializer.dump(log_message)).decode()
+            else:
+                debug_dict = {"type": "DEBUG", "message": record.getMessage(), "data": extras}
+                return filter_secrets(json.dumps(debug_dict))
         else:
             message = super().format(record)
             message = filter_secrets(message)

airbyte_cdk/sources/concurrent_source/concurrent_read_processor.py

@@ -95,14 +95,11 @@ def on_partition(self, partition: Partition) -> None:
         """
         stream_name = partition.stream_name()
         self._streams_to_running_partitions[stream_name].add(partition)
-        cursor = self._stream_name_to_instance[stream_name].cursor
         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._thread_pool_manager.submit(
-            self._partition_reader.process_partition, partition, cursor
-        )
+        self._thread_pool_manager.submit(self._partition_reader.process_partition, partition)
 
     def on_partition_complete_sentinel(
         self, sentinel: PartitionCompleteSentinel
@@ -115,16 +112,26 @@ def on_partition_complete_sentinel(
         """
         partition = sentinel.partition
 
-        partitions_running = self._streams_to_running_partitions[partition.stream_name()]
-        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
-            ):
-                yield from self._on_stream_is_done(partition.stream_name())
-        yield from self._message_repository.consume_queue()
+        try:
+            if sentinel.is_successful:
+                stream = self._stream_name_to_instance[partition.stream_name()]
+                stream.cursor.close_partition(partition)
+        except Exception as exception:
+            self._flag_exception(partition.stream_name(), exception)
+            yield AirbyteTracedException.from_exception(
+                exception, stream_descriptor=StreamDescriptor(name=partition.stream_name())
+            ).as_sanitized_airbyte_message()
+        finally:
+            partitions_running = self._streams_to_running_partitions[partition.stream_name()]
+            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
+                ):
+                    yield from self._on_stream_is_done(partition.stream_name())
+            yield from self._message_repository.consume_queue()
 
     def on_record(self, record: Record) -> Iterable[AirbyteMessage]:
         """