airbyte-cdk 0.72.1__py3-none-any.whl → 6.17.1.dev0__py3-none-any.whl

airbyte_cdk/sources/declarative/retrievers/simple_retriever.py

@@ -2,27 +2,38 @@
 # Copyright (c) 2023 Airbyte, Inc., all rights reserved.
 #
 
+import json
 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
 
 import requests
+
 from airbyte_cdk.models import AirbyteMessage
 from airbyte_cdk.sources.declarative.extractors.http_selector import HttpSelector
-from airbyte_cdk.sources.declarative.incremental.cursor import Cursor
+from airbyte_cdk.sources.declarative.incremental import ResumableFullRefreshCursor
+from airbyte_cdk.sources.declarative.incremental.declarative_cursor import DeclarativeCursor
 from airbyte_cdk.sources.declarative.interpolation import InterpolatedString
-from airbyte_cdk.sources.declarative.partition_routers.single_partition_router import SinglePartitionRouter
+from airbyte_cdk.sources.declarative.partition_routers.single_partition_router import (
+    SinglePartitionRouter,
+)
 from airbyte_cdk.sources.declarative.requesters.paginators.no_pagination import NoPagination
 from airbyte_cdk.sources.declarative.requesters.paginators.paginator import Paginator
+from airbyte_cdk.sources.declarative.requesters.request_options import (
+    DefaultRequestOptionsProvider,
+    RequestOptionsProvider,
+)
 from airbyte_cdk.sources.declarative.requesters.requester import Requester
 from airbyte_cdk.sources.declarative.retrievers.retriever import Retriever
 from airbyte_cdk.sources.declarative.stream_slicers.stream_slicer import StreamSlicer
-from airbyte_cdk.sources.declarative.types import Config, Record, StreamSlice, StreamState
 from airbyte_cdk.sources.http_logger import format_http_message
 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
 
+FULL_REFRESH_SYNC_COMPLETE_KEY = "__ab_full_refresh_sync_complete"
+
 
 @dataclass
 class SimpleRetriever(Retriever):
@@ -57,23 +68,34 @@ class SimpleRetriever(Retriever):
     primary_key: Optional[Union[str, List[str], List[List[str]]]]
     _primary_key: str = field(init=False, repr=False, default="")
     paginator: Optional[Paginator] = None
-    stream_slicer: StreamSlicer = SinglePartitionRouter(parameters={})
-    cursor: Optional[Cursor] = None
+    stream_slicer: StreamSlicer = field(
+        default_factory=lambda: SinglePartitionRouter(parameters={})
+    )
+    request_option_provider: RequestOptionsProvider = field(
+        default_factory=lambda: DefaultRequestOptionsProvider(parameters={})
+    )
+    cursor: Optional[DeclarativeCursor] = None
     ignore_stream_slicer_parameters_on_paginated_requests: bool = False
 
     def __post_init__(self, parameters: Mapping[str, Any]) -> None:
         self._paginator = self.paginator or NoPagination(parameters=parameters)
-        self._last_response: Optional[requests.Response] = None
-        self._records_from_last_response: List[Record] = []
         self._parameters = parameters
-        self._name = InterpolatedString(self._name, parameters=parameters) if isinstance(self._name, str) else self._name
+        self._name = (
+            InterpolatedString(self._name, parameters=parameters)
+            if isinstance(self._name, str)
+            else self._name
+        )
 
     @property  # type: ignore
     def name(self) -> str:
         """
         :return: Stream name
         """
-        return str(self._name.eval(self.config)) if isinstance(self._name, InterpolatedString) else self._name
+        return (
+            str(self._name.eval(self.config))
+            if isinstance(self._name, InterpolatedString)
+            else self._name
+        )
 
     @name.setter
     def name(self, value: str) -> None:
@@ -107,10 +129,20 @@ class SimpleRetriever(Retriever):
         """
         # FIXME we should eventually remove the usage of stream_state as part of the interpolation
         mappings = [
-            paginator_method(stream_state=stream_state, stream_slice=stream_slice, next_page_token=next_page_token),
+            paginator_method(
+                stream_state=stream_state,
+                stream_slice=stream_slice,
+                next_page_token=next_page_token,
+            ),
         ]
         if not next_page_token or not self.ignore_stream_slicer_parameters_on_paginated_requests:
-            mappings.append(stream_slicer_method(stream_state=stream_state, stream_slice=stream_slice, next_page_token=next_page_token))
+            mappings.append(
+                stream_slicer_method(
+                    stream_state=stream_state,
+                    stream_slice=stream_slice,
+                    next_page_token=next_page_token,
+                )
+            )
         return combine_mappings(mappings)
 
     def _request_headers(
@@ -128,7 +160,7 @@ class SimpleRetriever(Retriever):
             stream_slice,
             next_page_token,
             self._paginator.get_request_headers,
-            self.stream_slicer.get_request_headers,
+            self.request_option_provider.get_request_headers,
         )
         if isinstance(headers, str):
             raise ValueError("Request headers cannot be a string")
@@ -150,7 +182,7 @@ class SimpleRetriever(Retriever):
             stream_slice,
             next_page_token,
             self._paginator.get_request_params,
-            self.stream_slicer.get_request_params,
+            self.request_option_provider.get_request_params,
         )
         if isinstance(params, str):
             raise ValueError("Request params cannot be a string")
@@ -176,7 +208,7 @@ class SimpleRetriever(Retriever):
             stream_slice,
             next_page_token,
             self._paginator.get_request_body_data,
-            self.stream_slicer.get_request_body_data,
+            self.request_option_provider.get_request_body_data,
         )
 
     def _request_body_json(
@@ -195,23 +227,19 @@ class SimpleRetriever(Retriever):
             stream_slice,
             next_page_token,
             self._paginator.get_request_body_json,
-            self.stream_slicer.get_request_body_json,
+            self.request_option_provider.get_request_body_json,
         )
         if isinstance(body_json, str):
             raise ValueError("Request body json cannot be a string")
         return body_json
 
-    def _paginator_path(
-        self,
-    ) -> Optional[str]:
+    def _paginator_path(self, next_page_token: Optional[Mapping[str, Any]] = None) -> Optional[str]:
         """
         If the paginator points to a path, follow it, else return nothing so the requester is used.
-        :param stream_state:
-        :param stream_slice:
         :param next_page_token:
         :return:
         """
-        return self._paginator.path()
+        return self._paginator.path(next_page_token=next_page_token)
 
     def _parse_response(
         self,
@@ -222,20 +250,15 @@ class SimpleRetriever(Retriever):
         next_page_token: Optional[Mapping[str, Any]] = None,
     ) -> Iterable[Record]:
         if not response:
-            self._last_response = None
-            self._records_from_last_response = []
-            return []
-
-        self._last_response = response
-        records = self.record_selector.select_records(
-            response=response,
-            stream_state=stream_state,
-            records_schema=records_schema,
-            stream_slice=stream_slice,
-            next_page_token=next_page_token,
-        )
-        self._records_from_last_response = records
-        return records
+            yield from []
+        else:
+            yield from self.record_selector.select_records(
+                response=response,
+                stream_state=stream_state,
+                records_schema=records_schema,
+                stream_slice=stream_slice,
+                next_page_token=next_page_token,
+            )
 
     @property  # type: ignore
     def primary_key(self) -> Optional[Union[str, List[str], List[List[str]]]]:
@@ -247,7 +270,13 @@ class SimpleRetriever(Retriever):
         if not isinstance(value, property):
             self._primary_key = value
 
-    def _next_page_token(self, response: requests.Response) -> Optional[Mapping[str, Any]]:
+    def _next_page_token(
+        self,
+        response: requests.Response,
+        last_page_size: int,
+        last_record: Optional[Record],
+        last_page_token_value: Optional[Any],
+    ) -> Optional[Mapping[str, Any]]:
         """
         Specifies a pagination strategy.
 
@@ -255,49 +284,129 @@ class SimpleRetriever(Retriever):
 
         :return: The token for the next page from the input response object. Returning None means there are no more pages to read in this response.
         """
-        return self._paginator.next_page_token(response, self._records_from_last_response)
+        return self._paginator.next_page_token(
+            response=response,
+            last_page_size=last_page_size,
+            last_record=last_record,
+            last_page_token_value=last_page_token_value,
+        )
 
     def _fetch_next_page(
-        self, stream_state: Mapping[str, Any], stream_slice: StreamSlice, next_page_token: Optional[Mapping[str, Any]] = None
+        self,
+        stream_state: Mapping[str, Any],
+        stream_slice: StreamSlice,
+        next_page_token: Optional[Mapping[str, Any]] = None,
     ) -> Optional[requests.Response]:
         return self.requester.send_request(
-            path=self._paginator_path(),
+            path=self._paginator_path(next_page_token=next_page_token),
             stream_state=stream_state,
             stream_slice=stream_slice,
             next_page_token=next_page_token,
-            request_headers=self._request_headers(stream_state=stream_state, stream_slice=stream_slice, next_page_token=next_page_token),
-            request_params=self._request_params(stream_state=stream_state, stream_slice=stream_slice, next_page_token=next_page_token),
+            request_headers=self._request_headers(
+                stream_state=stream_state,
+                stream_slice=stream_slice,
+                next_page_token=next_page_token,
+            ),
+            request_params=self._request_params(
+                stream_state=stream_state,
+                stream_slice=stream_slice,
+                next_page_token=next_page_token,
+            ),
             request_body_data=self._request_body_data(
-                stream_state=stream_state, stream_slice=stream_slice, next_page_token=next_page_token
+                stream_state=stream_state,
+                stream_slice=stream_slice,
+                next_page_token=next_page_token,
             ),
             request_body_json=self._request_body_json(
-                stream_state=stream_state, stream_slice=stream_slice, next_page_token=next_page_token
+                stream_state=stream_state,
+                stream_slice=stream_slice,
+                next_page_token=next_page_token,
             ),
         )
 
     # 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[StreamData]],
+        records_generator_fn: Callable[[Optional[requests.Response]], Iterable[Record]],
         stream_state: Mapping[str, Any],
         stream_slice: StreamSlice,
-    ) -> Iterable[StreamData]:
+    ) -> Iterable[Record]:
         pagination_complete = False
-        next_page_token = None
+        initial_token = self._paginator.get_initial_token()
+        next_page_token: Optional[Mapping[str, Any]] = (
+            {"next_page_token": initial_token} if initial_token else None
+        )
         while not pagination_complete:
             response = self._fetch_next_page(stream_state, stream_slice, next_page_token)
-            yield from records_generator_fn(response)
+
+            last_page_size = 0
+            last_record: Optional[Record] = None
+            for record in records_generator_fn(response):
+                last_page_size += 1
+                last_record = record
+                yield record
 
             if not response:
                 pagination_complete = True
             else:
-                next_page_token = self._next_page_token(response)
+                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=response,
+                    last_page_size=last_page_size,
+                    last_record=last_record,
+                    last_page_token_value=last_page_token_value,
+                )
                 if not next_page_token:
                     pagination_complete = True
 
         # Always return an empty generator just in case no records were ever yielded
         yield from []
 
+    def _read_single_page(
+        self,
+        records_generator_fn: Callable[[Optional[requests.Response]], Iterable[Record]],
+        stream_state: Mapping[str, Any],
+        stream_slice: StreamSlice,
+    ) -> Iterable[StreamData]:
+        initial_token = stream_state.get("next_page_token")
+        if initial_token is None:
+            initial_token = self._paginator.get_initial_token()
+        next_page_token: Optional[Mapping[str, Any]] = (
+            {"next_page_token": initial_token} if initial_token else None
+        )
+
+        response = self._fetch_next_page(stream_state, stream_slice, next_page_token)
+
+        last_page_size = 0
+        last_record: Optional[Record] = None
+        for record in records_generator_fn(response):
+            last_page_size += 1
+            last_record = record
+            yield record
+
+        if not response:
+            next_page_token = {FULL_REFRESH_SYNC_COMPLETE_KEY: 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=response,
+                last_page_size=last_page_size,
+                last_record=last_record,
+                last_page_token_value=last_page_token_value,
+            ) or {FULL_REFRESH_SYNC_COMPLETE_KEY: True}
+
+        if self.cursor:
+            self.cursor.close_slice(
+                StreamSlice(cursor_slice=next_page_token, partition=stream_slice.partition)
+            )
+
+        # Always return an empty generator just in case no records were ever yielded
+        yield from []
+
     def read_records(
         self,
         records_schema: Mapping[str, Any],
@@ -311,8 +420,6 @@ class SimpleRetriever(Retriever):
         :return: The records read from the API source
         """
         _slice = stream_slice or StreamSlice(partition={}, cursor_slice={})  # None-check
-        # Fixing paginator types has a long tail of dependencies
-        self._paginator.reset()
 
         most_recent_record_from_slice = None
         record_generator = partial(
@@ -321,33 +428,56 @@ class SimpleRetriever(Retriever):
             stream_slice=_slice,
             records_schema=records_schema,
         )
-        for stream_data in self._read_pages(record_generator, self.state, _slice):
-            current_record = self._extract_record(stream_data, _slice)
-            if self.cursor and current_record:
-                self.cursor.observe(_slice, current_record)
 
-            # TODO this is just the most recent record *read*, not necessarily the most recent record *within slice boundaries*; once all
-            # cursors implement a meaningful `observe` method, it can be removed, both from here and the `Cursor.close_slice` method args
-            most_recent_record_from_slice = self._get_most_recent_record(most_recent_record_from_slice, current_record, _slice)
-            yield stream_data
+        if self.cursor and isinstance(self.cursor, ResumableFullRefreshCursor):
+            stream_state = self.state
 
-        if self.cursor:
-            self.cursor.close_slice(_slice, most_recent_record_from_slice)
+            # Before syncing the RFR stream, we check if the job's prior attempt was successful and don't need to
+            # fetch more records. The platform deletes stream state for full refresh streams before starting a
+            # new job, so we don't need to worry about this value existing for the initial attempt
+            if stream_state.get(FULL_REFRESH_SYNC_COMPLETE_KEY):
+                return
+
+            yield from self._read_single_page(record_generator, stream_state, _slice)
+        else:
+            for stream_data in self._read_pages(record_generator, self.state, _slice):
+                current_record = self._extract_record(stream_data, _slice)
+                if self.cursor and current_record:
+                    self.cursor.observe(_slice, current_record)
+
+                # Latest record read, not necessarily within slice boundaries.
+                # TODO Remove once all custom components implement `observe` method.
+                # https://github.com/airbytehq/airbyte-internal-issues/issues/6955
+                most_recent_record_from_slice = self._get_most_recent_record(
+                    most_recent_record_from_slice, current_record, _slice
+                )
+                yield stream_data
+
+            if self.cursor:
+                self.cursor.close_slice(_slice, most_recent_record_from_slice)
         return
 
     def _get_most_recent_record(
-        self, current_most_recent: Optional[Record], current_record: Optional[Record], stream_slice: StreamSlice
+        self,
+        current_most_recent: Optional[Record],
+        current_record: Optional[Record],
+        stream_slice: StreamSlice,
     ) -> Optional[Record]:
         if self.cursor and current_record:
             if not current_most_recent:
                 return current_record
             else:
-                return current_most_recent if self.cursor.is_greater_than_or_equal(current_most_recent, current_record) else current_record
+                return (
+                    current_most_recent
+                    if self.cursor.is_greater_than_or_equal(current_most_recent, current_record)
+                    else current_record
+                )
         else:
             return None
 
-    @staticmethod
-    def _extract_record(stream_data: StreamData, stream_slice: StreamSlice) -> Optional[Record]:
+    def _extract_record(
+        self, stream_data: StreamData, stream_slice: StreamSlice
+    ) -> Optional[Record]:
         """
         As we allow the output of _read_pages to be StreamData, it can be multiple things. Therefore, we need to filter out and normalize
         to data to streamline the rest of the process.
@@ -356,9 +486,15 @@ class SimpleRetriever(Retriever):
             # Record is not part of `StreamData` but is the most common implementation of `Mapping[str, Any]` which is part of `StreamData`
             return stream_data
         elif isinstance(stream_data, (dict, Mapping)):
-            return Record(dict(stream_data), stream_slice)
+            return Record(
+                data=dict(stream_data), associated_slice=stream_slice, stream_name=self.name
+            )
         elif isinstance(stream_data, AirbyteMessage) and stream_data.record:
-            return Record(stream_data.record.data, stream_slice)
+            return Record(
+                data=stream_data.record.data,  # type:ignore # AirbyteMessage always has record.data
+                associated_slice=stream_slice,
+                stream_name=self.name,
+            )
         return None
 
     # stream_slices is defined with arguments on http stream and fixing this has a long tail of dependencies. Will be resolved by the decoupling of http stream and simple retriever
@@ -389,7 +525,7 @@ class SimpleRetriever(Retriever):
         stream_state: Mapping[str, Any],
         records_schema: Mapping[str, Any],
         stream_slice: Optional[StreamSlice],
-    ) -> Iterable[StreamData]:
+    ) -> Iterable[Record]:
         yield from self._parse_response(
             response,
             stream_slice=stream_slice,
@@ -400,6 +536,11 @@ class SimpleRetriever(Retriever):
     def must_deduplicate_query_params(self) -> bool:
         return True
 
+    @staticmethod
+    def _to_partition_key(to_serialize: Any) -> str:
+        # separators have changed in Python 3.4. To avoid being impacted by further change, we explicitly specify our own value
+        return json.dumps(to_serialize, indent=None, separators=(",", ":"), sort_keys=True)
+
 
 @dataclass
 class SimpleRetrieverTestReadDecorator(SimpleRetriever):
@@ -422,20 +563,35 @@ class SimpleRetrieverTestReadDecorator(SimpleRetriever):
         return islice(super().stream_slices(), self.maximum_number_of_slices)
 
     def _fetch_next_page(
-        self, stream_state: Mapping[str, Any], stream_slice: StreamSlice, next_page_token: Optional[Mapping[str, Any]] = None
+        self,
+        stream_state: Mapping[str, Any],
+        stream_slice: StreamSlice,
+        next_page_token: Optional[Mapping[str, Any]] = None,
     ) -> Optional[requests.Response]:
         return self.requester.send_request(
-            path=self._paginator_path(),
+            path=self._paginator_path(next_page_token=next_page_token),
             stream_state=stream_state,
             stream_slice=stream_slice,
             next_page_token=next_page_token,
-            request_headers=self._request_headers(stream_state=stream_state, stream_slice=stream_slice, next_page_token=next_page_token),
-            request_params=self._request_params(stream_state=stream_state, stream_slice=stream_slice, next_page_token=next_page_token),
+            request_headers=self._request_headers(
+                stream_state=stream_state,
+                stream_slice=stream_slice,
+                next_page_token=next_page_token,
+            ),
+            request_params=self._request_params(
+                stream_state=stream_state,
+                stream_slice=stream_slice,
+                next_page_token=next_page_token,
+            ),
             request_body_data=self._request_body_data(
-                stream_state=stream_state, stream_slice=stream_slice, next_page_token=next_page_token
+                stream_state=stream_state,
+                stream_slice=stream_slice,
+                next_page_token=next_page_token,
             ),
             request_body_json=self._request_body_json(
-                stream_state=stream_state, stream_slice=stream_slice, next_page_token=next_page_token
+                stream_state=stream_state,
+                stream_slice=stream_slice,
+                next_page_token=next_page_token,
             ),
             log_formatter=lambda response: format_http_message(
                 response,

airbyte_cdk/sources/declarative/schema/__init__.py

@@ -3,8 +3,21 @@
 #
 
 from airbyte_cdk.sources.declarative.schema.default_schema_loader import DefaultSchemaLoader
+from airbyte_cdk.sources.declarative.schema.dynamic_schema_loader import (
+    DynamicSchemaLoader,
+    SchemaTypeIdentifier,
+    TypesMap,
+)
 from airbyte_cdk.sources.declarative.schema.inline_schema_loader import InlineSchemaLoader
 from airbyte_cdk.sources.declarative.schema.json_file_schema_loader import JsonFileSchemaLoader
 from airbyte_cdk.sources.declarative.schema.schema_loader import SchemaLoader
 
-__all__ = ["JsonFileSchemaLoader", "DefaultSchemaLoader", "SchemaLoader", "InlineSchemaLoader"]
+__all__ = [
+    "JsonFileSchemaLoader",
+    "DefaultSchemaLoader",
+    "SchemaLoader",
+    "InlineSchemaLoader",
+    "DynamicSchemaLoader",
+    "TypesMap",
+    "SchemaTypeIdentifier",
+]

airbyte_cdk/sources/declarative/schema/default_schema_loader.py

@@ -8,7 +8,7 @@ from typing import Any, Mapping
 
 from airbyte_cdk.sources.declarative.schema.json_file_schema_loader import JsonFileSchemaLoader
 from airbyte_cdk.sources.declarative.schema.schema_loader import SchemaLoader
-from airbyte_cdk.sources.declarative.types import Config
+from airbyte_cdk.sources.types import Config
 
 
 @dataclass
@@ -24,7 +24,7 @@ class DefaultSchemaLoader(SchemaLoader):
     config: Config
     parameters: InitVar[Mapping[str, Any]]
 
-    def __post_init__(self, parameters: Mapping[str, Any]):
+    def __post_init__(self, parameters: Mapping[str, Any]) -> None:
         self._parameters = parameters
         self.default_loader = JsonFileSchemaLoader(parameters=parameters, config=self.config)
 
@@ -41,5 +41,7 @@ class DefaultSchemaLoader(SchemaLoader):
             # A slight hack since we don't directly have the stream name. However, when building the default filepath we assume the
             # runtime options stores stream name 'name' so we'll do the same here
             stream_name = self._parameters.get("name", "")
-            logging.info(f"Could not find schema for stream {stream_name}, defaulting to the empty schema")
+            logging.info(
+                f"Could not find schema for stream {stream_name}, defaulting to the empty schema"
+            )
             return {}