From 846bf125b34cd35fe05491ff509b5c84984d2b50 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 16 Oct 2026 23:06:06 +0000 Subject: [PATCH 1/3] feat(declarative): add opaque token cursor support Co-Authored-By: bot_apk --- .../declarative_component_schema.yaml | 63 +++ .../declarative/incremental/__init__.py | 2 + .../incremental/opaque_token_cursor.py | 324 +++++++++++++++ .../models/declarative_component_schema.py | 40 +- .../parsers/model_to_component_factory.py | 213 +++++++++- .../paginators/default_paginator.py | 7 + .../requesters/request_options/__init__.py | 4 + .../opaque_token_request_options_provider.py | 85 ++++ .../retrievers/simple_retriever.py | 17 + .../streams/concurrent/default_stream.py | 13 + .../incremental/test_opaque_token_cursor.py | 314 ++++++++++++++ .../test_opaque_token_cursor_integration.py | 390 ++++++++++++++++++ 12 files changed, 1465 insertions(+), 7 deletions(-) create mode 100644 airbyte_cdk/sources/declarative/incremental/opaque_token_cursor.py create mode 100644 airbyte_cdk/sources/declarative/requesters/request_options/opaque_token_request_options_provider.py create mode 100644 unit_tests/sources/declarative/incremental/test_opaque_token_cursor.py create mode 100644 unit_tests/sources/declarative/incremental/test_opaque_token_cursor_integration.py diff --git a/airbyte_cdk/sources/declarative/declarative_component_schema.yaml b/airbyte_cdk/sources/declarative/declarative_component_schema.yaml index 6df783d1db..eadd756458 100644 --- a/airbyte_cdk/sources/declarative/declarative_component_schema.yaml +++ b/airbyte_cdk/sources/declarative/declarative_component_schema.yaml @@ -1428,6 +1428,68 @@ definitions: $parameters: type: object additionalProperties: true + OpaqueTokenCursor: + title: Opaque Token Cursor + description: >- + Incremental sync driven by an opaque change token that the API returns at the end of a + pagination pass, such as a sync token. The token is committed to state only after the + pass's records are emitted, and is sent back on the next sync. Tokens are never parsed, + compared or derived from records. + type: object + required: [type, token_path, token_option] + additionalProperties: true + properties: + type: + type: string + enum: [OpaqueTokenCursor] + token_path: + title: Token Path + description: >- + Path to the next token in the decoded response that ends a pagination pass. The + value must be a string or an integer, and is stored and sent as a string. A missing, + null or empty value means the response carried no token. + type: array + items: + type: string + interpolation_context: [config] + examples: + - ["syncToken"] + - ["nextSyncToken"] + - ["meta", "sync_token"] + 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" + completion_condition: + title: Completion Condition + description: >- + Must be true on the last response of a pass for its token to be committed. Defaults + to a token being present at token_path. Set it when the API repeats the old token on + intermediate pages or returns a position on every page. + type: string + interpolation_context: [config, headers, response, stream_partition, stream_slice] + examples: + - "{{ not response.moreDataAvailable }}" + initial_token: + title: Initial Token + description: >- + Value to send when no token is stored, for an API that needs one to start a full + pass, such as Todoist's "*". When unset, nothing is sent. + type: string + interpolation_context: [config] + examples: ["*"] + cursor_field: + title: Destination Cursor Field + description: >- + Optional top-level record field advertised as the stream's cursor, so destinations + can order versions of a record when deduplicating. It is never used to filter + records or to build state. + type: string + $parameters: + type: object + additionalProperties: true JwtAuthenticator: title: JWT Authenticator description: Authenticator for requests using JWT authentication flow. @@ -1845,6 +1907,7 @@ definitions: anyOf: - "$ref": "#/definitions/DatetimeBasedCursor" - "$ref": "#/definitions/IncrementingCountCursor" + - "$ref": "#/definitions/OpaqueTokenCursor" primary_key: title: Primary Key "$ref": "#/definitions/PrimaryKey" diff --git a/airbyte_cdk/sources/declarative/incremental/__init__.py b/airbyte_cdk/sources/declarative/incremental/__init__.py index ddefe0f8be..1fefbdd57d 100644 --- a/airbyte_cdk/sources/declarative/incremental/__init__.py +++ b/airbyte_cdk/sources/declarative/incremental/__init__.py @@ -6,8 +6,10 @@ ConcurrentCursorFactory, ConcurrentPerPartitionCursor, ) +from airbyte_cdk.sources.declarative.incremental.opaque_token_cursor import OpaqueTokenCursor __all__ = [ "ConcurrentCursorFactory", "ConcurrentPerPartitionCursor", + "OpaqueTokenCursor", ] diff --git a/airbyte_cdk/sources/declarative/incremental/opaque_token_cursor.py b/airbyte_cdk/sources/declarative/incremental/opaque_token_cursor.py new file mode 100644 index 0000000000..4beeddb845 --- /dev/null +++ b/airbyte_cdk/sources/declarative/incremental/opaque_token_cursor.py @@ -0,0 +1,324 @@ +# +# Copyright (c) 2025 Airbyte, Inc., all rights reserved. +# + +import copy +import datetime +import logging +import threading +import time +from collections import OrderedDict +from typing import Any, Iterable, Mapping, MutableMapping, Optional + +import dpath + +from airbyte_cdk.models import FailureType +from airbyte_cdk.sources.connector_state_manager import ConnectorStateManager +from airbyte_cdk.sources.declarative.decoders import Decoder +from airbyte_cdk.sources.declarative.incremental.concurrent_partition_cursor import ( + ConcurrentPerPartitionCursor, +) +from airbyte_cdk.sources.declarative.interpolation import InterpolatedString +from airbyte_cdk.sources.declarative.interpolation.interpolated_boolean import ( + InterpolatedBoolean, +) +from airbyte_cdk.sources.declarative.partition_routers.partition_router import PartitionRouter +from airbyte_cdk.sources.declarative.partition_routers.single_partition_router import ( + SinglePartitionRouter, +) +from airbyte_cdk.sources.declarative.retrievers.simple_retriever import PageCompleted +from airbyte_cdk.sources.message import MessageRepository +from airbyte_cdk.sources.streams.checkpoint.per_partition_key_serializer import ( + PerPartitionKeySerializer, +) +from airbyte_cdk.sources.streams.concurrent.cursor import Cursor, CursorField +from airbyte_cdk.sources.streams.concurrent.partitions.partition import Partition +from airbyte_cdk.sources.types import Config, Record, StreamSlice, StreamState +from airbyte_cdk.utils.traced_exception import AirbyteTracedException + +logger = logging.getLogger("airbyte") + + +class OpaqueTokenCursor(Cursor): + STATE_TYPE = "opaque_token" + + def __init__( + self, + *, + stream_name: str, + stream_namespace: Optional[str], + stream_state: Optional[StreamState], + message_repository: MessageRepository, + connector_state_manager: ConnectorStateManager, + partition_router: PartitionRouter, + token_path: list[str], + decoder: Decoder, + config: Config, + parameters: Mapping[str, Any], + completion_condition: Optional[str] = None, + cursor_field: Optional[str] = None, + max_partitions: int = ConcurrentPerPartitionCursor.DEFAULT_MAX_PARTITIONS_NUMBER, + ) -> None: + self._stream_name = stream_name + self._stream_namespace = stream_namespace + self._message_repository = message_repository + self._connector_state_manager = connector_state_manager + self._partition_router = partition_router + self._partitioned = not isinstance(partition_router, SinglePartitionRouter) + self._token_path = [ + InterpolatedString.create(path, parameters=parameters) for path in token_path + ] + self._decoder = decoder + self._config = config + self._completion_condition = ( + InterpolatedBoolean(completion_condition, parameters=parameters) + if completion_condition is not None + else None + ) + self.cursor_field = CursorField(cursor_field) if cursor_field else None + self._max_partitions = max_partitions + self._serializer = PerPartitionKeySerializer() + self._committed: OrderedDict[str, dict[str, Any]] = OrderedDict() + self._pending: dict[str, dict[str, Any]] = {} + self._lock = threading.RLock() + self._last_emission_time = 0.0 + self._warned_no_token = False + self._warned_cap = False + self._partitions_over_cap = 0 + self._load_state(stream_state or {}) + + @property + def state(self) -> MutableMapping[str, Any]: + with self._lock: + if not self._partitioned: + state: dict[str, Any] = {"state_type": self.STATE_TYPE} + if self._committed: + state.update(copy.deepcopy(next(iter(self._committed.values())))) + return state + return { + "state_type": self.STATE_TYPE, + "states": [ + { + "partition": self._serializer.to_partition(partition_key), + "cursor": copy.deepcopy(cursor), + } + for partition_key, cursor in self._committed.items() + ], + } + + def stream_slices(self) -> Iterable[StreamSlice]: + seen_partitions: set[str] = set() + for stream_slice in self._partition_router.stream_slices(): + partition = dict(stream_slice.partition or {}) + if "token" in partition: + raise ValueError( + f"Stream {self._stream_name}: partition router uses reserved key 'token'" + ) + partition_key = self._serializer.to_partition_key(partition) + if partition_key in seen_partitions: + continue + seen_partitions.add(partition_key) + with self._lock: + committed = self._committed.get(partition_key) + token = committed["token"] if committed else None + yield StreamSlice( + partition=partition, + cursor_slice={"token": token} if token else {}, + extra_fields=stream_slice.extra_fields, + ) + + if self._partitioned: + with self._lock: + vanished = set(self._committed) - seen_partitions + for partition_key in vanished: + del self._committed[partition_key] + if vanished: + logger.info( + "Stream %s: dropped the stored tokens of %s partitions the router no longer returns.", + self._stream_name, + len(vanished), + ) + + def on_page_completed(self, page: PageCompleted) -> None: + if not page.pagination_exhausted or page.response is None: + return + + body = next( + ( + decoded + for decoded in self._decoder.decode(page.response) + if isinstance(decoded, Mapping) + ), + {}, + ) + path = [str(part.eval(self._config)) for part in self._token_path] + token_value = dpath.get(body, path, default=None) if isinstance(body, Mapping) else None + token = self._coerce_token(token_value, path) + if token is None: + self._warn_missing_token() + return + + if self._completion_condition is not None: + complete = self._completion_condition.eval( + self._config, + response=body, + headers=page.response.headers, + stream_slice=page.stream_slice, + ) + else: + complete = True + if not complete: + self._warn_missing_token() + return + + partition_key = self._serializer.to_partition_key(page.stream_slice.partition or {}) + with self._lock: + self._pending[partition_key] = { + "token": token, + "committed_at": datetime.datetime.now(datetime.timezone.utc) + .isoformat(timespec="milliseconds") + .replace("+00:00", "Z"), + } + + @staticmethod + def _coerce_token(value: Any, path: list[str]) -> Optional[str]: + if value is None or value == "": + return None + if isinstance(value, bool): + value_type = "boolean" + elif isinstance(value, str): + return value + elif isinstance(value, int): + return str(value) + elif isinstance(value, float): + value_type = "float" + elif isinstance(value, Mapping): + value_type = "object" + elif isinstance(value, list): + value_type = "list" + else: + value_type = type(value).__name__ + raise AirbyteTracedException( + internal_message=f"Stream token_path {path} holds a {value_type}, not a string or integer.", + message=f"Token path {path} holds a {value_type}, not a string or integer. Check token_path.", + failure_type=FailureType.system_error, + ) + + def _warn_missing_token(self) -> None: + if not self._warned_no_token: + logger.warning( + "Stream %s: pagination ended without a new change token; kept the stored token. " + "The next sync re-reads from it.", + self._stream_name, + ) + self._warned_no_token = True + + def close_partition(self, partition: Partition) -> None: + stream_slice = partition.to_slice() + if stream_slice is None: + raise ValueError("stream_slice cannot be None") + partition_values = getattr(stream_slice, "partition", None) + if partition_values is None: + partition_values = stream_slice.get("partition", {}) + partition_key = self._serializer.to_partition_key(partition_values) + with self._lock: + pending = self._pending.pop(partition_key, None) + if pending is not None: + if partition_key in self._committed or len(self._committed) < self._max_partitions: + self._committed[partition_key] = pending + else: + self._partitions_over_cap += 1 + if not self._warned_cap: + logger.warning( + "Stream %s: %s partitions already store tokens; additional partitions read in full.", + self._stream_name, + len(self._committed), + ) + self._warned_cap = True + if not self._partitioned: + self._emit_state_message() + elif time.time() - self._last_emission_time > 600: + self._emit_state_message() + + def _emit_state_message(self) -> None: + self._last_emission_time = time.time() + self._connector_state_manager.update_state_for_stream( + self._stream_name, self._stream_namespace, self.state + ) + self._message_repository.emit_message( + self._connector_state_manager.create_state_message( + self._stream_name, self._stream_namespace + ) + ) + + def ensure_at_least_one_state_emitted(self) -> None: + self._emit_state_message() + + def observe(self, record: Record) -> None: + return + + def should_be_synced(self, record: Record) -> bool: + return True + + def _load_state(self, stream_state: StreamState) -> None: + if not stream_state: + return + if stream_state.get("state_type") != self.STATE_TYPE: + logger.info( + "Stream %s: ignoring foreign state and reading from the beginning.", + self._stream_name, + ) + return + if self._partitioned: + states = stream_state.get("states") + if not isinstance(states, list) or "token" in stream_state: + logger.info( + "Stream %s: ignoring state with an incompatible partition shape.", + self._stream_name, + ) + return + committed: dict[str, dict[str, Any]] = {} + for state in states: + if not isinstance(state, Mapping): + logger.info( + "Stream %s: ignoring state with an incompatible opaque token partition.", + self._stream_name, + ) + return + partition = state.get("partition") + cursor = state.get("cursor") + if ( + isinstance(partition, Mapping) + and isinstance(cursor, Mapping) + and isinstance(cursor.get("token"), str) + and cursor.get("token") + ): + committed[self._serializer.to_partition_key(partition)] = { + "token": cursor["token"], + "committed_at": cursor.get("committed_at", ""), + } + else: + logger.info( + "Stream %s: ignoring state with an incompatible opaque token partition.", + self._stream_name, + ) + return + self._committed.update(committed) + else: + if "states" in stream_state: + logger.info( + "Stream %s: ignoring state with an incompatible unpartitioned shape.", + self._stream_name, + ) + return + token = stream_state.get("token") + if isinstance(token, str) and token: + self._committed["{}"] = { + "token": token, + "committed_at": stream_state.get("committed_at", ""), + } + elif token is not None: + logger.info( + "Stream %s: ignoring an incompatible opaque token state.", + self._stream_name, + ) diff --git a/airbyte_cdk/sources/declarative/models/declarative_component_schema.py b/airbyte_cdk/sources/declarative/models/declarative_component_schema.py index 0317d52d65..01ed2f3bd8 100644 --- a/airbyte_cdk/sources/declarative/models/declarative_component_schema.py +++ b/airbyte_cdk/sources/declarative/models/declarative_component_schema.py @@ -2952,6 +2952,42 @@ class FileUploader(BaseModel): parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters") +class OpaqueTokenCursor(BaseModel): + class Config: + extra = Extra.allow + + type: Literal["OpaqueTokenCursor"] + token_path: List[str] = Field( + ..., + description="Path to the next token in the decoded response that ends a pagination pass. The value must be a string or an integer, and is stored and sent as a string. A missing, null or empty value means the response carried no token.", + examples=[["syncToken"], ["nextSyncToken"], ["meta", "sync_token"]], + title="Token Path", + ) + token_option: RequestOption = 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.", + title="Send Token In", + ) + completion_condition: Optional[str] = Field( + None, + description="Must be true on the last response of a pass for its token to be committed. Defaults to a token being present at token_path. Set it when the API repeats the old token on intermediate pages or returns a position on every page.", + examples=["{{ not response.moreDataAvailable }}"], + title="Completion Condition", + ) + initial_token: Optional[str] = Field( + None, + description='Value to send when no token is stored, for an API that needs one to start a full pass, such as Todoist\'s "*". When unset, nothing is sent.', + examples=["*"], + title="Initial Token", + ) + cursor_field: Optional[str] = Field( + None, + description="Optional top-level record field advertised as the stream's cursor, so destinations can order versions of a record when deduplicating. It is never used to filter records or to build state.", + title="Destination Cursor Field", + ) + parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters") + + class DeclarativeStream(BaseModel): class Config: extra = Extra.allow @@ -2963,7 +2999,9 @@ class Config: description="Component used to coordinate how records are extracted across stream slices and request pages.", title="Retriever", ) - incremental_sync: Optional[Union[DatetimeBasedCursor, IncrementingCountCursor]] = Field( + incremental_sync: Optional[ + Union[DatetimeBasedCursor, IncrementingCountCursor, OpaqueTokenCursor] + ] = Field( None, description="Component used to fetch data incrementally based on a time field in the data.", title="Incremental Sync", 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 219e0f70ad..23902e1484 100644 --- a/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py +++ b/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py @@ -128,6 +128,7 @@ from airbyte_cdk.sources.declarative.incremental import ( ConcurrentCursorFactory, ConcurrentPerPartitionCursor, + OpaqueTokenCursor, ) from airbyte_cdk.sources.declarative.interpolation import InterpolatedString from airbyte_cdk.sources.declarative.interpolation.interpolated_mapping import InterpolatedMapping @@ -391,6 +392,9 @@ from airbyte_cdk.sources.declarative.models.declarative_component_schema import ( OffsetIncrement as OffsetIncrementModel, ) +from airbyte_cdk.sources.declarative.models.declarative_component_schema import ( + OpaqueTokenCursor as OpaqueTokenCursorModel, +) from airbyte_cdk.sources.declarative.models.declarative_component_schema import ( PageIncrement as PageIncrementModel, ) @@ -572,6 +576,7 @@ DatetimeBasedRequestOptionsProvider, DefaultRequestOptionsProvider, InterpolatedRequestOptionsProvider, + OpaqueTokenRequestOptionsProvider, RequestOptionsProvider, ) from airbyte_cdk.sources.declarative.requesters.request_options.per_partition_request_option_provider import ( @@ -2127,6 +2132,7 @@ def create_default_stream( self, model: DeclarativeStreamModel, config: Config, is_parent: bool = False, **kwargs: Any ) -> AbstractStream: primary_key = model.primary_key.__root__ if model.primary_key else None + self._validate_opaque_token_stream(model) self._migrate_state(model, config) self._warn_on_ineffective_incremental_dependency(model) @@ -2137,6 +2143,7 @@ def create_default_stream( **kwargs, ) concurrent_cursor = self._build_concurrent_cursor(model, partition_router, config) + request_options_provider: Optional[RequestOptionsProvider] if model.incremental_sync and isinstance(model.incremental_sync, DatetimeBasedCursorModel): cursor_model: DatetimeBasedCursorModel = model.incremental_sync @@ -2170,6 +2177,27 @@ def create_default_stream( partition_router, datetime_request_options_provider ) ) + elif isinstance(model.incremental_sync, OpaqueTokenCursorModel): + opaque_cursor_model = model.incremental_sync + token_option = self._create_component_from_model( + opaque_cursor_model.token_option, + config, + parameters=opaque_cursor_model.parameters or {}, + ) + opaque_request_options_provider = OpaqueTokenRequestOptionsProvider( + token_option=token_option, + config=config, + parameters=opaque_cursor_model.parameters or {}, + initial_token=opaque_cursor_model.initial_token, + ) + request_options_provider = ( + PerPartitionRequestOptionsProvider( + partition_router, opaque_request_options_provider + ) + if isinstance(concurrent_cursor, OpaqueTokenCursor) + and concurrent_cursor._partitioned + else opaque_request_options_provider + ) elif model.incremental_sync and isinstance( model.incremental_sync, IncrementingCountCursorModel ): @@ -2283,8 +2311,120 @@ def create_default_stream( logger=logging.getLogger(f"airbyte.{stream_name}"), cursor=concurrent_cursor, supports_file_transfer=hasattr(model, "file_uploader") and bool(model.file_uploader), + source_defined_incremental=isinstance(model.incremental_sync, OpaqueTokenCursorModel), ) + @staticmethod + def _validate_opaque_token_stream(model: DeclarativeStreamModel) -> None: + cursor_model = model.incremental_sync + stream_name = model.name or "" + retriever_model = model.retriever + router_model = getattr(retriever_model, "partition_router", None) + routers = router_model if isinstance(router_model, list) else [router_model] + for router in routers: + if isinstance(router, SubstreamPartitionRouterModel): + for parent in router.parent_stream_configs: + parent_sync = getattr(parent.stream, "incremental_sync", None) + if isinstance(parent.stream, StateDelegatingStreamModel): + parent_sync = parent.stream.incremental_stream.incremental_sync + if parent.incremental_dependency and isinstance( + parent_sync, OpaqueTokenCursorModel + ): + raise ValueError( + f"Stream {stream_name}: an OpaqueTokenCursor parent cannot be used with " + "`incremental_dependency`." + ) + if not isinstance(cursor_model, OpaqueTokenCursorModel): + return + datetime_only_keys = { + "start_datetime", + "end_datetime", + "step", + "cursor_granularity", + "datetime_format", + "cursor_datetime_formats", + "start_time_option", + "end_time_option", + "lookback_window", + "partition_field_start", + "partition_field_end", + "is_data_feed", + "is_client_side_incremental", + "is_compare_strictly", + "global_substream_cursor", + "clamping", + } + unsupported = sorted(datetime_only_keys.intersection(cursor_model.__dict__)) + if unsupported: + raise ValueError( + f"Stream {stream_name}: OpaqueTokenCursor does not support datetime-only setting(s): " + f"{', '.join(unsupported)}." + ) + if isinstance(model.retriever, AsyncRetrieverModel): + raise ValueError( + f"Stream {stream_name}: OpaqueTokenCursor is not supported with AsyncRetriever; " + "use SimpleRetriever so response pages can be observed." + ) + if isinstance(retriever_model, CustomRetrieverModel): + LOGGER.warning( + "Stream %s: CustomRetriever does not expose response pages to OpaqueTokenCursor; " + "no change token will be captured.", + stream_name, + ) + + for router in routers: + if isinstance(router, GroupingPartitionRouterModel): + raise ValueError( + f"Stream {stream_name}: OpaqueTokenCursor does not support GroupingPartitionRouter; " + "use an ungrouped partition router." + ) + if isinstance(router, ListPartitionRouterModel) and router.cursor_field == "token": + raise ValueError( + f"Stream {stream_name}: partition router field 'token' conflicts with OpaqueTokenCursor." + ) + if isinstance(router, SubstreamPartitionRouterModel): + for parent in router.parent_stream_configs: + if parent.partition_field == "token": + raise ValueError( + f"Stream {stream_name}: parent partition field 'token' conflicts with OpaqueTokenCursor." + ) + if parent.lazy_read_pointer: + raise ValueError( + f"Stream {stream_name}: OpaqueTokenCursor is not supported with parent " + "`lazy_read_pointer`; remove the lazy pointer." + ) + if parent.incremental_dependency: + raise ValueError( + f"Stream {stream_name}: OpaqueTokenCursor is not supported with " + "`incremental_dependency` on a SubstreamPartitionRouter." + ) + pagination_reset = getattr(retriever_model, "pagination_reset", None) + if pagination_reset and pagination_reset.limits: + raise ValueError( + f"Stream {stream_name}: pagination reset `limits` are not supported with OpaqueTokenCursor." + ) + requester_model = getattr(retriever_model, "requester", None) + query_properties = [getattr(requester_model, "query_properties", None)] + request_parameters = getattr(requester_model, "request_parameters", None) + if isinstance(request_parameters, Mapping): + query_properties.extend( + value + for value in request_parameters.values() + if isinstance(value, QueryPropertiesModel) + or (isinstance(value, Mapping) and value.get("type") == "QueryProperties") + ) + for properties in query_properties: + property_chunking = ( + properties.get("property_chunking") + if isinstance(properties, Mapping) + else getattr(properties, "property_chunking", None) + ) + if property_chunking: + raise ValueError( + f"Stream {stream_name}: query property chunking is not supported with " + "OpaqueTokenCursor." + ) + def _warn_on_ineffective_incremental_dependency(self, model: DeclarativeStreamModel) -> None: """ `incremental_dependency: true` only takes effect when the substream defines its own @@ -2405,6 +2545,16 @@ def _build_concurrent_cursor( stream_name = model.name or "" stream_state = self._connector_state_manager.get_stream_state(stream_name, None) + if isinstance(model.incremental_sync, OpaqueTokenCursorModel): + return self.create_opaque_token_cursor( + model.incremental_sync, + config=config, + stream_name=stream_name, + stream_state=stream_state, + partition_router=stream_slicer or SinglePartitionRouter(parameters={}), + retriever_model=model.retriever, + ) + if ( model.incremental_sync and stream_slicer @@ -2455,6 +2605,37 @@ def _build_concurrent_cursor( ) return FinalStateCursor(stream_name, None, self._message_repository) + def create_opaque_token_cursor( + self, + model: OpaqueTokenCursorModel, + *, + config: Config, + stream_name: str, + stream_state: Optional[Mapping[str, Any]], + partition_router: PartitionRouter, + retriever_model: Any, + ) -> OpaqueTokenCursor: + decoder_model = getattr(retriever_model, "decoder", None) + decoder = ( + self._create_component_from_model(model=decoder_model, config=config) + if decoder_model + else JsonDecoder(parameters={}) + ) + return OpaqueTokenCursor( + stream_name=stream_name, + stream_namespace=None, + stream_state=stream_state, + message_repository=self._message_repository, + connector_state_manager=self._connector_state_manager, + partition_router=partition_router, + token_path=model.token_path, + decoder=decoder, + config=config or {}, + parameters=model.parameters or {}, + completion_condition=model.completion_condition, + cursor_field=model.cursor_field, + ) + def create_default_error_handler( self, model: DefaultErrorHandlerModel, config: Config, **kwargs: Any ) -> DefaultErrorHandler: @@ -3708,7 +3889,11 @@ def create_simple_retriever( transformations: List[RecordTransformation], file_uploader: Optional[DefaultFileUploader] = None, incremental_sync: Optional[ - Union[IncrementingCountCursorModel, DatetimeBasedCursorModel] + Union[ + IncrementingCountCursorModel, + DatetimeBasedCursorModel, + OpaqueTokenCursorModel, + ] ] = None, use_cache: Optional[bool] = None, log_formatter: Optional[Callable[[Response], Any]] = None, @@ -3765,9 +3950,13 @@ def _get_url(req: Requester) -> str: ) decoder = ( - self._create_component_from_model(model=model.decoder, config=config) - if model.decoder - else JsonDecoder(parameters={}) + cursor._decoder + if isinstance(cursor, OpaqueTokenCursor) + else ( + self._create_component_from_model(model=model.decoder, config=config) + if model.decoder + else JsonDecoder(parameters={}) + ) ) record_selector = self._create_component_from_model( model=model.record_selector, @@ -3843,6 +4032,15 @@ def _get_url(req: Requester) -> str: config=config, ) + if ( + isinstance(cursor, OpaqueTokenCursor) + and query_properties + and query_properties.property_chunking + ): + raise ValueError( + f"Stream {name}: query property chunking is not supported with OpaqueTokenCursor." + ) + requester = self._create_component_from_model( model=model.requester, decoder=decoder, @@ -3981,6 +4179,9 @@ def _get_url(req: Requester) -> str: else None ), post_pagination_filter=post_pagination_filter, + page_observer=( + cursor.on_page_completed if isinstance(cursor, OpaqueTokenCursor) else None + ), parameters=model.parameters or {}, ) @@ -4291,7 +4492,7 @@ def _create_request_window_splitting( name: str, cursor: Optional[Cursor], incremental_sync: Optional[ - Union[IncrementingCountCursorModel, DatetimeBasedCursorModel] + Union[IncrementingCountCursorModel, DatetimeBasedCursorModel, OpaqueTokenCursorModel] ] = None, query_properties: Optional[QueryProperties] = None, file_uploader: Optional[DefaultFileUploader] = None, @@ -4342,7 +4543,7 @@ def _validate_request_window_splitting_is_supported( name: str, cursor: Optional[Cursor], incremental_sync: Optional[ - Union[IncrementingCountCursorModel, DatetimeBasedCursorModel] + Union[IncrementingCountCursorModel, DatetimeBasedCursorModel, OpaqueTokenCursorModel] ] = None, min_split_window: Optional[str] = None, query_properties: Optional[QueryProperties] = None, diff --git a/airbyte_cdk/sources/declarative/requesters/paginators/default_paginator.py b/airbyte_cdk/sources/declarative/requesters/paginators/default_paginator.py index 053d609193..8245c66f84 100644 --- a/airbyte_cdk/sources/declarative/requesters/paginators/default_paginator.py +++ b/airbyte_cdk/sources/declarative/requesters/paginators/default_paginator.py @@ -284,11 +284,17 @@ def __init__(self, decorated: Paginator, maximum_number_of_pages: int = 5) -> No self._maximum_number_of_pages = maximum_number_of_pages self._decorated = decorated self._page_count = self._PAGE_COUNT_BEFORE_FIRST_NEXT_CALL + self._did_reach_page_limit = False def get_initial_token(self) -> Optional[Any]: self._page_count = self._PAGE_COUNT_BEFORE_FIRST_NEXT_CALL + self._did_reach_page_limit = False return self._decorated.get_initial_token() + @property + def did_reach_page_limit(self) -> bool: + return self._did_reach_page_limit + def get_page_size(self) -> Optional[int]: return self._decorated.get_page_size() @@ -302,6 +308,7 @@ def next_page_token( stream_slice: Optional[StreamSlice] = None, ) -> Optional[Mapping[str, Any]]: if self._page_count >= self._maximum_number_of_pages: + self._did_reach_page_limit = True return None self._page_count += 1 diff --git a/airbyte_cdk/sources/declarative/requesters/request_options/__init__.py b/airbyte_cdk/sources/declarative/requesters/request_options/__init__.py index a63705832d..cdd9bb8562 100644 --- a/airbyte_cdk/sources/declarative/requesters/request_options/__init__.py +++ b/airbyte_cdk/sources/declarative/requesters/request_options/__init__.py @@ -11,6 +11,9 @@ from airbyte_cdk.sources.declarative.requesters.request_options.interpolated_request_options_provider import ( InterpolatedRequestOptionsProvider, ) +from airbyte_cdk.sources.declarative.requesters.request_options.opaque_token_request_options_provider import ( + OpaqueTokenRequestOptionsProvider, +) from airbyte_cdk.sources.declarative.requesters.request_options.request_options_provider import ( RequestOptionsProvider, ) @@ -19,5 +22,6 @@ "DatetimeBasedRequestOptionsProvider", "DefaultRequestOptionsProvider", "InterpolatedRequestOptionsProvider", + "OpaqueTokenRequestOptionsProvider", "RequestOptionsProvider", ] diff --git a/airbyte_cdk/sources/declarative/requesters/request_options/opaque_token_request_options_provider.py b/airbyte_cdk/sources/declarative/requesters/request_options/opaque_token_request_options_provider.py new file mode 100644 index 0000000000..e536f2141f --- /dev/null +++ b/airbyte_cdk/sources/declarative/requesters/request_options/opaque_token_request_options_provider.py @@ -0,0 +1,85 @@ +# +# Copyright (c) 2025 Airbyte, Inc., all rights reserved. +# + +from dataclasses import InitVar, dataclass +from typing import Any, Mapping, Optional, Union + +from airbyte_cdk.sources.declarative.interpolation import InterpolatedString +from airbyte_cdk.sources.declarative.requesters.request_option import ( + RequestOption, + RequestOptionType, +) +from airbyte_cdk.sources.declarative.requesters.request_options.request_options_provider import ( + RequestOptionsProvider, +) +from airbyte_cdk.sources.types import Config, StreamSlice, StreamState + + +@dataclass +class OpaqueTokenRequestOptionsProvider(RequestOptionsProvider): + token_option: RequestOption + config: Config + parameters: InitVar[Mapping[str, Any]] + initial_token: Optional[str] = None + + def __post_init__(self, parameters: Mapping[str, Any]) -> None: + self._initial_token = ( + InterpolatedString.create(self.initial_token, parameters=parameters) + if self.initial_token is not None + else None + ) + + def _get_request_options( + self, option_type: RequestOptionType, stream_slice: Optional[StreamSlice] + ) -> Mapping[str, Any]: + if not self.token_option or self.token_option.inject_into != option_type: + return {} + token = ( + stream_slice.cursor_slice.get("token") + if stream_slice and stream_slice.cursor_slice + else None + ) + 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 + + def get_request_params( + self, + *, + stream_state: Optional[StreamState] = None, + stream_slice: Optional[StreamSlice] = None, + next_page_token: Optional[Mapping[str, Any]] = None, + ) -> Mapping[str, Any]: + return self._get_request_options(RequestOptionType.request_parameter, stream_slice) + + def get_request_headers( + self, + *, + stream_state: Optional[StreamState] = None, + stream_slice: Optional[StreamSlice] = None, + next_page_token: Optional[Mapping[str, Any]] = None, + ) -> Mapping[str, Any]: + return self._get_request_options(RequestOptionType.header, stream_slice) + + def get_request_body_data( + self, + *, + stream_state: Optional[StreamState] = None, + stream_slice: Optional[StreamSlice] = None, + next_page_token: Optional[Mapping[str, Any]] = None, + ) -> Union[Mapping[str, Any], str]: + return self._get_request_options(RequestOptionType.body_data, stream_slice) + + def get_request_body_json( + self, + *, + stream_state: Optional[StreamState] = None, + stream_slice: Optional[StreamSlice] = None, + next_page_token: Optional[Mapping[str, Any]] = None, + ) -> Mapping[str, Any]: + return self._get_request_options(RequestOptionType.body_json, stream_slice) diff --git a/airbyte_cdk/sources/declarative/retrievers/simple_retriever.py b/airbyte_cdk/sources/declarative/retrievers/simple_retriever.py index 19b9cce00b..091fa7cc4c 100644 --- a/airbyte_cdk/sources/declarative/retrievers/simple_retriever.py +++ b/airbyte_cdk/sources/declarative/retrievers/simple_retriever.py @@ -80,6 +80,13 @@ _MAX_REQUEST_WINDOW_SPLIT_DEPTH = 10 +@dataclass(frozen=True) +class PageCompleted: + stream_slice: StreamSlice + response: requests.Response + pagination_exhausted: bool + + @dataclass class SimpleRetriever(Retriever): """ @@ -146,6 +153,7 @@ class SimpleRetriever(Retriever): request_window_splitter: Optional[ Callable[[StreamSlice, Optional[datetime.timedelta]], Optional[List[StreamSlice]]] ] = None + page_observer: Optional[Callable[["PageCompleted"], None]] = None def __post_init__(self, parameters: Mapping[str, Any]) -> None: self._paginator = self.paginator or NoPagination(parameters=parameters) @@ -547,6 +555,15 @@ def _read_pages( **page_size_override_kwargs(page_size_override), **stream_slice_kwargs(self._next_page_token, stream_slice), ) + if response is not None and self.page_observer: + self.page_observer( + PageCompleted( + stream_slice=stream_slice, + response=response, + pagination_exhausted=not next_page_token + and not getattr(self._paginator, "did_reach_page_limit", False), + ) + ) if not next_page_token: break diff --git a/airbyte_cdk/sources/streams/concurrent/default_stream.py b/airbyte_cdk/sources/streams/concurrent/default_stream.py index 6cc6e44d44..c94db6959d 100644 --- a/airbyte_cdk/sources/streams/concurrent/default_stream.py +++ b/airbyte_cdk/sources/streams/concurrent/default_stream.py @@ -9,6 +9,7 @@ from airbyte_cdk.sources.declarative.incremental.concurrent_partition_cursor import ( ConcurrentPerPartitionCursor, ) +from airbyte_cdk.sources.declarative.incremental.opaque_token_cursor import OpaqueTokenCursor from airbyte_cdk.sources.declarative.partition_routers.partition_router import PartitionRouter from airbyte_cdk.sources.declarative.stream_slicers.declarative_partition_generator import ( StreamSlicerPartitionGenerator, @@ -34,6 +35,7 @@ def __init__( namespace: Optional[str] = None, supports_file_transfer: bool = False, block_simultaneous_read: str = "", + source_defined_incremental: bool = False, ) -> None: self._stream_partition_generator = partition_generator self._name = name @@ -45,6 +47,7 @@ def __init__( self._namespace = namespace self._supports_file_transfer = supports_file_transfer self._block_simultaneous_read = block_simultaneous_read + self._source_defined_incremental = source_defined_incremental def generate_partitions(self) -> Iterable[Partition]: yield from self._stream_partition_generator.generate() @@ -83,6 +86,14 @@ def as_airbyte_stream(self) -> AirbyteStream: stream.is_resumable = True stream.supported_sync_modes.append(SyncMode.incremental) stream.default_cursor_field = [self._cursor_field.cursor_field_key] + elif self._source_defined_incremental: + stream.source_defined_cursor = True + stream.is_resumable = True + stream.supported_sync_modes.append(SyncMode.incremental) + stream.default_cursor_field = [] + if self._source_defined_incremental and self._cursor_field: + stream.source_defined_cursor = True + stream.is_resumable = True keys = self._primary_key if keys and len(keys) > 0: @@ -117,6 +128,8 @@ def get_partition_router(self) -> PartitionRouter | None: if not isinstance(self._stream_partition_generator, StreamSlicerPartitionGenerator): return None stream_slicer = self._stream_partition_generator._stream_slicer + if isinstance(stream_slicer, OpaqueTokenCursor): + return stream_slicer._partition_router if not isinstance(stream_slicer, ConcurrentPerPartitionCursor): return None return stream_slicer._partition_router diff --git a/unit_tests/sources/declarative/incremental/test_opaque_token_cursor.py b/unit_tests/sources/declarative/incremental/test_opaque_token_cursor.py new file mode 100644 index 0000000000..9f5687ff87 --- /dev/null +++ b/unit_tests/sources/declarative/incremental/test_opaque_token_cursor.py @@ -0,0 +1,314 @@ +# +# Copyright (c) 2025 Airbyte, Inc., all rights reserved. +# + +import json +from unittest.mock import MagicMock + +import pytest +import requests +from requests import Response + +from airbyte_cdk.models import FailureType +from airbyte_cdk.sources.connector_state_manager import ConnectorStateManager +from airbyte_cdk.sources.declarative.decoders.json_decoder import JsonDecoder +from airbyte_cdk.sources.declarative.incremental.opaque_token_cursor import OpaqueTokenCursor +from airbyte_cdk.sources.declarative.partition_routers.single_partition_router import ( + SinglePartitionRouter, +) +from airbyte_cdk.sources.declarative.requesters.paginators.default_paginator import ( + PaginatorTestReadDecorator, +) +from airbyte_cdk.sources.declarative.retrievers.simple_retriever import PageCompleted +from airbyte_cdk.sources.message.repository import InMemoryMessageRepository +from airbyte_cdk.sources.types import StreamSlice +from airbyte_cdk.utils import AirbyteTracedException + + +class Router: + def __init__(self, slices): + self.slices = slices + + def stream_slices(self): + yield from self.slices + + +def _response(body): + response = requests.Response() + response.status_code = 200 + response._content = json.dumps(body).encode() + response.headers = {} + return response + + +def _cursor(*, state=None, router=None, max_partitions=1000, completion_condition=None): + return OpaqueTokenCursor( + stream_name="events", + stream_namespace=None, + stream_state=state, + message_repository=InMemoryMessageRepository(), + connector_state_manager=ConnectorStateManager(), + partition_router=router or SinglePartitionRouter(parameters={}), + token_path=["syncToken"], + decoder=JsonDecoder(parameters={}), + config={}, + parameters={}, + completion_condition=completion_condition, + max_partitions=max_partitions, + ) + + +def _partition(partition): + result = MagicMock() + result.to_slice.return_value = {"partition": partition, "cursor_slice": {}} + return result + + +def test_t_token_preserves_string_and_converts_integer_but_rejects_other_types(): + cursor = _cursor() + + for value in ("001", "000123", '"quoted"', "None"): + assert cursor._coerce_token(value, ["syncToken"]) == value + for value in (0, 17, 2**63): + assert cursor._coerce_token(value, ["syncToken"]) == str(value) + assert cursor._coerce_token(None, ["syncToken"]) is None + assert cursor._coerce_token("", ["syncToken"]) is None + + for value in (True, 1.5, {}, []): + with pytest.raises(AirbyteTracedException) as exc: + cursor._coerce_token(value, ["syncToken"]) + assert exc.value.failure_type == FailureType.system_error + + +def test_t_boot_and_t_inc_commit_only_after_partition_close_and_resume_from_state(): + cursor = _cursor() + stream_slice = StreamSlice(partition={}, cursor_slice={}) + cursor.on_page_completed( + PageCompleted(stream_slice, _response({"items": [], "syncToken": "token-1"}), True) + ) + + assert cursor.state == {"state_type": "opaque_token"} + cursor.close_partition(_partition({})) + assert cursor.state["state_type"] == "opaque_token" + assert cursor.state["token"] == "token-1" + assert cursor.state["committed_at"].endswith("Z") + + resumed = _cursor(state=cursor.state) + slices = list(resumed.stream_slices()) + assert len(slices) == 1 + assert slices[0].cursor_slice == {"token": "token-1"} + + +def test_t_noecho_and_t_empty_leave_committed_token_unchanged(): + cursor = _cursor( + state={ + "state_type": "opaque_token", + "token": "previous", + "committed_at": "2026-01-01T00:00:00Z", + } + ) + stream_slice = StreamSlice(partition={}, cursor_slice={}) + cursor.on_page_completed( + PageCompleted(stream_slice, _response({"syncToken": "intermediate"}), False) + ) + cursor.on_page_completed(PageCompleted(stream_slice, _response({"syncToken": ""}), True)) + cursor.close_partition(_partition({})) + + assert cursor.state == { + "state_type": "opaque_token", + "token": "previous", + "committed_at": "2026-01-01T00:00:00Z", + } + + +def test_t_completion_condition_must_pass_before_a_token_is_committed(): + cursor = _cursor(completion_condition="{{ response.get('complete') }}") + stream_slice = StreamSlice(partition={}, cursor_slice={}) + + cursor.on_page_completed( + PageCompleted( + stream_slice, + _response({"syncToken": "incomplete", "complete": False}), + True, + ) + ) + cursor.close_partition(_partition({})) + assert cursor.state == {"state_type": "opaque_token"} + + cursor.on_page_completed( + PageCompleted( + stream_slice, + _response({"syncToken": "complete", "complete": True}), + True, + ) + ) + cursor.close_partition(_partition({})) + assert cursor.state["token"] == "complete" + + +def test_t_fail_discards_pending_token_when_pass_does_not_complete(): + cursor = _cursor() + stream_slice = StreamSlice(partition={}, cursor_slice={}) + cursor.on_page_completed( + PageCompleted(stream_slice, _response({"syncToken": "uncommitted"}), False) + ) + cursor.close_partition(_partition({})) + + assert cursor.state == {"state_type": "opaque_token"} + + +def test_t_builder_page_limit_is_not_treated_as_a_completed_pass(): + paginator = PaginatorTestReadDecorator(MagicMock(), maximum_number_of_pages=1) + paginator.get_initial_token() + next_page_token = paginator.next_page_token(Response(), last_page_size=1, last_record=None) + + cursor = _cursor() + stream_slice = StreamSlice(partition={}, cursor_slice={}) + cursor.on_page_completed( + PageCompleted( + stream_slice, + _response({"syncToken": "truncated"}), + pagination_exhausted=not next_page_token and not paginator.did_reach_page_limit, + ) + ) + cursor.close_partition(_partition({})) + + assert paginator.did_reach_page_limit + assert cursor.state == {"state_type": "opaque_token"} + + +def test_t_par_keeps_tokens_partition_scoped(): + router = Router( + [ + StreamSlice(partition={"account": "a"}, cursor_slice={}), + StreamSlice(partition={"account": "b"}, cursor_slice={}), + StreamSlice(partition={"account": "c"}, cursor_slice={}), + ] + ) + cursor = _cursor( + router=router, + state={ + "state_type": "opaque_token", + "states": [ + { + "partition": {"account": "c"}, + "cursor": {"token": "token-c", "committed_at": "2026-01-01T00:00:00Z"}, + } + ], + }, + ) + for stream_slice, token in ( + (router.slices[1], "token-b-new"), + (router.slices[0], "token-a-new"), + ): + cursor.on_page_completed(PageCompleted(stream_slice, _response({"syncToken": token}), True)) + cursor.close_partition(_partition(stream_slice.partition)) + + assert cursor.state["state_type"] == "opaque_token" + assert { + (item["partition"]["account"], item["cursor"]["token"]) for item in cursor.state["states"] + } == {("a", "token-a-new"), ("b", "token-b-new"), ("c", "token-c")} + + +def test_t_stream_slices_deduplicates_partitions_and_uses_only_the_cursor_token(): + router = Router( + [ + StreamSlice( + partition={"account": "a"}, + cursor_slice={"router": "value"}, + extra_fields={"query": "first"}, + ), + StreamSlice( + partition={"account": "a"}, + cursor_slice={"router": "other"}, + extra_fields={"query": "duplicate"}, + ), + ] + ) + cursor = _cursor(router=router) + + slices = list(cursor.stream_slices()) + + assert len(slices) == 1 + assert slices[0].cursor_slice == {} + assert slices[0].extra_fields == {"query": "first"} + + +def test_t_cap_retains_tokens_after_failed_router_enumeration_and_drops_vanished_partitions(): + state = { + "state_type": "opaque_token", + "states": [ + { + "partition": {"account": "a"}, + "cursor": {"token": "token-a", "committed_at": "2026-01-01T00:00:00Z"}, + }, + { + "partition": {"account": "b"}, + "cursor": {"token": "token-b", "committed_at": "2026-01-01T00:00:00Z"}, + }, + ], + } + + class FailingRouter(Router): + def stream_slices(self): + yield self.slices[0] + raise RuntimeError("router failed") + + failing_cursor = _cursor( + router=FailingRouter([StreamSlice(partition={"account": "a"}, cursor_slice={})]), + state=state, + ) + with pytest.raises(RuntimeError, match="router failed"): + list(failing_cursor.stream_slices()) + assert {item["cursor"]["token"] for item in failing_cursor.state["states"]} == { + "token-a", + "token-b", + } + + current_cursor = _cursor( + router=Router([StreamSlice(partition={"account": "a"}, cursor_slice={})]), + state=state, + ) + list(current_cursor.stream_slices()) + assert {item["partition"]["account"] for item in current_cursor.state["states"]} == {"a"} + + +def test_t_cap_does_not_evict_existing_partitions(): + router = Router( + [ + StreamSlice(partition={"account": "a"}, cursor_slice={}), + StreamSlice(partition={"account": "b"}, cursor_slice={}), + ] + ) + cursor = _cursor(router=router, max_partitions=1) + for stream_slice, token in zip(router.slices, ("token-a", "token-b")): + cursor.on_page_completed(PageCompleted(stream_slice, _response({"syncToken": token}), True)) + cursor.close_partition(_partition(stream_slice.partition)) + + assert len(cursor.state["states"]) == 1 + assert cursor.state["states"][0]["partition"] == {"account": "a"} + assert cursor.state["states"][0]["cursor"]["token"] == "token-a" + + +def test_t_foreign_ignores_state_from_another_cursor_type(): + cursor = _cursor(state={"state_type": "datetime", "token": "not-a-token"}) + + assert cursor.state == {"state_type": "opaque_token"} + + +def test_t_schema_ignores_incompatible_partition_state_shape(): + cursor = _cursor( + router=Router([]), + state={"state_type": "opaque_token", "token": "wrong-shape"}, + ) + + assert cursor.state == {"state_type": "opaque_token", "states": []} + + +def test_t_valid_and_t_init_use_partition_router_and_initial_token_slice(): + router = Router([StreamSlice(partition={"account": "a"}, cursor_slice={})]) + cursor = _cursor(router=router) + + slices = list(cursor.stream_slices()) + assert slices[0].partition == {"account": "a"} + assert slices[0].cursor_slice == {} diff --git a/unit_tests/sources/declarative/incremental/test_opaque_token_cursor_integration.py b/unit_tests/sources/declarative/incremental/test_opaque_token_cursor_integration.py new file mode 100644 index 0000000000..08e5ca894a --- /dev/null +++ b/unit_tests/sources/declarative/incremental/test_opaque_token_cursor_integration.py @@ -0,0 +1,390 @@ +# +# Copyright (c) 2025 Airbyte, Inc., all rights reserved. +# + +import json +import logging +from unittest.mock import MagicMock + +import pytest + +from airbyte_cdk.models import ( + AirbyteStateType, + ConfiguredAirbyteCatalog, + ConfiguredAirbyteStream, + DestinationSyncMode, + SyncMode, + Type, +) +from airbyte_cdk.sources.declarative.concurrent_declarative_source import ( + ConcurrentDeclarativeSource, +) +from airbyte_cdk.sources.declarative.models.declarative_component_schema import ( + AsyncRetriever as AsyncRetrieverModel, +) +from airbyte_cdk.sources.declarative.models.declarative_component_schema import ( + CustomRetriever as CustomRetrieverModel, +) +from airbyte_cdk.sources.declarative.models.declarative_component_schema import ( + DeclarativeStream as DeclarativeStreamModel, +) +from airbyte_cdk.sources.declarative.models.declarative_component_schema import ( + GroupingPartitionRouter as GroupingPartitionRouterModel, +) +from airbyte_cdk.sources.declarative.models.declarative_component_schema import ( + ListPartitionRouter as ListPartitionRouterModel, +) +from airbyte_cdk.sources.declarative.models.declarative_component_schema import ( + OpaqueTokenCursor as OpaqueTokenCursorModel, +) +from airbyte_cdk.sources.declarative.models.declarative_component_schema import ( + PaginationReset as PaginationResetModel, +) +from airbyte_cdk.sources.declarative.models.declarative_component_schema import ( + PaginationResetLimits as PaginationResetLimitsModel, +) +from airbyte_cdk.sources.declarative.models.declarative_component_schema import ( + ParentStreamConfig as ParentStreamConfigModel, +) +from airbyte_cdk.sources.declarative.models.declarative_component_schema import ( + SubstreamPartitionRouter as SubstreamPartitionRouterModel, +) +from airbyte_cdk.sources.declarative.parsers.model_to_component_factory import ( + ModelToComponentFactory, +) +from airbyte_cdk.sources.declarative.retrievers.simple_retriever import SimpleRetriever +from airbyte_cdk.sources.declarative.stream_slicers.declarative_partition_generator import ( + StreamSlicerPartitionGenerator, +) +from airbyte_cdk.sources.streams.concurrent.default_stream import DefaultStream +from airbyte_cdk.test.mock_http import HttpMocker, HttpRequest, HttpResponse + +_URL = "https://api.example.com/events" + + +def _manifest(*, initial_token=None, partitioned=False): + stream = { + "type": "DeclarativeStream", + "name": "events", + "primary_key": [], + "schema_loader": { + "type": "InlineSchemaLoader", + "schema": { + "$schema": "http://json-schema.org/schema#", + "properties": {"id": {"type": "string"}}, + "type": "object", + }, + }, + "retriever": { + "type": "SimpleRetriever", + "requester": { + "type": "HttpRequester", + "url_base": "https://api.example.com", + "path": "/events", + "http_method": "POST", + }, + "record_selector": { + "type": "RecordSelector", + "extractor": {"type": "DpathExtractor", "field_path": ["items"]}, + }, + "paginator": { + "type": "DefaultPaginator", + "page_token_option": { + "type": "RequestOption", + "field_name": "nextCursor", + "inject_into": "body_json", + }, + "pagination_strategy": { + "type": "CursorPagination", + "cursor_value": "{{ response.get('nextCursor') }}", + "stop_condition": "{{ not response.get('nextCursor') }}", + }, + }, + }, + "incremental_sync": { + "type": "OpaqueTokenCursor", + "token_path": ["syncToken"], + "token_option": { + "type": "RequestOption", + "field_name": "syncToken", + "inject_into": "body_json", + }, + }, + } + if initial_token is not None: + stream["incremental_sync"]["initial_token"] = initial_token + if partitioned: + stream["retriever"]["partition_router"] = { + "type": "ListPartitionRouter", + "cursor_field": "account", + "values": ["a", "b"], + "request_option": { + "type": "RequestOption", + "field_name": "account", + "inject_into": "body_json", + }, + } + return { + "version": "0.34.2", + "type": "DeclarativeSource", + "check": {"type": "CheckStream", "stream_names": ["events"]}, + "streams": [stream], + "spec": { + "connection_specification": { + "$schema": "http://json-schema.org/draft-07/schema#", + "type": "object", + "required": [], + "properties": {}, + "additionalProperties": True, + }, + "documentation_url": "https://example.org", + "type": "Spec", + }, + } + + +def _source_and_catalog(manifest, state=None): + source = ConcurrentDeclarativeSource( + source_config=manifest, config={}, catalog=None, state=state + ) + discovered = source.discover(logger=MagicMock(), config={}) + stream = discovered.streams[0] + catalog = ConfiguredAirbyteCatalog( + streams=[ + ConfiguredAirbyteStream( + stream=stream, + sync_mode=SyncMode.incremental, + destination_sync_mode=DestinationSyncMode.append, + cursor_field=stream.default_cursor_field, + primary_key=stream.source_defined_primary_key, + ) + ] + ) + return source, catalog + + +def _read(source, catalog, state=None): + return list(source.read(logger=MagicMock(), config={}, catalog=catalog, state=state)) + + +def _state(messages): + return [ + message.state + for message in messages + if message.type == Type.STATE and message.state.type == AirbyteStateType.STREAM + ] + + +def test_t_boot_inc_builder_and_resume_send_exact_request_bodies(): + with HttpMocker() as http_mocker: + http_mocker.post( + HttpRequest(_URL), + HttpResponse(json.dumps({"items": [{"id": "1"}], "nextCursor": "p2"})), + ) + http_mocker.post( + HttpRequest(_URL, body={"nextCursor": "p2"}), + HttpResponse(json.dumps({"items": [{"id": "2"}], "syncToken": "sync-1"})), + ) + http_mocker.post( + HttpRequest(_URL, body={"syncToken": "sync-1"}), + HttpResponse(json.dumps({"items": [{"id": "3"}], "syncToken": "sync-2"})), + ) + + source, catalog = _source_and_catalog(_manifest()) + first_read = _read(source, catalog) + first_state = _state(first_read) + assert any(message.type == Type.RECORD for message in first_read) + first_stream_state = first_state[0].stream.stream_state.__dict__ + assert first_stream_state["state_type"] == "opaque_token" + assert first_stream_state["token"] == "sync-1" + assert first_stream_state["committed_at"].endswith("Z") + + resumed_source, resumed_catalog = _source_and_catalog(_manifest(), first_state) + resumed = _read(resumed_source, resumed_catalog, first_state) + assert _state(resumed)[0].stream.stream_state.__dict__["token"] == "sync-2" + + +def test_t_init_sends_initial_token_and_t_par_keeps_partition_state_separate(): + with HttpMocker() as http_mocker: + http_mocker.post( + HttpRequest(_URL, body={"syncToken": "*"}), + HttpResponse(json.dumps({"items": [], "syncToken": "initialized"})), + ) + source, catalog = _source_and_catalog(_manifest(initial_token="*")) + messages = _read(source, catalog) + assert _state(messages)[0].stream.stream_state.__dict__["token"] == "initialized" + + with HttpMocker() as http_mocker: + http_mocker.post( + HttpRequest(_URL, body={"account": "a"}), + HttpResponse(json.dumps({"items": [], "syncToken": "token-a"})), + ) + http_mocker.post( + HttpRequest(_URL, body={"account": "b"}), + HttpResponse(json.dumps({"items": [], "syncToken": "token-b"})), + ) + source, catalog = _source_and_catalog(_manifest(partitioned=True)) + messages = _read(source, catalog) + states = _state(messages)[-1].stream.stream_state.__dict__["states"] + + assert {(item["partition"]["account"], item["cursor"]["token"]) for item in states} == { + ("a", "token-a"), + ("b", "token-b"), + } + + +def test_t_disc_advertises_source_defined_incremental_without_cursor_field(): + source, _ = _source_and_catalog(_manifest()) + + discovered = source.discover(logger=MagicMock(), config={}) + stream = discovered.streams[0] + assert stream.source_defined_cursor is True + assert stream.default_cursor_field == [] + + +def test_t_valid_rejects_datetime_only_extra_settings(): + manifest = _manifest() + manifest["streams"][0]["incremental_sync"]["start_datetime"] = "2024-01-01" + source = ConcurrentDeclarativeSource( + source_config=manifest, config={}, catalog=None, state=None + ) + + with pytest.raises(ValueError, match="datetime-only setting.*start_datetime"): + source.discover(logger=MagicMock(), config={}) + + +def _stream_model(): + return DeclarativeStreamModel.parse_obj(_manifest()["streams"][0]) + + +def test_t_valid_rejects_grouping_async_lazy_dependency_and_reserved_partition(): + stream = _stream_model() + stream.retriever.partition_router = GroupingPartitionRouterModel.construct( + type="GroupingPartitionRouter" + ) + with pytest.raises(ValueError, match="GroupingPartitionRouter"): + ModelToComponentFactory._validate_opaque_token_stream(stream) + + stream = _stream_model() + stream.retriever = AsyncRetrieverModel.construct(type="AsyncRetriever") + with pytest.raises(ValueError, match="AsyncRetriever"): + ModelToComponentFactory._validate_opaque_token_stream(stream) + + for parent_config, error in ( + ({"lazy_read_pointer": ["items"]}, "lazy_read_pointer"), + ({"incremental_dependency": True}, "incremental_dependency"), + ({"partition_field": "token"}, "partition field 'token'"), + ): + stream = _stream_model() + parent_fields = { + "type": "ParentStreamConfig", + "stream": None, + "parent_key": "id", + "partition_field": "account", + "incremental_dependency": False, + "lazy_read_pointer": [], + } + parent_fields.update(parent_config) + parent = ParentStreamConfigModel.construct(**parent_fields) + stream.retriever.partition_router = SubstreamPartitionRouterModel.construct( + type="SubstreamPartitionRouter", parent_stream_configs=[parent] + ) + with pytest.raises(ValueError, match=error): + ModelToComponentFactory._validate_opaque_token_stream(stream) + + stream = _stream_model() + stream.retriever.partition_router = ListPartitionRouterModel.construct( + type="ListPartitionRouter", cursor_field="token" + ) + with pytest.raises(ValueError, match="partition router field 'token'"): + ModelToComponentFactory._validate_opaque_token_stream(stream) + + +def test_t_valid_rejects_opaque_token_parent_dependency(): + stream = _stream_model() + parent_stream = DeclarativeStreamModel.construct( + type="DeclarativeStream", + incremental_sync=OpaqueTokenCursorModel.construct(type="OpaqueTokenCursor"), + ) + parent = ParentStreamConfigModel.construct( + type="ParentStreamConfig", + stream=parent_stream, + parent_key="id", + partition_field="account", + incremental_dependency=True, + ) + stream.retriever.partition_router = SubstreamPartitionRouterModel.construct( + type="SubstreamPartitionRouter", parent_stream_configs=[parent] + ) + + with pytest.raises(ValueError, match="OpaqueTokenCursor parent"): + ModelToComponentFactory._validate_opaque_token_stream(stream) + + +def test_t_valid_warns_for_custom_retriever(caplog): + stream = _stream_model() + stream.retriever = CustomRetrieverModel.construct(type="CustomRetriever") + + with caplog.at_level(logging.WARNING, logger="airbyte"): + ModelToComponentFactory._validate_opaque_token_stream(stream) + + assert "CustomRetriever does not expose response pages" in caplog.text + + +def test_t_valid_rejects_pagination_reset_limits_and_query_property_chunking(): + stream = _stream_model() + stream.retriever.pagination_reset = PaginationResetModel.construct( + type="PaginationReset", + limits=PaginationResetLimitsModel.construct(type="PaginationResetLimits"), + ) + with pytest.raises(ValueError, match="pagination reset `limits`"): + ModelToComponentFactory._validate_opaque_token_stream(stream) + + stream = _stream_model() + stream.retriever.requester.query_properties = { + "type": "QueryProperties", + "property_chunking": {"type": "PropertyChunking"}, + } + with pytest.raises(ValueError, match="query property chunking"): + ModelToComponentFactory._validate_opaque_token_stream(stream) + + stream = _stream_model() + stream.retriever.requester.request_parameters = { + "properties": { + "type": "QueryProperties", + "property_chunking": {"type": "PropertyChunking"}, + } + } + with pytest.raises(ValueError, match="query property chunking"): + ModelToComponentFactory._validate_opaque_token_stream(stream) + + +def test_t_inert_preserves_initial_token_when_response_has_no_token(): + with HttpMocker() as http_mocker: + http_mocker.post( + HttpRequest(_URL, body={"syncToken": "*"}), + HttpResponse(json.dumps({"items": []})), + ) + source, catalog = _source_and_catalog(_manifest(initial_token="*")) + messages = _read(source, catalog) + + assert all("token" not in message.stream.stream_state.__dict__ for message in _state(messages)) + + +def test_t_inert_leaves_page_observer_unset_for_datetime_cursor(): + manifest = _manifest() + manifest["streams"][0]["incremental_sync"] = { + "type": "DatetimeBasedCursor", + "cursor_field": "created_at", + "start_datetime": "2020-01-01T00:00:00Z", + "datetime_format": "%Y-%m-%dT%H:%M:%SZ", + } + source = ConcurrentDeclarativeSource(source_config=manifest, config={}, catalog=None) + + stream = source.streams(config={})[0] + assert isinstance(stream, DefaultStream) + generator = stream._stream_partition_generator + assert isinstance(generator, StreamSlicerPartitionGenerator) + retriever = generator._partition_factory._retriever + assert isinstance(retriever, SimpleRetriever) + assert retriever.page_observer is None From a1642dde79ac3394392bbd8026933aee0bde3535 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 16 Oct 2026 23:23:51 +0000 Subject: [PATCH 2/3] Address opaque token cursor review feedback Co-Authored-By: bot_apk --- .../concurrent_declarative_source.py | 1 + .../incremental/opaque_token_cursor.py | 157 ++++++-- .../parsers/model_to_component_factory.py | 3 + .../paginators/default_paginator.py | 7 - .../retrievers/simple_retriever.py | 3 +- .../incremental/test_opaque_token_cursor.py | 170 +++++++-- .../test_opaque_token_cursor_integration.py | 361 +++++++++++++++++- 7 files changed, 619 insertions(+), 83 deletions(-) diff --git a/airbyte_cdk/sources/declarative/concurrent_declarative_source.py b/airbyte_cdk/sources/declarative/concurrent_declarative_source.py index 055576f45f..b015e1ddec 100644 --- a/airbyte_cdk/sources/declarative/concurrent_declarative_source.py +++ b/airbyte_cdk/sources/declarative/concurrent_declarative_source.py @@ -198,6 +198,7 @@ def __init__( max_concurrent_async_job_count=source_config.get("max_concurrent_async_job_count"), limit_pages_fetched_per_slice=limits.max_pages_per_slice if limits else None, limit_slices_fetched=limits.max_slices if limits else None, + max_records_limit=limits.max_records if limits else None, disable_retries=True if limits else False, disable_cache=True if limits else False, ) diff --git a/airbyte_cdk/sources/declarative/incremental/opaque_token_cursor.py b/airbyte_cdk/sources/declarative/incremental/opaque_token_cursor.py index 4beeddb845..cd75943f7e 100644 --- a/airbyte_cdk/sources/declarative/incremental/opaque_token_cursor.py +++ b/airbyte_cdk/sources/declarative/incremental/opaque_token_cursor.py @@ -4,6 +4,7 @@ import copy import datetime +import json import logging import threading import time @@ -80,11 +81,15 @@ def __init__( self._serializer = PerPartitionKeySerializer() self._committed: OrderedDict[str, dict[str, Any]] = OrderedDict() self._pending: dict[str, dict[str, Any]] = {} + self._page_counts: dict[str, int] = {} self._lock = threading.RLock() self._last_emission_time = 0.0 + self._logged_bootstrap_reasons: set[str] = set() + self._state_is_foreign = False + self._state_emitted = False self._warned_no_token = False self._warned_cap = False - self._partitions_over_cap = 0 + self._partitions_over_cap: set[str] = set() self._load_state(stream_state or {}) @property @@ -121,6 +126,14 @@ def stream_slices(self) -> Iterable[StreamSlice]: with self._lock: committed = self._committed.get(partition_key) token = committed["token"] if committed else None + at_partition_cap = len(self._committed) >= self._max_partitions + if token: + self._log_replay(partition, committed or {}) + elif at_partition_cap: + self._partitions_over_cap.add(partition_key) + self._log_bootstrap("partition cap") + elif not self._state_is_foreign: + self._log_bootstrap("none") yield StreamSlice( partition=partition, cursor_slice={"token": token} if token else {}, @@ -140,6 +153,10 @@ def stream_slices(self) -> Iterable[StreamSlice]: ) def on_page_completed(self, page: PageCompleted) -> None: + partition_key = self._serializer.to_partition_key(page.stream_slice.partition or {}) + with self._lock: + self._page_counts[partition_key] = self._page_counts.get(partition_key, 0) + 1 + if not page.pagination_exhausted or page.response is None: return @@ -180,8 +197,7 @@ def on_page_completed(self, page: PageCompleted) -> None: .replace("+00:00", "Z"), } - @staticmethod - def _coerce_token(value: Any, path: list[str]) -> Optional[str]: + def _coerce_token(self, value: Any, path: list[str]) -> Optional[str]: if value is None or value == "": return None if isinstance(value, bool): @@ -193,17 +209,82 @@ def _coerce_token(value: Any, path: list[str]) -> Optional[str]: elif isinstance(value, float): value_type = "float" elif isinstance(value, Mapping): - value_type = "object" + value_type = "dict" elif isinstance(value, list): value_type = "list" else: value_type = type(value).__name__ + message = ( + f"Stream {self._stream_name}: token_path {json.dumps(path)} holds a {value_type}, " + "not a string or integer. Check token_path." + ) raise AirbyteTracedException( - internal_message=f"Stream token_path {path} holds a {value_type}, not a string or integer.", - message=f"Token path {path} holds a {value_type}, not a string or integer. Check token_path.", + internal_message=message, + message=message, failure_type=FailureType.system_error, ) + def _log_bootstrap(self, reason: str) -> None: + if reason not in self._logged_bootstrap_reasons: + logger.info( + "Stream %s: no stored change token (reason: %s); reading all records.", + self._stream_name, + reason, + ) + self._logged_bootstrap_reasons.add(reason) + + def _log_foreign_state(self) -> None: + self._state_is_foreign = True + self._log_bootstrap("foreign state") + + def _log_replay(self, partition: Mapping[str, Any], committed: Mapping[str, Any]) -> None: + age = self._human_age(committed.get("committed_at")) + if self._partitioned: + logger.debug( + "Stream %s: replaying the change token committed %s ago for partition %s.", + self._stream_name, + age, + dict(partition), + ) + else: + logger.info( + "Stream %s: replaying the change token committed %s ago.", + self._stream_name, + age, + ) + + @staticmethod + def _human_age(committed_at: Any) -> str: + if not isinstance(committed_at, str) or not committed_at: + return "an unknown time" + try: + committed_time = datetime.datetime.fromisoformat(committed_at.replace("Z", "+00:00")) + if committed_time.tzinfo is None: + committed_time = committed_time.replace(tzinfo=datetime.timezone.utc) + age_seconds = max( + 0, + int( + (datetime.datetime.now(datetime.timezone.utc) - committed_time).total_seconds() + ), + ) + except (OverflowError, ValueError): + return "an unknown time" + + units = ( + ("day", 24 * 60 * 60), + ("hour", 60 * 60), + ("minute", 60), + ("second", 1), + ) + parts = [] + for unit, unit_seconds in units: + count, age_seconds = divmod(age_seconds, unit_seconds) + if count: + parts.append(f"{count} {unit}{'' if count == 1 else 's'}") + if len(parts) == 2: + break + return " ".join(parts) if parts else "0 seconds" + def _warn_missing_token(self) -> None: if not self._warned_no_token: logger.warning( @@ -223,18 +304,20 @@ def close_partition(self, partition: Partition) -> None: partition_key = self._serializer.to_partition_key(partition_values) with self._lock: 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: self._committed[partition_key] = pending + self._partitions_over_cap.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, + ) else: - self._partitions_over_cap += 1 - if not self._warned_cap: - logger.warning( - "Stream %s: %s partitions already store tokens; additional partitions read in full.", - self._stream_name, - len(self._committed), - ) - self._warned_cap = True + self._partitions_over_cap.add(partition_key) if not self._partitioned: self._emit_state_message() elif time.time() - self._last_emission_time > 600: @@ -250,9 +333,19 @@ def _emit_state_message(self) -> None: self._stream_name, self._stream_namespace ) ) + self._state_emitted = True def ensure_at_least_one_state_emitted(self) -> None: - self._emit_state_message() + 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 def observe(self, record: Record) -> None: return @@ -262,29 +355,21 @@ def should_be_synced(self, record: Record) -> bool: def _load_state(self, stream_state: StreamState) -> None: if not stream_state: + self._log_bootstrap("none") return if stream_state.get("state_type") != self.STATE_TYPE: - logger.info( - "Stream %s: ignoring foreign state and reading from the beginning.", - self._stream_name, - ) + self._log_foreign_state() return if self._partitioned: states = stream_state.get("states") if not isinstance(states, list) or "token" in stream_state: - logger.info( - "Stream %s: ignoring state with an incompatible partition shape.", - self._stream_name, - ) + self._log_foreign_state() return committed: dict[str, dict[str, Any]] = {} for state in states: if not isinstance(state, Mapping): - logger.info( - "Stream %s: ignoring state with an incompatible opaque token partition.", - self._stream_name, - ) - return + self._log_bootstrap("foreign state") + continue partition = state.get("partition") cursor = state.get("cursor") if ( @@ -298,18 +383,11 @@ def _load_state(self, stream_state: StreamState) -> None: "committed_at": cursor.get("committed_at", ""), } else: - logger.info( - "Stream %s: ignoring state with an incompatible opaque token partition.", - self._stream_name, - ) - return + self._log_bootstrap("foreign state") self._committed.update(committed) else: if "states" in stream_state: - logger.info( - "Stream %s: ignoring state with an incompatible unpartitioned shape.", - self._stream_name, - ) + self._log_foreign_state() return token = stream_state.get("token") if isinstance(token, str) and token: @@ -318,7 +396,4 @@ def _load_state(self, stream_state: StreamState) -> None: "committed_at": stream_state.get("committed_at", ""), } elif token is not None: - logger.info( - "Stream %s: ignoring an incompatible opaque token state.", - self._stream_name, - ) + self._log_foreign_state() 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 23902e1484..95f532c918 100644 --- a/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py +++ b/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py @@ -760,6 +760,7 @@ def __init__( self, limit_pages_fetched_per_slice: Optional[int] = None, limit_slices_fetched: Optional[int] = None, + max_records_limit: Optional[int] = None, emit_connector_builder_messages: bool = False, disable_retries: bool = False, disable_cache: bool = False, @@ -777,6 +778,7 @@ def __init__( self._custom_components_trusted = custom_components_trusted self._limit_pages_fetched_per_slice = limit_pages_fetched_per_slice self._limit_slices_fetched = limit_slices_fetched + self._max_records_limit = max_records_limit self._emit_connector_builder_messages = emit_connector_builder_messages self._disable_retries = disable_retries self._disable_cache = disable_cache @@ -2296,6 +2298,7 @@ def create_default_stream( schema_loader, retriever, self._message_repository, + max_records_limit=self._max_records_limit, ), stream_slicer, slice_limit=self._limit_slices_fetched, diff --git a/airbyte_cdk/sources/declarative/requesters/paginators/default_paginator.py b/airbyte_cdk/sources/declarative/requesters/paginators/default_paginator.py index 8245c66f84..053d609193 100644 --- a/airbyte_cdk/sources/declarative/requesters/paginators/default_paginator.py +++ b/airbyte_cdk/sources/declarative/requesters/paginators/default_paginator.py @@ -284,17 +284,11 @@ def __init__(self, decorated: Paginator, maximum_number_of_pages: int = 5) -> No self._maximum_number_of_pages = maximum_number_of_pages self._decorated = decorated self._page_count = self._PAGE_COUNT_BEFORE_FIRST_NEXT_CALL - self._did_reach_page_limit = False def get_initial_token(self) -> Optional[Any]: self._page_count = self._PAGE_COUNT_BEFORE_FIRST_NEXT_CALL - self._did_reach_page_limit = False return self._decorated.get_initial_token() - @property - def did_reach_page_limit(self) -> bool: - return self._did_reach_page_limit - def get_page_size(self) -> Optional[int]: return self._decorated.get_page_size() @@ -308,7 +302,6 @@ def next_page_token( stream_slice: Optional[StreamSlice] = None, ) -> Optional[Mapping[str, Any]]: if self._page_count >= self._maximum_number_of_pages: - self._did_reach_page_limit = True return None self._page_count += 1 diff --git a/airbyte_cdk/sources/declarative/retrievers/simple_retriever.py b/airbyte_cdk/sources/declarative/retrievers/simple_retriever.py index 091fa7cc4c..d7b63e85fc 100644 --- a/airbyte_cdk/sources/declarative/retrievers/simple_retriever.py +++ b/airbyte_cdk/sources/declarative/retrievers/simple_retriever.py @@ -560,8 +560,7 @@ def _read_pages( PageCompleted( stream_slice=stream_slice, response=response, - pagination_exhausted=not next_page_token - and not getattr(self._paginator, "did_reach_page_limit", False), + pagination_exhausted=not next_page_token, ) ) if not next_page_token: diff --git a/unit_tests/sources/declarative/incremental/test_opaque_token_cursor.py b/unit_tests/sources/declarative/incremental/test_opaque_token_cursor.py index 9f5687ff87..5cae3df7af 100644 --- a/unit_tests/sources/declarative/incremental/test_opaque_token_cursor.py +++ b/unit_tests/sources/declarative/incremental/test_opaque_token_cursor.py @@ -3,11 +3,11 @@ # import json +import logging from unittest.mock import MagicMock import pytest import requests -from requests import Response from airbyte_cdk.models import FailureType from airbyte_cdk.sources.connector_state_manager import ConnectorStateManager @@ -16,9 +16,6 @@ from airbyte_cdk.sources.declarative.partition_routers.single_partition_router import ( SinglePartitionRouter, ) -from airbyte_cdk.sources.declarative.requesters.paginators.default_paginator import ( - PaginatorTestReadDecorator, -) from airbyte_cdk.sources.declarative.retrievers.simple_retriever import PageCompleted from airbyte_cdk.sources.message.repository import InMemoryMessageRepository from airbyte_cdk.sources.types import StreamSlice @@ -78,6 +75,13 @@ def test_t_token_preserves_string_and_converts_integer_but_rejects_other_types() with pytest.raises(AirbyteTracedException) as exc: cursor._coerce_token(value, ["syncToken"]) assert exc.value.failure_type == FailureType.system_error + if isinstance(value, list): + expected = ( + 'Stream events: token_path ["syncToken"] holds a list, ' + "not a string or integer. Check token_path." + ) + assert exc.value.message == expected + assert exc.value.internal_message == expected def test_t_boot_and_t_inc_commit_only_after_partition_close_and_resume_from_state(): @@ -157,26 +161,6 @@ def test_t_fail_discards_pending_token_when_pass_does_not_complete(): assert cursor.state == {"state_type": "opaque_token"} -def test_t_builder_page_limit_is_not_treated_as_a_completed_pass(): - paginator = PaginatorTestReadDecorator(MagicMock(), maximum_number_of_pages=1) - paginator.get_initial_token() - next_page_token = paginator.next_page_token(Response(), last_page_size=1, last_record=None) - - cursor = _cursor() - stream_slice = StreamSlice(partition={}, cursor_slice={}) - cursor.on_page_completed( - PageCompleted( - stream_slice, - _response({"syncToken": "truncated"}), - pagination_exhausted=not next_page_token and not paginator.did_reach_page_limit, - ) - ) - cursor.close_partition(_partition({})) - - assert paginator.did_reach_page_limit - assert cursor.state == {"state_type": "opaque_token"} - - def test_t_par_keeps_tokens_partition_scoped(): router = Router( [ @@ -296,6 +280,138 @@ def test_t_foreign_ignores_state_from_another_cursor_type(): assert cursor.state == {"state_type": "opaque_token"} +@pytest.mark.parametrize( + "state", + [ + {"state_type": "datetime", "token": "not-a-token"}, + {"__ab_no_cursor_state_message": True}, + { + "state_type": "datetime", + "states": [ + { + "partition": {"account": "a"}, + "cursor": {"created_at": "2024-01-01T00:00:00Z"}, + } + ], + }, + ], +) +def test_t_foreign_datetime_and_missing_cursor_states_bootstrap(state): + router = Router([StreamSlice(partition={"account": "a"}, cursor_slice={})]) + cursor = _cursor(router=router, state=state) + + assert list(cursor.stream_slices())[0].cursor_slice == {} + assert cursor.state["state_type"] == "opaque_token" + assert cursor.state["states"] == [] + + +def test_partition_state_load_skips_only_malformed_entries(caplog): + router = Router( + [ + StreamSlice(partition={"account": "a"}, cursor_slice={}), + StreamSlice(partition={"account": "b"}, cursor_slice={}), + ] + ) + cursor = _cursor( + router=router, + state={ + "state_type": "opaque_token", + "states": [ + { + "partition": {"account": "a"}, + "cursor": {"token": "A1", "committed_at": "2026-01-01T00:00:00Z"}, + }, + None, + { + "partition": {"account": "b"}, + "cursor": {"token": ["invalid"]}, + }, + ], + }, + ) + + slices = list(cursor.stream_slices()) + + assert [stream_slice.cursor_slice for stream_slice in slices] == [ + {"token": "A1"}, + {}, + ] + assert len(cursor.state["states"]) == 1 + assert "reason: foreign state" in caplog.text + + +def test_t_cap_warns_once_at_final_emission_with_over_cap_partition_count(caplog): + router = Router( + [ + StreamSlice(partition={"account": account}, cursor_slice={}) + for account in ("a", "b", "c", "d", "e") + ] + ) + state = { + "state_type": "opaque_token", + "states": [ + { + "partition": {"account": account}, + "cursor": {"token": account.upper(), "committed_at": "2026-01-01T00:00:00Z"}, + } + for account in ("a", "b", "c") + ], + } + cursor = _cursor(router=router, state=state, max_partitions=3) + + with caplog.at_level(logging.INFO): + list(cursor.stream_slices()) + assert "partitions already store tokens" not in caplog.text + cursor.ensure_at_least_one_state_emitted() + cursor.ensure_at_least_one_state_emitted() + + assert ( + caplog.text.count( + "Stream events: 3 partitions already store tokens; 2 more read in full on every sync." + ) + == 1 + ) + assert caplog.text.count("reason: partition cap") == 1 + + +def test_bootstrap_and_replay_observability_are_deduplicated_and_aged(caplog): + router = Router( + [StreamSlice(partition={"account": account}, cursor_slice={}) for account in ("a", "b")] + ) + cursor = _cursor(router=router) + with caplog.at_level(logging.INFO): + list(cursor.stream_slices()) + assert caplog.text.count("reason: none") == 1 + + caplog.clear() + replay = _cursor( + state={ + "state_type": "opaque_token", + "token": "T1", + "committed_at": "2024-01-01T00:00:00Z", + } + ) + with caplog.at_level(logging.INFO): + list(replay.stream_slices()) + assert "Stream events: replaying the change token committed " in caplog.text + assert " ago." in caplog.text + + +def test_commit_observability_counts_completed_pages_per_partition(caplog): + cursor = _cursor() + stream_slice = StreamSlice(partition={}, cursor_slice={}) + with caplog.at_level(logging.DEBUG, logger="airbyte"): + for exhausted, body in ( + (False, {"nextCursor": "p2"}), + (False, {"nextCursor": "p3"}), + (True, {"syncToken": "T1"}), + ): + cursor.on_page_completed(PageCompleted(stream_slice, _response(body), exhausted)) + cursor.close_partition(_partition({})) + + assert "Stream events: committed change token T1 for partition {} after 3 pages." in caplog.text + + def test_t_schema_ignores_incompatible_partition_state_shape(): cursor = _cursor( router=Router([]), @@ -305,6 +421,12 @@ def test_t_schema_ignores_incompatible_partition_state_shape(): assert cursor.state == {"state_type": "opaque_token", "states": []} +def test_t_schema_ignores_partition_state_shape_on_unpartitioned_stream(): + cursor = _cursor(state={"state_type": "opaque_token", "states": []}) + + assert cursor.state == {"state_type": "opaque_token"} + + def test_t_valid_and_t_init_use_partition_router_and_initial_token_slice(): router = Router([StreamSlice(partition={"account": "a"}, cursor_slice={})]) cursor = _cursor(router=router) diff --git a/unit_tests/sources/declarative/incremental/test_opaque_token_cursor_integration.py b/unit_tests/sources/declarative/incremental/test_opaque_token_cursor_integration.py index 08e5ca894a..893315a86e 100644 --- a/unit_tests/sources/declarative/incremental/test_opaque_token_cursor_integration.py +++ b/unit_tests/sources/declarative/incremental/test_opaque_token_cursor_integration.py @@ -4,10 +4,11 @@ import json import logging -from unittest.mock import MagicMock +from unittest.mock import MagicMock, patch import pytest +from airbyte_cdk.connector_builder.test_reader.reader import TestReader from airbyte_cdk.models import ( AirbyteStateType, ConfiguredAirbyteCatalog, @@ -16,8 +17,10 @@ SyncMode, Type, ) +from airbyte_cdk.sources.connector_state_manager import ConnectorStateManager from airbyte_cdk.sources.declarative.concurrent_declarative_source import ( ConcurrentDeclarativeSource, + TestLimits, ) from airbyte_cdk.sources.declarative.models.declarative_component_schema import ( AsyncRetriever as AsyncRetrieverModel, @@ -58,11 +61,19 @@ ) from airbyte_cdk.sources.streams.concurrent.default_stream import DefaultStream from airbyte_cdk.test.mock_http import HttpMocker, HttpRequest, HttpResponse +from airbyte_cdk.utils.traced_exception import AirbyteTracedException _URL = "https://api.example.com/events" -def _manifest(*, initial_token=None, partitioned=False): +def _manifest( + *, + initial_token=None, + partitioned=False, + partition_values=("a", "b"), + completion_condition=None, + error_handler=None, +): stream = { "type": "DeclarativeStream", "name": "events", @@ -111,13 +122,17 @@ def _manifest(*, initial_token=None, partitioned=False): }, }, } + if completion_condition is not None: + stream["incremental_sync"]["completion_condition"] = completion_condition + if error_handler is not None: + stream["retriever"]["requester"]["error_handler"] = error_handler if initial_token is not None: stream["incremental_sync"]["initial_token"] = initial_token if partitioned: stream["retriever"]["partition_router"] = { "type": "ListPartitionRouter", "cursor_field": "account", - "values": ["a", "b"], + "values": list(partition_values), "request_option": { "type": "RequestOption", "field_name": "account", @@ -143,9 +158,16 @@ def _manifest(*, initial_token=None, partitioned=False): } -def _source_and_catalog(manifest, state=None): +def _source_and_catalog( + manifest, + state=None, + *, + limits=None, + limit_pages_fetched_per_slice=None, + sync_mode=SyncMode.incremental, +): source = ConcurrentDeclarativeSource( - source_config=manifest, config={}, catalog=None, state=state + source_config=manifest, config={}, catalog=None, state=state, limits=limits ) discovered = source.discover(logger=MagicMock(), config={}) stream = discovered.streams[0] @@ -153,13 +175,20 @@ def _source_and_catalog(manifest, state=None): streams=[ ConfiguredAirbyteStream( stream=stream, - sync_mode=SyncMode.incremental, + sync_mode=sync_mode, destination_sync_mode=DestinationSyncMode.append, cursor_field=stream.default_cursor_field, primary_key=stream.source_defined_primary_key, ) ] ) + if limit_pages_fetched_per_slice is not None: + source._constructor = ModelToComponentFactory( + message_repository=source._message_repository, + connector_state_manager=source._connector_state_manager, + configured_catalog=catalog, + limit_pages_fetched_per_slice=limit_pages_fetched_per_slice, + ) return source, catalog @@ -167,6 +196,14 @@ def _read(source, catalog, state=None): return list(source.read(logger=MagicMock(), config={}, catalog=catalog, state=state)) +def _read_until_error(source, catalog, state): + messages = [] + with pytest.raises(AirbyteTracedException): + for message in source.read(logger=MagicMock(), config={}, catalog=catalog, state=state): + messages.append(message) + return messages + + def _state(messages): return [ message.state @@ -175,6 +212,238 @@ def _state(messages): ] +def _incoming_state(stream_state): + state_manager = ConnectorStateManager() + state_manager.update_state_for_stream("events", None, stream_state) + return [state_manager.create_state_message("events", None).state] + + +def test_t_boot_commits_after_three_pages_and_records(): + with HttpMocker() as http_mocker: + http_mocker.post( + HttpRequest(_URL, body=None), + HttpResponse(json.dumps({"items": [{"id": "1"}], "nextCursor": "p2"})), + ) + http_mocker.post( + HttpRequest(_URL, body={"nextCursor": "p2"}), + HttpResponse(json.dumps({"items": [{"id": "2"}], "nextCursor": "p3"})), + ) + http_mocker.post( + HttpRequest(_URL, body={"nextCursor": "p3"}), + HttpResponse(json.dumps({"items": [{"id": "3"}], "syncToken": "T1"})), + ) + source, catalog = _source_and_catalog(_manifest()) + messages = _read(source, catalog) + + records = [index for index, message in enumerate(messages) if message.type == Type.RECORD] + state_messages = _state(messages) + assert len(state_messages) == 1 + assert state_messages[0].stream.stream_state.__dict__["token"] == "T1" + state_index = next( + index for index, message in enumerate(messages) if message.type == Type.STATE + ) + assert records and max(records) < state_index + + +@pytest.mark.parametrize("echo_initial_token", [True, False], ids=["incremental", "no-echo"]) +def test_t_inc_and_t_noecho_never_commit_intermediate_token(echo_initial_token): + first_response = {"items": [{"id": "1"}], "nextCursor": "p2"} + if echo_initial_token: + first_response["syncToken"] = "T1" + state = _incoming_state({"state_type": "opaque_token", "token": "T1"}) + with HttpMocker() as http_mocker: + http_mocker.post( + HttpRequest(_URL, body={"syncToken": "T1"}), + HttpResponse(json.dumps(first_response)), + ) + http_mocker.post( + HttpRequest(_URL, body={"syncToken": "T1", "nextCursor": "p2"}), + HttpResponse(json.dumps({"items": [{"id": "2"}], "syncToken": "T2"})), + ) + source, catalog = _source_and_catalog(_manifest(), state) + messages = _read(source, catalog, state) + + assert _state(messages)[-1].stream.stream_state.__dict__["token"] == "T2" + + +def test_t_empty_pages_can_commit_terminal_token(): + with HttpMocker() as http_mocker: + http_mocker.post( + HttpRequest(_URL, body=None), + HttpResponse(json.dumps({"items": [], "nextCursor": "p2"})), + ) + http_mocker.post( + HttpRequest(_URL, body={"nextCursor": "p2"}), + HttpResponse(json.dumps({"items": [], "syncToken": "T2"})), + ) + source, catalog = _source_and_catalog(_manifest()) + messages = _read(source, catalog) + + assert not any(message.type == Type.RECORD for message in messages) + assert _state(messages)[-1].stream.stream_state.__dict__["token"] == "T2" + + +def test_t_fail_preserves_committed_token_and_emits_incomplete_status(): + state = _incoming_state({"state_type": "opaque_token", "token": "T1"}) + manifest = _manifest( + error_handler={ + "type": "DefaultErrorHandler", + "max_retries": 0, + "backoff_strategies": [], + } + ) + with HttpMocker() as http_mocker: + http_mocker.post( + HttpRequest(_URL, body={"syncToken": "T1"}), + HttpResponse(json.dumps({"items": [{"id": "1"}], "nextCursor": "p2"})), + ) + http_mocker.post( + HttpRequest(_URL, body={"syncToken": "T1", "nextCursor": "p2"}), + HttpResponse("{}", status_code=500), + ) + source, catalog = _source_and_catalog(manifest, state) + messages = _read_until_error(source, catalog, state) + + assert any(message.type == Type.RECORD for message in messages) + assert _state(messages)[-1].stream.stream_state.__dict__["token"] == "T1" + assert any( + message.type == Type.TRACE + and message.trace.stream_status + and message.trace.stream_status.status.value == "INCOMPLETE" + for message in messages + ) + + +def test_t_builder_page_limit_and_max_records_do_not_commit_tokens(): + with HttpMocker() as http_mocker: + http_mocker.post( + HttpRequest(_URL, body=None), + HttpResponse(json.dumps({"items": [{"id": "1"}], "nextCursor": "p2"})), + ) + source, catalog = _source_and_catalog(_manifest(), limit_pages_fetched_per_slice=1) + messages = _read(source, catalog) + assert not any("token" in message.stream.stream_state.__dict__ for message in _state(messages)) + + with HttpMocker() as http_mocker: + http_mocker.post( + HttpRequest(_URL, body=None), + HttpResponse( + json.dumps( + { + "items": [{"id": "1"}, {"id": "2"}], + "syncToken": "must-not-commit", + "nextCursor": "p2", + } + ) + ), + ) + source, catalog = _source_and_catalog(_manifest(), limits=TestLimits(max_records=1)) + with patch.object( + source._message_repository, + "emit_message", + wraps=source._message_repository.emit_message, + ) as emit_message: + test_read = TestReader( + max_pages_per_slice=5, max_slices=5, max_record_limit=1 + ).run_test_read(source, {}, catalog, "events", [], record_limit=1) + assert test_read.test_read_limit_reached + assert len(test_read.slices) == 1 + emitted_states = [ + call.args[0] + for call in emit_message.call_args_list + if call.args and call.args[0].type == Type.STATE + ] + assert not any( + ( + message.state.stream.stream_state.get("token") + if isinstance(message.state.stream.stream_state, dict) + else message.state.stream.stream_state.__dict__.get("token") + ) + == "must-not-commit" + for message in emitted_states + if message.state and message.state.stream and message.state.stream.stream_state + ) + + +def test_t_resume_replays_existing_partition_and_bootstraps_missing_partition(): + state = _incoming_state( + { + "state_type": "opaque_token", + "states": [ + { + "partition": {"account": "a"}, + "cursor": {"token": "A1", "committed_at": "2024-01-01T00:00:00Z"}, + } + ], + } + ) + with HttpMocker() as http_mocker: + http_mocker.post( + HttpRequest(_URL, body={"account": "a", "syncToken": "A1"}), + HttpResponse(json.dumps({"items": [], "syncToken": "A2"})), + ) + http_mocker.post( + HttpRequest(_URL, body={"account": "b"}), + HttpResponse(json.dumps({"items": [], "syncToken": "B1"})), + ) + source, catalog = _source_and_catalog(_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": "A2", "b": "B1"} + + +def test_t_par_keeps_successful_partition_commits_when_another_fails(): + state = _incoming_state( + { + "state_type": "opaque_token", + "states": [ + { + "partition": {"account": "c"}, + "cursor": {"token": "C0", "committed_at": "2024-01-01T00:00:00Z"}, + } + ], + } + ) + manifest = _manifest( + partitioned=True, + partition_values=("b", "a", "c"), + error_handler={ + "type": "DefaultErrorHandler", + "max_retries": 0, + "backoff_strategies": [], + }, + ) + with HttpMocker() as http_mocker: + http_mocker.post( + HttpRequest(_URL, body={"account": "b"}), + HttpResponse(json.dumps({"items": [], "syncToken": "B1"})), + ) + http_mocker.post( + HttpRequest(_URL, body={"account": "a"}), + HttpResponse(json.dumps({"items": [], "syncToken": "A1"})), + ) + http_mocker.post( + HttpRequest(_URL, body={"account": "c", "syncToken": "C0"}), + HttpResponse("{}", status_code=500), + ) + source, catalog = _source_and_catalog(manifest, state) + messages = _read_until_error(source, catalog, state) + committed = { + item["partition"]["account"]: item["cursor"]["token"] + for item in _state(messages)[-1].stream.stream_state.__dict__["states"] + } + assert committed == {"a": "A1", "b": "B1", "c": "C0"} + assert any( + message.type == Type.TRACE + and message.trace.stream_status + and message.trace.stream_status.status.value == "INCOMPLETE" + for message in messages + ) + + def test_t_boot_inc_builder_and_resume_send_exact_request_bodies(): with HttpMocker() as http_mocker: http_mocker.post( @@ -242,17 +511,42 @@ def test_t_disc_advertises_source_defined_incremental_without_cursor_field(): assert stream.default_cursor_field == [] -def test_t_valid_rejects_datetime_only_extra_settings(): +def test_t_disc_datetime_cursor_has_non_none_default_cursor_field(): manifest = _manifest() - manifest["streams"][0]["incremental_sync"]["start_datetime"] = "2024-01-01" + manifest["streams"][0]["incremental_sync"] = { + "type": "DatetimeBasedCursor", + "cursor_field": "created_at", + "start_datetime": "2020-01-01T00:00:00Z", + "datetime_format": "%Y-%m-%dT%H:%M:%SZ", + } + source = ConcurrentDeclarativeSource(source_config=manifest, config={}, catalog=None) + + stream = source.discover(logger=MagicMock(), config={}).streams[0] + + assert stream.default_cursor_field is not None + assert stream.default_cursor_field == ["created_at"] + + +@pytest.mark.parametrize("setting", ["start_datetime", "lookback_window"]) +def test_t_schema_rejects_datetime_only_extra_settings(setting): + manifest = _manifest() + manifest["streams"][0]["incremental_sync"][setting] = "2024-01-01" source = ConcurrentDeclarativeSource( source_config=manifest, config={}, catalog=None, state=None ) - with pytest.raises(ValueError, match="datetime-only setting.*start_datetime"): + with pytest.raises(ValueError, match=setting): source.discover(logger=MagicMock(), config={}) +def test_t_schema_accepts_stream_parameters_and_builds_stream(): + manifest = _manifest() + manifest["streams"][0]["$parameters"] = {"parameter": "value"} + source, _ = _source_and_catalog(manifest) + + assert isinstance(source.streams(config={})[0], DefaultStream) + + def _stream_model(): return DeclarativeStreamModel.parse_obj(_manifest()["streams"][0]) @@ -388,3 +682,52 @@ def test_t_inert_leaves_page_observer_unset_for_datetime_cursor(): retriever = generator._partition_factory._retriever assert isinstance(retriever, SimpleRetriever) assert retriever.page_observer is None + + +@pytest.mark.parametrize("cursor_kind", ["datetime", "count", "none"]) +def test_t_inert_cursor_types_read_without_opaque_page_observer(cursor_kind): + manifest = _manifest() + if cursor_kind == "datetime": + manifest["streams"][0]["incremental_sync"] = { + "type": "DatetimeBasedCursor", + "cursor_field": "created_at", + "start_datetime": "2020-01-01T00:00:00Z", + "datetime_format": "%Y-%m-%dT%H:%M:%SZ", + } + elif cursor_kind == "count": + manifest["streams"][0]["incremental_sync"] = { + "type": "IncrementingCountCursor", + "cursor_field": "id", + "start_value": 0, + } + else: + manifest["streams"][0].pop("incremental_sync") + + sync_mode = SyncMode.full_refresh if cursor_kind == "none" else SyncMode.incremental + with HttpMocker() as http_mocker: + http_mocker.post( + HttpRequest(_URL, body=None), + HttpResponse( + json.dumps( + { + "items": [ + { + "id": 1 if cursor_kind == "count" else "1", + "created_at": "2020-01-02T00:00:00Z", + } + ] + } + ) + ), + ) + source, catalog = _source_and_catalog(manifest, sync_mode=sync_mode) + messages = _read(source, catalog) + + assert any(message.type == Type.RECORD for message in messages) + stream = source.streams(config={})[0] + assert isinstance(stream, DefaultStream) + generator = stream._stream_partition_generator + assert isinstance(generator, StreamSlicerPartitionGenerator) + retriever = generator._partition_factory._retriever + assert isinstance(retriever, SimpleRetriever) + assert retriever.page_observer is None From 9808890d0ecc4136509e9b1e16f811ff10d8a138 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 16 Oct 2026 23:26:18 +0000 Subject: [PATCH 3/3] Revert Builder max record cursor wiring Co-Authored-By: bot_apk --- .../concurrent_declarative_source.py | 1 - .../parsers/model_to_component_factory.py | 3 -- .../test_opaque_token_cursor_integration.py | 42 +++++-------------- 3 files changed, 11 insertions(+), 35 deletions(-) diff --git a/airbyte_cdk/sources/declarative/concurrent_declarative_source.py b/airbyte_cdk/sources/declarative/concurrent_declarative_source.py index b015e1ddec..055576f45f 100644 --- a/airbyte_cdk/sources/declarative/concurrent_declarative_source.py +++ b/airbyte_cdk/sources/declarative/concurrent_declarative_source.py @@ -198,7 +198,6 @@ def __init__( max_concurrent_async_job_count=source_config.get("max_concurrent_async_job_count"), limit_pages_fetched_per_slice=limits.max_pages_per_slice if limits else None, limit_slices_fetched=limits.max_slices if limits else None, - max_records_limit=limits.max_records if limits else None, disable_retries=True if limits else False, disable_cache=True if limits else False, ) 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 95f532c918..23902e1484 100644 --- a/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py +++ b/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py @@ -760,7 +760,6 @@ def __init__( self, limit_pages_fetched_per_slice: Optional[int] = None, limit_slices_fetched: Optional[int] = None, - max_records_limit: Optional[int] = None, emit_connector_builder_messages: bool = False, disable_retries: bool = False, disable_cache: bool = False, @@ -778,7 +777,6 @@ def __init__( self._custom_components_trusted = custom_components_trusted self._limit_pages_fetched_per_slice = limit_pages_fetched_per_slice self._limit_slices_fetched = limit_slices_fetched - self._max_records_limit = max_records_limit self._emit_connector_builder_messages = emit_connector_builder_messages self._disable_retries = disable_retries self._disable_cache = disable_cache @@ -2298,7 +2296,6 @@ def create_default_stream( schema_loader, retriever, self._message_repository, - max_records_limit=self._max_records_limit, ), stream_slicer, slice_limit=self._limit_slices_fetched, diff --git a/unit_tests/sources/declarative/incremental/test_opaque_token_cursor_integration.py b/unit_tests/sources/declarative/incremental/test_opaque_token_cursor_integration.py index 893315a86e..652d45f851 100644 --- a/unit_tests/sources/declarative/incremental/test_opaque_token_cursor_integration.py +++ b/unit_tests/sources/declarative/incremental/test_opaque_token_cursor_integration.py @@ -4,11 +4,10 @@ import json import logging -from unittest.mock import MagicMock, patch +from unittest.mock import MagicMock import pytest -from airbyte_cdk.connector_builder.test_reader.reader import TestReader from airbyte_cdk.models import ( AirbyteStateType, ConfiguredAirbyteCatalog, @@ -20,7 +19,6 @@ from airbyte_cdk.sources.connector_state_manager import ConnectorStateManager from airbyte_cdk.sources.declarative.concurrent_declarative_source import ( ConcurrentDeclarativeSource, - TestLimits, ) from airbyte_cdk.sources.declarative.models.declarative_component_schema import ( AsyncRetriever as AsyncRetrieverModel, @@ -162,12 +160,11 @@ def _source_and_catalog( manifest, state=None, *, - limits=None, limit_pages_fetched_per_slice=None, sync_mode=SyncMode.incremental, ): source = ConcurrentDeclarativeSource( - source_config=manifest, config={}, catalog=None, state=state, limits=limits + source_config=manifest, config={}, catalog=None, state=state ) discovered = source.discover(logger=MagicMock(), config={}) stream = discovered.streams[0] @@ -337,32 +334,15 @@ def test_t_builder_page_limit_and_max_records_do_not_commit_tokens(): ) ), ) - source, catalog = _source_and_catalog(_manifest(), limits=TestLimits(max_records=1)) - with patch.object( - source._message_repository, - "emit_message", - wraps=source._message_repository.emit_message, - ) as emit_message: - test_read = TestReader( - max_pages_per_slice=5, max_slices=5, max_record_limit=1 - ).run_test_read(source, {}, catalog, "events", [], record_limit=1) - assert test_read.test_read_limit_reached - assert len(test_read.slices) == 1 - emitted_states = [ - call.args[0] - for call in emit_message.call_args_list - if call.args and call.args[0].type == Type.STATE - ] - assert not any( - ( - message.state.stream.stream_state.get("token") - if isinstance(message.state.stream.stream_state, dict) - else message.state.stream.stream_state.__dict__.get("token") - ) - == "must-not-commit" - for message in emitted_states - if message.state and message.state.stream and message.state.stream.stream_state - ) + source, _ = _source_and_catalog(_manifest()) + stream = source.streams(config={})[0] + partition = next(iter(stream.generate_partitions())) + record_iterator = iter(partition.read()) + assert next(record_iterator).data == {"id": "1"} + record_iterator.close() + stream.cursor.close_partition(partition) + + assert "token" not in stream.cursor.state def test_t_resume_replays_existing_partition_and_bootstraps_missing_partition():