airbyte-cdk 6.61.6__py3-none-any.whl → 6.62.0.dev1__py3-none-any.whl

airbyte_cdk/sources/declarative/partition_routers/cartesian_product_stream_slicer.py

@@ -162,12 +162,6 @@ class CartesianProductStreamSlicer(PartitionRouter):
                 partition=partition, cursor_slice=cursor_slice, extra_fields=extra_fields
             )
 
-    def set_initial_state(self, stream_state: StreamState) -> None:
-        """
-        Parent stream states are not supported for cartesian product stream slicer
-        """
-        pass
-
     def get_stream_state(self) -> Optional[Mapping[str, StreamState]]:
         """
         Parent stream states are not supported for cartesian product stream slicer

airbyte_cdk/sources/declarative/partition_routers/grouping_partition_router.py

@@ -140,11 +140,6 @@ class GroupingPartitionRouter(PartitionRouter):
     ) -> Mapping[str, Any]:
         return {}
 
-    def set_initial_state(self, stream_state: StreamState) -> None:
-        """Delegate state initialization to the underlying partition router."""
-        self.underlying_partition_router.set_initial_state(stream_state)
-        self._state = self.underlying_partition_router.get_stream_state()
-
     def get_stream_state(self) -> Optional[Mapping[str, StreamState]]:
         """Delegate state retrieval to the underlying partition router."""
         return self._state

airbyte_cdk/sources/declarative/partition_routers/list_partition_router.py

@@ -108,12 +108,6 @@ class ListPartitionRouter(PartitionRouter):
         else:
             return {}
 
-    def set_initial_state(self, stream_state: StreamState) -> None:
-        """
-        ListPartitionRouter doesn't have parent streams
-        """
-        pass
-
     def get_stream_state(self) -> Optional[Mapping[str, StreamState]]:
         """
         ListPartitionRouter doesn't have parent streams

airbyte_cdk/sources/declarative/partition_routers/partition_router.py

@@ -15,31 +15,9 @@ class PartitionRouter(StreamSlicer):
     """
     Base class for partition routers.
     Methods:
-        set_parent_state(stream_state): Set the state of the parent streams.
-        get_parent_state(): Get the state of the parent streams.
+        get_stream_state(): Get the state of the parent streams.
     """
 
-    @abstractmethod
-    def set_initial_state(self, stream_state: StreamState) -> None:
-        """
-        Set the state of the parent streams.
-
-        This method should only be implemented if the slicer is based on some parent stream and needs to read this stream
-        incrementally using the state.
-
-        Args:
-            stream_state (StreamState): The state of the streams to be set. The expected format is a dictionary that includes
-                                        'parent_state' which is a dictionary of parent state names to their corresponding state.
-                Example:
-                {
-                    "parent_state": {
-                        "parent_stream_name_1": { ... },
-                        "parent_stream_name_2": { ... },
-                        ...
-                    }
-                }
-        """
-
     @abstractmethod
     def get_stream_state(self) -> Optional[Mapping[str, StreamState]]:
         """

airbyte_cdk/sources/declarative/partition_routers/single_partition_router.py

@@ -50,12 +50,6 @@ class SinglePartitionRouter(PartitionRouter):
     def stream_slices(self) -> Iterable[StreamSlice]:
         yield StreamSlice(partition={}, cursor_slice={})
 
-    def set_initial_state(self, stream_state: StreamState) -> None:
-        """
-        SinglePartitionRouter doesn't have parent streams
-        """
-        pass
-
     def get_stream_state(self) -> Optional[Mapping[str, StreamState]]:
         """
         SinglePartitionRouter doesn't have parent streams

airbyte_cdk/sources/declarative/partition_routers/substream_partition_router.py

@@ -7,24 +7,51 @@ import copy
 import json
 import logging
 from dataclasses import InitVar, dataclass
-from typing import TYPE_CHECKING, Any, Iterable, List, Mapping, MutableMapping, Optional, Union
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Iterable,
+    List,
+    Mapping,
+    MutableMapping,
+    Optional,
+    TypeVar,
+    Union,
+)
 
 import dpath
 import requests
 
 from airbyte_cdk.models import AirbyteMessage
-from airbyte_cdk.models import Type as MessageType
 from airbyte_cdk.sources.declarative.interpolation.interpolated_string import InterpolatedString
 from airbyte_cdk.sources.declarative.partition_routers.partition_router import PartitionRouter
 from airbyte_cdk.sources.declarative.requesters.request_option import (
     RequestOption,
     RequestOptionType,
 )
-from airbyte_cdk.sources.types import Config, Record, StreamSlice, StreamState
-from airbyte_cdk.utils import AirbyteTracedException
+from airbyte_cdk.sources.streams.concurrent.partitions.partition import Partition
+from airbyte_cdk.sources.types import Config, StreamSlice, StreamState
 
 if TYPE_CHECKING:
-    from airbyte_cdk.sources.declarative.declarative_stream import DeclarativeStream
+    from airbyte_cdk.sources.streams.concurrent.abstract_stream import AbstractStream
+
+
+T = TypeVar("T")
+
+
+def iterate_with_last_flag(generator: Iterable[T]) -> Iterable[tuple[T, bool]]:
+    iterator = iter(generator)
+
+    try:
+        current = next(iterator)
+    except StopIteration:
+        return  # Return an empty iterator
+
+    for next_item in iterator:
+        yield current, False
+        current = next_item
+
+    yield current, True
 
 
 @dataclass
@@ -40,7 +67,7 @@ class ParentStreamConfig:
     incremental_dependency (bool): Indicates if the parent stream should be read incrementally.
     """
 
-    stream: "DeclarativeStream"  # Parent streams must be DeclarativeStream because we can't know which part of the stream slice is a partition for regular Stream
+    stream: "AbstractStream"
     parent_key: Union[InterpolatedString, str]
     partition_field: Union[InterpolatedString, str]
     config: Config
@@ -176,59 +203,65 @@ class SubstreamPartitionRouter(PartitionRouter):
                         for field_path in parent_stream_config.extra_fields
                     ]
 
-                # read_stateless() assumes the parent is not concurrent. This is currently okay since the concurrent CDK does
-                # not support either substreams or RFR, but something that needs to be considered once we do
-                for parent_record in parent_stream.read_only_records():
-                    parent_partition = None
-                    # Skip non-records (eg AirbyteLogMessage)
-                    if isinstance(parent_record, AirbyteMessage):
-                        self.logger.warning(
-                            f"Parent stream {parent_stream.name} returns records of type AirbyteMessage. This SubstreamPartitionRouter is not able to checkpoint incremental parent state."
-                        )
-                        if parent_record.type == MessageType.RECORD:
-                            parent_record = parent_record.record.data  # type: ignore[union-attr, assignment]  # record is always a Record
-                        else:
-                            continue
-                    elif isinstance(parent_record, Record):
+                for partition, is_last_slice in iterate_with_last_flag(
+                    parent_stream.generate_partitions()
+                ):
+                    for parent_record, is_last_record_in_slice in iterate_with_last_flag(
+                        partition.read()
+                    ):
+                        # In the previous CDK implementation, state management was done internally by the stream.
+                        # However, this could cause issues when doing availability check for example as the availability
+                        # check would progress the state so state management was moved outside of the read method.
+                        # Hence, we need to call the cursor here.
+                        # Note that we call observe and close_partition before emitting the associated record as the
+                        # ConcurrentPerPartitionCursor will associate a record with the state of the stream after the
+                        # record was consumed.
+                        parent_stream.cursor.observe(parent_record)
                         parent_partition = (
                             parent_record.associated_slice.partition
                             if parent_record.associated_slice
                             else {}
                         )
-                        parent_record = parent_record.data
-                    elif not isinstance(parent_record, Mapping):
-                        # The parent_record should only take the form of a Record, AirbyteMessage, or Mapping. Anything else is invalid
-                        raise AirbyteTracedException(
-                            message=f"Parent stream returned records as invalid type {type(parent_record)}"
+                        record_data = parent_record.data
+
+                        try:
+                            partition_value = dpath.get(
+                                record_data,  # type: ignore [arg-type]
+                                parent_field,
+                            )
+                        except KeyError:
+                            # FIXME a log here would go a long way for debugging
+                            continue
+
+                        # Add extra fields
+                        extracted_extra_fields = self._extract_extra_fields(
+                            record_data, extra_fields
                         )
-                    try:
-                        partition_value = dpath.get(
-                            parent_record,  # type: ignore [arg-type]
-                            parent_field,
+
+                        if parent_stream_config.lazy_read_pointer:
+                            extracted_extra_fields = {
+                                "child_response": self._extract_child_response(
+                                    record_data,
+                                    parent_stream_config.lazy_read_pointer,  # type: ignore[arg-type]  # lazy_read_pointer type handeled in __post_init__ of parent_stream_config
+                                ),
+                                **extracted_extra_fields,
+                            }
+
+                        if is_last_record_in_slice:
+                            parent_stream.cursor.close_partition(partition)
+                            if is_last_slice:
+                                parent_stream.cursor.ensure_at_least_one_state_emitted()
+
+                        yield StreamSlice(
+                            partition={
+                                partition_field: partition_value,
+                                "parent_slice": parent_partition or {},
+                            },
+                            cursor_slice={},
+                            extra_fields=extracted_extra_fields,
                         )
-                    except KeyError:
-                        continue
-
-                    # Add extra fields
-                    extracted_extra_fields = self._extract_extra_fields(parent_record, extra_fields)
-
-                    if parent_stream_config.lazy_read_pointer:
-                        extracted_extra_fields = {
-                            "child_response": self._extract_child_response(
-                                parent_record,
-                                parent_stream_config.lazy_read_pointer,  # type: ignore[arg-type]  # lazy_read_pointer type handeled in __post_init__ of parent_stream_config
-                            ),
-                            **extracted_extra_fields,
-                        }
-
-                    yield StreamSlice(
-                        partition={
-                            partition_field: partition_value,
-                            "parent_slice": parent_partition or {},
-                        },
-                        cursor_slice={},
-                        extra_fields=extracted_extra_fields,
-                    )
+
+                yield from []
 
     def _extract_child_response(
         self, parent_record: Mapping[str, Any] | AirbyteMessage, pointer: List[InterpolatedString]
@@ -278,60 +311,6 @@ class SubstreamPartitionRouter(PartitionRouter):
                 extracted_extra_fields[".".join(extra_field_path)] = extra_field_value
         return extracted_extra_fields
 
-    def set_initial_state(self, stream_state: StreamState) -> None:
-        """
-        Set the state of the parent streams.
-
-        If the `parent_state` key is missing from `stream_state`, migrate the child stream state to the parent stream's state format.
-        This migration applies only to parent streams with incremental dependencies.
-
-        Args:
-            stream_state (StreamState): The state of the streams to be set.
-
-        Example of state format:
-        {
-            "parent_state": {
-                "parent_stream_name1": {
-                    "last_updated": "2023-05-27T00:00:00Z"
-                },
-                "parent_stream_name2": {
-                    "last_updated": "2023-05-27T00:00:00Z"
-                }
-            }
-        }
-
-        Example of migrating to parent state format:
-        - Initial state:
-        {
-            "updated_at": "2023-05-27T00:00:00Z"
-        }
-        - After migration:
-        {
-            "updated_at": "2023-05-27T00:00:00Z",
-            "parent_state": {
-                "parent_stream_name": {
-                    "parent_stream_cursor": "2023-05-27T00:00:00Z"
-                }
-            }
-        }
-        """
-        if not stream_state:
-            return
-
-        parent_state = stream_state.get("parent_state", {})
-
-        # Set state for each parent stream with an incremental dependency
-        for parent_config in self.parent_stream_configs:
-            if (
-                not parent_state.get(parent_config.stream.name, {})
-                and parent_config.incremental_dependency
-            ):
-                # Migrate child state to parent state format
-                parent_state = self._migrate_child_state_to_parent_state(stream_state)
-
-            if parent_config.incremental_dependency:
-                parent_config.stream.state = parent_state.get(parent_config.stream.name, {})
-
     def _migrate_child_state_to_parent_state(self, stream_state: StreamState) -> StreamState:
         """
         Migrate the child or global stream state into the parent stream's state format.
@@ -414,7 +393,9 @@ class SubstreamPartitionRouter(PartitionRouter):
         parent_state = {}
         for parent_config in self.parent_stream_configs:
             if parent_config.incremental_dependency:
-                parent_state[parent_config.stream.name] = copy.deepcopy(parent_config.stream.state)
+                parent_state[parent_config.stream.name] = copy.deepcopy(
+                    parent_config.stream.cursor.state
+                )
         return parent_state
 
     @property

airbyte_cdk/sources/declarative/requesters/request_options/per_partition_request_option_provider.py

@@ -0,0 +1,95 @@
+from typing import Any, Mapping, Optional, Union
+
+from airbyte_cdk.sources.declarative.partition_routers import PartitionRouter
+from airbyte_cdk.sources.declarative.requesters.request_options import RequestOptionsProvider
+from airbyte_cdk.sources.types import StreamSlice, StreamState
+
+
+class PerPartitionRequestOptionsProvider(RequestOptionsProvider):
+    def __init__(self, partition_router: PartitionRouter, cursor_provider: RequestOptionsProvider):
+        self._partition_router = partition_router
+        self._cursor_provider = cursor_provider
+
+    def get_request_params(
+        self,
+        *,
+        stream_state: Optional[StreamState] = None,
+        stream_slice: Optional[StreamSlice] = None,
+        next_page_token: Optional[Mapping[str, Any]] = None,
+    ) -> Mapping[str, Any]:
+        return self._partition_router.get_request_params(  # type: ignore # this always returns a mapping
+            stream_state=stream_state,
+            stream_slice=StreamSlice(partition=stream_slice.partition, cursor_slice={})
+            if stream_slice
+            else StreamSlice(partition={}, cursor_slice={}),
+            next_page_token=next_page_token,
+        ) | self._cursor_provider.get_request_params(
+            stream_state=stream_state,
+            stream_slice=StreamSlice(partition={}, cursor_slice=stream_slice.cursor_slice)
+            if stream_slice
+            else StreamSlice(partition={}, cursor_slice={}),
+            next_page_token=next_page_token,
+        )
+
+    def get_request_headers(
+        self,
+        *,
+        stream_state: Optional[StreamState] = None,
+        stream_slice: Optional[StreamSlice] = None,
+        next_page_token: Optional[Mapping[str, Any]] = None,
+    ) -> Mapping[str, Any]:
+        return self._partition_router.get_request_headers(  # type: ignore # this always returns a mapping
+            stream_state=stream_state,
+            stream_slice=StreamSlice(partition=stream_slice.partition, cursor_slice={})
+            if stream_slice
+            else stream_slice,
+            next_page_token=next_page_token,
+        ) | self._cursor_provider.get_request_headers(
+            stream_state=stream_state,
+            stream_slice=StreamSlice(partition={}, cursor_slice=stream_slice.cursor_slice)
+            if stream_slice
+            else stream_slice,
+            next_page_token=next_page_token,
+        )
+
+    def get_request_body_data(
+        self,
+        *,
+        stream_state: Optional[StreamState] = None,
+        stream_slice: Optional[StreamSlice] = None,
+        next_page_token: Optional[Mapping[str, Any]] = None,
+    ) -> Union[Mapping[str, Any], str]:
+        return self._partition_router.get_request_body_data(  # type: ignore # this always returns a mapping
+            stream_state=stream_state,
+            stream_slice=StreamSlice(partition=stream_slice.partition, cursor_slice={})
+            if stream_slice
+            else stream_slice,
+            next_page_token=next_page_token,
+        ) | self._cursor_provider.get_request_body_data(
+            stream_state=stream_state,
+            stream_slice=StreamSlice(partition={}, cursor_slice=stream_slice.cursor_slice)
+            if stream_slice
+            else stream_slice,
+            next_page_token=next_page_token,
+        )
+
+    def get_request_body_json(
+        self,
+        *,
+        stream_state: Optional[StreamState] = None,
+        stream_slice: Optional[StreamSlice] = None,
+        next_page_token: Optional[Mapping[str, Any]] = None,
+    ) -> Mapping[str, Any]:
+        return self._partition_router.get_request_body_json(  # type: ignore # this always returns a mapping
+            stream_state=stream_state,
+            stream_slice=StreamSlice(partition=stream_slice.partition, cursor_slice={})
+            if stream_slice
+            else stream_slice,
+            next_page_token=next_page_token,
+        ) | self._cursor_provider.get_request_body_json(
+            stream_state=stream_state,
+            stream_slice=StreamSlice(partition={}, cursor_slice=stream_slice.cursor_slice)
+            if stream_slice
+            else stream_slice,
+            next_page_token=next_page_token,
+        )

airbyte_cdk/sources/declarative/resolvers/http_components_resolver.py

@@ -17,6 +17,7 @@ from airbyte_cdk.sources.declarative.resolvers.components_resolver import (
 )
 from airbyte_cdk.sources.declarative.retrievers.retriever import Retriever
 from airbyte_cdk.sources.source import ExperimentalClassWarning
+from airbyte_cdk.sources.streams.concurrent.partitions.stream_slicer import StreamSlicer
 from airbyte_cdk.sources.types import Config
 
 
@@ -28,12 +29,14 @@ class HttpComponentsResolver(ComponentsResolver):
 
     Attributes:
         retriever (Retriever): The retriever used to fetch data from an API.
+        stream_slicer (StreamSlicer): The how the data is sliced.
         config (Config): Configuration object for the resolver.
         components_mapping (List[ComponentMappingDefinition]): List of mappings to resolve.
         parameters (InitVar[Mapping[str, Any]]): Additional parameters for interpolation.
     """
 
     retriever: Retriever
+    stream_slicer: StreamSlicer
     config: Config
     components_mapping: List[ComponentMappingDefinition]
     parameters: InitVar[Mapping[str, Any]]
@@ -88,7 +91,7 @@ class HttpComponentsResolver(ComponentsResolver):
         """
         kwargs = {"stream_template_config": stream_template_config}
 
-        for stream_slice in self.retriever.stream_slices():
+        for stream_slice in self.stream_slicer.stream_slices():
             for components_values in self.retriever.read_records(
                 records_schema={}, stream_slice=stream_slice
             ):

airbyte_cdk/sources/declarative/retrievers/retriever.py

@@ -5,6 +5,8 @@
 from abc import abstractmethod
 from typing import Any, Iterable, Mapping, Optional
 
+from typing_extensions import deprecated
+
 from airbyte_cdk.sources.streams.core import StreamData
 from airbyte_cdk.sources.types import StreamSlice, StreamState
 
@@ -29,11 +31,13 @@ class Retriever:
         """
 
     @abstractmethod
+    @deprecated("Stream slicing is being moved to the stream level.")
     def stream_slices(self) -> Iterable[Optional[StreamSlice]]:
         """Returns the stream slices"""
 
     @property
     @abstractmethod
+    @deprecated("State management is being moved to the stream level.")
     def state(self) -> StreamState:
         """State getter, should return state in form that can serialized to a string and send to the output
         as a STATE AirbyteMessage.
@@ -49,5 +53,6 @@ class Retriever:
 
     @state.setter
     @abstractmethod
+    @deprecated("State management is being moved to the stream level.")
     def state(self, value: StreamState) -> None:
         """State setter, accept state serialized by state getter."""

airbyte_cdk/sources/declarative/stream_slicers/declarative_partition_generator.py

@@ -14,21 +14,10 @@ from airbyte_cdk.sources.streams.concurrent.partitions.stream_slicer import Stre
 from airbyte_cdk.sources.types import Record, StreamSlice
 from airbyte_cdk.utils.slice_hasher import SliceHasher
 
-
 # For Connector Builder test read operations, we track the total number of records
-# read for the stream so that we can stop reading early if we exceed the record limit.
-class RecordCounter:
-    def __init__(self) -> None:
-        self.total_record_counter = 0
-
-    def increment(self) -> None:
-        self.total_record_counter += 1
-
-    def reset(self) -> None:
-        self.total_record_counter = 0
-
-    def get_total_records(self) -> int:
-        return self.total_record_counter
+# read for the stream at the global level so that we can stop reading early if we
+# exceed the record limit
+total_record_counter = 0
 
 
 class SchemaLoaderCachingDecorator(SchemaLoader):
@@ -62,7 +51,6 @@ class DeclarativePartitionFactory:
         self._retriever = retriever
         self._message_repository = message_repository
         self._max_records_limit = max_records_limit
-        self._record_counter = RecordCounter()
 
     def create(self, stream_slice: StreamSlice) -> Partition:
         return DeclarativePartition(
@@ -72,7 +60,6 @@ class DeclarativePartitionFactory:
             message_repository=self._message_repository,
             max_records_limit=self._max_records_limit,
             stream_slice=stream_slice,
-            record_counter=self._record_counter,
         )
 
 
@@ -85,7 +72,6 @@ class DeclarativePartition(Partition):
         message_repository: MessageRepository,
         max_records_limit: Optional[int],
         stream_slice: StreamSlice,
-        record_counter: RecordCounter,
     ):
         self._stream_name = stream_name
         self._schema_loader = schema_loader
@@ -94,17 +80,17 @@ class DeclarativePartition(Partition):
         self._max_records_limit = max_records_limit
         self._stream_slice = stream_slice
         self._hash = SliceHasher.hash(self._stream_name, self._stream_slice)
-        self._record_counter = record_counter
 
     def read(self) -> Iterable[Record]:
         if self._max_records_limit is not None:
-            if self._record_counter.get_total_records() >= self._max_records_limit:
+            global total_record_counter
+            if total_record_counter >= self._max_records_limit:
                 return
         for stream_data in self._retriever.read_records(
             self._schema_loader.get_json_schema(), self._stream_slice
         ):
             if self._max_records_limit is not None:
-                if self._record_counter.get_total_records() >= self._max_records_limit:
+                if total_record_counter >= self._max_records_limit:
                     break
 
             if isinstance(stream_data, Mapping):
@@ -122,7 +108,7 @@ class DeclarativePartition(Partition):
                 self._message_repository.emit_message(stream_data)
 
             if self._max_records_limit is not None:
-                self._record_counter.increment()
+                total_record_counter += 1
 
     def to_slice(self) -> Optional[Mapping[str, Any]]:
         return self._stream_slice

airbyte_cdk/sources/declarative/yaml_declarative_source.py

@@ -14,7 +14,7 @@ from airbyte_cdk.sources.declarative.concurrent_declarative_source import (
 from airbyte_cdk.sources.types import ConnectionDefinition
 
 
-class YamlDeclarativeSource(ConcurrentDeclarativeSource):
+class YamlDeclarativeSource(ConcurrentDeclarativeSource[List[AirbyteStateMessage]]):
     """Declarative source defined by a yaml file"""
 
     def __init__(

airbyte_cdk/sources/message/repository.py

@@ -95,6 +95,26 @@ class InMemoryMessageRepository(MessageRepository):
             yield self._message_queue.popleft()
 
 
+class StateFilteringMessageRepository(MessageRepository):
+    """
+    This message repository is used when creating parent streams for SubstreamPartitionRouter. As the child stream
+    manages the state for both the child and the parents, we want to prevent parents from emitting state messages.
+    """
+
+    def __init__(self, decorated: MessageRepository) -> None:
+        self._decorated = decorated
+
+    def emit_message(self, message: AirbyteMessage) -> None:
+        if message.type != Type.STATE:
+            self._decorated.emit_message(message)
+
+    def log_message(self, level: Level, message_provider: Callable[[], LogMessage]) -> None:
+        self._decorated.log_message(level, message_provider)
+
+    def consume_queue(self) -> Iterable[AirbyteMessage]:
+        yield from self._decorated.consume_queue()
+
+
 class LogAppenderMessageRepositoryDecorator(MessageRepository):
     def __init__(
         self,