feat(declarative): support RequestPath tokens and first-request checks in SyncTokenCursor

airbyte_cdk/sources/declarative/declarative_component_schema.yaml

@@ -1459,9 +1459,15 @@ definitions:
       token_option:
         title: Send Token In
         description: >-
-          Where to send the stored token on each request of a pass: the JSON body, the query
-          string, a header or the form body.
-        "$ref": "#/definitions/RequestOption"
+          How to send the stored token. A RequestOption puts it in the JSON body, the query
+          string, a header or the form body: on every request of a pass, or on the first request
+          only when the retriever sets ignore_stream_slicer_parameters_on_paginated_requests.
+          RequestPath makes it the URL of the first request of a pass, for APIs whose token is a
+          complete URL, such as a Microsoft Graph @odata.deltaLink; the paginator sends the
+          later requests.
+        anyOf:
+          - "$ref": "#/definitions/RequestOption"
+          - "$ref": "#/definitions/RequestPath"
       completion_condition:
         title: Completion Condition
         description: >-

airbyte_cdk/sources/declarative/incremental/sync_token_cursor.py

@@ -42,6 +42,7 @@
 
 class SyncTokenCursor(Cursor):
     STATE_TYPE = "sync_token"
+    DEFAULT_MAX_STATE_BYTES = 4 * 1024 * 1024
 
     def __init__(
         self,
@@ -59,6 +60,7 @@ def __init__(
         completion_condition: Optional[str] = None,
         cursor_field: Optional[str] = None,
         max_partitions: int = ConcurrentPerPartitionCursor.DEFAULT_MAX_PARTITIONS_NUMBER,
+        max_state_bytes: int = DEFAULT_MAX_STATE_BYTES,
     ) -> None:
         self._stream_name = stream_name
         self._stream_namespace = stream_namespace
@@ -78,8 +80,10 @@ def __init__(
         )
         self.cursor_field = CursorField(cursor_field) if cursor_field else None
         self._max_partitions = max_partitions
+        self._max_state_bytes = max_state_bytes
         self._serializer = PerPartitionKeySerializer()
         self._committed: OrderedDict[str, dict[str, Any]] = OrderedDict()
+        self._committed_bytes = 0
         self._pending: dict[str, dict[str, Any]] = {}
         self._page_counts: dict[str, int] = {}
         self._lock = threading.RLock()
@@ -88,8 +92,9 @@ def __init__(
         self._state_is_foreign = False
         self._state_emitted = False
         self._warned_no_token = False
-        self._warned_cap = False
+        self._warned_limit = False
         self._partitions_over_cap: set[str] = set()
+        self._partitions_over_state_size: set[str] = set()
         self._load_state(stream_state or {})
 
     @property
@@ -132,6 +137,9 @@ def stream_slices(self) -> Iterable[StreamSlice]:
                 elif at_partition_cap:
                     self._partitions_over_cap.add(partition_key)
                     self._log_bootstrap("partition cap")
+                elif self._committed_bytes >= self._max_state_bytes:
+                    self._partitions_over_state_size.add(partition_key)
+                    self._log_bootstrap("state size limit")
                 elif not self._state_is_foreign:
                     self._log_bootstrap("none")
             yield StreamSlice(
@@ -144,6 +152,9 @@ def stream_slices(self) -> Iterable[StreamSlice]:
             with self._lock:
                 vanished = set(self._committed) - seen_partitions
                 for partition_key in vanished:
+                    self._committed_bytes -= self._entry_bytes(
+                        partition_key, self._committed[partition_key]
+                    )
                     del self._committed[partition_key]
                 if vanished:
                     logger.info(
@@ -330,18 +341,39 @@ def close_partition(self, partition: Partition) -> None:
             pending = self._pending.pop(partition_key, None)
             page_count = self._page_counts.pop(partition_key, 0)
             if pending is not None:
-                if partition_key in self._committed or len(self._committed) < self._max_partitions:
+                existing = self._committed.get(partition_key)
+                if existing is not None:
+                    self._committed_bytes -= self._entry_bytes(partition_key, existing)
+                    self._committed_bytes += self._entry_bytes(partition_key, pending)
                     self._committed[partition_key] = pending
                     self._partitions_over_cap.discard(partition_key)
+                    self._partitions_over_state_size.discard(partition_key)
                     logger.debug(
                         "Stream %s: committed change token %s for partition %s after %s pages.",
                         self._stream_name,
                         pending["token"],
                         self._serializer.to_partition(partition_key),
                         page_count,
                     )
-                else:
+                elif (
+                    len(self._committed) < self._max_partitions
+                    and self._committed_bytes < self._max_state_bytes
+                ):
+                    self._committed[partition_key] = pending
+                    self._committed_bytes += self._entry_bytes(partition_key, pending)
+                    self._partitions_over_cap.discard(partition_key)
+                    self._partitions_over_state_size.discard(partition_key)
+                    logger.debug(
+                        "Stream %s: committed change token %s for partition %s after %s pages.",
+                        self._stream_name,
+                        pending["token"],
+                        self._serializer.to_partition(partition_key),
+                        page_count,
+                    )
+                elif len(self._committed) >= self._max_partitions:
                     self._partitions_over_cap.add(partition_key)
+                else:
+                    self._partitions_over_state_size.add(partition_key)
             if not self._partitioned:
                 self._emit_state_message()
             elif time.time() - self._last_emission_time > 600:
@@ -362,14 +394,25 @@ def _emit_state_message(self) -> None:
     def ensure_at_least_one_state_emitted(self) -> None:
         if self._partitioned or not self._state_emitted:
             self._emit_state_message()
-        if self._partitions_over_cap and not self._warned_cap:
-            logger.warning(
-                "Stream %s: %s partitions already store tokens; %s more read in full on every sync.",
-                self._stream_name,
-                len(self._committed),
-                len(self._partitions_over_cap),
-            )
-            self._warned_cap = True
+        if not self._warned_limit:
+            if self._partitions_over_state_size:
+                logger.warning(
+                    "Stream %s: stored tokens reached the %s state limit at %s partitions; "
+                    "%s more read in full on every sync.",
+                    self._stream_name,
+                    self._human_readable_bytes(self._max_state_bytes),
+                    f"{len(self._committed):,}",
+                    f"{len(self._partitions_over_state_size):,}",
+                )
+                self._warned_limit = True
+            elif self._partitions_over_cap:
+                logger.warning(
+                    "Stream %s: %s partitions already store tokens; %s more read in full on every sync.",
+                    self._stream_name,
+                    len(self._committed),
+                    len(self._partitions_over_cap),
+                )
+                self._warned_limit = True
 
     def observe(self, record: Record) -> None:
         return
@@ -409,6 +452,10 @@ def _load_state(self, stream_state: StreamState) -> None:
                 else:
                     self._log_bootstrap("foreign state")
             self._committed.update(committed)
+            self._committed_bytes += sum(
+                self._entry_bytes(partition_key, cursor)
+                for partition_key, cursor in committed.items()
+            )
         else:
             if "states" in stream_state:
                 self._log_foreign_state()
@@ -419,5 +466,19 @@ def _load_state(self, stream_state: StreamState) -> None:
                     "token": token,
                     "committed_at": stream_state.get("committed_at", ""),
                 }
+                self._committed_bytes += self._entry_bytes("{}", self._committed["{}"])
             elif token is not None:
                 self._log_foreign_state()
+
+    @staticmethod
+    def _entry_bytes(partition_key: str, cursor: Mapping[str, Any]) -> int:
+        return len(partition_key.encode()) + len(cursor["token"].encode())
+
+    @staticmethod
+    def _human_readable_bytes(size: int) -> str:
+        if size < 1024:
+            return f"{size:,} bytes"
+        size_in_kib = size / 1024
+        if size_in_kib < 1024:
+            return f"{size_in_kib:g} KB"
+        return f"{size_in_kib / 1024:g} MB"

airbyte_cdk/sources/declarative/models/declarative_component_schema.py

@@ -2965,9 +2965,9 @@ class Config:
         examples=[["syncToken"], ["nextSyncToken"], ["meta", "sync_token"]],
         title="Token Path",
     )
-    token_option: RequestOption = Field(
+    token_option: Union[RequestOption, RequestPath] = Field(
         ...,
-        description="Where to send the stored token on each request of a pass: the JSON body, the query string, a header or the form body.",
+        description="How to send the stored token. A RequestOption puts it in the JSON body, the query string, a header or the form body: on every request of a pass, or on the first request only when the retriever sets ignore_stream_slicer_parameters_on_paginated_requests. RequestPath makes it the URL of the first request of a pass, for APIs whose token is a complete URL, such as a Microsoft Graph @odata.deltaLink; the paginator sends the later requests.",
         title="Send Token In",
     )
     completion_condition: Optional[str] = Field(

airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py

@@ -2335,6 +2335,51 @@ def _validate_sync_token_stream(model: DeclarativeStreamModel) -> None:
                         )
         if not isinstance(cursor_model, SyncTokenCursorModel):
             return
+        if isinstance(retriever_model, SimpleRetrieverModel) and isinstance(
+            retriever_model.paginator, DefaultPaginatorModel
+        ):
+            token_option = cursor_model.token_option
+            ignore_stream_slicer_parameters = (
+                retriever_model.ignore_stream_slicer_parameters_on_paginated_requests is True
+            )
+            if isinstance(token_option, RequestOptionModel) and not ignore_stream_slicer_parameters:
+                for page_option in (
+                    retriever_model.paginator.page_token_option,
+                    retriever_model.paginator.page_size_option,
+                ):
+                    if (
+                        isinstance(page_option, RequestOptionModel)
+                        and token_option.inject_into == page_option.inject_into
+                    ):
+                        matching_field = None
+                        if token_option.field_name is not None and (
+                            token_option.field_name == page_option.field_name
+                        ):
+                            matching_field = token_option.field_name
+                        elif token_option.field_path is not None and (
+                            token_option.field_path == page_option.field_path
+                        ):
+                            matching_field = ".".join(token_option.field_path)
+                        if matching_field is not None:
+                            raise ValueError(
+                                f"Stream {stream_name}: token_option and the paginator both set "
+                                f'{token_option.inject_into.value} "{matching_field}". Set '
+                                "ignore_stream_slicer_parameters_on_paginated_requests: true on "
+                                "the retriever so the stored token is sent only on the first "
+                                "request of each pass."
+                            )
+
+            strategy = retriever_model.paginator.pagination_strategy
+            if (
+                isinstance(token_option, RequestPathModel) or ignore_stream_slicer_parameters
+            ) and isinstance(strategy, (PageIncrementModel, OffsetIncrementModel)):
+                if strategy.inject_on_first_request is True:
+                    raise ValueError(
+                        f"Stream {stream_name}: {strategy.type} with inject_on_first_request: "
+                        "true sends a page token on the first request, so the stored sync token "
+                        "would never be sent. Set inject_on_first_request: false or use "
+                        "CursorPagination."
+                    )
         datetime_only_keys = {
             "start_datetime",
             "end_datetime",
@@ -4055,6 +4100,18 @@ def _get_url(req: Requester) -> str:
             partition_router, PartitionRouter
         ):
             request_options_provider = partition_router
+        token_request_options_provider = (
+            request_options_provider._cursor_provider
+            if isinstance(request_options_provider, PerPartitionRequestOptionsProvider)
+            else request_options_provider
+        )
+        first_request_path = (
+            token_request_options_provider.get_request_path
+            if isinstance(incremental_sync, SyncTokenCursorModel)
+            and isinstance(token_request_options_provider, SyncTokenRequestOptionsProvider)
+            and isinstance(token_request_options_provider.token_option, RequestPath)
+            else None
+        )
 
         paginator = (
             self._create_component_from_model(
@@ -4140,6 +4197,7 @@ def _get_url(req: Requester) -> str:
                 config=config,
                 ignore_stream_slicer_parameters_on_paginated_requests=ignore_stream_slicer_parameters_on_paginated_requests,
                 post_pagination_filter=post_pagination_filter,
+                first_request_path=first_request_path,
                 parameters=model.parameters or {},
             )
 
@@ -4184,6 +4242,7 @@ def _get_url(req: Requester) -> str:
             token_reset_handler=(
                 cursor.restart_without_token if isinstance(cursor, SyncTokenCursor) else None
             ),
+            first_request_path=first_request_path,
             parameters=model.parameters or {},
         )
 

airbyte_cdk/sources/declarative/requesters/request_options/sync_token_request_options_provider.py

@@ -13,12 +13,13 @@
 from airbyte_cdk.sources.declarative.requesters.request_options.request_options_provider import (
     RequestOptionsProvider,
 )
+from airbyte_cdk.sources.declarative.requesters.request_path import RequestPath
 from airbyte_cdk.sources.types import Config, StreamSlice, StreamState
 
 
 @dataclass
 class SyncTokenRequestOptionsProvider(RequestOptionsProvider):
-    token_option: RequestOption
+    token_option: Union[RequestOption, RequestPath]
     config: Config
     parameters: InitVar[Mapping[str, Any]]
     initial_token: Optional[str] = None
@@ -33,20 +34,31 @@ def __post_init__(self, parameters: Mapping[str, Any]) -> None:
     def _get_request_options(
         self, option_type: RequestOptionType, stream_slice: Optional[StreamSlice]
     ) -> Mapping[str, Any]:
+        if isinstance(self.token_option, RequestPath):
+            return {}
         if not self.token_option or self.token_option.inject_into != option_type:
             return {}
+        token = self._get_token(stream_slice)
+        if token is None:
+            return {}
+        result: dict[str, Any] = {}
+        self.token_option.inject_into_request(result, token, self.config)
+        return result
+
+    def _get_token(self, stream_slice: Optional[StreamSlice]) -> Optional[Any]:
         token = (
             stream_slice.cursor_slice.get("token")
             if stream_slice and stream_slice.cursor_slice
             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
+        return token
+
+    def get_request_path(self, stream_slice: Optional[StreamSlice]) -> Optional[str]:
+        """Return the initial request URL; after reset, use `initial_token` or the requester path."""
+        token = self._get_token(stream_slice)
+        return str(token) if token is not None else None
 
     def get_request_params(
         self,