diff --git a/airbyte_cdk/sources/declarative/declarative_component_schema.yaml b/airbyte_cdk/sources/declarative/declarative_component_schema.yaml index f2e4d376ee..c36a992df3 100644 --- a/airbyte_cdk/sources/declarative/declarative_component_schema.yaml +++ b/airbyte_cdk/sources/declarative/declarative_component_schema.yaml @@ -1459,9 +1459,15 @@ definitions: token_option: title: Send Token In description: >- - Where to send the stored token on each request of a pass: the JSON body, the query - string, a header or the form body. - "$ref": "#/definitions/RequestOption" + How to send the stored token. A RequestOption puts it in the JSON body, the query + string, a header or the form body: on every request of a pass, or on the first request + only when the retriever sets ignore_stream_slicer_parameters_on_paginated_requests. + RequestPath makes it the URL of the first request of a pass, for APIs whose token is a + complete URL, such as a Microsoft Graph @odata.deltaLink; the paginator sends the + later requests. + anyOf: + - "$ref": "#/definitions/RequestOption" + - "$ref": "#/definitions/RequestPath" completion_condition: title: Completion Condition description: >- diff --git a/airbyte_cdk/sources/declarative/incremental/sync_token_cursor.py b/airbyte_cdk/sources/declarative/incremental/sync_token_cursor.py index caf76f85d4..1fc859cd31 100644 --- a/airbyte_cdk/sources/declarative/incremental/sync_token_cursor.py +++ b/airbyte_cdk/sources/declarative/incremental/sync_token_cursor.py @@ -42,6 +42,7 @@ class SyncTokenCursor(Cursor): STATE_TYPE = "sync_token" + DEFAULT_MAX_STATE_BYTES = 4 * 1024 * 1024 def __init__( self, @@ -59,6 +60,7 @@ def __init__( completion_condition: Optional[str] = None, cursor_field: Optional[str] = None, max_partitions: int = ConcurrentPerPartitionCursor.DEFAULT_MAX_PARTITIONS_NUMBER, + max_state_bytes: int = DEFAULT_MAX_STATE_BYTES, ) -> None: self._stream_name = stream_name self._stream_namespace = stream_namespace @@ -78,8 +80,10 @@ def __init__( ) self.cursor_field = CursorField(cursor_field) if cursor_field else None self._max_partitions = max_partitions + self._max_state_bytes = max_state_bytes self._serializer = PerPartitionKeySerializer() self._committed: OrderedDict[str, dict[str, Any]] = OrderedDict() + self._committed_bytes = 0 self._pending: dict[str, dict[str, Any]] = {} self._page_counts: dict[str, int] = {} self._lock = threading.RLock() @@ -88,8 +92,9 @@ def __init__( self._state_is_foreign = False self._state_emitted = False self._warned_no_token = False - self._warned_cap = False + self._warned_limit = False self._partitions_over_cap: set[str] = set() + self._partitions_over_state_size: set[str] = set() self._load_state(stream_state or {}) @property @@ -132,6 +137,9 @@ def stream_slices(self) -> Iterable[StreamSlice]: elif at_partition_cap: self._partitions_over_cap.add(partition_key) self._log_bootstrap("partition cap") + elif self._committed_bytes >= self._max_state_bytes: + self._partitions_over_state_size.add(partition_key) + self._log_bootstrap("state size limit") elif not self._state_is_foreign: self._log_bootstrap("none") yield StreamSlice( @@ -144,6 +152,9 @@ def stream_slices(self) -> Iterable[StreamSlice]: with self._lock: vanished = set(self._committed) - seen_partitions for partition_key in vanished: + self._committed_bytes -= self._entry_bytes( + partition_key, self._committed[partition_key] + ) del self._committed[partition_key] if vanished: logger.info( @@ -330,9 +341,13 @@ def close_partition(self, partition: Partition) -> None: pending = self._pending.pop(partition_key, None) page_count = self._page_counts.pop(partition_key, 0) if pending is not None: - if partition_key in self._committed or len(self._committed) < self._max_partitions: + existing = self._committed.get(partition_key) + if existing is not None: + self._committed_bytes -= self._entry_bytes(partition_key, existing) + self._committed_bytes += self._entry_bytes(partition_key, pending) self._committed[partition_key] = pending self._partitions_over_cap.discard(partition_key) + self._partitions_over_state_size.discard(partition_key) logger.debug( "Stream %s: committed change token %s for partition %s after %s pages.", self._stream_name, @@ -340,8 +355,25 @@ def close_partition(self, partition: Partition) -> None: self._serializer.to_partition(partition_key), page_count, ) - else: + elif ( + len(self._committed) < self._max_partitions + and self._committed_bytes < self._max_state_bytes + ): + self._committed[partition_key] = pending + self._committed_bytes += self._entry_bytes(partition_key, pending) + self._partitions_over_cap.discard(partition_key) + self._partitions_over_state_size.discard(partition_key) + logger.debug( + "Stream %s: committed change token %s for partition %s after %s pages.", + self._stream_name, + pending["token"], + self._serializer.to_partition(partition_key), + page_count, + ) + elif len(self._committed) >= self._max_partitions: self._partitions_over_cap.add(partition_key) + else: + self._partitions_over_state_size.add(partition_key) if not self._partitioned: self._emit_state_message() elif time.time() - self._last_emission_time > 600: @@ -362,14 +394,25 @@ def _emit_state_message(self) -> None: def ensure_at_least_one_state_emitted(self) -> None: if self._partitioned or not self._state_emitted: self._emit_state_message() - if self._partitions_over_cap and not self._warned_cap: - logger.warning( - "Stream %s: %s partitions already store tokens; %s more read in full on every sync.", - self._stream_name, - len(self._committed), - len(self._partitions_over_cap), - ) - self._warned_cap = True + if not self._warned_limit: + if self._partitions_over_state_size: + logger.warning( + "Stream %s: stored tokens reached the %s state limit at %s partitions; " + "%s more read in full on every sync.", + self._stream_name, + self._human_readable_bytes(self._max_state_bytes), + f"{len(self._committed):,}", + f"{len(self._partitions_over_state_size):,}", + ) + self._warned_limit = True + elif self._partitions_over_cap: + logger.warning( + "Stream %s: %s partitions already store tokens; %s more read in full on every sync.", + self._stream_name, + len(self._committed), + len(self._partitions_over_cap), + ) + self._warned_limit = True def observe(self, record: Record) -> None: return @@ -409,6 +452,10 @@ def _load_state(self, stream_state: StreamState) -> None: else: self._log_bootstrap("foreign state") self._committed.update(committed) + self._committed_bytes += sum( + self._entry_bytes(partition_key, cursor) + for partition_key, cursor in committed.items() + ) else: if "states" in stream_state: self._log_foreign_state() @@ -419,5 +466,19 @@ def _load_state(self, stream_state: StreamState) -> None: "token": token, "committed_at": stream_state.get("committed_at", ""), } + self._committed_bytes += self._entry_bytes("{}", self._committed["{}"]) elif token is not None: self._log_foreign_state() + + @staticmethod + def _entry_bytes(partition_key: str, cursor: Mapping[str, Any]) -> int: + return len(partition_key.encode()) + len(cursor["token"].encode()) + + @staticmethod + def _human_readable_bytes(size: int) -> str: + if size < 1024: + return f"{size:,} bytes" + size_in_kib = size / 1024 + if size_in_kib < 1024: + return f"{size_in_kib:g} KB" + return f"{size_in_kib / 1024:g} MB" diff --git a/airbyte_cdk/sources/declarative/models/declarative_component_schema.py b/airbyte_cdk/sources/declarative/models/declarative_component_schema.py index abca5f724d..32f3783b4b 100644 --- a/airbyte_cdk/sources/declarative/models/declarative_component_schema.py +++ b/airbyte_cdk/sources/declarative/models/declarative_component_schema.py @@ -2965,9 +2965,9 @@ class Config: examples=[["syncToken"], ["nextSyncToken"], ["meta", "sync_token"]], title="Token Path", ) - token_option: RequestOption = Field( + token_option: Union[RequestOption, RequestPath] = Field( ..., - description="Where to send the stored token on each request of a pass: the JSON body, the query string, a header or the form body.", + description="How to send the stored token. A RequestOption puts it in the JSON body, the query string, a header or the form body: on every request of a pass, or on the first request only when the retriever sets ignore_stream_slicer_parameters_on_paginated_requests. RequestPath makes it the URL of the first request of a pass, for APIs whose token is a complete URL, such as a Microsoft Graph @odata.deltaLink; the paginator sends the later requests.", title="Send Token In", ) completion_condition: Optional[str] = Field( diff --git a/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py b/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py index 121c573696..2775e71cb7 100644 --- a/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py +++ b/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py @@ -2335,6 +2335,51 @@ def _validate_sync_token_stream(model: DeclarativeStreamModel) -> None: ) if not isinstance(cursor_model, SyncTokenCursorModel): return + if isinstance(retriever_model, SimpleRetrieverModel) and isinstance( + retriever_model.paginator, DefaultPaginatorModel + ): + token_option = cursor_model.token_option + ignore_stream_slicer_parameters = ( + retriever_model.ignore_stream_slicer_parameters_on_paginated_requests is True + ) + if isinstance(token_option, RequestOptionModel) and not ignore_stream_slicer_parameters: + for page_option in ( + retriever_model.paginator.page_token_option, + retriever_model.paginator.page_size_option, + ): + if ( + isinstance(page_option, RequestOptionModel) + and token_option.inject_into == page_option.inject_into + ): + matching_field = None + if token_option.field_name is not None and ( + token_option.field_name == page_option.field_name + ): + matching_field = token_option.field_name + elif token_option.field_path is not None and ( + token_option.field_path == page_option.field_path + ): + matching_field = ".".join(token_option.field_path) + if matching_field is not None: + raise ValueError( + f"Stream {stream_name}: token_option and the paginator both set " + f'{token_option.inject_into.value} "{matching_field}". Set ' + "ignore_stream_slicer_parameters_on_paginated_requests: true on " + "the retriever so the stored token is sent only on the first " + "request of each pass." + ) + + strategy = retriever_model.paginator.pagination_strategy + if ( + isinstance(token_option, RequestPathModel) or ignore_stream_slicer_parameters + ) and isinstance(strategy, (PageIncrementModel, OffsetIncrementModel)): + if strategy.inject_on_first_request is True: + raise ValueError( + f"Stream {stream_name}: {strategy.type} with inject_on_first_request: " + "true sends a page token on the first request, so the stored sync token " + "would never be sent. Set inject_on_first_request: false or use " + "CursorPagination." + ) datetime_only_keys = { "start_datetime", "end_datetime", @@ -4055,6 +4100,18 @@ def _get_url(req: Requester) -> str: partition_router, PartitionRouter ): request_options_provider = partition_router + token_request_options_provider = ( + request_options_provider._cursor_provider + if isinstance(request_options_provider, PerPartitionRequestOptionsProvider) + else request_options_provider + ) + first_request_path = ( + token_request_options_provider.get_request_path + if isinstance(incremental_sync, SyncTokenCursorModel) + and isinstance(token_request_options_provider, SyncTokenRequestOptionsProvider) + and isinstance(token_request_options_provider.token_option, RequestPath) + else None + ) paginator = ( self._create_component_from_model( @@ -4140,6 +4197,7 @@ def _get_url(req: Requester) -> str: config=config, ignore_stream_slicer_parameters_on_paginated_requests=ignore_stream_slicer_parameters_on_paginated_requests, post_pagination_filter=post_pagination_filter, + first_request_path=first_request_path, parameters=model.parameters or {}, ) @@ -4184,6 +4242,7 @@ def _get_url(req: Requester) -> str: token_reset_handler=( cursor.restart_without_token if isinstance(cursor, SyncTokenCursor) else None ), + first_request_path=first_request_path, parameters=model.parameters or {}, ) diff --git a/airbyte_cdk/sources/declarative/requesters/request_options/sync_token_request_options_provider.py b/airbyte_cdk/sources/declarative/requesters/request_options/sync_token_request_options_provider.py index e621b88613..27e0bfe019 100644 --- a/airbyte_cdk/sources/declarative/requesters/request_options/sync_token_request_options_provider.py +++ b/airbyte_cdk/sources/declarative/requesters/request_options/sync_token_request_options_provider.py @@ -13,12 +13,13 @@ from airbyte_cdk.sources.declarative.requesters.request_options.request_options_provider import ( RequestOptionsProvider, ) +from airbyte_cdk.sources.declarative.requesters.request_path import RequestPath from airbyte_cdk.sources.types import Config, StreamSlice, StreamState @dataclass class SyncTokenRequestOptionsProvider(RequestOptionsProvider): - token_option: RequestOption + token_option: Union[RequestOption, RequestPath] config: Config parameters: InitVar[Mapping[str, Any]] initial_token: Optional[str] = None @@ -33,8 +34,18 @@ def __post_init__(self, parameters: Mapping[str, Any]) -> None: def _get_request_options( self, option_type: RequestOptionType, stream_slice: Optional[StreamSlice] ) -> Mapping[str, Any]: + if isinstance(self.token_option, RequestPath): + return {} if not self.token_option or self.token_option.inject_into != option_type: return {} + token = self._get_token(stream_slice) + if token is None: + return {} + result: dict[str, Any] = {} + self.token_option.inject_into_request(result, token, self.config) + return result + + def _get_token(self, stream_slice: Optional[StreamSlice]) -> Optional[Any]: token = ( stream_slice.cursor_slice.get("token") if stream_slice and stream_slice.cursor_slice @@ -42,11 +53,12 @@ def _get_request_options( ) if token is None and self._initial_token: token = self._initial_token.eval(self.config) - if token is None: - return {} - result: dict[str, Any] = {} - self.token_option.inject_into_request(result, token, self.config) - return result + return token + + def get_request_path(self, stream_slice: Optional[StreamSlice]) -> Optional[str]: + """Return the initial request URL; after reset, use `initial_token` or the requester path.""" + token = self._get_token(stream_slice) + return str(token) if token is not None else None def get_request_params( self, diff --git a/airbyte_cdk/sources/declarative/retrievers/simple_retriever.py b/airbyte_cdk/sources/declarative/retrievers/simple_retriever.py index 2cde0f8adb..1e3454ca4d 100644 --- a/airbyte_cdk/sources/declarative/retrievers/simple_retriever.py +++ b/airbyte_cdk/sources/declarative/retrievers/simple_retriever.py @@ -20,6 +20,7 @@ Tuple, Union, ) +from urllib.parse import urlsplit import requests from typing_extensions import deprecated @@ -159,6 +160,7 @@ class SimpleRetriever(Retriever): ] = None page_observer: Optional[Callable[["PageCompleted"], None]] = None token_reset_handler: Optional[Callable[[StreamSlice, Optional[str]], StreamSlice]] = None + first_request_path: Optional[Callable[[StreamSlice], Optional[str]]] = None def __post_init__(self, parameters: Mapping[str, Any]) -> None: self._paginator = self.paginator or NoPagination(parameters=parameters) @@ -326,11 +328,54 @@ def _paginator_path( :param next_page_token: :return: """ - return self._paginator.path( + path = self._paginator.path( next_page_token=next_page_token, stream_state={}, # stream_state as an interpolation context is deprecated stream_slice=stream_slice, ) + if ( + path is None + and not next_page_token + and self.first_request_path + and stream_slice is not None + ): + token_url = self.first_request_path(stream_slice) + if token_url: + self._check_token_host(token_url, stream_slice) + return token_url + return path + + def _check_token_host(self, token_url: str, stream_slice: StreamSlice) -> None: + token = urlsplit(token_url) + if not token.netloc: + return + + base_url = self.requester.get_url_base( + stream_state=None, stream_slice=stream_slice, next_page_token=None + ) or self.requester.get_url( + stream_state=None, stream_slice=stream_slice, next_page_token=None + ) + base = urlsplit(base_url) + same_scheme = token.scheme in ("", base.scheme) + same_origin = same_scheme and token.netloc.lower() == base.netloc.lower() + if same_origin: + return + + if same_scheme: + token_origin = token.netloc + base_origin = base.netloc + else: + token_origin = f"{token.scheme}://{token.netloc}" + base_origin = f"{base.scheme}://{base.netloc}" + message = ( + f"Stream {self.name}: the stored sync token points at {token_origin}, " + f"not the API host {base_origin}. Clear this stream's state to start over." + ) + raise AirbyteTracedException( + internal_message=message, + message=message, + failure_type=FailureType.config_error, + ) def _parse_response( self, diff --git a/unit_tests/sources/declarative/incremental/test_sync_token_cursor.py b/unit_tests/sources/declarative/incremental/test_sync_token_cursor.py index 0241c516dc..ce069f4bfd 100644 --- a/unit_tests/sources/declarative/incremental/test_sync_token_cursor.py +++ b/unit_tests/sources/declarative/incremental/test_sync_token_cursor.py @@ -38,7 +38,17 @@ def _response(body): return response -def _cursor(*, state=None, router=None, max_partitions=1000, completion_condition=None): +def _cursor( + *, + state=None, + router=None, + max_partitions=1000, + completion_condition=None, + max_state_bytes=None, +): + kwargs = {} + if max_state_bytes is not None: + kwargs["max_state_bytes"] = max_state_bytes return SyncTokenCursor( stream_name="events", stream_namespace=None, @@ -52,6 +62,7 @@ def _cursor(*, state=None, router=None, max_partitions=1000, completion_conditio parameters={}, completion_condition=completion_condition, max_partitions=max_partitions, + **kwargs, ) @@ -434,3 +445,41 @@ def test_t_valid_and_t_init_use_partition_router_and_initial_token_slice(): slices = list(cursor.stream_slices()) assert slices[0].partition == {"account": "a"} assert slices[0].cursor_slice == {} + + +def test_t_bytes_caps_new_partitions_but_updates_existing_partitions(caplog): + partitions = [{"account": f"account-{index}"} for index in range(6)] + router = Router([StreamSlice(partition=partition, cursor_slice={}) for partition in partitions]) + first_sync = _cursor(router=router, max_state_bytes=4096) + + for partition in partitions: + stream_slice = StreamSlice(partition=partition, cursor_slice={}) + first_sync.on_page_completed( + PageCompleted(stream_slice, _response({"syncToken": "A" * 1024}), True) + ) + first_sync.close_partition(_partition(partition)) + + assert len(first_sync.state["states"]) == 4 + + second_sync = _cursor(state=first_sync.state, router=router, max_state_bytes=4096) + for stream_slice in second_sync.stream_slices(): + partition = stream_slice.partition + token = "longer" * 300 if partition == partitions[0] else "B" * 1024 + second_sync.on_page_completed( + PageCompleted(stream_slice, _response({"syncToken": token}), True) + ) + second_sync.close_partition(_partition(partition)) + + with caplog.at_level(logging.WARNING): + second_sync.ensure_at_least_one_state_emitted() + second_sync.ensure_at_least_one_state_emitted() + + committed = { + item["partition"]["account"]: item["cursor"]["token"] + for item in second_sync.state["states"] + } + assert len(committed) == 4 + assert committed["account-0"] == "longer" * 300 + assert "account-4" not in committed + assert caplog.text.count("Stream events: stored tokens reached the 4 KB state limit") == 1 + assert "4 partitions; 2 more read in full on every sync." in caplog.text diff --git a/unit_tests/sources/declarative/incremental/test_sync_token_cursor_integration.py b/unit_tests/sources/declarative/incremental/test_sync_token_cursor_integration.py index 91f368f76d..d85923174d 100644 --- a/unit_tests/sources/declarative/incremental/test_sync_token_cursor_integration.py +++ b/unit_tests/sources/declarative/incremental/test_sync_token_cursor_integration.py @@ -36,6 +36,12 @@ from airbyte_cdk.sources.declarative.models.declarative_component_schema import ( ListPartitionRouter as ListPartitionRouterModel, ) +from airbyte_cdk.sources.declarative.models.declarative_component_schema import ( + OffsetIncrement as OffsetIncrementModel, +) +from airbyte_cdk.sources.declarative.models.declarative_component_schema import ( + PageIncrement as PageIncrementModel, +) from airbyte_cdk.sources.declarative.models.declarative_component_schema import ( PaginationReset as PaginationResetModel, ) @@ -45,6 +51,12 @@ from airbyte_cdk.sources.declarative.models.declarative_component_schema import ( ParentStreamConfig as ParentStreamConfigModel, ) +from airbyte_cdk.sources.declarative.models.declarative_component_schema import ( + RequestOption as RequestOptionModel, +) +from airbyte_cdk.sources.declarative.models.declarative_component_schema import ( + RequestPath as RequestPathModel, +) from airbyte_cdk.sources.declarative.models.declarative_component_schema import ( SubstreamPartitionRouter as SubstreamPartitionRouterModel, ) @@ -157,6 +169,39 @@ def _manifest( } +def _request_path_manifest( + *, + initial_token=None, + request_parameters=None, + partitioned=False, + partition_values=("a", "b", "c"), + error_handler=None, +): + manifest = _manifest( + initial_token=initial_token, + partitioned=partitioned, + partition_values=partition_values, + error_handler=error_handler, + ) + stream = manifest["streams"][0] + requester = stream["retriever"]["requester"] + requester["http_method"] = "GET" + if request_parameters is not None: + requester["request_parameters"] = request_parameters + stream["incremental_sync"]["token_path"] = ["@odata.deltaLink"] + stream["incremental_sync"]["token_option"] = {"type": "RequestPath"} + paginator = stream["retriever"]["paginator"] + paginator["page_token_option"] = {"type": "RequestPath"} + paginator["pagination_strategy"] = { + "type": "CursorPagination", + "cursor_value": "{{ response.get('@odata.nextLink') }}", + "stop_condition": "{{ not response.get('@odata.nextLink') }}", + } + if partitioned: + stream["retriever"]["partition_router"].pop("request_option") + return manifest + + def _source_and_catalog( manifest, state=None, @@ -1026,3 +1071,308 @@ def test_t_inert_cursor_types_read_without_sync_page_observer(cursor_kind): retriever = generator._partition_factory._retriever assert isinstance(retriever, SimpleRetriever) assert retriever.page_observer is None + assert retriever.first_request_path is None + + +def test_t_inert_request_option_sync_token_does_not_bind_first_request_path(): + source, _ = _source_and_catalog(_manifest()) + stream = source.streams(config={})[0] + retriever = stream._stream_partition_generator._partition_factory._retriever + + assert isinstance(retriever, SimpleRetriever) + assert retriever.first_request_path is None + + +@pytest.mark.parametrize( + ("initial_token", "expected_url"), + ( + (None, _URL), + ("https://api.example.com/initial?cursor=1", "https://api.example.com/initial?cursor=1"), + ), +) +def test_t_url_init_uses_initial_request_path_or_requester_path(initial_token, expected_url): + with HttpMocker() as http_mocker: + http_mocker.get( + HttpRequest(expected_url), + HttpResponse(json.dumps({"items": [], "@odata.deltaLink": "D2"})), + ) + source, catalog = _source_and_catalog(_request_path_manifest(initial_token=initial_token)) + messages = _read(source, catalog) + + assert _state(messages)[0].stream.stream_state.__dict__["token"] == "D2" + + +def test_t_url_resumes_from_stored_token_and_follows_next_link_exactly(): + token = "https://api.example.com/delta?cursor=stored" + next_link = "https://api.example.com/delta?cursor=page-2" + state = _incoming_state({"state_type": "sync_token", "token": token}) + with HttpMocker() as http_mocker: + http_mocker.get( + HttpRequest(token), + HttpResponse(json.dumps({"items": [{"id": "1"}], "@odata.nextLink": next_link})), + ) + http_mocker.get( + HttpRequest(next_link), + HttpResponse(json.dumps({"items": [{"id": "2"}], "@odata.deltaLink": "D2"})), + ) + source, catalog = _source_and_catalog(_request_path_manifest(), state) + messages = _read(source, catalog, state) + + assert [message.record.data for message in messages if message.type == Type.RECORD] == [ + {"id": "1"}, + {"id": "2"}, + ] + assert _state(messages)[-1].stream.stream_state.__dict__["token"] == "D2" + + +def test_t_url_params_keeps_configured_query_parameters_on_each_exact_url(): + token = "https://api.example.com/delta?cursor=stored" + next_link = "https://api.example.com/delta?cursor=page-2" + state = _incoming_state({"state_type": "sync_token", "token": token}) + with HttpMocker() as http_mocker: + http_mocker.get( + HttpRequest(f"{token}&$select=id"), + HttpResponse(json.dumps({"items": [], "@odata.nextLink": next_link})), + ) + http_mocker.get( + HttpRequest(f"{next_link}&$select=id"), + HttpResponse(json.dumps({"items": [], "@odata.deltaLink": "D2"})), + ) + source, catalog = _source_and_catalog( + _request_path_manifest(request_parameters={"$select": "id"}), state + ) + messages = _read(source, catalog, state) + + assert _state(messages)[-1].stream.stream_state.__dict__["token"] == "D2" + + +def test_t_url_reset_without_initial_token_restarts_at_requester_path(caplog): + token = "https://api.example.com/delta?cursor=expired" + state = _incoming_state({"state_type": "sync_token", "token": token}) + error_handler = _reset_error_handler() + for response_filter in error_handler["response_filters"]: + if response_filter["action"] == "RESET_SYNC_TOKEN": + response_filter["http_codes"] = [410] + + with HttpMocker() as http_mocker: + http_mocker.get( + HttpRequest(token), + HttpResponse(json.dumps({"success": False}), status_code=410), + ) + http_mocker.get( + HttpRequest(_URL), + HttpResponse(json.dumps({"items": [], "@odata.deltaLink": "D2"})), + ) + source, catalog = _source_and_catalog( + _request_path_manifest(error_handler=error_handler), state + ) + messages = _read(source, catalog, state) + + assert _state(messages)[-1].stream.stream_state.__dict__["token"] == "D2" + assert "sync_token_expired" in caplog.text or "RESET_SYNC_TOKEN" in caplog.text + + +def test_t_url_par_keeps_each_partition_token_independent(): + tokens = { + "a": "https://api.example.com/a?cursor=A1", + "b": "https://api.example.com/b?cursor=B1", + "c": "https://api.example.com/c?cursor=C1", + } + state = _incoming_state( + { + "state_type": "sync_token", + "states": [ + { + "partition": {"account": account}, + "cursor": {"token": token, "committed_at": "2026-01-01T00:00:00Z"}, + } + for account, token in tokens.items() + ], + } + ) + with HttpMocker() as http_mocker: + for account in ("a", "b", "c"): + http_mocker.get( + HttpRequest(tokens[account]), + HttpResponse(json.dumps({"items": [], "@odata.deltaLink": f"{account}-next"})), + ) + source, catalog = _source_and_catalog(_request_path_manifest(partitioned=True), state) + messages = _read(source, catalog, state) + + committed = { + item["partition"]["account"]: item["cursor"]["token"] + for item in _state(messages)[-1].stream.stream_state.__dict__["states"] + } + assert committed == {"a": "a-next", "b": "b-next", "c": "c-next"} + + +@pytest.mark.parametrize( + ("token", "expected_url"), + ( + ("/relative?cursor=1", "https://api.example.com/relative?cursor=1"), + ("https://API.EXAMPLE.COM/delta?cursor=1", "https://API.EXAMPLE.COM/delta?cursor=1"), + ), +) +def test_t_origin_accepts_relative_and_same_origin_urls(token, expected_url): + state = _incoming_state({"state_type": "sync_token", "token": token}) + with HttpMocker() as http_mocker: + http_mocker.get( + HttpRequest(expected_url), + HttpResponse(json.dumps({"items": [], "@odata.deltaLink": "D2"})), + ) + source, catalog = _source_and_catalog(_request_path_manifest(), state) + messages = _read(source, catalog, state) + + assert _state(messages)[-1].stream.stream_state.__dict__["token"] == "D2" + + +@pytest.mark.parametrize( + ("token", "expected_message"), + ( + ( + "https://evil.example/x", + "Stream events: the stored sync token points at evil.example, not the API host " + "api.example.com. Clear this stream's state to start over.", + ), + ( + "//evil.example/x", + "Stream events: the stored sync token points at evil.example, not the API host " + "api.example.com. Clear this stream's state to start over.", + ), + ( + "http://api.example.com/delta", + "Stream events: the stored sync token points at http://api.example.com, not the API " + "host https://api.example.com. Clear this stream's state to start over.", + ), + ( + "https://api.example.com:8443/delta", + "Stream events: the stored sync token points at api.example.com:8443, not the API " + "host api.example.com. Clear this stream's state to start over.", + ), + ), +) +def test_t_origin_rejects_other_host_or_scheme_before_http_request(token, expected_message): + state = _incoming_state({"state_type": "sync_token", "token": token}) + with HttpMocker() as http_mocker: + source, catalog = _source_and_catalog(_request_path_manifest(), state) + messages = _read_until_error(source, catalog, state) + + assert any( + message.type == Type.TRACE + and message.trace.error + and message.trace.error.failure_type == FailureType.config_error + and message.trace.error.message == expected_message + for message in messages + ) + assert http_mocker._mocker.request_history == [] + + +def test_t_pos_does_not_send_sync_token_with_paginator_on_later_pages(): + manifest = _manifest() + stream = manifest["streams"][0] + requester = stream["retriever"]["requester"] + requester["http_method"] = "GET" + stream["incremental_sync"]["token_path"] = ["next_stream_position"] + stream["incremental_sync"]["token_option"] = { + "type": "RequestOption", + "field_name": "stream_position", + "inject_into": "request_parameter", + } + stream["incremental_sync"]["completion_condition"] = "{{ response.chunk_size == 0 }}" + retriever = stream["retriever"] + retriever["ignore_stream_slicer_parameters_on_paginated_requests"] = True + paginator = retriever["paginator"] + paginator["page_token_option"] = { + "type": "RequestOption", + "field_name": "stream_position", + "inject_into": "request_parameter", + } + paginator["pagination_strategy"] = { + "type": "CursorPagination", + "cursor_value": "{{ response.get('next_stream_position') }}", + "stop_condition": "{{ response.get('chunk_size') == 0 }}", + } + state = _incoming_state({"state_type": "sync_token", "token": "P1"}) + + with HttpMocker() as http_mocker: + http_mocker.get( + HttpRequest(_URL, query_params={"stream_position": "P1"}), + HttpResponse( + json.dumps({"items": [{"id": "1"}], "chunk_size": 1, "next_stream_position": "P2"}) + ), + ) + http_mocker.get( + HttpRequest(_URL, query_params={"stream_position": "P2"}), + HttpResponse( + json.dumps({"items": [{"id": "2"}], "chunk_size": 1, "next_stream_position": "P3"}) + ), + ) + http_mocker.get( + HttpRequest(_URL, query_params={"stream_position": "P3"}), + HttpResponse(json.dumps({"items": [], "chunk_size": 0, "next_stream_position": "P4"})), + ) + source, catalog = _source_and_catalog(manifest, state) + messages = _read(source, catalog, state) + + assert _state(messages)[-1].stream.stream_state.__dict__["token"] == "P4" + + +@pytest.mark.parametrize("paginator_option", ("page_token_option", "page_size_option")) +@pytest.mark.parametrize( + ("field", "value"), + (("field_name", "stream_position"), ("field_path", ["stream_position"])), +) +def test_t_collide_requires_ignore_flag_for_matching_request_option(paginator_option, field, value): + stream = _stream_model() + token_option = RequestOptionModel.parse_obj( + {"type": "RequestOption", "inject_into": "request_parameter", field: value} + ) + paginator_request_option = RequestOptionModel.parse_obj( + {"type": "RequestOption", "inject_into": "request_parameter", field: value} + ) + stream.incremental_sync.token_option = token_option + setattr(stream.retriever.paginator, paginator_option, paginator_request_option) + + with pytest.raises(ValueError) as exc_info: + ModelToComponentFactory._validate_sync_token_stream(stream) + assert str(exc_info.value) == ( + "Stream events: token_option and the paginator both set request_parameter " + '"stream_position". Set ignore_stream_slicer_parameters_on_paginated_requests: ' + "true on the retriever so the stored token is sent only on the first request of each pass." + ) + + stream.retriever.ignore_stream_slicer_parameters_on_paginated_requests = True + ModelToComponentFactory._validate_sync_token_stream(stream) + + +@pytest.mark.parametrize("token_mode", ("request_path", "ignore_flag")) +@pytest.mark.parametrize( + ("strategy_model", "strategy_name"), + ( + (PageIncrementModel, "PageIncrement"), + (OffsetIncrementModel, "OffsetIncrement"), + ), +) +@pytest.mark.parametrize("inject_on_first_request", (True, False)) +def test_t_firstreq_requires_page_token_to_be_disabled_for_first_request( + token_mode, strategy_model, strategy_name, inject_on_first_request +): + stream = _stream_model() + if token_mode == "request_path": + stream.incremental_sync.token_option = RequestPathModel(type="RequestPath") + else: + stream.retriever.ignore_stream_slicer_parameters_on_paginated_requests = True + stream.retriever.paginator.pagination_strategy = strategy_model.parse_obj( + {"type": strategy_name, "inject_on_first_request": inject_on_first_request} + ) + + if inject_on_first_request: + with pytest.raises(ValueError) as exc_info: + ModelToComponentFactory._validate_sync_token_stream(stream) + assert str(exc_info.value) == ( + f"Stream events: {strategy_name} with inject_on_first_request: true sends a page " + "token on the first request, so the stored sync token would never be sent. Set " + "inject_on_first_request: false or use CursorPagination." + ) + else: + ModelToComponentFactory._validate_sync_token_stream(stream) diff --git a/unit_tests/sources/declarative/requesters/request_options/test_sync_token_request_options_provider.py b/unit_tests/sources/declarative/requesters/request_options/test_sync_token_request_options_provider.py new file mode 100644 index 0000000000..835a9bae32 --- /dev/null +++ b/unit_tests/sources/declarative/requesters/request_options/test_sync_token_request_options_provider.py @@ -0,0 +1,60 @@ +import pytest + +from airbyte_cdk.sources.declarative.requesters.request_options.sync_token_request_options_provider import ( + SyncTokenRequestOptionsProvider, +) +from airbyte_cdk.sources.declarative.requesters.request_path import RequestPath +from airbyte_cdk.sources.types import StreamSlice + + +@pytest.mark.parametrize( + "method_name", + ( + "get_request_params", + "get_request_headers", + "get_request_body_data", + "get_request_body_json", + ), +) +def test_request_path_token_option_does_not_inject_request_options(method_name): + provider = SyncTokenRequestOptionsProvider( + config={}, + token_option=RequestPath(parameters={}), + parameters={}, + ) + + method = getattr(provider, method_name) + assert method(stream_slice=StreamSlice(partition={}, cursor_slice={"token": "stored"})) == {} + + +@pytest.mark.parametrize( + ("cursor_slice", "initial_token", "config", "expected"), + ( + ( + {"token": "https://api.example.com/stored"}, + "{{ config['initial'] }}", + {"initial": "https://api.example.com/initial"}, + "https://api.example.com/stored", + ), + ( + {}, + "{{ config['initial'] }}", + {"initial": "https://api.example.com/initial"}, + "https://api.example.com/initial", + ), + ({}, None, {}, None), + ), +) +def test_get_request_path_prefers_stored_then_initial_token( + cursor_slice, initial_token, config, expected +): + provider = SyncTokenRequestOptionsProvider( + config=config, + token_option=RequestPath(parameters={}), + initial_token=initial_token, + parameters={}, + ) + + assert ( + provider.get_request_path(StreamSlice(partition={}, cursor_slice=cursor_slice)) == expected + )