Add extra_fields support

Author: lazebnyi

airbyte_cdk/sources/declarative/retrievers/simple_retriever.py

@@ -358,7 +358,7 @@ def _fetch_next_page(
     # This logic is similar to _read_pages in the HttpStream class. When making changes here, consider making changes there as well.
     def _read_pages(
         self,
-        records_generator_fn: Callable[[Optional[requests.Response]], Iterable[Record]],
+        records_generator_fn: Callable[[Optional[requests.Response], Optional[StreamSlice]], Iterable[Record]],
         stream_state: Mapping[str, Any],
         stream_slice: StreamSlice,
     ) -> Iterable[Record]:
@@ -372,7 +372,7 @@ def _read_pages(
 
             last_page_size = 0
             last_record: Optional[Record] = None
-            for record in records_generator_fn(response):
+            for record in records_generator_fn(response, stream_slice=stream_slice):  # type: ignore[call-arg] # only _parse_records expected as a func
                 last_page_size += 1
                 last_record = record
                 yield record
@@ -397,7 +397,7 @@ def _read_pages(
 
     def _read_single_page(
         self,
-        records_generator_fn: Callable[[Optional[requests.Response]], Iterable[Record]],
+        records_generator_fn: Callable[[Optional[requests.Response], Optional[StreamSlice]], Iterable[Record]],
         stream_state: Mapping[str, Any],
         stream_slice: StreamSlice,
     ) -> Iterable[StreamData]:
@@ -412,7 +412,7 @@ def _read_single_page(
 
         last_page_size = 0
         last_record: Optional[Record] = None
-        for record in records_generator_fn(response):
+        for record in records_generator_fn(response, stream_slice=stream_slice):  # type: ignore[call-arg] # only _parse_records expected as a func
             last_page_size += 1
             last_record = record
             yield record
@@ -456,7 +456,6 @@ def read_records(
         record_generator = partial(
             self._parse_records,
             stream_state=self.state or {},
-            stream_slice=_slice,
             records_schema=records_schema,
         )
 
@@ -665,7 +664,7 @@ class LazySimpleRetriever(SimpleRetriever):
 
     def _read_pages(
         self,
-        records_generator_fn: Callable[[Optional[requests.Response]], Iterable[Record]],
+        records_generator_fn: Callable[[Optional[requests.Response], Optional[StreamSlice]], Iterable[Record]],
         stream_state: Mapping[str, Any],
         stream_slice: StreamSlice,
     ) -> Iterable[Record]:
@@ -682,6 +681,20 @@ def _read_pages(
             child_records = self._extract_child_records(parent_record)
             response = self._create_response(child_records)
 
+            if parent_stream_config.extra_fields:
+                extra_fields = [
+                    [field_path_part.eval(self.config) for field_path_part in field_path]  # type: ignore [union-attr]
+                    for field_path in parent_stream_config.extra_fields
+                ]
+
+                extracted_extra_fields = self.partition_router._extract_extra_fields(parent_record, extra_fields)
+
+                stream_slice = StreamSlice(
+                    partition=stream_slice.partition,
+                    cursor_slice=stream_slice.cursor_slice,
+                    extra_fields=extracted_extra_fields
+                )
+
             yield from self._yield_records_with_pagination(
                 response,
                 records_generator_fn,
@@ -717,7 +730,7 @@ def _create_response(self, data: Mapping[str, Any]) -> SafeResponse:
     def _yield_records_with_pagination(
         self,
         response: requests.Response,
-        records_generator_fn: Callable[[Optional[requests.Response]], Iterable[Record]],
+        records_generator_fn: Callable[[Optional[requests.Response], Optional[StreamSlice]], Iterable[Record]],
         stream_state: Mapping[str, Any],
         stream_slice: StreamSlice,
         parent_record: MutableMapping[str, Any],
@@ -726,7 +739,7 @@ def _yield_records_with_pagination(
         """Yield records, handling pagination if needed."""
         last_page_size, last_record = 0, None
 
-        for record in records_generator_fn(response):
+        for record in records_generator_fn(response, stream_slice=stream_slice):  # type: ignore[call-arg] # only _parse_records expected as a func
             last_page_size += 1
             last_record = record
             yield record
@@ -745,7 +758,7 @@ def _yield_records_with_pagination(
     def _paginate(
         self,
         next_page_token: Any,
-        records_generator_fn: Callable[[Optional[requests.Response]], Iterable[Record]],
+        records_generator_fn: Callable[[Optional[requests.Response], Optional[StreamSlice]], Iterable[Record]],
         stream_state: Mapping[str, Any],
         stream_slice: StreamSlice,
         parent_record: MutableMapping[str, Any],
@@ -767,7 +780,7 @@ def _paginate(
             response = self._fetch_next_page(stream_state, stream_slice, next_page_token)
             last_page_size, last_record = 0, None
 
-            for record in records_generator_fn(response):
+            for record in records_generator_fn(response, stream_slice=stream_slice):  # type: ignore[call-arg] # only _parse_records expected as a func
                 last_page_size += 1
                 last_record = record
                 yield record