[WIP still tests failing] Remove DeclarativeStream instantiation

Author: maxime.c

airbyte_cdk/sources/declarative/incremental/concurrent_partition_cursor.py

@@ -11,6 +11,7 @@
 from datetime import timedelta
 from typing import Any, Callable, Iterable, List, Mapping, MutableMapping, Optional
 
+from airbyte_cdk.models import AirbyteStateMessage, AirbyteStateBlob, AirbyteStreamState, AirbyteStateType, StreamDescriptor
 from airbyte_cdk.sources.connector_state_manager import ConnectorStateManager
 from airbyte_cdk.sources.declarative.incremental.global_substream_cursor import (
     Timer,
@@ -128,6 +129,7 @@ def __init__(
 
         # FIXME this is a temporary field the time of the migration from declarative cursors to concurrent ones
         self._attempt_to_create_cursor_if_not_provided = attempt_to_create_cursor_if_not_provided
+        self._synced_some_data = False
 
     @property
     def cursor_field(self) -> CursorField:
@@ -168,8 +170,8 @@ def close_partition(self, partition: Partition) -> None:
         with self._lock:
             self._semaphore_per_partition[partition_key].acquire()
             if not self._use_global_cursor:
-                self._cursor_per_partition[partition_key].close_partition(partition=partition)
                 cursor = self._cursor_per_partition[partition_key]
+                cursor.close_partition(partition=partition)
                 if (
                     partition_key in self._partitions_done_generating_stream_slices
                     and self._semaphore_per_partition[partition_key]._value == 0
@@ -213,8 +215,10 @@ def ensure_at_least_one_state_emitted(self) -> None:
         if not any(
             semaphore_item[1]._value for semaphore_item in self._semaphore_per_partition.items()
         ):
-            self._global_cursor = self._new_global_cursor
-            self._lookback_window = self._timer.finish()
+            if self._synced_some_data:
+                # we only update those if we actually synced some data
+                self._global_cursor = self._new_global_cursor
+                self._lookback_window = self._timer.finish()
             self._parent_state = self._partition_router.get_stream_state()
         self._emit_state_message(throttle=False)
 
@@ -458,6 +462,7 @@ def observe(self, record: Record) -> None:
         except ValueError:
             return
 
+        self._synced_some_data = True
         record_cursor = self._connector_state_converter.output_format(
             self._connector_state_converter.parse_value(record_cursor_value)
         )
@@ -541,3 +546,23 @@ def _get_cursor(self, record: Record) -> ConcurrentCursor:
 
     def limit_reached(self) -> bool:
         return self._number_of_partitions > self.SWITCH_TO_GLOBAL_LIMIT
+
+    @staticmethod
+    def get_parent_state(stream_state: Optional[StreamState], parent_stream_name: str) -> Optional[AirbyteStateMessage]:
+        return AirbyteStateMessage(
+            type=AirbyteStateType.STREAM,
+            stream=AirbyteStreamState(
+                stream_descriptor=StreamDescriptor(parent_stream_name, None),
+                stream_state=AirbyteStateBlob(stream_state["parent_state"][parent_stream_name])
+            )
+        ) if stream_state and "parent_state" in stream_state else None
+
+    @staticmethod
+    def get_global_state(stream_state: Optional[StreamState], parent_stream_name: str) -> Optional[AirbyteStateMessage]:
+        return AirbyteStateMessage(
+            type=AirbyteStateType.STREAM,
+            stream=AirbyteStreamState(
+                stream_descriptor=StreamDescriptor(parent_stream_name, None),
+                stream_state=AirbyteStateBlob(stream_state["state"])
+            )
+        ) if stream_state and "state" in stream_state else None

airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py

@@ -7,24 +7,39 @@
 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, Union, TypeVar
 
 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
+
+
+def iterate_with_last_flag(generator: Iterable[Partition]) -> Iterable[tuple[Partition, 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 +55,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 +191,51 @@ def stream_slices(self) -> Iterable[StreamSlice]:
                         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()):
+                        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)}"
-                        )
-                    try:
-                        partition_value = dpath.get(
-                            parent_record,  # type: ignore [arg-type]
-                            parent_field,
+                        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)
+
+                        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)
+
+                        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,
-                    )
+
+                parent_stream.cursor.ensure_at_least_one_state_emitted()
+                yield from []
 
     def _extract_child_response(
         self, parent_record: Mapping[str, Any] | AirbyteMessage, pointer: List[InterpolatedString]
@@ -414,7 +421,7 @@ def get_stream_state(self) -> Optional[Mapping[str, StreamState]]:
         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/partition_routers/substream_partition_router.py

@@ -0,0 +1,79 @@
+from typing import Optional, Mapping, Any, 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={}),
+            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),
+            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={}),
+            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),
+            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={}),
+            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),
+            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={}),
+            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),
+            next_page_token=next_page_token,
+        )

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

@@ -17,6 +17,7 @@
 )
 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 @@ def resolve_components(
         """
         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/resolvers/http_components_resolver.py

@@ -4,6 +4,7 @@
 
 from abc import abstractmethod
 from typing import Any, Iterable, Mapping, Optional
+from typing_extensions import deprecated
 
 from airbyte_cdk.sources.declarative.incremental.per_partition_cursor import StreamSlice
 from airbyte_cdk.sources.streams.core import StreamData
@@ -30,11 +31,13 @@ def read_records(
         """
 
     @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.
@@ -50,5 +53,6 @@ def state(self) -> StreamState:
 
     @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/retrievers/retriever.py

@@ -95,6 +95,21 @@ def consume_queue(self) -> Iterable[AirbyteMessage]:
             yield self._message_queue.popleft()
 
 
+class StateFilteringMessageRepository(MessageRepository):
+    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,