airbyte-cdk 6.37.0.dev1__py3-none-any.whl → 6.37.2__py3-none-any.whl

airbyte_cdk/connector_builder/models.py

@@ -21,20 +21,6 @@ class HttpRequest:
     body: Optional[str] = None
 
 
-@dataclass
-class StreamReadPages:
-    records: List[object]
-    request: Optional[HttpRequest] = None
-    response: Optional[HttpResponse] = None
-
-
-@dataclass
-class StreamReadSlices:
-    pages: List[StreamReadPages]
-    slice_descriptor: Optional[Dict[str, Any]]
-    state: Optional[List[Dict[str, Any]]] = None
-
-
 @dataclass
 class LogMessage:
     message: str
@@ -46,11 +32,27 @@ class LogMessage:
 @dataclass
 class AuxiliaryRequest:
     title: str
+    type: str
     description: str
     request: HttpRequest
     response: HttpResponse
 
 
+@dataclass
+class StreamReadPages:
+    records: List[object]
+    request: Optional[HttpRequest] = None
+    response: Optional[HttpResponse] = None
+
+
+@dataclass
+class StreamReadSlices:
+    pages: List[StreamReadPages]
+    slice_descriptor: Optional[Dict[str, Any]]
+    state: Optional[List[Dict[str, Any]]] = None
+    auxiliary_requests: Optional[List[AuxiliaryRequest]] = None
+
+
 @dataclass
 class StreamRead(object):
     logs: List[LogMessage]

airbyte_cdk/connector_builder/test_reader/helpers.py

@@ -28,7 +28,7 @@ from airbyte_cdk.utils.schema_inferrer import (
     SchemaInferrer,
 )
 
-from .types import LOG_MESSAGES_OUTPUT_TYPE
+from .types import ASYNC_AUXILIARY_REQUEST_TYPES, LOG_MESSAGES_OUTPUT_TYPE
 
 # -------
 # Parsers
@@ -226,7 +226,8 @@ def should_close_page(
         at_least_one_page_in_group
         and is_log_message(message)
         and (
-            is_page_http_request(json_message) or message.log.message.startswith("slice:")  # type: ignore[union-attr] # AirbyteMessage with MessageType.LOG has log.message
+            is_page_http_request(json_message)
+            or message.log.message.startswith(SliceLogger.SLICE_LOG_PREFIX)  # type: ignore[union-attr] # AirbyteMessage with MessageType.LOG has log.message
         )
     )
 
@@ -330,6 +331,10 @@ def is_auxiliary_http_request(message: Optional[Dict[str, Any]]) -> bool:
     return is_http_log(message) and message.get("http", {}).get("is_auxiliary", False)
 
 
+def is_async_auxiliary_request(message: AuxiliaryRequest) -> bool:
+    return message.type in ASYNC_AUXILIARY_REQUEST_TYPES
+
+
 def is_log_message(message: AirbyteMessage) -> bool:
     """
     Determines whether the provided message is of type LOG.
@@ -413,6 +418,7 @@ def handle_current_slice(
     current_slice_pages: List[StreamReadPages],
     current_slice_descriptor: Optional[Dict[str, Any]] = None,
     latest_state_message: Optional[Dict[str, Any]] = None,
+    auxiliary_requests: Optional[List[AuxiliaryRequest]] = None,
 ) -> StreamReadSlices:
     """
     Handles the current slice by packaging its pages, descriptor, and state into a StreamReadSlices instance.
@@ -421,6 +427,7 @@ def handle_current_slice(
         current_slice_pages (List[StreamReadPages]): The pages to be included in the slice.
         current_slice_descriptor (Optional[Dict[str, Any]]): Descriptor for the current slice, optional.
         latest_state_message (Optional[Dict[str, Any]]): The latest state message, optional.
+        auxiliary_requests (Optional[List[AuxiliaryRequest]]): The auxiliary requests to include, optional.
 
     Returns:
         StreamReadSlices: An object containing the current slice's pages, descriptor, and state.
@@ -429,6 +436,7 @@ def handle_current_slice(
         pages=current_slice_pages,
         slice_descriptor=current_slice_descriptor,
         state=[latest_state_message] if latest_state_message else [],
+        auxiliary_requests=auxiliary_requests if auxiliary_requests else [],
     )
 
 
@@ -486,29 +494,24 @@ def handle_auxiliary_request(json_message: Dict[str, JsonType]) -> AuxiliaryRequ
     Raises:
         ValueError: If any of the "airbyte_cdk", "stream", or "http" fields is not a dictionary.
     """
-    airbyte_cdk = json_message.get("airbyte_cdk", {})
-
-    if not isinstance(airbyte_cdk, dict):
-        raise ValueError(
-            f"Expected airbyte_cdk to be a dict, got {airbyte_cdk} of type {type(airbyte_cdk)}"
-        )
-
-    stream = airbyte_cdk.get("stream", {})
 
-    if not isinstance(stream, dict):
-        raise ValueError(f"Expected stream to be a dict, got {stream} of type {type(stream)}")
+    airbyte_cdk = get_airbyte_cdk_from_message(json_message)
+    stream = get_stream_from_airbyte_cdk(airbyte_cdk)
+    title_prefix = get_auxiliary_request_title_prefix(stream)
+    http = get_http_property_from_message(json_message)
+    request_type = get_auxiliary_request_type(stream, http)
 
-    title_prefix = "Parent stream: " if stream.get("is_substream", False) else ""
-    http = json_message.get("http", {})
-
-    if not isinstance(http, dict):
-        raise ValueError(f"Expected http to be a dict, got {http} of type {type(http)}")
+    title = title_prefix + str(http.get("title", None))
+    description = str(http.get("description", None))
+    request = create_request_from_log_message(json_message)
+    response = create_response_from_log_message(json_message)
 
     return AuxiliaryRequest(
-        title=title_prefix + str(http.get("title", None)),
-        description=str(http.get("description", None)),
-        request=create_request_from_log_message(json_message),
-        response=create_response_from_log_message(json_message),
+        title=title,
+        type=request_type,
+        description=description,
+        request=request,
+        response=response,
     )
 
 
@@ -558,7 +561,8 @@ def handle_log_message(
         at_least_one_page_in_group,
         current_page_request,
         current_page_response,
-        auxiliary_request or log_message,
+        auxiliary_request,
+        log_message,
     )
 
 
@@ -589,3 +593,97 @@ def handle_record_message(
         datetime_format_inferrer.accumulate(message.record)  # type: ignore
 
     return records_count
+
+
+# -------
+# Reusable Getters
+# -------
+
+
+def get_airbyte_cdk_from_message(json_message: Dict[str, JsonType]) -> dict:  # type: ignore
+    """
+    Retrieves the "airbyte_cdk" dictionary from the provided JSON message.
+
+    This function validates that the extracted "airbyte_cdk" is of type dict,
+    raising a ValueError if the validation fails.
+
+    Parameters:
+        json_message (Dict[str, JsonType]): A dictionary representing the JSON message.
+
+    Returns:
+        dict: The "airbyte_cdk" dictionary extracted from the JSON message.
+
+    Raises:
+        ValueError: If the "airbyte_cdk" field is not a dictionary.
+    """
+    airbyte_cdk = json_message.get("airbyte_cdk", {})
+
+    if not isinstance(airbyte_cdk, dict):
+        raise ValueError(
+            f"Expected airbyte_cdk to be a dict, got {airbyte_cdk} of type {type(airbyte_cdk)}"
+        )
+
+    return airbyte_cdk
+
+
+def get_stream_from_airbyte_cdk(airbyte_cdk: dict) -> dict:  # type: ignore
+    """
+    Retrieves the "stream" dictionary from the provided "airbyte_cdk" dictionary.
+
+    This function ensures that the extracted "stream" is of type dict,
+    raising a ValueError if the validation fails.
+
+    Parameters:
+        airbyte_cdk (dict): The dictionary representing the Airbyte CDK data.
+
+    Returns:
+        dict: The "stream" dictionary extracted from the Airbyte CDK data.
+
+    Raises:
+        ValueError: If the "stream" field is not a dictionary.
+    """
+
+    stream = airbyte_cdk.get("stream", {})
+
+    if not isinstance(stream, dict):
+        raise ValueError(f"Expected stream to be a dict, got {stream} of type {type(stream)}")
+
+    return stream
+
+
+def get_auxiliary_request_title_prefix(stream: dict) -> str:  # type: ignore
+    """
+    Generates a title prefix based on the stream type.
+    """
+    return "Parent stream: " if stream.get("is_substream", False) else ""
+
+
+def get_http_property_from_message(json_message: Dict[str, JsonType]) -> dict:  # type: ignore
+    """
+    Retrieves the "http" dictionary from the provided JSON message.
+
+    This function validates that the extracted "http" is of type dict,
+    raising a ValueError if the validation fails.
+
+    Parameters:
+        json_message (Dict[str, JsonType]): A dictionary representing the JSON message.
+
+    Returns:
+        dict: The "http" dictionary extracted from the JSON message.
+
+    Raises:
+        ValueError: If the "http" field is not a dictionary.
+    """
+    http = json_message.get("http", {})
+
+    if not isinstance(http, dict):
+        raise ValueError(f"Expected http to be a dict, got {http} of type {type(http)}")
+
+    return http
+
+
+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))

airbyte_cdk/connector_builder/test_reader/message_grouper.py

@@ -6,6 +6,7 @@
 from typing import Any, Dict, Iterator, List, Mapping, Optional
 
 from airbyte_cdk.connector_builder.models import (
+    AuxiliaryRequest,
     HttpRequest,
     HttpResponse,
     StreamReadPages,
@@ -24,6 +25,7 @@ from .helpers import (
     handle_current_slice,
     handle_log_message,
     handle_record_message,
+    is_async_auxiliary_request,
     is_config_update_message,
     is_log_message,
     is_record_message,
@@ -89,6 +91,7 @@ def get_message_groups(
     current_page_request: Optional[HttpRequest] = None
     current_page_response: Optional[HttpResponse] = None
     latest_state_message: Optional[Dict[str, Any]] = None
+    slice_auxiliary_requests: List[AuxiliaryRequest] = []
 
     while records_count < limit and (message := next(messages, None)):
         json_message = airbyte_message_to_json(message)
@@ -106,6 +109,7 @@ def get_message_groups(
                 current_slice_pages,
                 current_slice_descriptor,
                 latest_state_message,
+                slice_auxiliary_requests,
             )
             current_slice_descriptor = parse_slice_description(message.log.message)  # type: ignore
             current_slice_pages = []
@@ -118,7 +122,8 @@ def get_message_groups(
                 at_least_one_page_in_group,
                 current_page_request,
                 current_page_response,
-                log_or_auxiliary_request,
+                auxiliary_request,
+                log_message,
             ) = handle_log_message(
                 message,
                 json_message,
@@ -126,8 +131,15 @@ def get_message_groups(
                 current_page_request,
                 current_page_response,
             )
-            if log_or_auxiliary_request:
-                yield log_or_auxiliary_request
+
+            if auxiliary_request:
+                if is_async_auxiliary_request(auxiliary_request):
+                    slice_auxiliary_requests.append(auxiliary_request)
+                else:
+                    yield auxiliary_request
+
+            if log_message:
+                yield log_message
         elif is_trace_with_error(message):
             if message.trace is not None:
                 yield message.trace
@@ -157,4 +169,5 @@ def get_message_groups(
                 current_slice_pages,
                 current_slice_descriptor,
                 latest_state_message,
+                slice_auxiliary_requests,
             )

airbyte_cdk/connector_builder/test_reader/types.py

@@ -71,5 +71,13 @@ LOG_MESSAGES_OUTPUT_TYPE = tuple[
     bool,
     HttpRequest | None,
     HttpResponse | None,
-    AuxiliaryRequest | AirbyteLogMessage | None,
+    AuxiliaryRequest | None,
+    AirbyteLogMessage | None,
+]
+
+ASYNC_AUXILIARY_REQUEST_TYPES = [
+    "ASYNC_CREATE",
+    "ASYNC_POLL",
+    "ASYNC_ABORT",
+    "ASYNC_DELETE",
 ]

airbyte_cdk/sources/declarative/auth/token_provider.py

@@ -58,6 +58,7 @@ class SessionTokenProvider(TokenProvider):
                 "Obtains session token",
                 None,
                 is_auxiliary=True,
+                type="AUTH",
             ),
         )
         if response is None:

airbyte_cdk/sources/declarative/concurrent_declarative_source.py

@@ -31,6 +31,9 @@ from airbyte_cdk.sources.declarative.models.declarative_component_schema import
 from airbyte_cdk.sources.declarative.models.declarative_component_schema import (
     DatetimeBasedCursor as DatetimeBasedCursorModel,
 )
+from airbyte_cdk.sources.declarative.models.declarative_component_schema import (
+    IncrementingCountCursor as IncrementingCountCursorModel,
+)
 from airbyte_cdk.sources.declarative.parsers.model_to_component_factory import (
     ModelToComponentFactory,
 )
@@ -44,6 +47,7 @@ from airbyte_cdk.sources.declarative.types import ConnectionDefinition
 from airbyte_cdk.sources.source import TState
 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.availability_strategy import (
     AlwaysAvailableAvailabilityStrategy,
 )
@@ -118,6 +122,12 @@ class ConcurrentDeclarativeSource(ManifestDeclarativeSource, Generic[TState]):
             message_repository=self.message_repository,
         )
 
+    # TODO: Remove this. This property is necessary to safely migrate Stripe during the transition state.
+    @property
+    def is_partially_declarative(self) -> bool:
+        """This flag used to avoid unexpected AbstractStreamFacade processing as concurrent streams."""
+        return False
+
     def read(
         self,
         logger: logging.Logger,
@@ -215,7 +225,7 @@ class ConcurrentDeclarativeSource(ManifestDeclarativeSource, Generic[TState]):
                     and not incremental_sync_component_definition
                 )
 
-                if self._is_datetime_incremental_without_partition_routing(
+                if self._is_concurrent_cursor_incremental_without_partition_routing(
                     declarative_stream, incremental_sync_component_definition
                 ):
                     stream_state = self._connector_state_manager.get_stream_state(
@@ -247,15 +257,26 @@ class ConcurrentDeclarativeSource(ManifestDeclarativeSource, Generic[TState]):
                             stream_slicer=declarative_stream.retriever.stream_slicer,
                         )
                     else:
-                        cursor = (
-                            self._constructor.create_concurrent_cursor_from_datetime_based_cursor(
+                        if (
+                            incremental_sync_component_definition
+                            and incremental_sync_component_definition.get("type")
+                            == IncrementingCountCursorModel.__name__
+                        ):
+                            cursor = self._constructor.create_concurrent_cursor_from_incrementing_count_cursor(
+                                model_type=IncrementingCountCursorModel,
+                                component_definition=incremental_sync_component_definition,  # type: ignore  # Not None because of the if condition above
+                                stream_name=declarative_stream.name,
+                                stream_namespace=declarative_stream.namespace,
+                                config=config or {},
+                            )
+                        else:
+                            cursor = self._constructor.create_concurrent_cursor_from_datetime_based_cursor(
                                 model_type=DatetimeBasedCursorModel,
                                 component_definition=incremental_sync_component_definition,  # type: ignore  # Not None because of the if condition above
                                 stream_name=declarative_stream.name,
                                 stream_namespace=declarative_stream.namespace,
                                 config=config or {},
                             )
-                        )
                         partition_generator = StreamSlicerPartitionGenerator(
                             partition_factory=DeclarativePartitionFactory(
                                 declarative_stream.name,
@@ -369,12 +390,20 @@ class ConcurrentDeclarativeSource(ManifestDeclarativeSource, Generic[TState]):
                     )
                 else:
                     synchronous_streams.append(declarative_stream)
+            # TODO: Remove this. This check is necessary to safely migrate Stripe during the transition state.
+            # Condition below needs to ensure that concurrent support is not lost for sources that already support
+            # it before migration, but now are only partially migrated to declarative implementation (e.g., Stripe).
+            elif (
+                isinstance(declarative_stream, AbstractStreamFacade)
+                and self.is_partially_declarative
+            ):
+                concurrent_streams.append(declarative_stream.get_underlying_stream())
             else:
                 synchronous_streams.append(declarative_stream)
 
         return concurrent_streams, synchronous_streams
 
-    def _is_datetime_incremental_without_partition_routing(
+    def _is_concurrent_cursor_incremental_without_partition_routing(
         self,
         declarative_stream: DeclarativeStream,
         incremental_sync_component_definition: Mapping[str, Any] | None,
@@ -382,11 +411,18 @@ class ConcurrentDeclarativeSource(ManifestDeclarativeSource, Generic[TState]):
         return (
             incremental_sync_component_definition is not None
             and bool(incremental_sync_component_definition)
-            and incremental_sync_component_definition.get("type", "")
-            == DatetimeBasedCursorModel.__name__
+            and (
+                incremental_sync_component_definition.get("type", "")
+                in (DatetimeBasedCursorModel.__name__, IncrementingCountCursorModel.__name__)
+            )
             and hasattr(declarative_stream.retriever, "stream_slicer")
             and (
                 isinstance(declarative_stream.retriever.stream_slicer, DatetimeBasedCursor)
+                # IncrementingCountCursorModel is hardcoded to be of type DatetimeBasedCursor
+                # add isintance check here if we want to create a Declarative IncrementingCountCursor
+                # or isinstance(
+                #     declarative_stream.retriever.stream_slicer, IncrementingCountCursor
+                # )
                 or isinstance(declarative_stream.retriever.stream_slicer, AsyncJobPartitionRouter)
             )
         )

airbyte_cdk/sources/declarative/datetime/datetime_parser.py

@@ -31,7 +31,8 @@ class DatetimeParser:
             return datetime.datetime.fromtimestamp(float(date), tz=datetime.timezone.utc)
         elif format == "%ms":
             return self._UNIX_EPOCH + datetime.timedelta(milliseconds=int(date))
-
+        elif "%_ms" in format:
+            format = format.replace("%_ms", "%f")
         parsed_datetime = datetime.datetime.strptime(str(date), format)
         if self._is_naive(parsed_datetime):
             return parsed_datetime.replace(tzinfo=datetime.timezone.utc)
@@ -48,6 +49,11 @@ class DatetimeParser:
         if format == "%ms":
             # timstamp() returns a float representing the number of seconds since the unix epoch
             return str(int(dt.timestamp() * 1000))
+        if "%_ms" in format:
+            _format = format.replace("%_ms", "%f")
+            milliseconds = int(dt.microsecond / 1000)
+            formatted_dt = dt.strftime(_format).replace(dt.strftime("%f"), "%03d" % milliseconds)
+            return formatted_dt
         else:
             return dt.strftime(format)