Fix mypy

Author: lazebnyi

airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py

@@ -2711,15 +2711,27 @@ def create_simple_retriever(
             model.ignore_stream_slicer_parameters_on_paginated_requests or False
         )
 
+        if model.lazy_read_pointer and not hasattr(model, "partition_router"):
+            raise ValueError(
+                "LazySimpleRetriever requires a 'partition_router' when 'lazy_read_pointer' is set. "
+                "Please either define 'partition_router' or remove 'lazy_read_pointer' from the model."
+            )
+
         if model.lazy_read_pointer and not bool(
-            self._connector_state_manager.get_stream_state(name, None)
+                self._connector_state_manager.get_stream_state(name, None)
         ):
+            if model.partition_router.type != "SubstreamPartitionRouterModel":  # type: ignore[union-attr] # model.partition_router has BaseModel type
+                raise ValueError(
+                    "LazySimpleRetriever only supports 'SubstreamPartitionRouterModel' as the 'partition_router' type. "  # type: ignore[union-attr] # model.partition_router has BaseModel type
+                    f"Found: '{model.partition_router.type}'."
+                )
+
             lazy_read_pointer = [
                 InterpolatedString.create(path, parameters=model.parameters or {})
                 for path in model.lazy_read_pointer
             ]
             partition_router = self._create_component_from_model(
-                model=model.partition_router, config=config
+                model=model.partition_router, config=config  # type: ignore[arg-type]
             )
             stream_slicer = (
                 self._create_component_from_model(model=incremental_sync, config=config)

airbyte_cdk/sources/declarative/partition_routers/substream_partition_router.py

@@ -144,8 +144,8 @@ def _get_request_option(
         return params
 
     def process_parent_record(
-        self, parent_record: Union[AirbyteMessage, Record, Mapping], parent_stream_name: str
-    ) -> Tuple[Optional[Mapping], Optional[Mapping]]:
+        self, parent_record: Union[AirbyteMessage, Record, Mapping[str, Any]], parent_stream_name: str
+    ) -> Tuple[Optional[MutableMapping[str, Any]], Optional[MutableMapping[str, Any]]]:
         """
         Processes and extracts data from a parent record, handling different record types
         and ensuring only valid types proceed.
@@ -161,23 +161,21 @@ def process_parent_record(
                 f"This SubstreamPartitionRouter is not able to checkpoint incremental parent state."
             )
             if parent_record.type == MessageType.RECORD:
-                return parent_record.record.data, {}
+                return parent_record.record.data, {}  # type: ignore[union-attr] # parent_record.record is always AirbyteRecordMessage
             return None, None  # Skip invalid or non-record data
 
-        # Handle Record type
         if isinstance(parent_record, Record):
             parent_partition = (
                 parent_record.associated_slice.partition if parent_record.associated_slice else {}
             )
-            return parent_record.data, parent_partition
+            return {**parent_record.data}, {**parent_partition}
 
-        # Validate the record type
-        if not isinstance(parent_record, Mapping):
-            raise AirbyteTracedException(
-                message=f"Parent stream returned records as invalid type {type(parent_record)}"
-            )
+        if isinstance(parent_record, Mapping):
+            return {**parent_record}, {}
 
-        return parent_record, {}
+        raise AirbyteTracedException(
+            message=f"Parent stream returned records as invalid type {type(parent_record)}"
+        )
 
     def stream_slices(self) -> Iterable[StreamSlice]:
         """
@@ -210,10 +208,10 @@ def stream_slices(self) -> Iterable[StreamSlice]:
 
                 # 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():
+                for raw_parent_record in parent_stream.read_only_records():
                     # Process the parent record
                     parent_record, parent_partition = self.process_parent_record(
-                        parent_record, parent_stream.name
+                        raw_parent_record, parent_stream.name
                     )
 
                     # Skip invalid or non-record data

airbyte_cdk/sources/declarative/retrievers/simple_retriever.py

@@ -6,10 +6,11 @@
 from dataclasses import InitVar, dataclass, field
 from functools import partial
 from itertools import islice
-from typing import Any, Callable, Iterable, List, Mapping, Optional, Set, Tuple, Union
+from typing import Any, Callable, Iterable, List, Mapping, Optional, Set, Tuple, Union, MutableMapping
 
 import dpath
 import requests
+from typing_extensions import deprecated
 
 from airbyte_cdk.models import AirbyteMessage
 from airbyte_cdk.sources.declarative.extractors.http_selector import HttpSelector
@@ -35,6 +36,7 @@
 from airbyte_cdk.sources.streams.core import StreamData
 from airbyte_cdk.sources.types import Config, Record, StreamSlice, StreamState
 from airbyte_cdk.utils.mapping_helpers import combine_mappings
+from airbyte_cdk.sources.source import ExperimentalClassWarning
 
 FULL_REFRESH_SYNC_COMPLETE_KEY = "__ab_full_refresh_sync_complete"
 
@@ -625,25 +627,29 @@ def _fetch_next_page(
 
 
 class SafeResponse(requests.Response):
-    def __getattr__(self, name):
+    def __getattr__(self, name: str) -> Any:
         return getattr(requests.Response, name, None)
 
     @property
-    def content(self):
+    def content(self) -> Optional[bytes]:
         return super().content
 
     @content.setter
-    def content(self, value):
+    def content(self, value: Union[str, bytes]) -> None:
         self._content = value.encode() if isinstance(value, str) else value
 
 
+@deprecated(
+    "This class is experimental. Use at your own risk.",
+    category=ExperimentalClassWarning,
+)
 @dataclass
 class LazySimpleRetriever(SimpleRetriever):
     """
     A retriever that supports lazy loading from parent streams.
     """
 
-    partition_router: SubstreamPartitionRouter = field(init=True, repr=False, default=None)
+    partition_router: SubstreamPartitionRouter = field(init=True, repr=False, default=None)  # type: ignore[assignment]  # 'partition_router' is required for LazySimpleRetriever and is validated in the constructor
     lazy_read_pointer: Optional[List[InterpolatedString]] = None
 
     def _read_pages(
@@ -655,9 +661,9 @@ def _read_pages(
         parent_stream_config = self.partition_router.parent_stream_configs[-1]
         parent_stream = parent_stream_config.stream
 
-        for parent_record in parent_stream.read_only_records():
+        for raw_parent_record in parent_stream.read_only_records():
             parent_record, parent_partition = self.partition_router.process_parent_record(
-                parent_record, parent_stream.name
+                raw_parent_record, parent_stream.name
             )
             if parent_record is None:
                 continue
@@ -676,19 +682,19 @@ def _read_pages(
 
         yield from []
 
-    def _extract_child_records(self, parent_record: Mapping) -> Mapping:
+    def _extract_child_records(self, parent_record: MutableMapping[str, Any]) -> MutableMapping[str, Any]:
         """Extract child records from a parent record based on lazy pointers."""
         if not self.lazy_read_pointer:
             return parent_record
 
         path = [path.eval(self.config) for path in self.lazy_read_pointer]
         return (
-            dpath.values(parent_record, path)
+            dpath.values(parent_record, path)  # type: ignore # return value will be a MutableMapping, given input data structure
             if "*" in path
             else dpath.get(parent_record, path, default=[])
         )
 
-    def _create_response(self, data: Mapping) -> SafeResponse:
+    def _create_response(self, data: Mapping[str, Any]) -> SafeResponse:
         """Create a SafeResponse with the given data."""
         response = SafeResponse()
         response.content = json.dumps(data).encode("utf-8")
@@ -701,7 +707,7 @@ def _yield_records_with_pagination(
         records_generator_fn: Callable[[Optional[requests.Response]], Iterable[Record]],
         stream_state: Mapping[str, Any],
         stream_slice: StreamSlice,
-        parent_record: Record,
+        parent_record: MutableMapping[str, Any],
         parent_stream_config: Any,
     ) -> Iterable[Record]:
         """Yield records, handling pagination if needed."""
@@ -729,7 +735,7 @@ def _paginate(
         records_generator_fn: Callable[[Optional[requests.Response]], Iterable[Record]],
         stream_state: Mapping[str, Any],
         stream_slice: StreamSlice,
-        parent_record: Record,
+        parent_record: MutableMapping[str, Any],
         parent_stream_config: Any,
     ) -> Iterable[Record]:
         """Handle pagination by fetching subsequent pages."""
@@ -742,7 +748,9 @@ def _paginate(
             cursor_slice=stream_slice.cursor_slice,
         )
 
-        while next_page_token:
+        pagination_complete = False
+
+        while not pagination_complete:
             response = self._fetch_next_page(stream_state, stream_slice, next_page_token)
             last_page_size, last_record = 0, None
 
@@ -751,9 +759,15 @@ def _paginate(
                 last_record = record
                 yield record
 
-            last_page_token_value = (
-                next_page_token.get("next_page_token") if next_page_token else None
-            )
-            next_page_token = self._next_page_token(
-                response, last_page_size, last_record, last_page_token_value
-            )
+            if not response:
+                pagination_complete = True
+            else:
+                last_page_token_value = (
+                    next_page_token.get("next_page_token") if next_page_token else None
+                )
+                next_page_token = self._next_page_token(
+                    response, last_page_size, last_record, last_page_token_value
+                )
+
+                if not next_page_token:
+                    pagination_complete = True