airbyte-cdk 6.36.1__tar.gz → 6.36.3__tar.gz

{airbyte_cdk-6.36.1 → airbyte_cdk-6.36.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: airbyte-cdk
-Version: 6.36.1
+Version: 6.36.3
 Summary: A framework for writing Airbyte Connectors.
 Home-page: https://airbyte.com
 License: MIT

{airbyte_cdk-6.36.1 → airbyte_cdk-6.36.3}/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-6.36.1 → airbyte_cdk-6.36.3}/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-6.36.1 → airbyte_cdk-6.36.3}/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-6.36.1 → airbyte_cdk-6.36.3}/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-6.36.1 → airbyte_cdk-6.36.3}/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-6.36.1 → airbyte_cdk-6.36.3}/airbyte_cdk/sources/declarative/declarative_component_schema.yaml

@@ -1779,6 +1779,9 @@ definitions:
           - stream_interval
           - stream_partition
           - stream_slice
+          - creation_response
+          - polling_response
+          - download_target
         examples:
           - "/products"
           - "/quotes/{{ stream_partition['id'] }}/quote_line_groups"
@@ -3223,7 +3226,7 @@ definitions:
       - polling_requester
       - download_requester
       - status_extractor
-      - urls_extractor
+      - download_target_extractor
     properties:
       type:
         type: string
@@ -3240,7 +3243,7 @@ definitions:
         anyOf:
           - "$ref": "#/definitions/CustomRecordExtractor"
           - "$ref": "#/definitions/DpathExtractor"
-      urls_extractor:
+      download_target_extractor:
         description: Responsible for fetching the final result `urls` provided by the completed / finished / ready async job.
         anyOf:
           - "$ref": "#/definitions/CustomRecordExtractor"
@@ -3261,7 +3264,7 @@ definitions:
         anyOf:
           - "$ref": "#/definitions/CustomRequester"
           - "$ref": "#/definitions/HttpRequester"
-      url_requester:
+      download_target_requester:
         description: Requester component that describes how to prepare HTTP requests to send to the source API to extract the url from polling response by the completed async job.
         anyOf:
           - "$ref": "#/definitions/CustomRequester"
@@ -3667,6 +3670,21 @@ interpolation:
             self: https://api.sendgrid.com/v3/marketing/lists?page_size=1&page_token=
             next: https://api.sendgrid.com/v3/marketing/lists?page_size=1&page_token=0236d6d2
             count: 82
+    - title: creation_response
+      description: The response received from the creation_requester in the AsyncRetriever component.
+      type: object
+      examples:
+        - id: "1234"
+    - title: polling_response
+      description: The response received from the polling_requester in the AsyncRetriever component.
+      type: object
+      examples:
+        - id: "1234"
+    - title: download_target
+      description: The `URL` received from the polling_requester in the AsyncRetriever with jobStatus as `COMPLETED`.
+      type: string
+      examples:
+        - "https://api.sendgrid.com/v3/marketing/lists?page_size=1&page_token=0236d6d2&filename=xxx_yyy_zzz.csv"
     - title: stream_interval
       description: The current stream interval being processed. The keys are defined by the incremental sync component. Default keys are `start_time` and `end_time`.
       type: object

{airbyte_cdk-6.36.1 → airbyte_cdk-6.36.3}/airbyte_cdk/sources/declarative/extractors/response_to_file_extractor.py

@@ -136,6 +136,7 @@ class ResponseToFileExtractor(RecordExtractor):
         """
 
         try:
+            # TODO: Add support for other file types, like `json`, with `pd.read_json()`
             with open(path, "r", encoding=file_encoding) as data:
                 chunks = pd.read_csv(
                     data, chunksize=chunk_size, iterator=True, dialect="unix", dtype=object

{airbyte_cdk-6.36.1 → airbyte_cdk-6.36.3}/airbyte_cdk/sources/declarative/incremental/concurrent_partition_cursor.py

@@ -95,6 +95,10 @@ class ConcurrentPerPartitionCursor(Cursor):
         # the oldest partitions can be efficiently removed, maintaining the most recent partitions.
         self._cursor_per_partition: OrderedDict[str, ConcurrentCursor] = OrderedDict()
         self._semaphore_per_partition: OrderedDict[str, threading.Semaphore] = OrderedDict()
+
+        # Parent-state tracking: store each partition’s parent state in creation order
+        self._partition_parent_state_map: OrderedDict[str, Mapping[str, Any]] = OrderedDict()
+
         self._finished_partitions: set[str] = set()
         self._lock = threading.Lock()
         self._timer = Timer()
@@ -155,11 +159,62 @@ class ConcurrentPerPartitionCursor(Cursor):
                     and self._semaphore_per_partition[partition_key]._value == 0
                 ):
                     self._update_global_cursor(cursor.state[self.cursor_field.cursor_field_key])
-                self._emit_state_message()
+
+            self._check_and_update_parent_state()
+
+            self._emit_state_message()
+
+    def _check_and_update_parent_state(self) -> None:
+        """
+        Pop the leftmost partition state from _partition_parent_state_map only if
+        *all partitions* up to (and including) that partition key in _semaphore_per_partition
+        are fully finished (i.e. in _finished_partitions and semaphore._value == 0).
+        Additionally, delete finished semaphores with a value of 0 to free up memory,
+        as they are only needed to track errors and completion status.
+        """
+        last_closed_state = None
+
+        while self._partition_parent_state_map:
+            # Look at the earliest partition key in creation order
+            earliest_key = next(iter(self._partition_parent_state_map))
+
+            # Verify ALL partitions from the left up to earliest_key are finished
+            all_left_finished = True
+            for p_key, sem in list(
+                self._semaphore_per_partition.items()
+            ):  # Use list to allow modification during iteration
+                # If any earlier partition is still not finished, we must stop
+                if p_key not in self._finished_partitions or sem._value != 0:
+                    all_left_finished = False
+                    break
+                # Once we've reached earliest_key in the semaphore order, we can stop checking
+                if p_key == earliest_key:
+                    break
+
+            # If the partitions up to earliest_key are not all finished, break the while-loop
+            if not all_left_finished:
+                break
+
+            # Pop the leftmost entry from parent-state map
+            _, closed_parent_state = self._partition_parent_state_map.popitem(last=False)
+            last_closed_state = closed_parent_state
+
+            # Clean up finished semaphores with value 0 up to and including earliest_key
+            for p_key in list(self._semaphore_per_partition.keys()):
+                sem = self._semaphore_per_partition[p_key]
+                if p_key in self._finished_partitions and sem._value == 0:
+                    del self._semaphore_per_partition[p_key]
+                    logger.debug(f"Deleted finished semaphore for partition {p_key} with value 0")
+                if p_key == earliest_key:
+                    break
+
+        # Update _parent_state if we popped at least one partition
+        if last_closed_state is not None:
+            self._parent_state = last_closed_state
 
     def ensure_at_least_one_state_emitted(self) -> None:
         """
-        The platform expect to have at least one state message on successful syncs. Hence, whatever happens, we expect this method to be
+        The platform expects at least one state message on successful syncs. Hence, whatever happens, we expect this method to be
         called.
         """
         if not any(
@@ -201,13 +256,19 @@ class ConcurrentPerPartitionCursor(Cursor):
 
         slices = self._partition_router.stream_slices()
         self._timer.start()
-        for partition in slices:
-            yield from self._generate_slices_from_partition(partition)
+        for partition, last, parent_state in iterate_with_last_flag_and_state(
+            slices, self._partition_router.get_stream_state
+        ):
+            yield from self._generate_slices_from_partition(partition, parent_state)
 
-    def _generate_slices_from_partition(self, partition: StreamSlice) -> Iterable[StreamSlice]:
+    def _generate_slices_from_partition(
+        self, partition: StreamSlice, parent_state: Mapping[str, Any]
+    ) -> Iterable[StreamSlice]:
         # Ensure the maximum number of partitions is not exceeded
         self._ensure_partition_limit()
 
+        partition_key = self._to_partition_key(partition.partition)
+
         cursor = self._cursor_per_partition.get(self._to_partition_key(partition.partition))
         if not cursor:
             cursor = self._create_cursor(
@@ -216,18 +277,26 @@ class ConcurrentPerPartitionCursor(Cursor):
             )
             with self._lock:
                 self._number_of_partitions += 1
-                self._cursor_per_partition[self._to_partition_key(partition.partition)] = cursor
-            self._semaphore_per_partition[self._to_partition_key(partition.partition)] = (
-                threading.Semaphore(0)
-            )
+                self._cursor_per_partition[partition_key] = cursor
+        self._semaphore_per_partition[partition_key] = threading.Semaphore(0)
+
+        with self._lock:
+            if (
+                len(self._partition_parent_state_map) == 0
+                or self._partition_parent_state_map[
+                    next(reversed(self._partition_parent_state_map))
+                ]
+                != parent_state
+            ):
+                self._partition_parent_state_map[partition_key] = deepcopy(parent_state)
 
         for cursor_slice, is_last_slice, _ in iterate_with_last_flag_and_state(
             cursor.stream_slices(),
             lambda: None,
         ):
-            self._semaphore_per_partition[self._to_partition_key(partition.partition)].release()
+            self._semaphore_per_partition[partition_key].release()
             if is_last_slice:
-                self._finished_partitions.add(self._to_partition_key(partition.partition))
+                self._finished_partitions.add(partition_key)
             yield StreamSlice(
                 partition=partition, cursor_slice=cursor_slice, extra_fields=partition.extra_fields
             )
@@ -257,9 +326,9 @@ class ConcurrentPerPartitionCursor(Cursor):
             while len(self._cursor_per_partition) > self.DEFAULT_MAX_PARTITIONS_NUMBER - 1:
                 # Try removing finished partitions first
                 for partition_key in list(self._cursor_per_partition.keys()):
-                    if (
-                        partition_key in self._finished_partitions
-                        and self._semaphore_per_partition[partition_key]._value == 0
+                    if partition_key in self._finished_partitions and (
+                        partition_key not in self._semaphore_per_partition
+                        or self._semaphore_per_partition[partition_key]._value == 0
                     ):
                         oldest_partition = self._cursor_per_partition.pop(
                             partition_key
@@ -338,9 +407,6 @@ class ConcurrentPerPartitionCursor(Cursor):
                 self._cursor_per_partition[self._to_partition_key(state["partition"])] = (
                     self._create_cursor(state["cursor"])
                 )
-                self._semaphore_per_partition[self._to_partition_key(state["partition"])] = (
-                    threading.Semaphore(0)
-                )
 
             # set default state for missing partitions if it is per partition with fallback to global
             if self._GLOBAL_STATE_KEY in stream_state:

{airbyte_cdk-6.36.1 → airbyte_cdk-6.36.3}/airbyte_cdk/sources/declarative/models/declarative_component_schema.py

@@ -2263,7 +2263,7 @@ class AsyncRetriever(BaseModel):
     status_extractor: Union[CustomRecordExtractor, DpathExtractor] = Field(
         ..., description="Responsible for fetching the actual status of the async job."
     )
-    urls_extractor: Union[CustomRecordExtractor, DpathExtractor] = Field(
+    download_target_extractor: Union[CustomRecordExtractor, DpathExtractor] = Field(
         ...,
         description="Responsible for fetching the final result `urls` provided by the completed / finished / ready async job.",
     )
@@ -2278,7 +2278,7 @@ class AsyncRetriever(BaseModel):
         ...,
         description="Requester component that describes how to prepare HTTP requests to send to the source API to fetch the status of the running async job.",
     )
-    url_requester: Optional[Union[CustomRequester, HttpRequester]] = Field(
+    download_target_requester: Optional[Union[CustomRequester, HttpRequester]] = Field(
         None,
         description="Requester component that describes how to prepare HTTP requests to send to the source API to extract the url from polling response by the completed async job.",
     )