Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -754,6 +754,9 @@ definitions:
- last_record
- page_size
- response
- stream_interval
- stream_partition
- stream_slice
examples:
- "{{ headers.link.next.cursor }}"
- "{{ last_record['key'] }}"
Expand Down Expand Up @@ -789,6 +792,9 @@ definitions:
- last_record
- page_size
- response
- stream_interval
- stream_partition
- stream_slice
examples:
- "{{ response.data.has_more is false }}"
- "{{ 'next' not in headers['link'] }}"
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@
from airbyte_cdk.sources.declarative.requesters.paginators.paginator import (
Paginator,
page_size_override_kwargs,
stream_slice_kwargs,
)
from airbyte_cdk.sources.declarative.requesters.paginators.strategies.pagination_strategy import (
PaginationStrategy,
Expand Down Expand Up @@ -146,13 +147,15 @@ def next_page_token(
last_record: Optional[Record],
last_page_token_value: Optional[Any] = None,
page_size_override: Optional[int] = None,
stream_slice: Optional[StreamSlice] = None,
) -> Optional[Mapping[str, Any]]:
next_page_token = self.pagination_strategy.next_page_token(
response=response,
last_page_size=last_page_size,
last_record=last_record,
last_page_token_value=last_page_token_value,
**page_size_override_kwargs(page_size_override),
**stream_slice_kwargs(self.pagination_strategy.next_page_token, stream_slice),
)
if next_page_token:
return {"next_page_token": next_page_token}
Expand Down Expand Up @@ -292,6 +295,7 @@ def next_page_token(
last_record: Optional[Record],
last_page_token_value: Optional[Any] = None,
page_size_override: Optional[int] = None,
stream_slice: Optional[StreamSlice] = None,
) -> Optional[Mapping[str, Any]]:
if self._page_count >= self._maximum_number_of_pages:
return None
Expand All @@ -303,6 +307,7 @@ def next_page_token(
last_record,
last_page_token_value,
**page_size_override_kwargs(page_size_override),
**stream_slice_kwargs(self._decorated.next_page_token, stream_slice),
)

def path(
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -73,5 +73,6 @@ def next_page_token(
last_record: Optional[Record],
last_page_token_value: Optional[Any],
page_size_override: Optional[int] = None,
stream_slice: Optional[StreamSlice] = None,
) -> Optional[Mapping[str, Any]]:
return {}
Original file line number Diff line number Diff line change
Expand Up @@ -2,9 +2,11 @@
# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
#

import inspect
from abc import ABC, abstractmethod
from dataclasses import dataclass
from typing import Any, Dict, Mapping, Optional
from functools import lru_cache
from typing import Any, Callable, Dict, Mapping, Optional

import requests

Expand All @@ -24,6 +26,39 @@ def page_size_override_kwargs(page_size_override: Optional[int]) -> Dict[str, An
return {"page_size_override": page_size_override} if page_size_override is not None else {}


def stream_slice_kwargs(
next_page_token: Callable[..., Any], stream_slice: Optional[StreamSlice]
) -> Dict[str, Any]:
"""
Build the `stream_slice` keyword argument only for a `next_page_token` that accepts it.

Unlike `page_size_override`, the slice is almost always present, so passing it only when it is set would still
break a paginator or pagination strategy defined outside of the CDK whose signature predates the argument. The
signature is checked instead, and a callee that does not declare `stream_slice` (or `**kwargs`) is called as
before.
"""
if stream_slice is None:
return {}
function = getattr(next_page_token, "__func__", next_page_token)
try:
accepts_stream_slice = _accepts_stream_slice(function)
except TypeError:
# Unhashable callables cannot be cached; they are rare enough to inspect on every call.
accepts_stream_slice = _accepts_stream_slice.__wrapped__(function)
return {"stream_slice": stream_slice} if accepts_stream_slice else {}


@lru_cache(maxsize=None)
def _accepts_stream_slice(function: Callable[..., Any]) -> bool:
try:
parameters = inspect.signature(function).parameters
except (TypeError, ValueError):
return False
return "stream_slice" in parameters or any(
parameter.kind == inspect.Parameter.VAR_KEYWORD for parameter in parameters.values()
)


@dataclass
class Paginator(ABC, RequestOptionsProvider):
"""
Expand All @@ -47,6 +82,7 @@ def next_page_token(
last_record: Optional[Record],
last_page_token_value: Optional[Any],
page_size_override: Optional[int] = None,
stream_slice: Optional[StreamSlice] = None,
) -> Optional[Mapping[str, Any]]:
"""
Returns the next_page_token to use to fetch the next page of records.
Expand All @@ -57,6 +93,8 @@ def next_page_token(
:param last_page_token_value: The current value of the page token made on the last request
:param page_size_override: the page size that was actually requested, when it differs from the configured
one because of a `REDUCE_PAGE_SIZE` response action
:param stream_slice: the slice the page was read for, so that a stop condition can compare the page against
the slice's own window or partition
:return: A mapping {"next_page_token": <token>} for the next page from the input response object. Returning None means there are no more pages to read in this response.
"""
pass
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@
from airbyte_cdk.sources.declarative.requesters.paginators.strategies.pagination_strategy import (
PaginationStrategy,
)
from airbyte_cdk.sources.types import Config, Record
from airbyte_cdk.sources.types import Config, Record, StreamSlice


@dataclass
Expand Down Expand Up @@ -115,6 +115,7 @@ def next_page_token(
last_record: Optional[Record],
last_page_token_value: Optional[Any] = None,
page_size_override: Optional[int] = None,
stream_slice: Optional[StreamSlice] = None,
) -> Optional[Any]:
# The next page is a cursor read from the response, so `page_size_override` does not change how the token
# is computed. It is still exposed to the interpolation context as `page_size` because a `stop_condition`
Expand All @@ -137,6 +138,7 @@ def next_page_token(
last_record=last_record,
last_page_size=last_page_size,
page_size=requested_page_size,
stream_slice=stream_slice,
)
if should_stop:
return None
Expand All @@ -147,6 +149,7 @@ def next_page_token(
last_record=last_record,
last_page_size=last_page_size,
page_size=requested_page_size,
stream_slice=stream_slice,
)
return token if token else None

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@
from airbyte_cdk.sources.declarative.requesters.paginators.strategies.pagination_strategy import (
PaginationStrategy,
)
from airbyte_cdk.sources.types import Config, Record
from airbyte_cdk.sources.types import Config, Record, StreamSlice


@dataclass
Expand Down Expand Up @@ -75,6 +75,7 @@ def next_page_token(
last_record: Optional[Record],
last_page_token_value: Optional[Any] = None,
page_size_override: Optional[int] = None,
stream_slice: Optional[StreamSlice] = None,
) -> Optional[Any]:
decoded_response = next(self.decoder.decode(response))

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@
from airbyte_cdk.sources.declarative.requesters.paginators.strategies.pagination_strategy import (
PaginationStrategy,
)
from airbyte_cdk.sources.types import Config, Record
from airbyte_cdk.sources.types import Config, Record, StreamSlice
from airbyte_cdk.utils.traced_exception import AirbyteTracedException


Expand Down Expand Up @@ -56,6 +56,7 @@ def next_page_token(
last_record: Optional[Record],
last_page_token_value: Optional[Any],
page_size_override: Optional[int] = None,
stream_slice: Optional[StreamSlice] = None,
) -> Optional[Any]:
if page_size_override is not None:
# Reachable only when the factory is bypassed: a manifest naming PageIncrement, and a
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@

import requests

from airbyte_cdk.sources.types import Record
from airbyte_cdk.sources.types import Record, StreamSlice


@dataclass
Expand All @@ -32,6 +32,7 @@ def next_page_token(
last_record: Optional[Record],
last_page_token_value: Optional[Any],
page_size_override: Optional[int] = None,
stream_slice: Optional[StreamSlice] = None,
) -> Optional[Any]:
"""
:param response: response to process
Expand All @@ -41,6 +42,8 @@ def next_page_token(
:param page_size_override: the page size that was actually requested, when it differs from the configured
one because of a `REDUCE_PAGE_SIZE` response action. Strategies that use their configured page size as
a stop condition must honor this value, else they end the pagination early and skip records.
:param stream_slice: the slice the page was read for. Only passed to strategies that declare it, so existing
strategies defined outside of the CDK keep working unchanged.
:return: next page token. Returns None if there are no more pages to fetch
"""
pass
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -9,12 +9,13 @@

from airbyte_cdk.sources.declarative.requesters.paginators.paginator import (
page_size_override_kwargs,
stream_slice_kwargs,
)
from airbyte_cdk.sources.declarative.requesters.paginators.strategies.pagination_strategy import (
PaginationStrategy,
)
from airbyte_cdk.sources.streams.concurrent.cursor import Cursor
from airbyte_cdk.sources.types import Record
from airbyte_cdk.sources.types import Record, StreamSlice


class PaginationStopCondition(ABC):
Expand Down Expand Up @@ -51,6 +52,7 @@ def next_page_token(
last_record: Optional[Record],
last_page_token_value: Optional[Any] = None,
page_size_override: Optional[int] = None,
stream_slice: Optional[StreamSlice] = None,
) -> Optional[Any]:
# We evaluate in reverse order because the assumption is that most of the APIs using data feed structure
# will return records in descending order. In terms of performance/memory, we return the records lazily
Expand All @@ -62,6 +64,7 @@ def next_page_token(
last_record,
last_page_token_value,
**page_size_override_kwargs(page_size_override),
**stream_slice_kwargs(self._delegate.next_page_token, stream_slice),
)

def get_page_size(self) -> Optional[int]:
Expand Down
18 changes: 16 additions & 2 deletions airbyte_cdk/sources/declarative/retrievers/simple_retriever.py
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,7 @@
from airbyte_cdk.sources.declarative.requesters.paginators.paginator import (
Paginator,
page_size_override_kwargs,
stream_slice_kwargs,
)
from airbyte_cdk.sources.declarative.requesters.query_properties import QueryProperties
from airbyte_cdk.sources.declarative.requesters.request_options import (
Expand Down Expand Up @@ -329,6 +330,7 @@ def _next_page_token(
last_record: Optional[Record],
last_page_token_value: Optional[Any],
page_size_override: Optional[int] = None,
stream_slice: Optional[StreamSlice] = None,
) -> Optional[Mapping[str, Any]]:
"""
Specifies a pagination strategy.
Expand All @@ -343,6 +345,7 @@ def _next_page_token(
last_record=last_record,
last_page_token_value=last_page_token_value,
**page_size_override_kwargs(page_size_override),
**stream_slice_kwargs(self._paginator.next_page_token, stream_slice),
)

def _fetch_next_page(
Expand Down Expand Up @@ -518,6 +521,7 @@ def _read_pages(
last_record=last_record,
last_page_token_value=last_page_token_value,
**page_size_override_kwargs(page_size_override),
**stream_slice_kwargs(self._next_page_token, stream_slice),
)
if not next_page_token:
break
Expand Down Expand Up @@ -629,7 +633,13 @@ def _read_pages(
last_record = record
yield record

next_page_token = self._next_page_token(response, last_page_size, last_record, None)
next_page_token = self._next_page_token(
response,
last_page_size,
last_record,
None,
**stream_slice_kwargs(self._next_page_token, stream_slice),
)
if next_page_token:
yield from self._paginate(
next_page_token,
Expand Down Expand Up @@ -669,7 +679,11 @@ def _paginate(
next_page_token.get("next_page_token") if next_page_token else None
)
next_page_token = self._next_page_token(
response, last_page_size, last_record, last_page_token_value
response,
last_page_size,
last_record,
last_page_token_value,
**stream_slice_kwargs(self._next_page_token, stream_slice),
)

if not next_page_token:
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@
from airbyte_cdk.sources.declarative.requesters.paginators.strategies.cursor_pagination_strategy import (
CursorPaginationStrategy,
)
from airbyte_cdk.sources.types import Record
from airbyte_cdk.sources.types import Record, StreamSlice


@pytest.mark.parametrize(
Expand Down Expand Up @@ -305,3 +305,61 @@ def test_given_a_page_size_when_an_expression_reads_page_size_then_accept(page_s
assert strategy.next_page_token(response, 100, None, None) == "a token"
assert strategy.next_page_token(response, 50, None, None, page_size_override=50) == "a token"
assert strategy.next_page_token(response, 49, None, None, page_size_override=50) is None


@pytest.mark.parametrize(
"oldest_created_at, expected_token",
[
pytest.param(
"2024-06-20T00:00:00Z", "page-2", id="test_oldest_run_inside_the_window_then_continue"
),
pytest.param(
"2024-05-30T00:00:00Z",
"page-2",
id="test_oldest_run_exactly_at_the_limit_then_continue",
),
pytest.param("2024-05-29T23:59:59Z", None, id="test_oldest_run_past_the_limit_then_stop"),
pytest.param(None, "page-2", id="test_oldest_run_without_created_at_then_continue"),
],
)
def test_stop_condition_can_compare_the_page_against_the_stream_slice(
oldest_created_at, expected_token
):
strategy = CursorPaginationStrategy(
cursor_value="{{ response.next }}",
stop_condition=(
"{{ str_to_datetime((response.get('runs') or [{}])[-1].get('created_at') or '9999-01-01T00:00:00Z')"
" < str_to_datetime(stream_interval['start_time']) - duration('P32D') }}"
),
config={},
parameters={},
)
response = requests.Response()
response._content = json.dumps(
{"next": "page-2", "runs": [{"created_at": oldest_created_at}]}
).encode("utf-8")
stream_slice = StreamSlice(
partition={"repository": "airbytehq/airbyte"},
cursor_slice={"start_time": "2024-07-01T00:00:00Z", "end_time": "2024-07-31T00:00:00Z"},
)

assert (
strategy.next_page_token(response, 1, None, None, stream_slice=stream_slice)
== expected_token
)


def test_cursor_value_can_read_the_stream_partition():
strategy = CursorPaginationStrategy(
cursor_value="{{ response.next ~ ':' ~ stream_partition['repository'] }}",
config={},
parameters={},
)
response = requests.Response()
response._content = json.dumps({"next": "page-2"}).encode("utf-8")
stream_slice = StreamSlice(partition={"repository": "airbytehq/airbyte"}, cursor_slice={})

assert (
strategy.next_page_token(response, 1, None, None, stream_slice=stream_slice)
== "page-2:airbytehq/airbyte"
)
Loading
Loading