airbyte_cdk.sources.declarative.retrievers.page_size_reducer

  1# Copyright (c) 2026 Airbyte, Inc., all rights reserved.
  2
  3import logging
  4import time
  5from dataclasses import dataclass
  6from enum import Enum
  7from typing import Callable, Optional
  8
  9from airbyte_cdk.models import FailureType
 10from airbyte_cdk.utils.traced_exception import AirbyteTracedException
 11
 12LOGGER = logging.getLogger("airbyte")
 13
 14
 15class PageSizeResetPolicy(Enum):
 16    NEVER = "NEVER"
 17    AFTER_SUCCESSFUL_PAGE = "AFTER_SUCCESSFUL_PAGE"
 18
 19
 20@dataclass(frozen=True)
 21class PageSizeReduction:
 22    """
 23    How much to shrink the page size when an error handler resolves to `ResponseAction.REDUCE_PAGE_SIZE`.
 24
 25    This is immutable configuration: it is created once per stream and shared, while the page size in effect
 26    lives in a `PageSizeReducer` created per partition read.
 27    """
 28
 29    reduction_factor: float = 2.0
 30    minimum_page_size: int = 1
 31    max_attempts: int = 5
 32    reset_policy: PageSizeResetPolicy = PageSizeResetPolicy.NEVER
 33    # Base wait before a page is re-issued, multiplied by the number of attempts made in a row. A
 34    # `REDUCE_PAGE_SIZE` response never reaches the HTTP retry budget - the exception raised for it is
 35    # deliberately not a backoff exception - so this is the only thing spacing these requests out, and an
 36    # API whose 502 means "we are briefly unwell" rather than "your page is too big" needs it to be more
 37    # than a token pause.
 38    backoff_seconds: float = 0.5
 39    # How many times the same page is re-issued unchanged once the page size cannot be shrunk any further,
 40    # before the read gives up. Zero keeps the strict behaviour: the first response that cannot be answered
 41    # with a smaller page fails the stream. It is what gives an API whose error is transient a budget at the
 42    # floor, where reducing is no longer an option but waiting still is.
 43    retries_at_minimum_page_size: int = 0
 44    # Appended to the two messages raised once the page size cannot be reduced any further. Those are
 45    # `transient_error`s the CDK has no remediation for - it only knows that the API rejected every page size
 46    # asked for - while the connector knows what narrows a query down on this particular API.
 47    failure_message: Optional[str] = None
 48
 49    def __post_init__(self) -> None:
 50        if self.reduction_factor <= 1:
 51            raise ValueError(
 52                f"The page size reduction factor needs to be greater than 1. Got {self.reduction_factor}"
 53            )
 54        if self.minimum_page_size < 1:
 55            raise ValueError(
 56                f"The minimum page size needs to be strictly positive. Got {self.minimum_page_size}"
 57            )
 58        if self.max_attempts < 1:
 59            raise ValueError(
 60                f"The maximum number of page size reductions needs to be strictly positive. Got {self.max_attempts}"
 61            )
 62        if self.backoff_seconds < 0:
 63            raise ValueError(
 64                f"The wait between page size reductions cannot be negative. Got {self.backoff_seconds}"
 65            )
 66        if self.retries_at_minimum_page_size < 0:
 67            raise ValueError(
 68                f"The number of retries at the minimum page size cannot be negative. Got {self.retries_at_minimum_page_size}"
 69            )
 70
 71
 72class PageSizeReducer:
 73    """
 74    Tracks the page size to use while reading one partition when the API asks for smaller pages.
 75
 76    One instance is created per `SimpleRetriever._read_pages` call. This is deliberate: a retriever and its
 77    paginator are shared by every partition of a stream and partitions are read concurrently, so the reduced
 78    page size must not be stored on the paginator or on the retriever.
 79    """
 80
 81    def __init__(
 82        self,
 83        config: PageSizeReduction,
 84        configured_page_size: Optional[int],
 85        stream_name: str = "",
 86        sleep: Optional[Callable[[float], None]] = None,
 87    ) -> None:
 88        self._config = config
 89        self._configured_page_size = configured_page_size
 90        self._stream_name = stream_name
 91        # Resolved on each call rather than bound here: a default of `time.sleep` would capture
 92        # the function object, and a test that patches `time.sleep` to keep a run of reductions
 93        # from taking its two minutes for real would have no effect on an already-bound default.
 94        self._sleep_override = sleep
 95        self._current_page_size: Optional[int] = None
 96        self._attempts = 0
 97        self._total_reductions = 0
 98        self._retries_at_minimum_page_size = 0
 99
100    def _sleep(self, seconds: float) -> None:
101        if self._sleep_override is not None:
102            self._sleep_override(seconds)
103            return
104        time.sleep(seconds)
105
106    @property
107    def page_size_override(self) -> Optional[int]:
108        """
109        :return: the reduced page size to request, or None while the configured page size is in effect
110        """
111        return self._current_page_size
112
113    def reduce(self, error_message: Optional[str] = None) -> None:
114        """
115        Shrink the page size used for the next request. Raises once the page size cannot be shrunk any further
116        so that an endpoint that keeps failing does not loop forever. Those errors include `error_message`, the
117        error handler's message for the rejected page.
118        """
119        if self._configured_page_size is None:
120            raise AirbyteTracedException(
121                internal_message=f"Stream {self._stream_name} received a REDUCE_PAGE_SIZE response action but its paginator does not inject a page size",
122                message=f"Stream {self._stream_name} is set up to reduce its page size on error but does not define "
123                f"one. Set `page_size` on the pagination strategy and `page_size_option` on the paginator.",
124                failure_type=FailureType.config_error,
125            )
126
127        current_page_size = (
128            self._current_page_size
129            if self._current_page_size is not None
130            else self._configured_page_size
131        )
132        if not isinstance(current_page_size, int) or isinstance(current_page_size, bool):
133            # A custom pagination strategy can return anything from `get_page_size`. Reducing
134            # is arithmetic, so a non-integer would otherwise fail with a bare TypeError in
135            # the middle of a sync.
136            raise AirbyteTracedException(
137                internal_message=f"Stream {self._stream_name} has a page size of type {type(current_page_size).__name__}: {current_page_size!r}",
138                message=f"The page size of stream {self._stream_name} is not a whole number, so the connector cannot "
139                f"reduce it. Make sure the pagination strategy's `page_size` is a number.",
140                failure_type=FailureType.config_error,
141            )
142
143        rejection = f"; the API rejected the page: {error_message}" if error_message else ""
144        reduced_page_size = max(
145            self._config.minimum_page_size,
146            int(current_page_size // self._config.reduction_factor),
147        )
148        if reduced_page_size >= current_page_size:
149            # Nothing left to give up on the page size. Whether that is the end of the read is
150            # `retries_at_minimum_page_size`'s call, not this branch's: the response may still be transient.
151            self._retry_at_minimum_page_size(current_page_size, rejection)
152            return
153
154        self._attempts += 1
155        self._total_reductions += 1
156        if self._attempts > self._config.max_attempts:
157            # The budget counts the reductions that did *not* get a page through, which is what separates a
158            # partition that is stuck from one that is merely expensive. A partition where pages keep
159            # succeeding restarts this counter on each of them, under either reset policy, and reads to the
160            # end however many pages it has; a partition where nothing gets through burns the budget here.
161            raise AirbyteTracedException(
162                internal_message=f"Stream {self._stream_name} reduced its page size {self._attempts - 1} times in a row without a single page succeeding, which is the configured maximum of {self._config.max_attempts} ({self._total_reductions - 1} reductions so far while reading this partition){rejection}",
163                # `transient_error`, so the only remediation is the connector's own, if it defined one.
164                message=self._with_failure_message(
165                    f"The source keeps rejecting pages of stream {self._stream_name} at every page size the "
166                    f"connector requested, down to {current_page_size} records per page."
167                ),
168                failure_type=FailureType.transient_error,
169            )
170
171        backoff = self._config.backoff_seconds * self._attempts
172        LOGGER.info(
173            f"Reducing the page size of stream {self._stream_name} from {current_page_size} to {reduced_page_size} "
174            f"and retrying the same page in {backoff}s."
175        )
176        self._current_page_size = reduced_page_size
177        self._sleep(backoff)
178
179    def _retry_at_minimum_page_size(self, current_page_size: int, rejection: str) -> None:
180        """
181        Handle a `REDUCE_PAGE_SIZE` response that arrives when the page size is already as small as the
182        connector is allowed to request.
183
184        Reducing is out of options here, but re-issuing the page is not: an API that answers 502 to a page it
185        considers too heavy answers the same 502 when it is briefly unwell, and the error handler cannot tell
186        the two apart. `REDUCE_PAGE_SIZE` bypasses the HTTP retry budget, so without this budget the second
187        kind of 502 ends the stream on the first response once the floor is reached - fewer attempts than the
188        same connector got before it adopted the reduction.
189
190        The budget applies however the page size arrived at the floor, whether by reduction or because the
191        configured page size was already there. Only once it is spent does the failure depend on that: a page
192        size the connector's own floor blocks from ever being reduced is a configuration error, while a floor
193        the reduction walked down to means the API kept rejecting every size, which is transient.
194        """
195        # The retries come first, including on a stream whose page size started at the floor. How the page size
196        # arrived there says nothing about the response, so letting the misconfiguration branch below decide it
197        # would give the same manifest a retry budget or none depending on the user's `page_size`.
198        if self._retries_at_minimum_page_size < self._config.retries_at_minimum_page_size:
199            self._retries_at_minimum_page_size += 1
200            backoff = self._config.backoff_seconds * self._retries_at_minimum_page_size
201            LOGGER.info(
202                f"Stream {self._stream_name} cannot request a page smaller than {current_page_size} records, "
203                f"so the same page is retried unchanged in {backoff}s "
204                f"({self._retries_at_minimum_page_size} of {self._config.retries_at_minimum_page_size})."
205            )
206            self._sleep(backoff)
207            return
208
209        if self._current_page_size is None and self._config.minimum_page_size > 1:
210            # No reduction was ever applied and the connector's own floor is what blocks it, so the page size
211            # can never be reduced on this stream however the API behaves. That is a configuration error, and
212            # it is actionable: both numbers in the message are the connector's to change. It is reported once
213            # the retries above are spent, so a manifest that asked for them still gets them.
214            raise AirbyteTracedException(
215                internal_message=f"Stream {self._stream_name} has a configured page size of {current_page_size} which is not greater than the configured minimum page size of {self._config.minimum_page_size}, so it can never be reduced"
216                + (
217                    f" ({self._retries_at_minimum_page_size} retries at that size were spent first)"
218                    if self._retries_at_minimum_page_size
219                    else ""
220                )
221                + rejection,
222                message=f"The page size of stream {self._stream_name} ({current_page_size}) is already at or below "
223                f"the configured minimum of {self._config.minimum_page_size}, so the connector cannot reduce it. "
224                f"Raise the page size of the stream, or lower `minimum_page_size`.",
225                failure_type=FailureType.config_error,
226            )
227
228        raise AirbyteTracedException(
229            internal_message=f"Stream {self._stream_name} still fails with a page size of {current_page_size}, which is the smallest page size allowed by the configured minimum of {self._config.minimum_page_size}"
230            + (
231                f", after {self._retries_at_minimum_page_size} retries at that size"
232                if self._retries_at_minimum_page_size
233                else ""
234            )
235            + rejection,
236            # `transient_error`, so the only remediation is the connector's own, if it defined one.
237            message=self._with_failure_message(
238                f"The source keeps rejecting pages of stream {self._stream_name} at the smallest page size "
239                f"the connector is allowed to request ({current_page_size} records per page)."
240            ),
241            failure_type=FailureType.transient_error,
242        )
243
244    def _with_failure_message(self, message: str) -> str:
245        """
246        :return: the message followed by the connector's `failure_message`, when it defined one
247        """
248        failure_message = (self._config.failure_message or "").strip()
249        if not failure_message:
250            return message
251        return f"{message} {failure_message}"
252
253    def on_successful_page(self) -> None:
254        """
255        Called after each page that did not require a reduction.
256
257        The `max_attempts` budget restarts under both policies. It counts the reductions made *in a row*
258        without a single page succeeding, which is what separates a partition that is stuck from one that is
259        merely expensive: on a stream whose per-page cost varies - the GraphQL case this feature exists for -
260        a handful of heavy pages spread over a long partition is a healthy read, and a budget spanning the
261        whole partition would fail it at the `max_attempts + 1`-th heavy page while every reduction so far had
262        been followed by a successful page. The terminal message says the source rejected every page size the
263        connector asked for, so the budget has to mean exactly that.
264
265        The budget still terminates the read, because only a page that succeeded restarts it and only
266        `_read_pages` calls this, once per page it consumed. So between any two restarts the partition made one
267        page of progress, and the reductions that make no progress are bounded by `max_attempts`. Under `NEVER`
268        the reduced page size is never restored either, so it strictly decreases and `minimum_page_size` bounds
269        the reductions of the whole partition on its own.
270
271        Only `AFTER_SUCCESSFUL_PAGE` restores the page size. That policy is for an API that rejects the
272        configured page size on every page, so every page legitimately costs one reduction; `NEVER` keeps the
273        reduced size for the rest of the partition, which is the right behaviour when reductions are one-off.
274        """
275        self._attempts = 0
276        self._retries_at_minimum_page_size = 0
277
278        if self._config.reset_policy != PageSizeResetPolicy.AFTER_SUCCESSFUL_PAGE:
279            return
280
281        if self._current_page_size is not None:
282            LOGGER.info(
283                f"Restoring the page size of stream {self._stream_name} from {self._current_page_size} to {self._configured_page_size}."
284            )
285            self._current_page_size = None