fix linting but break test_per_partition_cursor.py for now

Author: maxime.c

airbyte_cdk/sources/declarative/incremental/concurrent_partition_cursor.py

@@ -432,9 +432,6 @@ def _set_initial_state(self, stream_state: StreamState) -> None:
         if stream_state.get("parent_state"):
             self._parent_state = stream_state["parent_state"]
 
-        # Set parent state for partition routers based on parent streams
-        self._partition_router.set_initial_state(stream_state)
-
     def _set_global_state(self, stream_state: Mapping[str, Any]) -> None:
         """
         Initializes the global cursor state from the provided stream state.
@@ -557,6 +554,9 @@ def limit_reached(self) -> bool:
     def get_parent_state(
         stream_state: Optional[StreamState], parent_stream_name: str
     ) -> Optional[AirbyteStateMessage]:
+        if not stream_state:
+            return stream_state
+
         if "parent_state" not in stream_state:
             logger.warning(
                 f"Trying to get_parent_state for stream `{parent_stream_name}` when there are not parent state in the state"

airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py

@@ -98,6 +98,7 @@
     RecordSelector,
     ResponseToFileExtractor,
 )
+from airbyte_cdk.sources.declarative.extractors.record_extractor import RecordExtractor
 from airbyte_cdk.sources.declarative.extractors.record_filter import (
     ClientSideIncrementalRecordFilterDecorator,
 )
@@ -106,7 +107,6 @@
     ConcurrentPerPartitionCursor,
     CursorFactory,
     DatetimeBasedCursor,
-    DeclarativeCursor,
     GlobalSubstreamCursor,
     PerPartitionWithGlobalCursor,
 )
@@ -512,7 +512,7 @@
     PerPartitionRequestOptionsProvider,
 )
 from airbyte_cdk.sources.declarative.requesters.request_path import RequestPath
-from airbyte_cdk.sources.declarative.requesters.requester import HttpMethod
+from airbyte_cdk.sources.declarative.requesters.requester import HttpMethod, Requester
 from airbyte_cdk.sources.declarative.resolvers import (
     ComponentMappingDefinition,
     ConfigComponentsResolver,
@@ -529,6 +529,7 @@
 from airbyte_cdk.sources.declarative.retrievers.file_uploader import (
     ConnectorBuilderFileUploader,
     DefaultFileUploader,
+    FileUploader,
     LocalFileSystemFileWriter,
     NoopFileWriter,
 )
@@ -553,6 +554,7 @@
 )
 from airbyte_cdk.sources.declarative.transformations import (
     AddFields,
+    RecordTransformation,
     RemoveFields,
 )
 from airbyte_cdk.sources.declarative.transformations.add_fields import AddedFieldDefinition
@@ -3255,7 +3257,7 @@ def create_simple_retriever(
         log_formatter: Optional[Callable[[Response], Any]] = None,
         **kwargs: Any,
     ) -> SimpleRetriever:
-        def _get_url() -> str:
+        def _get_url(req: Requester) -> str:
             """
             Closure to get the URL from the requester. This is used to get the URL in the case of a lazy retriever.
             This is needed because the URL is not set until the requester is created.
@@ -3264,12 +3266,12 @@ def _get_url() -> str:
             _url: str = (
                 model.requester.url
                 if hasattr(model.requester, "url") and model.requester.url is not None
-                else requester.get_url()
+                else req.get_url(stream_state=None, stream_slice=None, next_page_token=None)
             )
             _url_base: str = (
                 model.requester.url_base
                 if hasattr(model.requester, "url_base") and model.requester.url_base is not None
-                else requester.get_url_base()
+                else req.get_url(stream_state=None, stream_slice=None, next_page_token=None)
             )
 
             return _url or _url_base
@@ -3355,7 +3357,7 @@ def _get_url() -> str:
             self._create_component_from_model(
                 model=model.paginator,
                 config=config,
-                url_base=_get_url(),
+                url_base=_get_url(requester),
                 extractor_model=model.record_selector.extractor,
                 decoder=decoder,
                 cursor_used_for_stop_condition=stop_condition_cursor or None,
@@ -3538,12 +3540,14 @@ def create_async_retriever(
         transformations: List[RecordTransformation],
         **kwargs: Any,
     ) -> AsyncRetriever:
-        def _get_download_retriever() -> SimpleRetriever:
+        def _get_download_retriever(
+            requester: Requester, extractor: RecordExtractor, _decoder: Decoder
+        ) -> SimpleRetriever:
             # We create a record selector for the download retriever
             # with no schema normalization and no transformations, neither record filter
             # as all this occurs in the record_selector of the AsyncRetriever
             record_selector = RecordSelector(
-                extractor=download_extractor,
+                extractor=extractor,
                 name=name,
                 record_filter=None,
                 transformations=[],
@@ -3554,7 +3558,7 @@ def _get_download_retriever() -> SimpleRetriever:
             paginator = (
                 self._create_component_from_model(
                     model=model.download_paginator,
-                    decoder=decoder,
+                    decoder=_decoder,
                     config=config,
                     url_base="",
                 )
@@ -3563,7 +3567,7 @@ def _get_download_retriever() -> SimpleRetriever:
             )
 
             return SimpleRetriever(
-                requester=download_requester,
+                requester=requester,
                 record_selector=record_selector,
                 primary_key=None,
                 name=name,
@@ -3657,7 +3661,9 @@ def _get_job_timeout() -> datetime.timedelta:
             config=config,
             name=job_download_components_name,
         )
-        download_retriever = _get_download_retriever()
+        download_retriever = _get_download_retriever(
+            download_requester, download_extractor, download_decoder
+        )
         abort_requester = (
             self._create_component_from_model(
                 model=model.abort_requester,
@@ -3840,7 +3846,9 @@ def _create_message_repository_substream_wrapper(
             model=model, config=config, has_parent_state=has_parent_state, **kwargs
         )
 
-    def _instantiate_parent_stream_state_manager(self, child_state, config, model):
+    def _instantiate_parent_stream_state_manager(
+        self, child_state: MutableMapping[str, Any], config: Config, model: ParentStreamConfigModel
+    ):
         """
         With DefaultStream, the state needs to be provided during __init__ of the cursor as opposed to the
         `set_initial_state` flow that existed for the declarative cursors. This state is taken from

airbyte_cdk/sources/declarative/partition_routers/cartesian_product_stream_slicer.py

@@ -159,12 +159,6 @@ def stream_slices(self) -> Iterable[StreamSlice]:
                 cursor_slice = {}
             yield StreamSlice(partition=partition, cursor_slice=cursor_slice)
 
-    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 @@ def get_request_body_json(
     ) -> 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 @@ def _get_request_option(
         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 @@ def get_request_body_json(
     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

@@ -36,7 +36,10 @@
     from airbyte_cdk.sources.streams.concurrent.abstract_stream import AbstractStream
 
 
-def iterate_with_last_flag(generator: Iterable[Partition]) -> Iterable[tuple[Partition, bool]]:
+T = TypeVar("T")
+
+
+def iterate_with_last_flag(generator: Iterable[T]) -> Iterable[tuple[T, bool]]:
     iterator = iter(generator)
 
     try:
@@ -307,60 +310,6 @@ def _extract_extra_fields(
                 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.

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

@@ -19,11 +19,15 @@ def get_request_params(
     ) -> 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={}),
+            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),
+            stream_slice=StreamSlice(partition={}, cursor_slice=stream_slice.cursor_slice)
+            if stream_slice
+            else StreamSlice(partition={}, cursor_slice={}),
             next_page_token=next_page_token,
         )
 
@@ -36,11 +40,15 @@ def get_request_headers(
     ) -> 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={}),
+            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),
+            stream_slice=StreamSlice(partition={}, cursor_slice=stream_slice.cursor_slice)
+            if stream_slice
+            else stream_slice,
             next_page_token=next_page_token,
         )
 
@@ -53,11 +61,15 @@ def get_request_body_data(
     ) -> 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={}),
+            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),
+            stream_slice=StreamSlice(partition={}, cursor_slice=stream_slice.cursor_slice)
+            if stream_slice
+            else stream_slice,
             next_page_token=next_page_token,
         )
 
@@ -70,10 +82,14 @@ def get_request_body_json(
     ) -> 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={}),
+            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),
+            stream_slice=StreamSlice(partition={}, cursor_slice=stream_slice.cursor_slice)
+            if stream_slice
+            else stream_slice,
             next_page_token=next_page_token,
         )