diff --git a/airbyte_cdk/sources/declarative/declarative_component_schema.yaml b/airbyte_cdk/sources/declarative/declarative_component_schema.yaml index 390cd793a2..192e0e72fb 100644 --- a/airbyte_cdk/sources/declarative/declarative_component_schema.yaml +++ b/airbyte_cdk/sources/declarative/declarative_component_schema.yaml @@ -2803,6 +2803,7 @@ definitions: - RATE_LIMITED - REFRESH_TOKEN_THEN_RETRY - REDUCE_PAGE_SIZE + - SPLIT_REQUEST_WINDOW examples: - SUCCESS - FAIL @@ -2812,6 +2813,7 @@ definitions: - RATE_LIMITED - REFRESH_TOKEN_THEN_RETRY - REDUCE_PAGE_SIZE + - SPLIT_REQUEST_WINDOW failure_type: title: Failure Type description: Failure type of traced exception if a response matches the filter. @@ -4379,6 +4381,13 @@ definitions: re-issuing the page would emit them twice. A page_token_option of type RequestPath is rejected as well, because the next page is then a URL built by the API which already carries the page size. "$ref": "#/definitions/PageSizeReduction" + request_window_splitting: + description: >- + Describes how the request window is split when an error handler resolves to the + SPLIT_REQUEST_WINDOW action. Requires an incremental_sync cursor that defines cursor_granularity. + Distinct from page_size_reduction: this splits the cursor's window rather than a paginator's page + size, and applies even when the API rejects the request before returning a single page. + "$ref": "#/definitions/RequestWindowSplitting" ignore_stream_slicer_parameters_on_paginated_requests: description: If true, the partition router and incremental request options will be ignored when paginating requests. Request options set directly on the requester will not be ignored. type: boolean @@ -4559,6 +4568,42 @@ definitions: - NEVER - AFTER_SUCCESSFUL_PAGE default: NEVER + RequestWindowSplitting: + title: Request Window Splitting + description: >- + Describes how the request window is split when an error handler resolves to the SPLIT_REQUEST_WINDOW + action. The connector replaces the failing window with two smaller, ordered windows covering the same + range and reads each one fully before checkpointing. If the failing window had already emitted records, + those may be re-emitted; downstream primary-key deduplication is expected to handle that. Only supported + on a stream whose incremental_sync cursor defines cursor_granularity. Distinct from page_size_reduction, + which shrinks a paginator's page size rather than the cursor's window. + type: object + required: + - type + properties: + type: + type: string + enum: [RequestWindowSplitting] + failure_message: + title: Failure Message + description: >- + Sentence appended to the error message when the connector can no longer split the window further. + Use it to tell the user what to do for this API, for instance which filter narrows the query down. + type: string + examples: + - Narrow the sync down by selecting fewer fields on this stream. + - Lower time_window so that each request covers less data. + min_split_window: + title: Minimum Split Window + description: >- + Approximate smallest window (ISO 8601 duration) the connector will split down to - a window at or + below this size is not split further, though the result can end up a bit smaller than this, not + larger. Independent of, and usually looser than, cursor_granularity - useful when a cursor could + split further but many small requests are worse for this API than a few large ones. + type: string + examples: + - "P1D" + - "PT1H" GzipDecoder: title: gzip description: Select 'gzip' for response data that is compressed with gzip. Requires specifying an inner data type/decoder to parse the decompressed data. diff --git a/airbyte_cdk/sources/declarative/models/declarative_component_schema.py b/airbyte_cdk/sources/declarative/models/declarative_component_schema.py index fcb7e66445..f33abaaffd 100644 --- a/airbyte_cdk/sources/declarative/models/declarative_component_schema.py +++ b/airbyte_cdk/sources/declarative/models/declarative_component_schema.py @@ -18,12 +18,6 @@ class AuthFlowType(Enum): oauth1_0 = "oauth1.0" -class ScopesJoinStrategy(Enum): - space = "space" - comma = "comma" - plus = "plus" - - class BasicHttpAuthenticator(BaseModel): type: Literal["BasicHttpAuthenticator"] username: str = Field( @@ -52,6 +46,42 @@ class BearerAuthenticator(BaseModel): parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters") +class HttpMethod(Enum): + GET = "GET" + POST = "POST" + + +class QuotaStatusSource(BaseModel): + type: Literal["QuotaStatusSource"] + url: str = Field( + ..., + description="The full URL of the quota status endpoint.", + examples=[ + "https://api.github.com/rate_limit", + "{{ config.get('api_url', 'https://api.github.com') }}/rate_limit", + ], + title="URL", + ) + http_method: Optional[HttpMethod] = Field( + HttpMethod.GET, + description="The HTTP method used to fetch the quota status.", + title="HTTP Method", + ) + request_headers: Optional[Dict[str, str]] = Field( + None, + description="Additional headers to send with the quota status request.", + title="Request Headers", + ) + unavailable_status_codes: Optional[List[int]] = Field( + None, + description="Status codes from the quota status endpoint that mean quota tracking is unavailable rather than broken, such as a self-hosted deployment with rate limiting turned off. Every pool of the token whose request returned that status is then treated as untracked, so the authenticator stops waiting for quota resets, stops throttling proactively and stops rotating on exhaustion for it, while still signing requests. A token untracked this way stays untracked for the rest of the sync, because the endpoint is never consulted for it again, so a status the endpoint can also return transiently costs quota tracking for the whole run. If only some tokens return that status the others stay tracked, but they are no longer refreshed either, because the authenticator stops waiting for quota resets as soon as one token is untracked; once their counters are locally spent all traffic moves onto the untracked tokens. Rate limiting reported by ordinary responses is still handled by the stream's error handler, so one that retries 429 or 403 keeps working, and a retry rotates onto the next token; it pays the backoff the response asks for rather than the shortened one a tracked pool would get, since an untracked pool has no counters with which to argue the rejection was about that credential. Any status not listed still fails the connection, and this field never excuses a quota path missing from a response the endpoint did answer, so list only the codes the endpoint uses to report that rate limiting is not enabled. Do not list authentication or authorization statuses, since a 401 or 403 from a revoked credential would then be read as quota tracking being unavailable rather than as a credentials failure.", + examples=[[404]], + title="Unavailable Status Codes", + unique_items=True, + ) + parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters") + + class DynamicStreamCheckConfig(BaseModel): type: Literal["DynamicStreamCheckConfig"] dynamic_stream_name: str = Field( @@ -503,7 +533,7 @@ class Config: ) weight: Optional[Union[int, str]] = Field( None, - description="The weight of a request matching this matcher when acquiring a call from the rate limiter. Different endpoints can consume different amounts from a shared budget by specifying different weights. If not set, each request counts as 1.", + description="The weight of a request matching this matcher when acquiring a call from the rate limiter. Different endpoints can consume different amounts from a shared budget by specifying different weights. If not set, each request counts as 1.\n", title="Weight", ) @@ -578,157 +608,6 @@ class SessionTokenRequestBearerAuthenticator(BaseModel): type: Literal["Bearer"] -class HttpMethod(Enum): - GET = "GET" - POST = "POST" - - -class QuotaStatusSource(BaseModel): - type: Literal["QuotaStatusSource"] - url: str = Field( - ..., - description="The full URL of the quota status endpoint.", - examples=[ - "https://api.github.com/rate_limit", - "{{ config.get('api_url', 'https://api.github.com') }}/rate_limit", - ], - title="URL", - ) - http_method: Optional[HttpMethod] = Field( - HttpMethod.GET, - description="The HTTP method used to fetch the quota status.", - title="HTTP Method", - ) - request_headers: Optional[Dict[str, str]] = Field( - None, - description="Additional headers to send with the quota status request.", - title="Request Headers", - ) - unavailable_status_codes: Optional[List[int]] = Field( - None, - description="Status codes from the quota status endpoint that mean quota tracking is unavailable rather than broken, such as a self-hosted deployment with rate limiting turned off. Every pool of the token whose request returned that status is then treated as untracked, so the authenticator stops waiting for quota resets, stops throttling proactively and stops rotating on exhaustion for it, while still signing requests. A token untracked this way stays untracked for the rest of the sync, because the endpoint is never consulted for it again, so a status the endpoint can also return transiently costs quota tracking for the whole run. If only some tokens return that status the others stay tracked, but they are no longer refreshed either, because the authenticator stops waiting for quota resets as soon as one token is untracked; once their counters are locally spent all traffic moves onto the untracked tokens. Rate limiting reported by ordinary responses is still handled by the stream's error handler, so one that retries 429 or 403 keeps working, and a retry rotates onto the next token; it pays the backoff the response asks for rather than the shortened one a tracked pool would get, since an untracked pool has no counters with which to argue the rejection was about that credential. Any status not listed still fails the connection, and this field never excuses a quota path missing from a response the endpoint did answer, so list only the codes the endpoint uses to report that rate limiting is not enabled. Do not list authentication or authorization statuses, since a 401 or 403 from a revoked credential would then be read as quota tracking being unavailable rather than as a credentials failure.", - examples=[[404]], - title="Unavailable Status Codes", - unique_items=True, - ) - parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters") - - -class TokenQuota(BaseModel): - type: Literal["TokenQuota"] - name: str = Field( - ..., - description="Name of the quota pool.", - examples=["rest", "graphql"], - title="Name", - ) - remaining_path: List[str] = Field( - ..., - description="Path to the remaining call count for this pool in the quota status response.", - examples=[["resources", "core", "remaining"]], - title="Remaining Path", - ) - reset_path: List[str] = Field( - ..., - description="Path to the quota reset timestamp for this pool in the quota status response.", - examples=[["resources", "core", "reset"]], - title="Reset Path", - ) - limit_path: Optional[List[str]] = Field( - None, - description="Optional path to the total call limit for this pool in the quota status response. Used to compute the proactive throttling reserve; falls back to the initially observed remaining count when not set. Setting it on every pool is recommended so the reserve does not shrink when a sync starts with the pool already partially consumed.", - examples=[["resources", "core", "limit"]], - title="Limit Path", - ) - matchers: Optional[List[HttpRequestRegexMatcher]] = Field( - None, - description="List of matchers that classify outgoing requests into this quota pool. The first pool whose matcher matches a request is used. A pool with no matchers acts as the default pool.", - title="Matchers", - ) - remaining_header: Optional[str] = Field( - None, - description="Optional response header carrying the remaining call count for this pool. When set, the pool's counter is reconciled against this header on every response, which corrects drift caused by sharing the token with other clients, by requests in flight concurrently, or by a sync running long enough for the initial quota status read to go stale. Without it the pool is only ever seeded from the quota status endpoint.", - examples=["X-RateLimit-Remaining"], - title="Remaining Header", - ) - reset_header: Optional[str] = Field( - None, - description="Optional response header carrying the quota reset timestamp for this pool. Parsed with the same rules as `reset_path`, so epoch seconds and ISO 8601 both work. Used to tell a rolled-over quota window from the current one; a response proving the window has rolled over restores the pool to its limit. Most useful alongside `remaining_header`.", - examples=["X-RateLimit-Reset"], - title="Reset Header", - ) - limit_header: Optional[str] = Field( - None, - description="Optional response header carrying the total call limit for this pool, used to keep the proactive throttling reserve accurate as the limit changes.", - examples=["X-RateLimit-Limit"], - title="Limit Header", - ) - exhaustion_status_codes: Optional[List[int]] = Field( - None, - description="Response status codes that mean this token's pool is spent. These have two effects. A response carrying one of them but no remaining count sets the pool to zero, so the next request rotates to another token instead of waiting out the reset window. They also mark which responses may report a zero for a quota window that has already elapsed, so a rate limit whose reset header trails the value being held still stops the token being used; a zero on any other response is treated as the last call of a finished window and ignored. Leaving this empty means such trailing rejections are ignored unless their reset is within the skew tolerance of the current window. Only list codes the API uses exclusively for rate limiting -- a code that also signals other failures would park a healthy token.", - examples=[[429]], - title="Exhaustion Status Codes", - ) - parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters") - - -class RateLimitedMultipleTokenAuthenticator(BaseModel): - type: Literal["RateLimitedMultipleTokenAuthenticator"] - tokens: Union[str, List[str]] = Field( - ..., - description="The tokens to rotate between. Either an explicit list of tokens, or a single string containing multiple tokens separated by `token_delimiter`.", - examples=[ - "{{ config['credentials']['personal_access_token'] }}", - ["{{ config['token_1'] }}", "{{ config['token_2'] }}"], - ], - title="Tokens", - ) - token_delimiter: Optional[str] = Field( - ",", - description="Delimiter used to split a single token string into multiple tokens.", - title="Token Delimiter", - ) - auth_method: Optional[str] = Field( - "Bearer", - description="The prefix to prepend to the token in the auth header value (e.g. `Authorization: Bearer `).", - examples=["Bearer", "token"], - title="Auth Method", - ) - header: Optional[str] = Field( - "Authorization", - description="The name of the HTTP header in which to inject the token.", - title="Header Name", - ) - quota_status_source: QuotaStatusSource = Field( - ..., - description="Defines where to fetch each token's current quota status. Called once per token at startup and after an exhaustion wait, not per data request.", - title="Quota Status Source", - ) - quotas: List[TokenQuota] = Field( - ..., - description="Quota pools tracked per token. Each outgoing request is classified into the first pool whose matchers match the request; a pool with no matchers acts as the default. The `remaining_path` and `reset_path` locate each pool's values in the quota status response.\n", - min_items=1, - title="Quota Pools", - ) - max_wait_time: Optional[str] = Field( - "PT2H", - description="ISO 8601 duration. When all tokens are exhausted, the maximum time to wait for a quota reset before raising a transient error.", - examples=["PT2H", "PT30M", "PT{{ config.get('max_waiting_time', 120) }}M"], - title="Maximum Wait Time", - ) - budget_reserve_fraction: Optional[float] = Field( - 0.1, - description="Fraction of each token's quota to keep in reserve. When every token drops below its reserve, requests are proactively throttled to spread the remaining calls until the quota reset. Set to 0 (along with `budget_min_reserve`) to disable throttling.", - title="Budget Reserve Fraction", - ) - budget_min_reserve: Optional[int] = Field( - 50, - description="Minimum number of calls to keep in reserve per token before proactive throttling kicks in.", - title="Budget Minimum Reserve", - ) - parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters") - - class Action(Enum): SUCCESS = "SUCCESS" FAIL = "FAIL" @@ -738,6 +617,7 @@ class Action(Enum): RATE_LIMITED = "RATE_LIMITED" REFRESH_TOKEN_THEN_RETRY = "REFRESH_TOKEN_THEN_RETRY" REDUCE_PAGE_SIZE = "REDUCE_PAGE_SIZE" + SPLIT_REQUEST_WINDOW = "SPLIT_REQUEST_WINDOW" class FailureType(Enum): @@ -760,6 +640,7 @@ class HttpResponseFilter(BaseModel): "RATE_LIMITED", "REFRESH_TOKEN_THEN_RETRY", "REDUCE_PAGE_SIZE", + "SPLIT_REQUEST_WINDOW", ], title="Action", ) @@ -860,12 +741,13 @@ class JsonItemsDecoder(BaseModel): type: Literal["JsonItemsDecoder"] items_path: str = Field( ..., - description="Dot-separated path to the JSON array whose elements should be yielded as records. Uses `ijson` path syntax (e.g. `data.users`), not JSONPath syntax \u2014 do not include leading `$.` or trailing `[*]`.", + description="Dot-separated path to the JSON array whose elements should be yielded as records. Uses `ijson` path syntax (e.g. `data.users`), not JSONPath syntax — do not include leading `$.` or trailing `[*]`.", + examples=["dataByDepartmentAndSearchTerm", "dataByAsin", "data.users"], title="Items Path", ) encoding: Optional[str] = Field( "utf-8", - description="The character encoding of the JSON data. Defaults to UTF-8.", + description="Text encoding used to decode the streamed bytes before JSON parsing.", title="Encoding", ) @@ -1028,22 +910,32 @@ class NoPagination(BaseModel): type: Literal["NoPagination"] -class State(BaseModel): +class Scope(BaseModel): class Config: extra = Extra.allow - min: int - max: int + scope: str = Field(..., description="The OAuth scope string to request from the provider.") -class OAuthScope(BaseModel): +class OptionalScope(BaseModel): class Config: extra = Extra.allow - scope: str = Field( - ..., - description="The OAuth scope string to request from the provider.", - ) + scope: str = Field(..., description="The OAuth scope string to request from the provider.") + + +class ScopesJoinStrategy(Enum): + space = "space" + comma = "comma" + plus = "plus" + + +class State(BaseModel): + class Config: + extra = Extra.allow + + min: int + max: int class OauthConnectorInputSpecification(BaseModel): @@ -1065,17 +957,13 @@ class Config: examples=["user:read user:read_orders workspaces:read"], title="Scopes", ) - # NOTE: scopes, optional_scopes, and scopes_join_strategy are processed by the - # platform OAuth handler (DeclarativeOAuthSpecHandler.kt), not by the CDK runtime. - # The CDK schema defines the manifest contract; the platform reads these fields - # during the OAuth consent flow to build the authorization URL. - scopes: Optional[List[OAuthScope]] = Field( + scopes: Optional[List[Scope]] = Field( None, description="List of OAuth scope objects. When present, takes precedence over the `scope` string property.\nThe scope values are joined using the `scopes_join_strategy` (default: space) before being\nsent to the OAuth provider.", examples=[[{"scope": "user:read"}, {"scope": "user:write"}]], title="Scopes", ) - optional_scopes: Optional[List[OAuthScope]] = Field( + optional_scopes: Optional[List[OptionalScope]] = Field( None, description="Optional OAuth scope objects that may or may not be granted.", examples=[[{"scope": "admin:read"}]], @@ -1446,7 +1334,7 @@ class PageSizeReduction(BaseModel): 2, description="Divisor applied to the page size on each reduction. The new page size is floor(current page size / reduction factor).", examples=[2, 4], - gt=1.0, + gt=1, title="Reduction Factor", ) minimum_page_size: Optional[int] = Field( @@ -1467,7 +1355,7 @@ class PageSizeReduction(BaseModel): 0.5, description='Base number of seconds to wait before the page is re-issued, multiplied by the number of attempts made in a row, so the second attempt waits twice as long as the first. A REDUCE_PAGE_SIZE response never reaches the error handler\'s retry budget, backoff_strategies or a Retry-After header, so this is the only thing spacing those requests out. Raise it on an API whose error also means "we are briefly unwell" rather than only "your page is too big", since the default spaces the whole run of attempts over a few seconds.', examples=[0.5, 5], - ge=0.0, + ge=0, title="Backoff Seconds", ) retries_at_minimum_page_size: Optional[int] = Field( @@ -1493,6 +1381,25 @@ class PageSizeReduction(BaseModel): ) +class RequestWindowSplitting(BaseModel): + type: Literal["RequestWindowSplitting"] + failure_message: Optional[str] = Field( + None, + description="Sentence appended to the error message when the connector can no longer split the window further. Use it to tell the user what to do for this API, for instance which filter narrows the query down.", + examples=[ + "Narrow the sync down by selecting fewer fields on this stream.", + "Lower time_window so that each request covers less data.", + ], + title="Failure Message", + ) + min_split_window: Optional[str] = Field( + None, + description="Approximate smallest window (ISO 8601 duration) the connector will split down to - a window at or below this size is not split further, though the result can end up a bit smaller than this, not larger. Independent of, and usually looser than, cursor_granularity - useful when a cursor could split further but many small requests are worse for this API than a few large ones.", + examples=["P1D", "PT1H"], + title="Minimum Split Window", + ) + + class CsvDecoder(BaseModel): type: Literal["CsvDecoder"] encoding: Optional[str] = "utf-8" @@ -1506,7 +1413,14 @@ class AsyncJobStatusMap(BaseModel): completed: List[str] failed: List[str] timeout: List[str] - skipped: Optional[List[str]] = None + skipped: Optional[List[str]] = Field( + None, + description="Statuses that indicate the job was skipped because there is no data to return. Jobs with these statuses will not be retried and no records will be fetched.", + ) + + +class BlockSimultaneousSyncsAction(BaseModel): + type: Literal["BlockSimultaneousSyncsAction"] class ValueType(Enum): @@ -1874,6 +1788,64 @@ class AuthFlow(BaseModel): oauth_config_specification: Optional[OAuthConfigSpecification] = None +class TokenQuota(BaseModel): + type: Literal["TokenQuota"] + name: str = Field( + ..., + description="Name of the quota pool.", + examples=["rest", "graphql"], + title="Name", + ) + remaining_path: List[str] = Field( + ..., + description="Path to the remaining call count for this pool in the quota status response.", + examples=[["resources", "core", "remaining"]], + title="Remaining Path", + ) + reset_path: List[str] = Field( + ..., + description="Path to the quota reset timestamp for this pool in the quota status response.", + examples=[["resources", "core", "reset"]], + title="Reset Path", + ) + limit_path: Optional[List[str]] = Field( + None, + description="Optional path to the total call limit for this pool in the quota status response. Used to compute the proactive throttling reserve; falls back to the initially observed remaining count when not set. Setting it on every pool is recommended so the reserve does not shrink when a sync starts with the pool already partially consumed.", + examples=[["resources", "core", "limit"]], + title="Limit Path", + ) + matchers: Optional[List[HttpRequestRegexMatcher]] = Field( + None, + description="List of matchers that classify outgoing requests into this quota pool. The first pool whose matcher matches a request is used. A pool with no matchers acts as the default pool.", + title="Matchers", + ) + remaining_header: Optional[str] = Field( + None, + description="Optional response header carrying the remaining call count for this pool. When set, the pool's counter is reconciled against this header on every response, which corrects drift caused by sharing the token with other clients, by requests in flight concurrently, or by a sync running long enough for the initial quota status read to go stale. Without it the pool is only ever seeded from the quota status endpoint.", + examples=["X-RateLimit-Remaining"], + title="Remaining Header", + ) + reset_header: Optional[str] = Field( + None, + description="Optional response header carrying the quota reset timestamp for this pool. Parsed with the same rules as `reset_path`, so epoch seconds and ISO 8601 both work. Used to tell a rolled-over quota window from the current one; a response proving the window has rolled over restores the pool to its limit. Most useful alongside `remaining_header`.", + examples=["X-RateLimit-Reset"], + title="Reset Header", + ) + limit_header: Optional[str] = Field( + None, + description="Optional response header carrying the total call limit for this pool, used to keep the proactive throttling reserve accurate as the limit changes.", + examples=["X-RateLimit-Limit"], + title="Limit Header", + ) + exhaustion_status_codes: Optional[List[int]] = Field( + None, + description="Response status codes that mean this token's pool is spent. These have two effects. A response carrying one of them but no remaining count sets the pool to zero, so the next request rotates to another token instead of waiting out the reset window. They also mark which responses may report a zero for a quota window that has already elapsed, so a rate limit whose reset header trails the value being held still stops the token being used; a zero on any other response is treated as the last call of a finished window and ignored. Leaving this empty means such trailing rejections are ignored unless their reset is within the skew tolerance of the current window. Only list codes the API uses exclusively for rate limiting -- a code that also signals other failures would park a healthy token.", + examples=[[429]], + title="Exhaustion Status Codes", + ) + parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters") + + class CheckStream(BaseModel): type: Literal["CheckStream"] stream_names: Optional[List[str]] = Field( @@ -2555,6 +2527,63 @@ class ConfigAddFields(BaseModel): ) +class RateLimitedMultipleTokenAuthenticator(BaseModel): + type: Literal["RateLimitedMultipleTokenAuthenticator"] + tokens: Union[str, List[str]] = Field( + ..., + description="The tokens to rotate between. Either an explicit list of tokens, or a single string containing multiple tokens separated by `token_delimiter`.", + examples=[ + "{{ config['credentials']['personal_access_token'] }}", + ["{{ config['token_1'] }}", "{{ config['token_2'] }}"], + ], + title="Tokens", + ) + token_delimiter: Optional[str] = Field( + ",", + description="Delimiter used to split a single token string into multiple tokens.", + title="Token Delimiter", + ) + auth_method: Optional[str] = Field( + "Bearer", + description="The prefix to prepend to the token in the auth header value (e.g. `Authorization: Bearer `).", + examples=["Bearer", "token"], + title="Auth Method", + ) + header: Optional[str] = Field( + "Authorization", + description="The name of the HTTP header in which to inject the token.", + title="Header Name", + ) + quota_status_source: QuotaStatusSource = Field( + ..., + description="Defines where to fetch each token's current quota status. Called once per token at startup and after an exhaustion wait, not per data request.", + title="Quota Status Source", + ) + quotas: List[TokenQuota] = Field( + ..., + description="Quota pools tracked per token. Each outgoing request is classified into the first pool whose matchers match the request; a pool with no matchers acts as the default. The `remaining_path` and `reset_path` locate each pool's values in the quota status response.\n", + min_items=1, + title="Quota Pools", + ) + max_wait_time: Optional[str] = Field( + "PT2H", + description="ISO 8601 duration. When all tokens are exhausted, the maximum time to wait for a quota reset before raising a transient error.", + examples=["PT2H", "PT30M", "PT{{ config.get('max_waiting_time', 120) }}M"], + title="Maximum Wait Time", + ) + budget_reserve_fraction: Optional[float] = Field( + 0.1, + description="Fraction of each token's quota to keep in reserve. When every token drops below its reserve, requests are proactively throttled to spread the remaining calls until the quota reset. Set to 0 (along with `budget_min_reserve`) to disable throttling.", + title="Budget Reserve Fraction", + ) + budget_min_reserve: Optional[int] = Field( + 50, + description="Minimum number of calls to keep in reserve per token before proactive throttling kicks in.", + title="Budget Minimum Reserve", + ) + parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters") + + class CompositeErrorHandler(BaseModel): type: Literal["CompositeErrorHandler"] error_handlers: List[Union[CompositeErrorHandler, DefaultErrorHandler, CustomErrorHandler]] = ( @@ -2766,7 +2795,7 @@ class Config: api_budget: Optional[HTTPAPIBudget] = None stream_groups: Optional[Dict[str, StreamGroup]] = Field( None, - description="Groups of streams that share a common resource and should not be read simultaneously. Each group defines a set of stream references and an action that controls how concurrent reads are managed. Only applies to ConcurrentDeclarativeSource.", + description="Groups of streams that share a common resource and should not be read simultaneously. Each group defines a set of stream references and an action that controls how concurrent reads are managed. Only applies to ConcurrentDeclarativeSource.\n", title="Stream Groups", ) max_concurrent_async_job_count: Optional[Union[int, str]] = Field( @@ -2806,7 +2835,7 @@ class Config: api_budget: Optional[HTTPAPIBudget] = None stream_groups: Optional[Dict[str, StreamGroup]] = Field( None, - description="Groups of streams that share a common resource and should not be read simultaneously. Each group defines a set of stream references and an action that controls how concurrent reads are managed. Only applies to ConcurrentDeclarativeSource.", + description="Groups of streams that share a common resource and should not be read simultaneously. Each group defines a set of stream references and an action that controls how concurrent reads are managed. Only applies to ConcurrentDeclarativeSource.\n", title="Stream Groups", ) max_concurrent_async_job_count: Optional[Union[int, str]] = Field( @@ -3306,7 +3335,7 @@ class StateDelegatingStream(BaseModel): ) api_retention_period: Optional[str] = Field( None, - description="The data retention period of the incremental API (ISO8601 duration). If the cursor value is older than this retention period, the connector will automatically fall back to a full refresh to avoid data loss.\nThis is useful for APIs like Stripe Events API which only retain data for 30 days.\n * **PT1H**: 1 hour\n * **P1D**: 1 day\n * **P1W**: 1 week\n * **P1M**: 1 month\n * **P1Y**: 1 year\n * **P30D**: 30 days\n", + description="The data retention period of the incremental API (ISO8601 duration). If the cursor value is older than this retention period, the connector will automatically fall back to a full refresh to avoid data loss.\nThis is useful for APIs like Stripe Events API which only retain data for 30 days.\n* **PT1H**: 1 hour\n* **P1D**: 1 day\n* **P1W**: 1 week\n* **P1M**: 1 month\n* **P1Y**: 1 year\n* **P30D**: 30 days\n", examples=["P30D", "P90D", "P1Y"], title="API Retention Period", ) @@ -3352,6 +3381,10 @@ class SimpleRetriever(BaseModel): None, description="Describes how the page size is reduced when an error handler resolves to the REDUCE_PAGE_SIZE action. Requires a DefaultPaginator that defines both page_size_option and a pagination strategy with a page_size. Cannot be combined with query properties, a file uploader, or a parent stream read lazily through lazy_read_pointer, because in those cases records of the failing page have already been emitted and re-issuing the page would emit them twice. A page_token_option of type RequestPath is rejected as well, because the next page is then a URL built by the API which already carries the page size.", ) + request_window_splitting: Optional[RequestWindowSplitting] = Field( + None, + description="Describes how the request window is split when an error handler resolves to the SPLIT_REQUEST_WINDOW action. Requires an incremental_sync cursor that defines cursor_granularity. Distinct from page_size_reduction: this splits the cursor's window rather than a paginator's page size, and applies even when the API rejects the request before returning a single page.", + ) ignore_stream_slicer_parameters_on_paginated_requests: Optional[bool] = Field( False, description="If true, the partition router and incremental request options will be ignored when paginating requests. Request options set directly on the requester will not be ignored.", @@ -3498,20 +3531,14 @@ class AsyncRetriever(BaseModel): parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters") -class BlockSimultaneousSyncsAction(BaseModel): - type: Literal["BlockSimultaneousSyncsAction"] - - class StreamGroup(BaseModel): - streams: List[str] = Field( + streams: List[DeclarativeStream] = Field( ..., - description='List of references to streams that belong to this group. Use JSON references to stream definitions (e.g., "#/definitions/my_stream").', + description="List of references to streams that belong to this group.\n", title="Streams", ) action: BlockSimultaneousSyncsAction = Field( - ..., - description="The action to apply to streams in this group.", - title="Action", + ..., description="The action to apply to streams in this group.", title="Action" ) @@ -3536,7 +3563,7 @@ class GroupingPartitionRouter(BaseModel): underlying_partition_router: Union[ ListPartitionRouter, SubstreamPartitionRouter, - "UnionPartitionRouter", + UnionPartitionRouter, CustomPartitionRouter, ] = Field( ..., @@ -3617,8 +3644,9 @@ class DynamicDeclarativeStream(BaseModel): ConditionalStreams.update_forward_refs() FileUploader.update_forward_refs() DeclarativeStream.update_forward_refs() -SessionTokenAuthenticator.update_forward_refs() +DpathExtractor.update_forward_refs() RecordExpander.update_forward_refs() +SessionTokenAuthenticator.update_forward_refs() HttpRequester.update_forward_refs() DynamicSchemaLoader.update_forward_refs() ParentStreamConfig.update_forward_refs() @@ -3626,4 +3654,3 @@ class DynamicDeclarativeStream(BaseModel): SimpleRetriever.update_forward_refs() AsyncRetriever.update_forward_refs() GroupingPartitionRouter.update_forward_refs() -UnionPartitionRouter.update_forward_refs() 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 338e062020..a44dc22b75 100644 --- a/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py +++ b/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py @@ -604,6 +604,9 @@ PageSizeResetPolicy, ) from airbyte_cdk.sources.declarative.retrievers.pagination_tracker import PaginationTracker +from airbyte_cdk.sources.declarative.retrievers.request_window_splitting import ( + RequestWindowSplitting, +) from airbyte_cdk.sources.declarative.schema import ( ComplexFieldType, DefaultSchemaLoader, @@ -725,6 +728,23 @@ # this would be a circular import MAX_SLICES = 5 +# Smallest duration each `DatetimeParser`/`strftime` directive can distinguish, used by +# `_smallest_datetime_format_unit` to check `datetime_format` can represent `cursor_granularity`. A format with +# no entry here (e.g. a bare `%Y`/`%m`) is treated as "cannot determine" rather than rejected. +_DATETIME_FORMAT_DIRECTIVE_GRANULARITY: List[Tuple[str, datetime.timedelta]] = [ + ("%epoch_microseconds", datetime.timedelta(microseconds=1)), + ("%s_as_float", datetime.timedelta(microseconds=1)), + ("%f", datetime.timedelta(microseconds=1)), + ("%_ms", datetime.timedelta(milliseconds=1)), + ("%ms", datetime.timedelta(milliseconds=1)), + ("%S", datetime.timedelta(seconds=1)), + ("%s", datetime.timedelta(seconds=1)), + ("%M", datetime.timedelta(minutes=1)), + ("%H", datetime.timedelta(hours=1)), + ("%I", datetime.timedelta(hours=1)), + ("%d", datetime.timedelta(days=1)), +] + LOGGER = logging.getLogger(f"airbyte.model_to_component_factory") @@ -1227,6 +1247,9 @@ def create_session_token_authenticator( self._reject_reduce_page_size_action( model.login_requester, f"`login_requester` of the SessionTokenAuthenticator of {name}" ) + self._reject_split_request_window_action( + model.login_requester, f"`login_requester` of the SessionTokenAuthenticator of {name}" + ) decoder = ( self._create_component_from_model(model=model.decoder, config=config) if model.decoder @@ -3875,6 +3898,21 @@ def _get_url(req: Requester) -> str: f"stream's `lazy_read_pointer` for stream {name}." ) + if reads_parent_stream_lazily and ( + model.request_window_splitting + or self._uses_split_request_window_action( + getattr(model.requester, "error_handler", None) + ) + ): + # Same reasoning as the page_size_reduction check above: LazySimpleRetriever reads records + # embedded in the parent's own pages rather than requesting a window of its own, so there is + # nothing here for a split child window to be read from. + raise ValueError( + f"`request_window_splitting` and the SPLIT_REQUEST_WINDOW response action are not " + f"supported when reading a parent stream lazily. Remove either the request window " + f"splitting or the parent stream's `lazy_read_pointer` for stream {name}." + ) + if reads_parent_stream_lazily and not bool( self._connector_state_manager.get_stream_state(name, None) ): @@ -3915,6 +3953,9 @@ def _get_url(req: Requester) -> str: ): raise ValueError("PaginationResetLimits are not supported while having record filter.") + request_window_splitting = self._create_request_window_splitting( + model, name, cursor, incremental_sync, query_properties, file_uploader + ) return SimpleRetriever( name=name, paginator=paginator, @@ -3933,6 +3974,12 @@ def _get_url(req: Requester) -> str: page_size_reduction=self._create_page_size_reduction( model, name, query_properties, file_uploader ), + request_window_splitting=request_window_splitting, + request_window_splitter=( + cursor.split_request_window + if request_window_splitting and hasattr(cursor, "split_request_window") + else None + ), post_pagination_filter=post_pagination_filter, parameters=model.parameters or {}, ) @@ -4238,6 +4285,209 @@ def _uses_reduce_page_size_action(self, error_handler: Any) -> bool: # A CustomErrorHandler can return any action and we cannot inspect it, so we do not validate it. return False + def _create_request_window_splitting( + self, + model: SimpleRetrieverModel, + name: str, + cursor: Optional[Cursor], + incremental_sync: Optional[ + Union[IncrementingCountCursorModel, DatetimeBasedCursorModel] + ] = None, + query_properties: Optional[QueryProperties] = None, + file_uploader: Optional[DefaultFileUploader] = None, + ) -> Optional[RequestWindowSplitting]: + # A CustomRequester does not necessarily define an error handler. A CustomRequester that does define + # one keeps it as a raw dict rather than a typed model, so this returns False for it as well. + error_handler = getattr(model.requester, "error_handler", None) + uses_action = self._uses_split_request_window_action(error_handler) + if uses_action and not model.request_window_splitting: + raise ValueError( + f"Stream {name} has a response filter with the SPLIT_REQUEST_WINDOW action but the " + f"retriever does not define `request_window_splitting`. Add a `request_window_splitting` " + f"block to the retriever." + ) + + if not model.request_window_splitting: + return None + + if not uses_action: + # Not raised: a CustomErrorHandler can resolve to SPLIT_REQUEST_WINDOW without us being able to + # see it, and custom code can raise RequestWindowSplitRequiredException directly without + # going through any error handler at all - so the only safe reaction to a block we cannot tie to + # a known trigger is a warning. + LOGGER.warning( + f"Stream {name} defines `request_window_splitting` but no response filter with the " + f"SPLIT_REQUEST_WINDOW action was found on its requester. The request window will never " + f"be split unless a custom error handler or custom code resolves to that action." + ) + + self._validate_request_window_splitting_is_supported( + name, + cursor, + incremental_sync, + model.request_window_splitting.min_split_window, + query_properties, + file_uploader, + ) + + return RequestWindowSplitting( + failure_message=model.request_window_splitting.failure_message, + min_split_window=parse_duration(model.request_window_splitting.min_split_window) + if model.request_window_splitting.min_split_window + else None, + ) + + def _validate_request_window_splitting_is_supported( + self, + name: str, + cursor: Optional[Cursor], + incremental_sync: Optional[ + Union[IncrementingCountCursorModel, DatetimeBasedCursorModel] + ] = None, + min_split_window: Optional[str] = None, + query_properties: Optional[QueryProperties] = None, + file_uploader: Optional[DefaultFileUploader] = None, + ) -> None: + """ + Request window splitting replaces a failing slice with smaller children derived from the stream's own + cursor, so it only makes sense on a stream whose cursor exposes a `split_request_window` method and it + is rejected for the same structural reasons `page_size_reduction` is: `additional_query_properties` and + a `file_uploader` both mean records of the failing window may already have been emitted by the time the + split is requested, from inside the record generator rather than from a page boundary this retriever + controls. + """ + if not hasattr(cursor, "split_request_window"): + raise ValueError( + f"`request_window_splitting` requires an incremental cursor that supports window splitting on " + f"stream {name}. Found {type(cursor).__name__ if cursor is not None else 'no cursor'}: this is " + f"only supported today for a `DatetimeBasedCursor` with `cursor_granularity` set and no " + f"partition router combining multiple cursors." + ) + + if not ( + isinstance(incremental_sync, DatetimeBasedCursorModel) + and incremental_sync.cursor_granularity + ): + # `cursor_granularity` is both the smallest window the connector will ever request and what keeps + # two child windows from overlapping at their shared edge, so `split_request_window` cannot split + # at all without it - checked here, on the manifest model, rather than by reaching into the + # cursor's private state from outside the class it belongs to. + raise ValueError( + f"`request_window_splitting` requires `cursor_granularity` on the `DatetimeBasedCursor` of " + f"stream {name}: without it there is no smallest window to stop splitting at, and no way to " + f"keep two child windows from overlapping at their shared edge." + ) + + if incremental_sync.is_client_side_incremental: + # A client-side-incremental cursor filters records after they are read rather than sending the + # window as request parameters, so every child window sends the exact same request as its parent: + # splitting changes nothing, and the sync would keep splitting all the way to `min_split_window` + # or the safety-net depth before failing. + raise ValueError( + f"`request_window_splitting` cannot be used together with `is_client_side_incremental` on " + f"stream {name}: the window is filtered client-side rather than sent to the API, so splitting " + f"it would not change the request and could not resolve a SPLIT_REQUEST_WINDOW response." + ) + + parsed_cursor_granularity = parse_duration(incremental_sync.cursor_granularity) + if isinstance( + parsed_cursor_granularity, datetime.timedelta + ) and parsed_cursor_granularity <= datetime.timedelta(0): + # Passes the truthy check above (the string is non-empty), but a zero-or-negative granularity can + # never produce a child strictly smaller than the parent. + raise ValueError( + f"`request_window_splitting` requires a `cursor_granularity` greater than zero on stream " + f"{name}: `{incremental_sync.cursor_granularity}` parses to a zero-length duration, which can " + f"never produce a smaller child window." + ) + + smallest_format_unit = self._smallest_datetime_format_unit(incremental_sync.datetime_format) + if ( + smallest_format_unit is not None + and isinstance(parsed_cursor_granularity, datetime.timedelta) + and parsed_cursor_granularity < smallest_format_unit + ): + # Two children split closer together than `datetime_format` can render would format to the same + # boundary value instead of splitting. + raise ValueError( + f"`request_window_splitting` requires `cursor_granularity` to be no finer than what " + f"`datetime_format` can represent on stream {name}: `datetime_format` " + f"{incremental_sync.datetime_format!r} cannot represent a value finer than " + f"{smallest_format_unit}, but `cursor_granularity` {incremental_sync.cursor_granularity!r} is " + f"finer than that. Use a `datetime_format` precise enough to represent `cursor_granularity`, " + f"or a coarser `cursor_granularity`." + ) + + if min_split_window: + parsed_min_split_window = parse_duration(min_split_window) + if not isinstance( + parsed_min_split_window, datetime.timedelta + ) or parsed_min_split_window <= datetime.timedelta(0): + # A duration with years or months (e.g. `P1M`) parses to an `isodate.Duration`, which has no + # fixed length to compare a window's span against - `split_request_window` would raise + # mid-sync trying to. A zero-or-negative duration would never stop splitting. + raise ValueError( + f"`min_split_window` must be a positive duration expressible as a fixed number of " + f"days/hours/minutes/seconds on stream {name}: `{min_split_window}` is not." + ) + + if query_properties: + raise ValueError( + f"`request_window_splitting` cannot be used together with query properties on stream {name}. " + f"Records from the earlier property chunks may have already been emitted when a chunk asks to " + f"split the window, so splitting and re-reading it could duplicate them." + ) + + if file_uploader: + raise ValueError( + f"`request_window_splitting` cannot be used together with a `file_uploader` on stream {name}. " + f"The file uploader sends one request per record from inside the window's record generator, so " + f"a split requested partway through could re-emit records already yielded by it." + ) + + @staticmethod + def _smallest_datetime_format_unit(datetime_format: str) -> Optional[datetime.timedelta]: + """ + :return: the finest duration `datetime_format` can distinguish, per + `_DATETIME_FORMAT_DIRECTIVE_GRANULARITY`, or `None` if it contains none of those directives. + """ + matches = [ + granularity + for directive, granularity in _DATETIME_FORMAT_DIRECTIVE_GRANULARITY + if directive in datetime_format + ] + return min(matches) if matches else None + + def _reject_split_request_window_action(self, requester: Any, description: str) -> None: + """ + `error_handler` is defined on `HttpRequester`, which is referenced by requesters that have no window of + their own to split, so SPLIT_REQUEST_WINDOW is schema-legal in places where nothing can honor it. Only + the main requester of a `SimpleRetriever` reads a cursor-sliced window, so every other requester is + rejected here rather than surfacing the exception mid-sync as a generic failure. + """ + if requester is None: + return + if self._uses_split_request_window_action(getattr(requester, "error_handler", None)): + raise ValueError( + f"The SPLIT_REQUEST_WINDOW response action is not supported on the {description}: only the " + f"main requester of a SimpleRetriever reads a window that can be split. Use a different " + f"action on that error handler." + ) + + def _uses_split_request_window_action(self, error_handler: Any) -> bool: + if isinstance(error_handler, CompositeErrorHandlerModel): + return any( + self._uses_split_request_window_action(nested) + for nested in error_handler.error_handlers + ) + if isinstance(error_handler, DefaultErrorHandlerModel): + return any( + response_filter.action == HttpResponseFilterActionModel.SPLIT_REQUEST_WINDOW + for response_filter in error_handler.response_filters or [] + ) + # A CustomErrorHandler can return any action and we cannot inspect it, so we do not validate it. + return False + def _create_pagination_tracker_factory( self, model: Optional[PaginationResetModel], cursor: Cursor ) -> Callable[[], PaginationTracker]: @@ -4517,6 +4767,10 @@ def create_async_retriever( getattr(model, requester_field, None), f"`{requester_field}` of the AsyncRetriever of stream {name}", ) + self._reject_split_request_window_action( + getattr(model, requester_field, None), + f"`{requester_field}` of the AsyncRetriever of stream {name}", + ) def _get_download_retriever( requester: Requester, extractor: RecordExtractor, _decoder: Decoder @@ -5142,6 +5396,7 @@ def create_file_uploader( self, model: FileUploaderModel, config: Config, **kwargs: Any ) -> FileUploader: self._reject_reduce_page_size_action(model.requester, "requester of a `file_uploader`") + self._reject_split_request_window_action(model.requester, "requester of a `file_uploader`") name = "File Uploader" requester = self._create_component_from_model( model=model.requester, diff --git a/airbyte_cdk/sources/declarative/requesters/error_handlers/composite_error_handler.py b/airbyte_cdk/sources/declarative/requesters/error_handlers/composite_error_handler.py index 5cf9532ec8..2ef35a06ff 100644 --- a/airbyte_cdk/sources/declarative/requesters/error_handlers/composite_error_handler.py +++ b/airbyte_cdk/sources/declarative/requesters/error_handlers/composite_error_handler.py @@ -72,6 +72,7 @@ def interpret_response( ResponseAction.IGNORE, ResponseAction.RESET_PAGINATION, ResponseAction.REDUCE_PAGE_SIZE, + ResponseAction.SPLIT_REQUEST_WINDOW, ]: return matched_error_resolution diff --git a/airbyte_cdk/sources/declarative/requesters/paginators/default_paginator.py b/airbyte_cdk/sources/declarative/requesters/paginators/default_paginator.py index 839a4a678a..053d609193 100644 --- a/airbyte_cdk/sources/declarative/requesters/paginators/default_paginator.py +++ b/airbyte_cdk/sources/declarative/requesters/paginators/default_paginator.py @@ -268,6 +268,10 @@ class PaginatorTestReadDecorator(Paginator): class (`connector_builder/test_reader/reader.py::_has_reached_limit`). The two counts agree because `HttpClient` logs a response resolving to `REDUCE_PAGE_SIZE` as an auxiliary request, which the Builder does not turn into a page. + + `get_initial_token()` also resets this count, and `SimpleRetriever._read_pages` calls it again from scratch + for every child window a `SPLIT_REQUEST_WINDOW` response produces. A split slice can therefore fetch up to + `maximum_number_of_pages` per child rather than sharing one budget across the whole original partition. """ _PAGE_COUNT_BEFORE_FIRST_NEXT_CALL = 1 diff --git a/airbyte_cdk/sources/declarative/retrievers/request_window_splitting.py b/airbyte_cdk/sources/declarative/retrievers/request_window_splitting.py new file mode 100644 index 0000000000..1431928b92 --- /dev/null +++ b/airbyte_cdk/sources/declarative/retrievers/request_window_splitting.py @@ -0,0 +1,21 @@ +# +# Copyright (c) 2026 Airbyte, Inc., all rights reserved. +# + +import datetime +from dataclasses import dataclass +from typing import Optional + + +@dataclass(frozen=True) +class RequestWindowSplitting: + """ + Configuration for the `request_window_splitting` retriever field - created once per stream and shared, + mirroring `PageSizeReduction`. Unlike `PageSizeReducer`, there is no per-partition mutable counterpart: a + split window is either split again (recursion) or the sync fails, so a plain config object is enough. + """ + + failure_message: Optional[str] = None + # Independent of, and typically looser than, the cursor's own granularity floor - useful when a cursor + # could split further but the API's rate limit makes many small requests worse than a few large ones. + min_split_window: Optional[datetime.timedelta] = None diff --git a/airbyte_cdk/sources/declarative/retrievers/simple_retriever.py b/airbyte_cdk/sources/declarative/retrievers/simple_retriever.py index 21f56347a9..26f0d616c3 100644 --- a/airbyte_cdk/sources/declarative/retrievers/simple_retriever.py +++ b/airbyte_cdk/sources/declarative/retrievers/simple_retriever.py @@ -2,6 +2,7 @@ # Copyright (c) 2025 Airbyte, Inc., all rights reserved. # +import datetime import json import logging from collections import defaultdict @@ -49,6 +50,9 @@ PageSizeReduction, ) from airbyte_cdk.sources.declarative.retrievers.pagination_tracker import PaginationTracker +from airbyte_cdk.sources.declarative.retrievers.request_window_splitting import ( + RequestWindowSplitting, +) from airbyte_cdk.sources.declarative.retrievers.retriever import Retriever from airbyte_cdk.sources.declarative.stream_slicers.stream_slicer import StreamSlicer from airbyte_cdk.sources.source import ExperimentalClassWarning @@ -60,6 +64,10 @@ from airbyte_cdk.sources.streams.http.pagination_reset_exception import ( PaginationResetRequiredException, ) +from airbyte_cdk.sources.streams.http.request_window_split_exception import ( + RequestWindowSplitNotSupportedException, + RequestWindowSplitRequiredException, +) from airbyte_cdk.sources.types import Config, Record, StreamSlice from airbyte_cdk.utils.mapping_helpers import combine_mappings from airbyte_cdk.utils.traced_exception import AirbyteTracedException @@ -67,6 +75,10 @@ FULL_REFRESH_SYNC_COMPLETE_KEY = "__ab_full_refresh_sync_complete" LOGGER = logging.getLogger("airbyte") +# A defense-in-depth bound independent of any `request_window_splitter`'s own no-progress guard: protects +# against a misbehaving custom cursor whose children don't actually shrink the window. Not user-configurable. +_MAX_REQUEST_WINDOW_SPLIT_DEPTH = 10 + @dataclass class SimpleRetriever(Retriever): @@ -97,6 +109,14 @@ class SimpleRetriever(Retriever): `_read_pages` creates per call, so the retriever and its paginator - both shared by every partition of the stream, read concurrently - stay stateless. When `page_size_reduction` is `None` no reducer is created and `_read_pages` keeps its previous behaviour + request_window_splitting (Optional[RequestWindowSplitting]): Policy applied when an error handler + resolves to `ResponseAction.SPLIT_REQUEST_WINDOW`, or when `RequestWindowSplitRequiredException` + is raised directly by custom code. `None` disables request-window splitting entirely: `read_records` + converts the exception into `RequestWindowSplitNotSupportedException` instead of attempting a + split it cannot perform. + request_window_splitter: Bound method - normally the stream's own cursor's `split_request_window` - + asked to replace a failing `StreamSlice` with smaller children. `None` has the same effect as + `request_window_splitting` being `None`. """ requester: Requester @@ -122,6 +142,10 @@ class SimpleRetriever(Retriever): ) post_pagination_filter: Optional[ClientSideIncrementalRecordFilterDecorator] = None page_size_reduction: Optional[PageSizeReduction] = None + request_window_splitting: Optional[RequestWindowSplitting] = None + request_window_splitter: Optional[ + Callable[[StreamSlice, Optional[datetime.timedelta]], Optional[List[StreamSlice]]] + ] = None def __post_init__(self, parameters: Mapping[str, Any]) -> None: self._paginator = self.paginator or NoPagination(parameters=parameters) @@ -547,27 +571,158 @@ def read_records( :return: The records read from the API source """ _slice = stream_slice or StreamSlice(partition={}, cursor_slice={}) # None-check + yield from self._read_records_or_split_request_window( + records_schema, _slice, _slice, depth=0 + ) + def _read_records_or_split_request_window( + self, + records_schema: Mapping[str, Any], + stream_slice: StreamSlice, + original_slice: StreamSlice, + depth: int, + ) -> Iterable[StreamData]: + """ + Read `stream_slice` to completion, replacing it with smaller children and recursing into each of them in + turn when a `RequestWindowSplitRequiredException` is raised while reading it. + + Each recursive call re-enters this method - and, through it, `_read_pages` - from scratch for the child + slice it is given, so every child gets its own paginator token, `PaginationTracker`, and + `PageSizeReducer` for free: nothing here needs to reset that state explicitly. Because the whole + recursion lives inside the single generator `DeclarativePartition.read()` consumes, + `PartitionReader.process_partition()` only calls `cursor.close_partition()` once, after this generator is + fully exhausted - so a failure anywhere in the recursion (a child that cannot be read, or a window that + cannot be split any further) propagates out without ever checkpointing the original partition. + + `original_slice` is threaded through the recursion unchanged so every record yielded, however deep the + recursion went to produce it, is re-stamped with the partition's own slice rather than the child slice it + was actually read against - see `_reassociate_with_original_slice`. `depth` bounds the recursion + independently of `request_window_splitter`'s own no-progress guard: it is enforced here so a + misbehaving custom cursor cannot recurse indefinitely regardless of what that guard does or does not + catch. + """ record_generator = partial( self._parse_records, stream_slice=stream_slice, records_schema=records_schema, ) - records: Iterable[Mapping[str, Any]] = self._read_pages(record_generator, _slice) - if self.post_pagination_filter: - # A data feed paginates until it reaches a record older than the cursor, so the page that triggers the stop - # condition still holds already-synced records. Those are filtered here rather than in the record selector - # so that the paginator keeps seeing the whole page: the stop condition is evaluated on the last record of - # the page, which is precisely one of the records being dropped. Two consequences of filtering this late: - # the pagination tracker observes the dropped records, and a `file_uploader` on the record selector has - # already uploaded their files by the time they are dropped. - records = self.post_pagination_filter.filter_records( - records, - # the filter is only used for its cursor comparison, which does not read the stream state - stream_state={}, - stream_slice=_slice, + + emitted_count = 0 + try: + records: Iterable[Mapping[str, Any]] = self._read_pages(record_generator, stream_slice) + if self.post_pagination_filter: + # A data feed paginates until it reaches a record older than the cursor, so the page that triggers the stop + # condition still holds already-synced records. Those are filtered here rather than in the record selector + # so that the paginator keeps seeing the whole page: the stop condition is evaluated on the last record of + # the page, which is precisely one of the records being dropped. Two consequences of filtering this late: + # the pagination tracker observes the dropped records, and a `file_uploader` on the record selector has + # already uploaded their files by the time they are dropped. + records = self.post_pagination_filter.filter_records( + records, + # the filter is only used for its cursor comparison, which does not read the stream state + stream_state={}, + stream_slice=stream_slice, + ) + for record in records: + emitted_count += 1 + yield self._reassociate_with_original_slice(record, original_slice) + except RequestWindowSplitRequiredException as exception: + if self.request_window_splitting is None or self.request_window_splitter is None: + raise RequestWindowSplitNotSupportedException(stream_name=self.name) from exception + + if emitted_count: + # Splitting and re-reading the window re-emits these records: a failed partition is never + # checkpointed, so the next attempt would re-emit them anyway. Always replaying makes forward + # progress instead of retrying the same oversized window forever; primary-key dedup downstream + # handles the duplicates either way. + LOGGER.warning( + f"Stream {self.name} already emitted {emitted_count} record(s) from {stream_slice} before " + f"it was rejected; splitting and re-reading the window may re-emit them." + ) + + if depth >= _MAX_REQUEST_WINDOW_SPLIT_DEPTH: + # Logged separately from the exception below so it stays visible even if the trace message's + # internal_message isn't surfaced by whatever catches it. + LOGGER.warning( + f"Stream {self.name} hit the maximum request window split depth " + f"({_MAX_REQUEST_WINDOW_SPLIT_DEPTH}) while splitting {original_slice}." + ) + raise AirbyteTracedException( + internal_message=f"Stream {self.name} exceeded the maximum request window split depth of {_MAX_REQUEST_WINDOW_SPLIT_DEPTH} while splitting {original_slice}", + # `transient_error`, so the only remediation is the connector's own, if it defined one. + message=self._with_request_window_failure_message( + f"Stream {self.name} could not split its request window to a size the API accepts " + f"within {_MAX_REQUEST_WINDOW_SPLIT_DEPTH} splits. This usually means the stream's " + f"cursor is not actually shrinking the window on each split; if it uses a custom " + f"cursor, check its `split_request_window` implementation." + ), + failure_type=FailureType.transient_error, + ) from exception + + children = self.request_window_splitter( + stream_slice, self.request_window_splitting.min_split_window + ) + if children is None: + min_split_window_note = ( + f", its configured `min_split_window` ({self.request_window_splitting.min_split_window})" + if self.request_window_splitting.min_split_window + else "" + ) + raise AirbyteTracedException( + internal_message=f"Stream {self.name} could not split its request window {stream_slice} any further", + # `transient_error` regardless of how the triggering response was classified: exhaustion is + # never the user's fault, matching `PageSizeReducer`'s equivalent exhaustion branch. + message=self._with_request_window_failure_message( + f"The API kept rejecting stream {self.name}'s request window even at the smallest window " + f"its cursor allows{min_split_window_note}, or splitting the window further would not " + f"make progress." + ), + failure_type=FailureType.transient_error, + ) from exception + + LOGGER.info( + f"Stream {self.name}: the API rejected request window {stream_slice} (split depth {depth}); " + f"reducing it to {children} and reading each in turn." + ) + for child in children: + yield from self._read_records_or_split_request_window( + records_schema, child, original_slice, depth + 1 + ) + + def _with_request_window_failure_message(self, message: str) -> str: + """ + :return: the message followed by `request_window_splitting`'s configured `failure_message`, when it + defined one - mirroring `PageSizeReducer._with_failure_message` + """ + assert self.request_window_splitting is not None + failure_message = (self.request_window_splitting.failure_message or "").strip() + if not failure_message: + return message + return f"{message} {failure_message}" + + @staticmethod + def _reassociate_with_original_slice( + record: StreamData, original_slice: StreamSlice + ) -> StreamData: + """ + Records read while recursing into a split child window are parsed against that child's `StreamSlice`, + so `RecordSelector` stamps them with the child as `associated_slice`. `DeclarativePartition.read()` + passes an already-built `Record` through unchanged rather than re-wrapping it, so left uncorrected the + child slice - not the partition's own slice - is what `ConcurrentCursor.observe()` would key its + per-partition bookkeeping by. `close_partition()` always looks that bookkeeping up by the partition's own + slice, so it would never find it, silently losing the precise most-recently-observed cursor value for a + split partition (recoverable in practice today only because callers of that lookup fall back to the + slice's end boundary when it is missing). Re-stamping every record with `original_slice` here restores + the same tracking a non-split read already gets, rather than relying on that fallback. + """ + if isinstance(record, Record) and record.associated_slice is not original_slice: + return Record( + data=record.data, + stream_name=record.stream_name, + associated_slice=original_slice, + file_reference=record.file_reference, ) - yield from records + return record def _parse_records( self, diff --git a/airbyte_cdk/sources/streams/concurrent/cursor.py b/airbyte_cdk/sources/streams/concurrent/cursor.py index 11eaad235c..ce4e3f644e 100644 --- a/airbyte_cdk/sources/streams/concurrent/cursor.py +++ b/airbyte_cdk/sources/streams/concurrent/cursor.py @@ -19,6 +19,7 @@ Union, ) +from airbyte_cdk.models import FailureType from airbyte_cdk.sources.connector_state_manager import ConnectorStateManager from airbyte_cdk.sources.message import MessageRepository, NoopMessageRepository from airbyte_cdk.sources.streams import NO_CURSOR_STATE_KEY @@ -30,6 +31,7 @@ AbstractStreamStateConverter, ) from airbyte_cdk.sources.types import Record, StreamSlice +from airbyte_cdk.utils.traced_exception import AirbyteTracedException LOGGER = logging.getLogger("airbyte") @@ -607,6 +609,96 @@ def reduce_slice_range(self, stream_slice: StreamSlice) -> StreamSlice: else: return stream_slice + def split_request_window( + self, stream_slice: StreamSlice, min_split_window: Optional[datetime.timedelta] = None + ) -> Optional[List[StreamSlice]]: + """ + Split `stream_slice` in half along this cursor's own boundary fields, granularity, and output format. + + Returns `None` when the slice cannot be split any further: either it spans at most one + `cursor_granularity` unit already (the floor), it is already at or below `min_split_window` (a + connector-configured floor independent of and typically looser than `cursor_granularity`), the + boundaries are reversed, or no `cursor_granularity` was configured at all (there is then no way to know + the smallest addressable unit, or how to keep two children from overlapping at their shared edge). + Raises `AirbyteTracedException` instead if a boundary is missing or unparseable: that is a bug in how + the slice was produced, not the API rejecting a window, so it must not be reported as one. This is the + implementation `SimpleRetriever.request_window_splitter` binds to (see + `airbyte_cdk.sources.declarative.retrievers.simple_retriever`) that reuses the same + parsing/formatting/comparison operations `_split_per_slice_range` already relies on, so it works for + any `CursorValueType`/`GapType` pair this cursor was built with (datetime/timedelta, or int/int), + without a second, duplicate implementation of that logic living outside the cursor. + """ + if not self._cursor_granularity: + return None + + start_field, end_field = self._slice_boundary_fields_wrapper + try: + start_value = self._connector_state_converter.parse_value( + stream_slice.cursor_slice[start_field] + ) + end_value = self._connector_state_converter.parse_value( + stream_slice.cursor_slice[end_field] + ) + except (KeyError, ValueError, TypeError) as exception: + # A missing or unparseable boundary means this slice was produced (or passed in) wrong - a bug in + # slice generation, not something the API did. Raising here, rather than returning `None` (which + # `SimpleRetriever` reports as "the API kept rejecting..."), keeps that distinction visible instead + # of blaming the API for our own state. + raise AirbyteTracedException( + internal_message=( + f"Cannot split {stream_slice}: boundary field {start_field!r} or {end_field!r} is missing " + f"or not parseable: {exception}" + ), + message=( + f"Stream {self._stream_name} produced a request window with a missing or invalid " + f"{start_field!r}/{end_field!r} boundary and could not split it further." + ), + failure_type=FailureType.system_error, + ) from exception + + if start_value >= end_value: + # Reversed or degenerate boundaries: nothing safe to split. + return None + + span = end_value - start_value # type: ignore[operator] # concretely a GapType (timedelta/int) at runtime + if min_split_window is not None and span <= min_split_window: # type: ignore[operator] + return None + + total_units = span // self._cursor_granularity + 1 # type: ignore[operator] # inclusive count of granularity units + if total_units <= 1: + # Already at the minimum granularity and still rejected; nothing left to split. + return None + + half_units = total_units // 2 + first_end = ( + start_value + (half_units * self._cursor_granularity) - self._cursor_granularity # type: ignore[operator] + ) + second_start = first_end + self._cursor_granularity + + if not (start_value <= first_end < second_start <= end_value): + # No-progress guard, independent of the granularity check above: refuse to return children that + # are not strictly ordered and non-degenerate, rather than risk an identical or overlapping split. + return None + + return [ + StreamSlice( + partition=stream_slice.partition, + cursor_slice={ + start_field: self._connector_state_converter.output_format(start_value), + end_field: self._connector_state_converter.output_format(first_end), + }, + extra_fields=stream_slice.extra_fields, + ), + StreamSlice( + partition=stream_slice.partition, + cursor_slice={ + start_field: self._connector_state_converter.output_format(second_start), + end_field: self._connector_state_converter.output_format(end_value), + }, + extra_fields=stream_slice.extra_fields, + ), + ] + def get_cursor_datetime_from_state( self, stream_state: Mapping[str, Any] ) -> datetime.datetime | None: diff --git a/airbyte_cdk/sources/streams/http/error_handlers/response_models.py b/airbyte_cdk/sources/streams/http/error_handlers/response_models.py index 10416c1a7e..4cda91ffb8 100644 --- a/airbyte_cdk/sources/streams/http/error_handlers/response_models.py +++ b/airbyte_cdk/sources/streams/http/error_handlers/response_models.py @@ -20,6 +20,7 @@ class ResponseAction(Enum): RATE_LIMITED = "RATE_LIMITED" REFRESH_TOKEN_THEN_RETRY = "REFRESH_TOKEN_THEN_RETRY" REDUCE_PAGE_SIZE = "REDUCE_PAGE_SIZE" + SPLIT_REQUEST_WINDOW = "SPLIT_REQUEST_WINDOW" @dataclass diff --git a/airbyte_cdk/sources/streams/http/http_client.py b/airbyte_cdk/sources/streams/http/http_client.py index f79276a7aa..12f792c47a 100644 --- a/airbyte_cdk/sources/streams/http/http_client.py +++ b/airbyte_cdk/sources/streams/http/http_client.py @@ -53,6 +53,9 @@ rate_limit_default_backoff_handler, user_defined_backoff_handler, ) +from airbyte_cdk.sources.streams.http.request_window_split_exception import ( + RequestWindowSplitRequiredException, +) # Imported from the leaf module rather than the package: `protocols` pulls in nothing from the # CDK, so this import cannot cycle no matter what else lands in `requests_native_auth` -- an @@ -475,20 +478,36 @@ def _send( and self._message_repository is not None ): formatter = log_formatter - # A response resolving to REDUCE_PAGE_SIZE is not a page of the stream: the retriever discards it - # and re-issues the same page with a smaller page size. Logging it as an auxiliary request keeps it - # visible in the Connector Builder while keeping it out of the per-slice page count, which would - # otherwise report "limit reached" on a read that only retried. - log_as_auxiliary = error_resolution.response_action == ResponseAction.REDUCE_PAGE_SIZE + # A response resolving to REDUCE_PAGE_SIZE or SPLIT_REQUEST_WINDOW is not a page of the stream: + # the retriever discards it and re-issues the request with a smaller page size or a narrower + # window. Logging it as an auxiliary request keeps it visible in the Connector Builder while + # keeping it out of the per-slice page count, which would otherwise report "limit reached" on a + # read that only retried. + log_as_auxiliary = error_resolution.response_action in ( + ResponseAction.REDUCE_PAGE_SIZE, + ResponseAction.SPLIT_REQUEST_WINDOW, + ) self._message_repository.log_message( Level.DEBUG, lambda: _as_auxiliary_request_log( formatter(response), - title=f"Stream '{self._name}' page rejected, retrying with a smaller page size", + title=( + f"Stream '{self._name}' page rejected, retrying with a smaller page size" + if error_resolution.response_action == ResponseAction.REDUCE_PAGE_SIZE + else f"Stream '{self._name}' request window rejected, retrying with a smaller window" + ), description=( - f"Request for stream '{self._name}' whose response asked for a smaller page. The " - f"same page is requested again with a reduced page size, so this request produced " - f"no records." + ( + f"Request for stream '{self._name}' whose response asked for a smaller page. The " + f"same page is requested again with a reduced page size, so this request produced " + f"no records." + ) + if error_resolution.response_action == ResponseAction.REDUCE_PAGE_SIZE + else ( + f"Request for stream '{self._name}' whose response asked for a smaller request " + f"window. The window is split and re-read as smaller children, so this request " + f"produced no records." + ) ), ) if log_as_auxiliary @@ -578,6 +597,13 @@ def _handle_error_resolution( stream_name=self._name, error_message=error_resolution.error_message ) + if error_resolution.response_action == ResponseAction.SPLIT_REQUEST_WINDOW: + raise RequestWindowSplitRequiredException( + stream_name=self._name, + error_message=error_resolution.error_message, + failure_type=error_resolution.failure_type, + ) + # Emit stream status RUNNING with the reason RATE_LIMITED to log that the rate limit has been reached if error_resolution.response_action == ResponseAction.RATE_LIMITED: # TODO: Update to handle with message repository when concurrent message repository is ready diff --git a/airbyte_cdk/sources/streams/http/request_window_split_exception.py b/airbyte_cdk/sources/streams/http/request_window_split_exception.py new file mode 100644 index 0000000000..0ee937fa00 --- /dev/null +++ b/airbyte_cdk/sources/streams/http/request_window_split_exception.py @@ -0,0 +1,62 @@ +# +# Copyright (c) 2026 Airbyte, Inc., all rights reserved. +# + +from typing import Optional + +from airbyte_cdk.models import FailureType +from airbyte_cdk.utils.traced_exception import AirbyteTracedException + + +class RequestWindowSplitRequiredException(AirbyteTracedException): + """ + Signals that the request window backing the current read must be split and retried as smaller children. + + This is control flow rather than a failure: `SimpleRetriever.read_records` catches it and replaces the + failing window with ordered, non-overlapping child windows produced by the stream's cursor. Unlike + `PageSizeReductionRequiredException`, which is only ever raised by `HttpClient`, this exception is public + and is meant to be raised directly by custom code - a `CustomRetriever`, a `CustomExtractor`, or code + handling a transport or streamed-decoding failure - so a connector can opt into window splitting for a + failure the built-in `HttpResponseFilter`/`ResponseAction.SPLIT_REQUEST_WINDOW` path cannot see, without + the CDK needing to know about that connector's specific exception types. + + The classifying error's `error_message` is preserved on this exception so it can be surfaced if the window + is never split any further. `failure_type` is not: exhausting every window size the API will accept is + never the user's fault, so `SimpleRetriever` always reports that as `transient_error` regardless of how the + triggering response was classified. Callers with an underlying exception should chain it the normal Python + way, `raise RequestWindowSplitRequiredException(...) from original_exception`, so the original traceback is + not lost. + """ + + def __init__( + self, + stream_name: Optional[str] = None, + error_message: Optional[str] = None, + failure_type: Optional[FailureType] = None, + ) -> None: + stream = f" of stream {stream_name}" if stream_name else "" + detail = f": {error_message}" if error_message else "" + super().__init__( + internal_message=f"A request window split was requested{stream}{detail}", + message=f"The API rejected the current request window{stream} and requires a smaller one. If this message ends a sync, the stream is not set up to split its window: add `request_window_splitting` to its retriever, or remove the `SPLIT_REQUEST_WINDOW` action from its error handler.", + failure_type=failure_type or FailureType.config_error, + ) + + +class RequestWindowSplitNotSupportedException(AirbyteTracedException): + """ + Raised when a window split is requested but the retriever cannot honor it. + + The factory rejects this at config time wherever it can see the error handler, so reaching this means the + signal came from somewhere it cannot inspect - a custom error handler, a custom requester/retriever raising + `RequestWindowSplitRequiredException` directly, or a stream whose cursor has no `split_request_window` + method (for example, a non-datetime or non-incremental slicer). + """ + + def __init__(self, stream_name: Optional[str] = None) -> None: + stream = f"Stream {stream_name}" if stream_name else "The stream" + super().__init__( + internal_message=f"A request window split was requested for stream {stream_name} but its retriever defines no `request_window_splitting`, or its cursor does not support window splitting. The action is only supported on a SimpleRetriever whose cursor has a `split_request_window` method and that defines `request_window_splitting`.", + message=f"{stream} resolves to a request window split but is not set up to split its window. Add `request_window_splitting` to the stream's retriever (this requires a datetime-based incremental cursor with `cursor_granularity` set), or remove the `SPLIT_REQUEST_WINDOW` action from its error handler.", + failure_type=FailureType.config_error, + ) diff --git a/unit_tests/sources/declarative/parsers/test_model_to_component_factory.py b/unit_tests/sources/declarative/parsers/test_model_to_component_factory.py index d6bc80a736..2933b05754 100644 --- a/unit_tests/sources/declarative/parsers/test_model_to_component_factory.py +++ b/unit_tests/sources/declarative/parsers/test_model_to_component_factory.py @@ -233,6 +233,9 @@ PageSizeReduction, PageSizeResetPolicy, ) +from airbyte_cdk.sources.declarative.retrievers.request_window_splitting import ( + RequestWindowSplitting, +) from airbyte_cdk.sources.declarative.schema import InlineSchemaLoader, JsonFileSchemaLoader from airbyte_cdk.sources.declarative.schema.caching_schema_loader_decorator import ( CachingSchemaLoaderDecorator, @@ -7859,6 +7862,550 @@ def test_given_file_uploader_and_page_size_reduction_then_raise(): assert "file_uploader" in str(exception.value) +_REQUEST_WINDOW_REDUCTION_STREAM = """ +type: DeclarativeStream +name: Test +primary_key: id +schema_loader: + type: InlineSchemaLoader + schema: + type: object +{incremental_sync} +retriever: + type: SimpleRetriever + {request_window_splitting} + requester: + type: HttpRequester + url_base: "https://airbyte.io" + path: "/items" + http_method: GET + error_handler: + type: DefaultErrorHandler + response_filters: + - type: HttpResponseFilter + http_codes: [400] + action: {action} + record_selector: + type: RecordSelector + extractor: + type: DpathExtractor + field_path: ["items"] +""" + +_DATETIME_BASED_CURSOR_WITH_GRANULARITY = """incremental_sync: + type: DatetimeBasedCursor + start_datetime: "2024-01-01T00:00:00Z" + end_datetime: "2024-01-31T00:00:00Z" + step: "P7D" + cursor_field: "updated_at" + cursor_granularity: "PT1S" + datetime_format: "%Y-%m-%dT%H:%M:%SZ\"""" + + +def _request_window_splitting_stream( + request_window_splitting="request_window_splitting:\n type: RequestWindowSplitting", + action="SPLIT_REQUEST_WINDOW", + incremental_sync=_DATETIME_BASED_CURSOR_WITH_GRANULARITY, +): + content = _REQUEST_WINDOW_REDUCTION_STREAM.format( + incremental_sync=incremental_sync, + request_window_splitting=request_window_splitting, + action=action, + ) + stream_manifest = transformer.propagate_types_and_parameters( + "", resolver.preprocess_manifest(YamlDeclarativeSource._parse(content)), {} + ) + return factory.create_component( + model_type=DeclarativeStreamModel, component_definition=stream_manifest, config={} + ) + + +def test_given_request_window_splitting_then_create_retriever_with_defaults(): + retriever = get_retriever(_request_window_splitting_stream()) + + assert retriever.request_window_splitting == RequestWindowSplitting( + failure_message=None, + ) + assert callable(retriever.request_window_splitter) + assert retriever.request_window_splitter.__name__ == "split_request_window" + + +def test_given_request_window_splitting_values_then_create_retriever_with_those_values(): + retriever = get_retriever( + _request_window_splitting_stream( + request_window_splitting=( + "request_window_splitting:\n" + " type: RequestWindowSplitting\n" + " failure_message: Lower time_window so that each request covers less data.\n" + " min_split_window: P1D" + ) + ) + ) + + assert retriever.request_window_splitting == RequestWindowSplitting( + failure_message="Lower time_window so that each request covers less data.", + min_split_window=timedelta(days=1), + ) + + +def test_given_no_request_window_splitting_then_retriever_has_none(): + retriever = get_retriever( + _request_window_splitting_stream(request_window_splitting="", action="RETRY") + ) + + assert retriever.request_window_splitting is None + assert retriever.request_window_splitter is None + + +def test_given_split_request_window_action_without_request_window_splitting_then_raise(): + with pytest.raises(ValueError) as exception: + _request_window_splitting_stream(request_window_splitting="") + + assert "SPLIT_REQUEST_WINDOW" in str(exception.value) + + +def test_given_request_window_splitting_without_split_request_window_action_then_warn(caplog): + """ + A `CustomErrorHandler` can resolve to the action without being inspectable, and custom code can raise + `RequestWindowSplitRequiredException` directly without going through any error handler at all - so this + cannot raise. It must not stay silent either: the feature would be dead on a stream that only defines the + block because it expects to need it. + """ + with caplog.at_level(logging.WARNING, logger="airbyte.model_to_component_factory"): + retriever = get_retriever(_request_window_splitting_stream(action="RETRY")) + + assert retriever.request_window_splitting == RequestWindowSplitting() + assert "SPLIT_REQUEST_WINDOW" in caplog.text + + +def test_given_composite_error_handler_with_split_request_window_action_then_require_request_window_splitting(): + """The action can be nested in a CompositeErrorHandler, which the guard has to recurse into.""" + stream_definition = { + "type": "DeclarativeStream", + "name": "Test", + "primary_key": "id", + "schema_loader": {"type": "InlineSchemaLoader", "schema": {"type": "object"}}, + "incremental_sync": { + "type": "DatetimeBasedCursor", + "start_datetime": "2024-01-01T00:00:00Z", + "end_datetime": "2024-01-31T00:00:00Z", + "step": "P7D", + "cursor_field": "updated_at", + "cursor_granularity": "PT1S", + "datetime_format": "%Y-%m-%dT%H:%M:%SZ", + }, + "retriever": { + "type": "SimpleRetriever", + "requester": { + "type": "HttpRequester", + "url_base": "https://airbyte.io", + "path": "/items", + "http_method": "GET", + "error_handler": { + "type": "CompositeErrorHandler", + "error_handlers": [ + { + "type": "DefaultErrorHandler", + "response_filters": [ + { + "type": "HttpResponseFilter", + "http_codes": [429], + "action": "RATE_LIMITED", + } + ], + }, + { + "type": "DefaultErrorHandler", + "response_filters": [ + { + "type": "HttpResponseFilter", + "http_codes": [400], + "action": "SPLIT_REQUEST_WINDOW", + } + ], + }, + ], + }, + }, + "record_selector": { + "type": "RecordSelector", + "extractor": {"type": "DpathExtractor", "field_path": ["items"]}, + }, + }, + } + + with pytest.raises(ValueError) as exception: + factory.create_component( + model_type=DeclarativeStreamModel, + component_definition=stream_definition, + config={}, + ) + + assert "SPLIT_REQUEST_WINDOW" in str(exception.value) + + # and the same manifest with the block present builds + stream_definition["retriever"]["request_window_splitting"] = {"type": "RequestWindowSplitting"} + retriever = get_retriever( + factory.create_component( + model_type=DeclarativeStreamModel, + component_definition=stream_definition, + config={}, + ) + ) + assert retriever.request_window_splitting == RequestWindowSplitting() + + +def test_given_no_incremental_sync_and_request_window_splitting_then_raise(): + with pytest.raises(ValueError) as exception: + _request_window_splitting_stream(incremental_sync="") + + assert "cursor_granularity" in str(exception.value) + + +def test_given_datetime_based_cursor_without_cursor_granularity_then_raise(): + with pytest.raises(ValueError) as exception: + _request_window_splitting_stream( + incremental_sync=( + 'incremental_sync:\n type: DatetimeBasedCursor\n start_datetime: "2024-01-01T00:00:00Z"\n' + ' end_datetime: "2024-01-31T00:00:00Z"\n step: "P7D"\n cursor_field: "updated_at"\n' + ' datetime_format: "%Y-%m-%dT%H:%M:%SZ"' + ) + ) + + assert "cursor_granularity" in str(exception.value) + + +def test_given_zero_length_cursor_granularity_and_request_window_splitting_then_raise(): + with pytest.raises(ValueError) as exception: + _request_window_splitting_stream( + incremental_sync=( + 'incremental_sync:\n type: DatetimeBasedCursor\n start_datetime: "2024-01-01T00:00:00Z"\n' + ' end_datetime: "2024-01-31T00:00:00Z"\n step: "P7D"\n cursor_field: "updated_at"\n' + ' cursor_granularity: "PT0S"\n datetime_format: "%Y-%m-%dT%H:%M:%SZ"' + ) + ) + + assert "cursor_granularity" in str(exception.value) + assert "zero-length" in str(exception.value) + + +def test_given_datetime_format_coarser_than_cursor_granularity_and_request_window_splitting_then_raise(): + """ + `cursor_granularity: PT1S` claims second-level precision, but `%Y-%m-%d` can only render day boundaries: two + children split less than a day apart would format to the same date and the retriever would re-request the + same (or a reversed) window instead of a smaller one. + """ + with pytest.raises(ValueError) as exception: + _request_window_splitting_stream( + incremental_sync=( + 'incremental_sync:\n type: DatetimeBasedCursor\n start_datetime: "2024-01-01"\n' + ' end_datetime: "2024-01-31"\n step: "P7D"\n cursor_field: "updated_at"\n' + ' cursor_granularity: "PT1S"\n datetime_format: "%Y-%m-%d"' + ) + ) + + assert "cursor_granularity" in str(exception.value) + assert "datetime_format" in str(exception.value) + + +def test_given_client_side_incremental_and_request_window_splitting_then_raise(): + """ + A client-side-incremental cursor filters records after they're read instead of sending the window as + request parameters, so every split child would send the exact same request as its parent - splitting would + never resolve a SPLIT_REQUEST_WINDOW response. + """ + with pytest.raises(ValueError) as exception: + _request_window_splitting_stream( + incremental_sync=( + 'incremental_sync:\n type: DatetimeBasedCursor\n start_datetime: "2024-01-01T00:00:00Z"\n' + ' end_datetime: "2024-01-31T00:00:00Z"\n step: "P7D"\n cursor_field: "updated_at"\n' + ' cursor_granularity: "PT1S"\n datetime_format: "%Y-%m-%dT%H:%M:%SZ"\n' + " is_client_side_incremental: true" + ) + ) + + assert "is_client_side_incremental" in str(exception.value) + + +def test_given_month_or_year_min_split_window_and_request_window_splitting_then_raise(): + """ + `P1M` parses to an `isodate.Duration`, which has no fixed length to compare a window's span against - + `ConcurrentCursor.split_request_window` would raise a `TypeError` mid-sync trying to. This must be rejected + at config time instead. + """ + with pytest.raises(ValueError) as exception: + _request_window_splitting_stream( + request_window_splitting=( + "request_window_splitting:\n type: RequestWindowSplitting\n min_split_window: P1M" + ) + ) + + assert "min_split_window" in str(exception.value) + + +def test_given_negative_min_split_window_and_request_window_splitting_then_raise(): + with pytest.raises(ValueError) as exception: + _request_window_splitting_stream( + request_window_splitting=( + "request_window_splitting:\n type: RequestWindowSplitting\n min_split_window: -PT1H" + ) + ) + + assert "min_split_window" in str(exception.value) + + +def test_given_incrementing_count_cursor_and_request_window_splitting_then_raise(): + """ + `IncrementingCountCursor` builds a `ConcurrentCursor` too, but not one carrying the datetime boundary + fields, granularity, or formatting `split_request_window` needs, so it must be rejected the same way a + stream with no incremental_sync at all is. + """ + with pytest.raises(ValueError) as exception: + _request_window_splitting_stream( + incremental_sync=( + "incremental_sync:\n type: IncrementingCountCursor\n cursor_field: id" + ) + ) + + assert "cursor_granularity" in str(exception.value) + + +def test_given_partition_router_and_request_window_splitting_then_raise(): + """ + A partition router combines multiple cursors (one per partition) behind a `ConcurrentPerPartitionCursor`, + which does not implement `split_request_window` itself - there is no longer one cursor whose own + granularity/boundary fields splitting could rely on. Rejected the same way a stream with no incremental + cursor at all is. + """ + content = _REQUEST_WINDOW_REDUCTION_STREAM.format( + incremental_sync=_DATETIME_BASED_CURSOR_WITH_GRANULARITY, + request_window_splitting="request_window_splitting:\n type: RequestWindowSplitting", + action="SPLIT_REQUEST_WINDOW", + ).replace( + " requester:\n", + " partition_router:\n" + " type: ListPartitionRouter\n" + ' values: ["a", "b"]\n' + " cursor_field: partition_field\n" + " requester:\n", + ) + stream_manifest = transformer.propagate_types_and_parameters( + "", resolver.preprocess_manifest(YamlDeclarativeSource._parse(content)), {} + ) + + with pytest.raises(ValueError) as exception: + factory.create_component( + model_type=DeclarativeStreamModel, component_definition=stream_manifest, config={} + ) + + assert "request_window_splitting" in str(exception.value) + assert "partition router" in str(exception.value) + + +def test_given_query_properties_and_request_window_splitting_then_raise(): + content = _REQUEST_WINDOW_REDUCTION_STREAM.format( + incremental_sync=_DATETIME_BASED_CURSOR_WITH_GRANULARITY, + request_window_splitting="request_window_splitting:\n type: RequestWindowSplitting", + action="SPLIT_REQUEST_WINDOW", + ).replace( + " http_method: GET\n", + " http_method: GET\n" + " query_properties:\n" + " type: QueryProperties\n" + ' property_list: ["a", "b"]\n', + ) + stream_manifest = transformer.propagate_types_and_parameters( + "", resolver.preprocess_manifest(YamlDeclarativeSource._parse(content)), {} + ) + + with pytest.raises(ValueError) as exception: + factory.create_component( + model_type=DeclarativeStreamModel, component_definition=stream_manifest, config={} + ) + + assert "query properties" in str(exception.value) + + +def test_given_file_uploader_and_request_window_splitting_then_raise(): + content = ( + _REQUEST_WINDOW_REDUCTION_STREAM.format( + incremental_sync=_DATETIME_BASED_CURSOR_WITH_GRANULARITY, + request_window_splitting="request_window_splitting:\n type: RequestWindowSplitting", + action="SPLIT_REQUEST_WINDOW", + ) + + """file_uploader: + type: FileUploader + requester: + type: HttpRequester + url_base: "https://airbyte.io" + path: "/download" + download_target_extractor: + type: DpathExtractor + field_path: ["url"] +""" + ) + stream_manifest = transformer.propagate_types_and_parameters( + "", resolver.preprocess_manifest(YamlDeclarativeSource._parse(content)), {} + ) + + with pytest.raises(ValueError) as exception: + factory.create_component( + model_type=DeclarativeStreamModel, component_definition=stream_manifest, config={} + ) + + assert "file_uploader" in str(exception.value) + + +def test_given_lazy_read_pointer_and_request_window_splitting_then_raise(): + """ + A LazySimpleRetriever reads records embedded in the parent's own pages rather than requesting a window of + its own, so there is nothing for a reduced child window to be read from. + """ + stream_definition = { + "type": "DeclarativeStream", + "name": "items", + "primary_key": [], + "schema_loader": { + "type": "InlineSchemaLoader", + "schema": {"type": "object", "properties": {}}, + }, + "retriever": { + "type": "SimpleRetriever", + "request_window_splitting": {"type": "RequestWindowSplitting"}, + "requester": { + "type": "HttpRequester", + "url_base": "https://api.test.com", + "path": "parent/{{ stream_partition.parent_id }}/items", + "http_method": "GET", + "error_handler": { + "type": "DefaultErrorHandler", + "response_filters": [ + { + "type": "HttpResponseFilter", + "http_codes": [400], + "action": "SPLIT_REQUEST_WINDOW", + } + ], + }, + }, + "record_selector": { + "type": "RecordSelector", + "extractor": {"type": "DpathExtractor", "field_path": ["data"]}, + }, + "partition_router": { + "type": "SubstreamPartitionRouter", + "parent_stream_configs": [ + { + "type": "ParentStreamConfig", + "parent_key": "id", + "partition_field": "parent_id", + "lazy_read_pointer": ["items"], + "stream": { + "type": "DeclarativeStream", + "name": "parent", + "schema_loader": { + "type": "InlineSchemaLoader", + "schema": {"type": "object", "properties": {}}, + }, + "retriever": { + "type": "SimpleRetriever", + "requester": { + "type": "HttpRequester", + "url_base": "https://api.test.com", + "path": "/parents", + "http_method": "GET", + }, + "record_selector": { + "type": "RecordSelector", + "extractor": { + "type": "DpathExtractor", + "field_path": ["data"], + }, + }, + }, + }, + } + ], + }, + }, + } + + with pytest.raises(ValueError) as exception: + factory.create_component( + model_type=DeclarativeStreamModel, + component_definition=stream_definition, + config={}, + ) + + assert "request_window_splitting" in str(exception.value) + assert "lazily" in str(exception.value) + + +def test_given_login_requester_with_split_request_window_action_then_raise(): + """ + Only the main requester of a SimpleRetriever reads a cursor-sliced window, so the action is rejected on a + SessionTokenAuthenticator's login_requester even though it is schema-legal there. + """ + stream_definition = { + "type": "DeclarativeStream", + "name": "Test", + "primary_key": "id", + "schema_loader": {"type": "InlineSchemaLoader", "schema": {"type": "object"}}, + "retriever": { + "type": "SimpleRetriever", + "requester": { + "type": "HttpRequester", + "url_base": "https://airbyte.io", + "path": "/items", + "http_method": "GET", + "authenticator": { + "type": "SessionTokenAuthenticator", + "login_requester": { + "type": "HttpRequester", + "url_base": "https://airbyte.io", + "path": "/session", + "http_method": "POST", + "error_handler": { + "type": "DefaultErrorHandler", + "response_filters": [ + { + "type": "HttpResponseFilter", + "http_codes": [400], + "action": "SPLIT_REQUEST_WINDOW", + } + ], + }, + }, + "session_token_path": ["session_token"], + "expiration_duration": "PT1H", + "request_authentication": { + "type": "ApiKey", + "inject_into": { + "type": "RequestOption", + "inject_into": "header", + "field_name": "X-Session-Token", + }, + }, + }, + }, + "record_selector": { + "type": "RecordSelector", + "extractor": {"type": "DpathExtractor", "field_path": ["items"]}, + }, + }, + } + + with pytest.raises(ValueError) as exception: + factory.create_component( + model_type=DeclarativeStreamModel, + component_definition=stream_definition, + config={}, + ) + + assert "SPLIT_REQUEST_WINDOW" in str(exception.value) + assert "login_requester" in str(exception.value) + + _REDUCE_PAGE_SIZE_ERROR_HANDLER = { "type": "DefaultErrorHandler", "response_filters": [ @@ -7995,6 +8542,114 @@ def test_given_reduce_page_size_action_on_an_async_retriever_requester_then_rais assert "REDUCE_PAGE_SIZE" in str(exception.value) +_SPLIT_REQUEST_WINDOW_ERROR_HANDLER = { + "type": "DefaultErrorHandler", + "response_filters": [ + {"type": "HttpResponseFilter", "http_codes": [400], "action": "SPLIT_REQUEST_WINDOW"} + ], +} + + +def test_given_split_request_window_action_on_the_file_uploader_requester_then_raise(): + """ + `error_handler` is defined on `HttpRequester`, so the action is manifest-legal on requesters that have no + window of their own. Nothing there can honor it, so it is rejected rather than raised mid-sync. + """ + with pytest.raises(ValueError) as exception: + factory.create_component( + model_type=FileUploaderModel, + component_definition={ + "type": "FileUploader", + "requester": { + "type": "HttpRequester", + "url_base": "https://airbyte.io", + "path": "/download", + "error_handler": _SPLIT_REQUEST_WINDOW_ERROR_HANDLER, + }, + "download_target_extractor": { + "type": "DpathExtractor", + "field_path": ["url"], + }, + }, + config={}, + ) + + assert "SPLIT_REQUEST_WINDOW" in str(exception.value) + + +@pytest.mark.parametrize( + "requester_field", + [ + "creation_requester", + "polling_requester", + "download_requester", + "download_target_requester", + "abort_requester", + "delete_requester", + ], +) +def test_given_split_request_window_action_on_an_async_retriever_requester_then_raise( + requester_field, +): + definition = { + "type": "AsyncRetriever", + "status_mapping": { + "type": "AsyncJobStatusMap", + "running": ["running"], + "completed": ["ready"], + "failed": ["failed"], + "timeout": ["timeout"], + }, + "status_extractor": {"type": "DpathExtractor", "field_path": ["status"]}, + "record_selector": { + "type": "RecordSelector", + "extractor": {"type": "DpathExtractor", "field_path": ["items"]}, + }, + "creation_requester": { + "type": "HttpRequester", + "url_base": "https://airbyte.io", + "path": "/jobs", + "http_method": "POST", + }, + "polling_requester": { + "type": "HttpRequester", + "url_base": "https://airbyte.io", + "path": "/jobs/{{ creation_response.id }}", + }, + "download_requester": { + "type": "HttpRequester", + "url_base": "https://airbyte.io", + "path": "/jobs/{{ creation_response.id }}/download", + }, + } + definition[requester_field] = { + "type": "HttpRequester", + "url_base": "https://airbyte.io", + "path": "/jobs", + "error_handler": _SPLIT_REQUEST_WINDOW_ERROR_HANDLER, + } + if requester_field == "download_target_requester": + # Without it the `download_target_extractor` guard fires first and the assertion below would pass for + # the wrong reason. + definition["download_target_extractor"] = { + "type": "DpathExtractor", + "field_path": ["url"], + } + + with pytest.raises(ValueError) as exception: + factory.create_component( + model_type=AsyncRetrieverModel, + component_definition=definition, + config={}, + name="a_stream", + primary_key=None, + stream_slicer=None, + transformations=[], + ) + + assert "SPLIT_REQUEST_WINDOW" in str(exception.value) + + def get_retriever(stream: Union[DeclarativeStream, DefaultStream]): return ( stream.retriever diff --git a/unit_tests/sources/declarative/requesters/error_handlers/test_composite_error_handler.py b/unit_tests/sources/declarative/requesters/error_handlers/test_composite_error_handler.py index afc767e5ea..78d3b7ee8a 100644 --- a/unit_tests/sources/declarative/requesters/error_handlers/test_composite_error_handler.py +++ b/unit_tests/sources/declarative/requesters/error_handlers/test_composite_error_handler.py @@ -372,3 +372,51 @@ def test_given_reduce_page_size_when_interpret_response_then_stop_at_the_matchin assert error_resolution.response_action == ResponseAction.REDUCE_PAGE_SIZE assert failing_handler.interpret_response.call_count == 0 + + +def test_given_split_request_window_when_interpret_response_then_stop_at_the_matching_handler(): + """ + Mirrors test_given_reduce_page_size_when_interpret_response_then_stop_at_the_matching_handler: without the + short circuit, a later handler resolving to FAIL would win over the request window reduction. + """ + reducing_handler = MagicMock() + reducing_handler.interpret_response.return_value = ErrorResolution( + response_action=ResponseAction.SPLIT_REQUEST_WINDOW, + failure_type=FailureType.config_error, + ) + failing_handler = MagicMock() + failing_handler.interpret_response.return_value = ErrorResolution( + response_action=ResponseAction.FAIL, failure_type=FailureType.system_error + ) + error_handler = CompositeErrorHandler( + error_handlers=[reducing_handler, failing_handler], parameters={} + ) + + error_resolution = error_handler.interpret_response(MagicMock()) + + assert error_resolution.response_action == ResponseAction.SPLIT_REQUEST_WINDOW + assert failing_handler.interpret_response.call_count == 0 + + +def test_given_unrelated_response_error_when_interpret_response_then_behavior_is_unaffected(): + """ + Adding SPLIT_REQUEST_WINDOW to the short-circuit list must not change behavior for actions that were + never in it: RATE_LIMITED still does not short-circuit, so a later handler's FAIL still wins over an + earlier handler's RATE_LIMITED, exactly as before this action existed. + """ + rate_limited_handler = MagicMock() + rate_limited_handler.interpret_response.return_value = ErrorResolution( + response_action=ResponseAction.RATE_LIMITED, failure_type=FailureType.transient_error + ) + failing_handler = MagicMock() + failing_handler.interpret_response.return_value = ErrorResolution( + response_action=ResponseAction.FAIL, failure_type=FailureType.system_error + ) + error_handler = CompositeErrorHandler( + error_handlers=[rate_limited_handler, failing_handler], parameters={} + ) + + error_resolution = error_handler.interpret_response(MagicMock()) + + assert error_resolution.response_action == ResponseAction.FAIL + assert failing_handler.interpret_response.call_count == 1 diff --git a/unit_tests/sources/declarative/retrievers/test_simple_retriever.py b/unit_tests/sources/declarative/retrievers/test_simple_retriever.py index c0fa483d78..4b87f0fed5 100644 --- a/unit_tests/sources/declarative/retrievers/test_simple_retriever.py +++ b/unit_tests/sources/declarative/retrievers/test_simple_retriever.py @@ -3,11 +3,13 @@ # import json +import logging import threading from collections import defaultdict from concurrent.futures import ThreadPoolExecutor +from datetime import timedelta from functools import partial -from typing import Any, Iterable, Mapping, Optional +from typing import Any, Callable, Iterable, List, Mapping, Optional from unittest.mock import MagicMock, Mock, patch import pytest @@ -52,13 +54,23 @@ PageSizeResetPolicy, ) from airbyte_cdk.sources.declarative.retrievers.pagination_tracker import PaginationTracker -from airbyte_cdk.sources.declarative.retrievers.simple_retriever import SimpleRetriever +from airbyte_cdk.sources.declarative.retrievers.request_window_splitting import ( + RequestWindowSplitting, +) +from airbyte_cdk.sources.declarative.retrievers.simple_retriever import ( + _MAX_REQUEST_WINDOW_SPLIT_DEPTH, + SimpleRetriever, +) from airbyte_cdk.sources.streams.http.page_size_reduction_exception import ( PageSizeReductionRequiredException, ) from airbyte_cdk.sources.streams.http.pagination_reset_exception import ( PaginationResetRequiredException, ) +from airbyte_cdk.sources.streams.http.request_window_split_exception import ( + RequestWindowSplitNotSupportedException, + RequestWindowSplitRequiredException, +) from airbyte_cdk.sources.types import Record, StreamSlice from airbyte_cdk.sources.utils.transform import TransformConfig, TypeTransformer from airbyte_cdk.utils.traced_exception import AirbyteTracedException @@ -2058,3 +2070,661 @@ def test_given_no_data_feed_when_read_records_then_emit_every_record(): ) assert actual_records == page + + +# --- request_window_splitting --- + +A_WINDOW_SLICE = StreamSlice( + cursor_slice={"start_time": "2024-01-01T00:00:00Z", "end_time": "2024-01-01T23:59:59Z"}, + partition={}, +) +A_FIRST_HALF_SLICE = StreamSlice( + cursor_slice={"start_time": "2024-01-01T00:00:00Z", "end_time": "2024-01-01T11:59:59Z"}, + partition={}, +) +A_SECOND_HALF_SLICE = StreamSlice( + cursor_slice={"start_time": "2024-01-01T12:00:00Z", "end_time": "2024-01-01T23:59:59Z"}, + partition={}, +) + + +def _request_window_splitting_retriever( + requester: Requester, + paginator: Paginator, + record_selector: HttpSelector, + request_window_splitter: Callable[ + [StreamSlice, Optional[timedelta]], Optional[List[StreamSlice]] + ], + request_window_splitting: Optional[RequestWindowSplitting] = None, +) -> SimpleRetriever: + return SimpleRetriever( + name=A_STREAM_NAME, + primary_key=primary_key, + requester=requester, + record_selector=record_selector, + paginator=paginator, + request_window_splitting=request_window_splitting or RequestWindowSplitting(), + request_window_splitter=request_window_splitter, + parameters={}, + config={}, + ) + + +def test_given_request_window_splitting_when_read_records_then_split_and_read_both_children(): + requester = Mock(spec=Requester) + requester.send_request.side_effect = [ + RequestWindowSplitRequiredException(), + [{"id": 1}], + [{"id": 2}], + ] + record_selector = Mock(spec=HttpSelector) + record_selector.select_records.side_effect = [[{"id": 1}], [{"id": 2}]] + paginator = _mock_paginator() + paginator.get_initial_token.return_value = None + paginator.next_page_token.return_value = None + + request_window_splitter = Mock() + request_window_splitter.return_value = [ + A_FIRST_HALF_SLICE, + A_SECOND_HALF_SLICE, + ] + + retriever = _request_window_splitting_retriever( + requester, paginator, record_selector, request_window_splitter + ) + + records = list(retriever.read_records(A_RECORD_SCHEMA, A_WINDOW_SLICE)) + + assert records == [{"id": 1}, {"id": 2}] + request_window_splitter.assert_called_once_with(A_WINDOW_SLICE, None) + assert requester.send_request.call_count == 3 + assert [call.kwargs["stream_slice"] for call in requester.send_request.call_args_list] == [ + A_WINDOW_SLICE, + A_FIRST_HALF_SLICE, + A_SECOND_HALF_SLICE, + ] + + +def test_given_min_split_window_configured_when_read_records_then_pass_it_to_the_splitter(): + """ + `min_split_window` lives on `request_window_splitting`, not on the splitter itself, so the retriever must + thread it through on every call rather than the splitter reading it from somewhere else. + """ + requester = Mock(spec=Requester) + requester.send_request.side_effect = RequestWindowSplitRequiredException() + record_selector = Mock(spec=HttpSelector) + paginator = _mock_paginator() + paginator.get_initial_token.return_value = None + + request_window_splitter = Mock() + request_window_splitter.return_value = None + + retriever = _request_window_splitting_retriever( + requester, + paginator, + record_selector, + request_window_splitter, + request_window_splitting=RequestWindowSplitting(min_split_window=timedelta(days=1)), + ) + + with pytest.raises(AirbyteTracedException): + list(retriever.read_records(A_RECORD_SCHEMA, A_WINDOW_SLICE)) + + request_window_splitter.assert_called_once_with(A_WINDOW_SLICE, timedelta(days=1)) + + +def test_given_only_one_child_needs_further_reduction_when_read_records_then_recurse_asymmetrically(): + """ + Nested/asymmetric reduction: the first child is itself rejected and split again, while the second + child (read after the first child's full recursive read completes) succeeds on the first try. + """ + requester = Mock(spec=Requester) + requester.send_request.side_effect = [ + RequestWindowSplitRequiredException(), # top-level window + RequestWindowSplitRequiredException(), # first half, rejected again + [{"id": 1}], # first quarter of the first half + [{"id": 2}], # second quarter of the first half + [{"id": 3}], # second half, succeeds directly + ] + record_selector = Mock(spec=HttpSelector) + record_selector.select_records.side_effect = [[{"id": 1}], [{"id": 2}], [{"id": 3}]] + paginator = _mock_paginator() + paginator.get_initial_token.return_value = None + paginator.next_page_token.return_value = None + + a_quarter_slice = StreamSlice( + cursor_slice={"start_time": "2024-01-01T00:00:00Z", "end_time": "2024-01-01T05:59:59Z"}, + partition={}, + ) + another_quarter_slice = StreamSlice( + cursor_slice={"start_time": "2024-01-01T06:00:00Z", "end_time": "2024-01-01T11:59:59Z"}, + partition={}, + ) + request_window_splitter = Mock() + request_window_splitter.side_effect = [ + [A_FIRST_HALF_SLICE, A_SECOND_HALF_SLICE], + [a_quarter_slice, another_quarter_slice], + ] + + retriever = _request_window_splitting_retriever( + requester, paginator, record_selector, request_window_splitter + ) + + records = list(retriever.read_records(A_RECORD_SCHEMA, A_WINDOW_SLICE)) + + assert records == [{"id": 1}, {"id": 2}, {"id": 3}] + assert request_window_splitter.call_args_list == [ + ((A_WINDOW_SLICE, None),), + ((A_FIRST_HALF_SLICE, None),), + ] + assert [call.kwargs["stream_slice"] for call in requester.send_request.call_args_list] == [ + A_WINDOW_SLICE, + A_FIRST_HALF_SLICE, + a_quarter_slice, + another_quarter_slice, + A_SECOND_HALF_SLICE, + ] + + +def test_given_no_request_window_splitting_configured_when_reduction_required_then_raise_not_supported(): + requester = Mock(spec=Requester) + requester.send_request.side_effect = RequestWindowSplitRequiredException() + record_selector = Mock(spec=HttpSelector) + paginator = _mock_paginator() + paginator.get_initial_token.return_value = None + + retriever = SimpleRetriever( + name=A_STREAM_NAME, + primary_key=primary_key, + requester=requester, + record_selector=record_selector, + paginator=paginator, + parameters={}, + config={}, + ) + + with pytest.raises(RequestWindowSplitNotSupportedException): + list(retriever.read_records(A_RECORD_SCHEMA, A_WINDOW_SLICE)) + + +def test_given_no_request_window_splitter_when_reduction_required_then_raise_not_supported(): + requester = Mock(spec=Requester) + requester.send_request.side_effect = RequestWindowSplitRequiredException() + record_selector = Mock(spec=HttpSelector) + paginator = _mock_paginator() + paginator.get_initial_token.return_value = None + + retriever = SimpleRetriever( + name=A_STREAM_NAME, + primary_key=primary_key, + requester=requester, + record_selector=record_selector, + paginator=paginator, + request_window_splitting=RequestWindowSplitting(), + request_window_splitter=None, + parameters={}, + config={}, + ) + + with pytest.raises(RequestWindowSplitNotSupportedException): + list(retriever.read_records(A_RECORD_SCHEMA, A_WINDOW_SLICE)) + + +def test_given_records_already_emitted_when_reduction_requested_then_split_and_replay_anyway( + caplog, +): + """ + A later page of the same window fails after an earlier page already emitted a record: since a failed + partition is never checkpointed, refusing to split here wouldn't avoid the duplicate anyway - the next + attempt would just re-read the same window and re-emit it. So this always splits and re-reads, logging a + warning with how many records were already emitted. + """ + requester = Mock(spec=Requester) + requester.send_request.side_effect = [ + Mock(), + RequestWindowSplitRequiredException(), + [{"id": 2}], + [{"id": 3}], + ] + record_selector = Mock(spec=HttpSelector) + record_selector.select_records.side_effect = [[{"id": 1}], [{"id": 2}], [{"id": 3}]] + paginator = _mock_paginator() + paginator.get_initial_token.return_value = None + paginator.next_page_token.side_effect = [{"next_page_token": "page2"}, None, None] + + request_window_splitter = Mock() + request_window_splitter.return_value = [ + A_FIRST_HALF_SLICE, + A_SECOND_HALF_SLICE, + ] + + retriever = _request_window_splitting_retriever( + requester, paginator, record_selector, request_window_splitter + ) + + with caplog.at_level(logging.WARNING, logger="airbyte"): + records = list(retriever.read_records(A_RECORD_SCHEMA, A_WINDOW_SLICE)) + + # the record from the page emitted before the split signal is included once more: splitting and + # re-reading accepts this at-least-once duplication rather than silently dropping it or failing the sync. + assert records == [{"id": 1}, {"id": 2}, {"id": 3}] + assert "already emitted 1 record" in caplog.text + request_window_splitter.assert_called_once_with(A_WINDOW_SLICE, None) + + +def test_given_request_window_splitter_returns_none_when_reduction_requested_then_raise_transient_error(): + """ + `request_window_splitter` returning `None` means the window cannot be split any further - either the + granularity floor or `min_split_window` was reached, or splitting would not make progress. This is never + the user's fault - the API kept rejecting every window size tried - so it is always `transient_error`, + matching `PageSizeReducer`'s equivalent exhaustion branch, regardless of how the triggering response was + classified (here, deliberately the opposite - `config_error` - to prove it is overridden). + """ + requester = Mock(spec=Requester) + requester.send_request.side_effect = RequestWindowSplitRequiredException( + failure_type=FailureType.config_error + ) + record_selector = Mock(spec=HttpSelector) + paginator = _mock_paginator() + paginator.get_initial_token.return_value = None + + request_window_splitter = Mock() + request_window_splitter.return_value = None + + retriever = _request_window_splitting_retriever( + requester, + paginator, + record_selector, + request_window_splitter, + request_window_splitting=RequestWindowSplitting( + failure_message="Lower time_window so that each request covers less data." + ), + ) + + with pytest.raises(AirbyteTracedException) as exception: + list(retriever.read_records(A_RECORD_SCHEMA, A_WINDOW_SLICE)) + + assert exception.value.failure_type == FailureType.transient_error + assert "Lower time_window so that each request covers less data." in exception.value.message + request_window_splitter.assert_called_once_with(A_WINDOW_SLICE, None) + + +def test_given_request_window_splitter_returns_none_and_min_split_window_configured_then_name_it_in_the_message(): + """ + When `min_split_window` is configured, it may be the actual reason splitting stopped rather than the + cursor's own granularity floor - the terminal message should name it so the user knows which knob to check. + """ + requester = Mock(spec=Requester) + requester.send_request.side_effect = RequestWindowSplitRequiredException() + record_selector = Mock(spec=HttpSelector) + paginator = _mock_paginator() + paginator.get_initial_token.return_value = None + + request_window_splitter = Mock() + request_window_splitter.return_value = None + + retriever = _request_window_splitting_retriever( + requester, + paginator, + record_selector, + request_window_splitter, + request_window_splitting=RequestWindowSplitting(min_split_window=timedelta(days=1)), + ) + + with pytest.raises(AirbyteTracedException) as exception: + list(retriever.read_records(A_RECORD_SCHEMA, A_WINDOW_SLICE)) + + assert "min_split_window" in exception.value.message + assert str(timedelta(days=1)) in exception.value.message + + +def test_given_empty_children_list_when_reduction_requested_then_emit_no_records_and_do_not_loop(): + requester = Mock(spec=Requester) + requester.send_request.side_effect = RequestWindowSplitRequiredException() + record_selector = Mock(spec=HttpSelector) + paginator = _mock_paginator() + paginator.get_initial_token.return_value = None + + request_window_splitter = Mock() + request_window_splitter.return_value = [] + + retriever = _request_window_splitting_retriever( + requester, paginator, record_selector, request_window_splitter + ) + + records = list(retriever.read_records(A_RECORD_SCHEMA, A_WINDOW_SLICE)) + + assert records == [] + request_window_splitter.assert_called_once_with(A_WINDOW_SLICE, None) + + +def test_given_a_later_child_fails_when_read_records_then_the_failure_propagates(): + """ + The first child reads fully and its records are emitted before the second child fails; the failure of a + later child must still surface as a real error, not be silently swallowed after the first child's success. + """ + requester = Mock(spec=Requester) + requester.send_request.side_effect = [ + RequestWindowSplitRequiredException(), + [{"id": 1}], + RequestWindowSplitRequiredException(), + ] + record_selector = Mock(spec=HttpSelector) + record_selector.select_records.return_value = [{"id": 1}] + paginator = _mock_paginator() + paginator.get_initial_token.return_value = None + paginator.next_page_token.return_value = None + + request_window_splitter = Mock() + request_window_splitter.side_effect = [ + [A_FIRST_HALF_SLICE, A_SECOND_HALF_SLICE], + None, + ] + + retriever = _request_window_splitting_retriever( + requester, paginator, record_selector, request_window_splitter + ) + + emitted = [] + with pytest.raises(AirbyteTracedException): + for record in retriever.read_records(A_RECORD_SCHEMA, A_WINDOW_SLICE): + emitted.append(record) + + # the first child's records were already emitted to the caller before the second child failed + assert emitted == [{"id": 1}] + assert request_window_splitter.call_args_list == [ + ((A_WINDOW_SLICE, None),), + ((A_SECOND_HALF_SLICE, None),), + ] + + +def test_given_request_window_splitting_when_read_records_then_each_child_gets_fresh_pagination_state(): + """ + Each recursive call re-enters read_records/_read_pages from scratch for its child slice, so the + paginator's initial token is requested again for every child rather than resuming from wherever the + failed parent window left off. + """ + requester = Mock(spec=Requester) + requester.send_request.side_effect = [ + RequestWindowSplitRequiredException(), + [{"id": 1}], + [{"id": 2}], + ] + record_selector = Mock(spec=HttpSelector) + record_selector.select_records.side_effect = [[{"id": 1}], [{"id": 2}]] + paginator = _mock_paginator() + paginator.get_initial_token.return_value = "initial token" + paginator.next_page_token.return_value = None + + request_window_splitter = Mock() + request_window_splitter.return_value = [ + A_FIRST_HALF_SLICE, + A_SECOND_HALF_SLICE, + ] + + retriever = _request_window_splitting_retriever( + requester, paginator, record_selector, request_window_splitter + ) + + list(retriever.read_records(A_RECORD_SCHEMA, A_WINDOW_SLICE)) + + assert [call.kwargs["next_page_token"] for call in requester.send_request.call_args_list] == [ + {"next_page_token": "initial token"}, + {"next_page_token": "initial token"}, + {"next_page_token": "initial token"}, + ] + + +def test_reassociate_with_original_slice_rewraps_a_record_stamped_with_a_different_slice(): + child_slice = StreamSlice(cursor_slice={"start_time": "child"}, partition={}) + record = Record( + data={"id": 1}, + stream_name=A_STREAM_NAME, + associated_slice=child_slice, + ) + + result = SimpleRetriever._reassociate_with_original_slice(record, A_WINDOW_SLICE) + + assert result.associated_slice is A_WINDOW_SLICE + assert result.data == {"id": 1} + assert result.stream_name == A_STREAM_NAME + + +def test_reassociate_with_original_slice_is_a_no_op_when_already_associated(): + record = Record( + data={"id": 1}, + stream_name=A_STREAM_NAME, + associated_slice=A_WINDOW_SLICE, + ) + + result = SimpleRetriever._reassociate_with_original_slice(record, A_WINDOW_SLICE) + + assert result is record + + +def test_reassociate_with_original_slice_passes_through_non_record_stream_data(): + message = {"raw": "not a Record instance"} + + result = SimpleRetriever._reassociate_with_original_slice(message, A_WINDOW_SLICE) + + assert result is message + + +def test_given_request_window_splitting_when_split_then_records_are_associated_with_original_slice(): + """ + Regression test: `ConcurrentCursor.observe()` keys its per-partition bookkeeping by `Record.associated_slice`, + but `ConcurrentCursor.close_partition()` always looks that bookkeeping up by the partition's own (original) + slice. A record read from a reduced child window - via a real `RecordSelector`, not the mocks the other + tests in this module use - would otherwise carry the child as its `associated_slice`, a key + `close_partition()` can never find, silently degrading the most-recently-observed cursor value tracked for a + split partition. + """ + requester = Mock(spec=Requester) + requester.send_request.side_effect = [ + RequestWindowSplitRequiredException(), + Mock(), + Mock(), + ] + + def select_records(*, response, stream_slice, **kwargs): + return [ + Record( + data={"start": stream_slice.cursor_slice["start_time"]}, + stream_name=A_STREAM_NAME, + associated_slice=stream_slice, + ) + ] + + record_selector = Mock(spec=HttpSelector) + record_selector.select_records.side_effect = select_records + paginator = _mock_paginator() + paginator.get_initial_token.return_value = None + paginator.next_page_token.return_value = None + + request_window_splitter = Mock() + request_window_splitter.return_value = [ + A_FIRST_HALF_SLICE, + A_SECOND_HALF_SLICE, + ] + + retriever = _request_window_splitting_retriever( + requester, paginator, record_selector, request_window_splitter + ) + + records = list(retriever.read_records(A_RECORD_SCHEMA, A_WINDOW_SLICE)) + + assert len(records) == 2 + assert all(record.associated_slice is A_WINDOW_SLICE for record in records) + # the underlying data still reflects which child the record actually came from + assert {record.data["start"] for record in records} == { + A_FIRST_HALF_SLICE.cursor_slice["start_time"], + A_SECOND_HALF_SLICE.cursor_slice["start_time"], + } + + +def test_given_exception_raised_during_record_extraction_when_read_then_still_split_request_window(): + """ + Google Ads/Iterable-style exception-driven cases fail during response streaming/decoding, not at the point + the request is sent - unlike PayPal, where the API rejects the request outright before any body is read. + The catch boundary must cover the entire read, not only the requester call, so custom code raising + `RequestWindowSplitRequiredException` from anywhere in the pipeline is still honored. + """ + requester = Mock(spec=Requester) + requester.send_request.side_effect = [Mock(), Mock(), Mock()] + record_selector = Mock(spec=HttpSelector) + record_selector.select_records.side_effect = [ + RequestWindowSplitRequiredException(), + [{"id": 1}], + [{"id": 2}], + ] + paginator = _mock_paginator() + paginator.get_initial_token.return_value = None + paginator.next_page_token.return_value = None + + request_window_splitter = Mock() + request_window_splitter.return_value = [ + A_FIRST_HALF_SLICE, + A_SECOND_HALF_SLICE, + ] + + retriever = _request_window_splitting_retriever( + requester, paginator, record_selector, request_window_splitter + ) + + records = list(retriever.read_records(A_RECORD_SCHEMA, A_WINDOW_SLICE)) + + assert records == [{"id": 1}, {"id": 2}] + + +def test_given_max_split_depth_exceeded_when_read_records_then_raise_terminal_error(caplog): + """ + Defense-in-depth: even if a (misbehaving, e.g. custom) `request_window_splitter` never itself signals "cannot split + any further", the retriever's own depth counter must still terminate the read deterministically rather than + recursing without bound. + """ + requester = Mock(spec=Requester) + requester.send_request.side_effect = RequestWindowSplitRequiredException() + record_selector = Mock(spec=HttpSelector) + paginator = _mock_paginator() + paginator.get_initial_token.return_value = None + + request_window_splitter = Mock() + # never signals "no further progress possible" - always splits into two children identical to the parent + request_window_splitter.side_effect = lambda stream_slice, min_split_window=None: [ + stream_slice, + stream_slice, + ] + + retriever = _request_window_splitting_retriever( + requester, + paginator, + record_selector, + request_window_splitter, + request_window_splitting=RequestWindowSplitting( + failure_message="Lower time_window so that each request covers less data.", + ), + ) + + with caplog.at_level(logging.WARNING, logger="airbyte"): + with pytest.raises(AirbyteTracedException) as exception: + list(retriever.read_records(A_RECORD_SCHEMA, A_WINDOW_SLICE)) + + # exhausting the split depth is not the user's fault - the API kept rejecting every window size tried - + # so it is transient_error, matching PageSizeReducer's equivalent exhaustion branch, not config_error + assert exception.value.failure_type == FailureType.transient_error + assert "maximum request window split depth" in exception.value.internal_message + assert f"within {_MAX_REQUEST_WINDOW_SPLIT_DEPTH} splits" in exception.value.message + assert "Lower time_window so that each request covers less data." in exception.value.message + # the safety net is expected to be rare, so it's logged separately from the exception (per review feedback + # to make it visible for later review even if the trace message's internal_message isn't surfaced) + assert ( + f"hit the maximum request window split depth ({_MAX_REQUEST_WINDOW_SPLIT_DEPTH})" + in caplog.text + ) + # the single leftmost recursion path is explored to the full depth (one split_request_window call per + # depth, from 0 up to the cap) before the cap stops the next call - the exception then propagates + # immediately, so sibling branches at shallower depths are never explored + assert request_window_splitter.call_count == _MAX_REQUEST_WINDOW_SPLIT_DEPTH + + +def test_given_partitions_read_concurrently_then_window_reduction_isolates_between_partitions(): + """ + One retriever instance - and one shared `request_window_splitter` (the stream's cursor) - is used for every partition + of a stream, so a reduction triggered while reading one partition must not affect a concurrently-running + read of another partition at all: the other partition must complete normally, and the reducer must only ever + be asked to split the slice that actually needed it. + """ + slice_a = StreamSlice( + cursor_slice={"start_time": "2024-01-01T00:00:00Z", "end_time": "2024-01-01T23:59:59Z"}, + partition={"id": "a"}, + ) + slice_b = StreamSlice( + cursor_slice={"start_time": "2024-02-01T00:00:00Z", "end_time": "2024-02-01T23:59:59Z"}, + partition={"id": "b"}, + ) + b_requested = threading.Event() + + def send_request(*, stream_slice, **kwargs): + # only the original top-level slice_a is rejected (its children must succeed, or the recursion would + # never terminate); identity, not partition id, distinguishes it from the children split_request_window builds + if stream_slice is slice_a: + # only raise once b's own, unrelated request is known to be in flight, so the two reads + # genuinely overlap rather than running one after the other + assert b_requested.wait(timeout=10) + raise RequestWindowSplitRequiredException() + if stream_slice is slice_b: + b_requested.set() + return Mock() + + requester = Mock(spec=Requester) + requester.send_request.side_effect = send_request + + def select_records(*, response, stream_slice, **kwargs): + return [ + { + "partition": stream_slice.partition["id"], + "start": stream_slice.cursor_slice["start_time"], + } + ] + + record_selector = Mock(spec=HttpSelector) + record_selector.select_records.side_effect = select_records + paginator = _mock_paginator() + paginator.get_initial_token.return_value = None + paginator.next_page_token.return_value = None + + def split_request_window(stream_slice, min_split_window=None): + # preserve the parent's own partition, matching what a real request_window_splitter would do + return [ + StreamSlice( + cursor_slice={ + "start_time": stream_slice.cursor_slice["start_time"], + "end_time": "midpoint", + }, + partition=stream_slice.partition, + ), + StreamSlice( + cursor_slice={ + "start_time": "midpoint", + "end_time": stream_slice.cursor_slice["end_time"], + }, + partition=stream_slice.partition, + ), + ] + + request_window_splitter = Mock() + request_window_splitter.side_effect = split_request_window + + retriever = _request_window_splitting_retriever( + requester, paginator, record_selector, request_window_splitter + ) + + with ThreadPoolExecutor(max_workers=2) as executor: + future_a = executor.submit(lambda: list(retriever.read_records(A_RECORD_SCHEMA, slice_a))) + future_b = executor.submit(lambda: list(retriever.read_records(A_RECORD_SCHEMA, slice_b))) + records_a = future_a.result(timeout=10) + records_b = future_b.result(timeout=10) + + assert records_b == [{"partition": "b", "start": "2024-02-01T00:00:00Z"}] + assert len(records_a) == 2 + request_window_splitter.assert_called_once_with(slice_a, None) diff --git a/unit_tests/sources/declarative/test_concurrent_declarative_source.py b/unit_tests/sources/declarative/test_concurrent_declarative_source.py index 9d4e438940..1a65aa0b94 100644 --- a/unit_tests/sources/declarative/test_concurrent_declarative_source.py +++ b/unit_tests/sources/declarative/test_concurrent_declarative_source.py @@ -5167,6 +5167,171 @@ def test_given_reductions_exhausted_when_read_then_emit_a_transient_error(): ) +def _request_window_splitting_manifest(): + return { + "version": "0.34.2", + "type": "DeclarativeSource", + "check": {"type": "CheckStream", "stream_names": ["Test"]}, + "streams": [ + { + "type": "DeclarativeStream", + "name": "Test", + "schema_loader": { + "type": "InlineSchemaLoader", + "schema": {"type": "object"}, + }, + "incremental_sync": { + "type": "DatetimeBasedCursor", + "start_datetime": "2024-01-01T00:00:00Z", + "end_datetime": "2024-01-01T23:59:59Z", + "step": "P1D", + "cursor_field": "updated_at", + "cursor_granularity": "PT1S", + "datetime_format": "%Y-%m-%dT%H:%M:%SZ", + "start_time_option": { + "type": "RequestOption", + "inject_into": "request_parameter", + "field_name": "start", + }, + "end_time_option": { + "type": "RequestOption", + "inject_into": "request_parameter", + "field_name": "end", + }, + }, + "retriever": { + "type": "SimpleRetriever", + "request_window_splitting": {"type": "RequestWindowSplitting"}, + "requester": { + "type": "HttpRequester", + "url_base": "https://example.org", + "path": "/test", + "authenticator": {"type": "NoAuth"}, + "error_handler": { + "type": "DefaultErrorHandler", + "response_filters": [ + { + "type": "HttpResponseFilter", + "http_codes": [400], + "action": "SPLIT_REQUEST_WINDOW", + }, + ], + }, + }, + "record_selector": { + "type": "RecordSelector", + "extractor": {"type": "DpathExtractor", "field_path": ["items"]}, + }, + }, + } + ], + "spec": { + "type": "Spec", + "documentation_url": "https://example.org", + "connection_specification": {}, + }, + } + + +def _read_request_window_splitting_source(manifest): + catalog = create_catalog("Test") + source = ConcurrentDeclarativeSource( + source_config=manifest, + config={}, + catalog=catalog, + state=None, + ) + yield from source.read(logger=source.logger, config={}, catalog=catalog, state=[]) + + +def test_given_split_request_window_action_when_read_then_split_and_read_both_halves(): + """ + Mirrors the real PayPal `RESULTSET_TOO_LARGE` incident this feature exists to generalize: a request for + the full day is rejected outright, and the connector reads it back as two half-day requests instead, with + a complete, gap-free, non-duplicated record set. + """ + full_window_request = HttpRequest( + "https://example.org/test", + query_params={"start": "2024-01-01T00:00:00Z", "end": "2024-01-01T23:59:59Z"}, + ) + first_half_request = HttpRequest( + "https://example.org/test", + query_params={"start": "2024-01-01T00:00:00Z", "end": "2024-01-01T11:59:59Z"}, + ) + second_half_request = HttpRequest( + "https://example.org/test", + query_params={"start": "2024-01-01T12:00:00Z", "end": "2024-01-01T23:59:59Z"}, + ) + + with HttpMocker() as http_mocker: + http_mocker.get(full_window_request, HttpResponse("", 400)) + http_mocker.get(first_half_request, HttpResponse(json.dumps({"items": [{"id": 1}]}), 200)) + http_mocker.get(second_half_request, HttpResponse(json.dumps({"items": [{"id": 2}]}), 200)) + + messages = list(_read_request_window_splitting_source(_request_window_splitting_manifest())) + + http_mocker.assert_number_of_calls(full_window_request, 1) + http_mocker.assert_number_of_calls(first_half_request, 1) + http_mocker.assert_number_of_calls(second_half_request, 1) + + assert sorted( + message.record.data["id"] for message in messages if message.type == Type.RECORD + ) == [1, 2] + + +def test_given_split_request_window_action_when_read_then_final_state_reflects_full_original_window(): + """ + Regression coverage for a bug where records read from a split child window carried the child's own slice + as `Record.associated_slice`, which `ConcurrentCursor.observe()` keys its bookkeeping by - a key + `close_partition()` (which always looks up by the *original* partition's slice) could never find. Left + unfixed, `_get_latest_complete_time`'s `first_interval.get("most_recent_cursor_value") or + first_interval[START_KEY]` fallback would silently commit the window's *start* boundary as the emitted + state instead of the actual highest cursor value observed across the two split children - reverting the + cursor all the way back to the beginning of the (already fully synced) day on the very next sync, rather + than a merely imprecise-but-safe value. + """ + full_window_request = HttpRequest( + "https://example.org/test", + query_params={"start": "2024-01-01T00:00:00Z", "end": "2024-01-01T23:59:59Z"}, + ) + first_half_request = HttpRequest( + "https://example.org/test", + query_params={"start": "2024-01-01T00:00:00Z", "end": "2024-01-01T11:59:59Z"}, + ) + second_half_request = HttpRequest( + "https://example.org/test", + query_params={"start": "2024-01-01T12:00:00Z", "end": "2024-01-01T23:59:59Z"}, + ) + + with HttpMocker() as http_mocker: + http_mocker.get(full_window_request, HttpResponse("", 400)) + http_mocker.get( + first_half_request, + HttpResponse( + json.dumps({"items": [{"id": 1, "updated_at": "2024-01-01T05:00:00Z"}]}), 200 + ), + ) + http_mocker.get( + second_half_request, + # the higher cursor value lives in the *second* half, so a fallback to the slice's end boundary + # instead of the true observed maximum would be easy to miss if it happened to match by accident + HttpResponse( + json.dumps({"items": [{"id": 2, "updated_at": "2024-01-01T20:00:00Z"}]}), 200 + ), + ) + + messages = list(_read_request_window_splitting_source(_request_window_splitting_manifest())) + + states = get_states_for_stream(stream_name="Test", messages=messages) + assert states + # the declarative DatetimeBasedCursor emits sequential-format state (a flat {cursor_field: value} dict) by + # default; every emitted state must show the true observed maximum, never the window's start boundary + assert all( + state.stream.stream_state.__dict__ == {"updated_at": "2024-01-01T20:00:00Z"} + for state in states + ) + + def test_given_pagination_limit_reached_when_read_then_reset_pagination(): input_config = {} manifest = { diff --git a/unit_tests/sources/streams/concurrent/test_cursor.py b/unit_tests/sources/streams/concurrent/test_cursor.py index 13fe1df875..fb86f175d3 100644 --- a/unit_tests/sources/streams/concurrent/test_cursor.py +++ b/unit_tests/sources/streams/concurrent/test_cursor.py @@ -11,6 +11,7 @@ import freezegun import pytest +from airbyte_cdk.models import FailureType from airbyte_cdk.sources.connector_state_manager import ConnectorStateManager from airbyte_cdk.sources.message import MessageRepository from airbyte_cdk.sources.streams import NO_CURSOR_STATE_KEY @@ -36,6 +37,7 @@ EpochValueConcurrentStreamStateConverter, ) from airbyte_cdk.sources.types import Record, StreamSlice +from airbyte_cdk.utils.traced_exception import AirbyteTracedException _A_STREAM_NAME = "a stream name" _A_STREAM_NAMESPACE = "a stream namespace" @@ -1411,3 +1413,264 @@ def test_final_state_cursor_get_cursor_datetime_from_state_returns_now_for_no_cu result_with_empty_state = cursor.get_cursor_datetime_from_state({}) assert result_with_empty_state is None + + +class ConcurrentCursorReduceWindowTest(TestCase): + """ + Covers ConcurrentCursor.split_request_window(), the implementation `SimpleRetriever.request_window_splitter` + binds to for declarative `request_window_splitting`. See + airbyte_cdk.sources.declarative.retrievers.simple_retriever. + """ + + def setUp(self) -> None: + self._message_repository = Mock(spec=MessageRepository) + self._state_manager = Mock(spec=ConnectorStateManager) + + def _cursor( + self, + granularity: Optional[timedelta] = timedelta(seconds=1), + datetime_format: str = "%Y-%m-%dT%H:%M:%SZ", + slice_boundary_fields: Optional[tuple] = _SLICE_BOUNDARY_FIELDS, + ) -> ConcurrentCursor: + return ConcurrentCursor( + _A_STREAM_NAME, + _A_STREAM_NAMESPACE, + {}, + self._message_repository, + self._state_manager, + CustomFormatConcurrentStreamStateConverter( + datetime_format, is_sequential_state=_NOT_SEQUENTIAL + ), + CursorField(_A_CURSOR_FIELD_KEY), + slice_boundary_fields, + datetime(2024, 1, 1, tzinfo=timezone.utc), + lambda: datetime(2024, 2, 1, tzinfo=timezone.utc), + cursor_granularity=granularity, + ) + + def _slice(self, start: str, end: str) -> StreamSlice: + return StreamSlice( + partition={}, + cursor_slice={_LOWER_SLICE_BOUNDARY_FIELD: start, _UPPER_SLICE_BOUNDARY_FIELD: end}, + ) + + def test_given_window_when_reduce_then_split_in_half_at_cursor_granularity_boundary( + self, + ) -> None: + cursor = self._cursor() + + children = cursor.split_request_window( + self._slice("2024-01-01T00:00:00Z", "2024-01-01T23:59:59Z") + ) + + assert children == [ + self._slice("2024-01-01T00:00:00Z", "2024-01-01T11:59:59Z"), + self._slice("2024-01-01T12:00:00Z", "2024-01-01T23:59:59Z"), + ] + + def test_given_odd_number_of_granularity_units_when_reduce_then_children_still_cover_parent_exactly( + self, + ) -> None: + cursor = self._cursor() + + # 3 whole seconds: [00:00:00, 00:00:02] + children = cursor.split_request_window( + self._slice("2024-01-01T00:00:00Z", "2024-01-01T00:00:02Z") + ) + + assert children == [ + self._slice("2024-01-01T00:00:00Z", "2024-01-01T00:00:00Z"), + self._slice("2024-01-01T00:00:01Z", "2024-01-01T00:00:02Z"), + ] + + def test_given_single_granularity_unit_window_when_reduce_then_return_none(self) -> None: + cursor = self._cursor() + + assert ( + cursor.split_request_window(self._slice("2024-01-01T00:00:00Z", "2024-01-01T00:00:00Z")) + is None + ) + + def test_given_window_at_or_below_min_split_window_when_reduce_then_return_none(self) -> None: + """ + `min_split_window` is a connector-configured floor independent of, and typically looser than, + `cursor_granularity`: a window already at or below it must not be split further even though + `cursor_granularity` alone would still allow it. + """ + cursor = self._cursor() + + # spans exactly 1 day - at the floor when min_split_window is 1 day, even though cursor_granularity + # (1 second) would otherwise allow splitting this down much further + assert ( + cursor.split_request_window( + self._slice("2024-01-01T00:00:00Z", "2024-01-01T23:59:59Z"), + min_split_window=timedelta(days=1), + ) + is None + ) + + def test_given_window_above_min_split_window_when_reduce_then_split_normally(self) -> None: + cursor = self._cursor() + + children = cursor.split_request_window( + self._slice("2024-01-01T00:00:00Z", "2024-01-02T23:59:59Z"), + min_split_window=timedelta(days=1), + ) + + assert children == [ + self._slice("2024-01-01T00:00:00Z", "2024-01-01T23:59:59Z"), + self._slice("2024-01-02T00:00:00Z", "2024-01-02T23:59:59Z"), + ] + + def test_given_repeated_reduction_when_reduce_then_eventually_reaches_granularity_floor( + self, + ) -> None: + cursor = self._cursor() + current = self._slice("2024-01-01T00:00:00Z", "2024-01-01T23:59:59Z") + + depth = 0 + while True: + children = cursor.split_request_window(current) + if children is None: + break + # exactly at the boundary + assert ( + children[0].cursor_slice[_UPPER_SLICE_BOUNDARY_FIELD] + < current.cursor_slice[_UPPER_SLICE_BOUNDARY_FIELD] + ) + current = children[0] + depth += 1 + assert depth < 30, "split_request_window did not converge to the granularity floor" + + assert ( + current.cursor_slice[_LOWER_SLICE_BOUNDARY_FIELD] + == current.cursor_slice[_UPPER_SLICE_BOUNDARY_FIELD] + == "2024-01-01T00:00:00Z" + ) + + def test_given_reversed_boundaries_when_reduce_then_return_none(self) -> None: + cursor = self._cursor() + + assert ( + cursor.split_request_window(self._slice("2024-01-01T00:00:05Z", "2024-01-01T00:00:00Z")) + is None + ) + + def test_given_no_cursor_granularity_when_reduce_then_return_none(self) -> None: + cursor = self._cursor(granularity=None) + + assert ( + cursor.split_request_window(self._slice("2024-01-01T00:00:00Z", "2024-01-01T23:59:59Z")) + is None + ) + + def test_given_malformed_boundary_when_reduce_then_raise(self) -> None: + """ + A malformed boundary means the slice was produced wrong - a bug, not the API rejecting a window - so + this must not be silently swallowed as `None` (which `SimpleRetriever` reports as "the API kept + rejecting...", misattributing the cause). + """ + cursor = self._cursor() + + with pytest.raises(AirbyteTracedException) as exception: + cursor.split_request_window( + StreamSlice( + partition={}, + cursor_slice={ + _LOWER_SLICE_BOUNDARY_FIELD: "not-a-date", + _UPPER_SLICE_BOUNDARY_FIELD: "2024-01-01T23:59:59Z", + }, + ) + ) + assert exception.value.failure_type == FailureType.system_error + assert _LOWER_SLICE_BOUNDARY_FIELD in exception.value.internal_message + + def test_given_missing_boundary_field_when_reduce_then_raise(self) -> None: + cursor = self._cursor() + + with pytest.raises(AirbyteTracedException) as exception: + cursor.split_request_window( + StreamSlice(partition={}, cursor_slice={"only_one_field": "x"}) + ) + assert exception.value.failure_type == FailureType.system_error + + def test_split_request_window_preserves_partition_and_extra_fields(self) -> None: + cursor = self._cursor() + original = StreamSlice( + partition={"parent_id": "123"}, + cursor_slice={ + _LOWER_SLICE_BOUNDARY_FIELD: "2024-01-01T00:00:00Z", + _UPPER_SLICE_BOUNDARY_FIELD: "2024-01-01T23:59:59Z", + }, + extra_fields={"some_extra": "value"}, + ) + + children = cursor.split_request_window(original) + + for child in children: + assert child.partition == {"parent_id": "123"} + assert child.extra_fields == {"some_extra": "value"} + + def test_given_epoch_cursor_when_reduce_then_split_using_epoch_output_format(self) -> None: + cursor = ConcurrentCursor( + _A_STREAM_NAME, + _A_STREAM_NAMESPACE, + {}, + self._message_repository, + self._state_manager, + EpochValueConcurrentStreamStateConverter(is_sequential_state=_NOT_SEQUENTIAL), + CursorField(_A_CURSOR_FIELD_KEY), + _SLICE_BOUNDARY_FIELDS, + 0, + lambda: 1000, + cursor_granularity=timedelta(seconds=1), + ) + + children = cursor.split_request_window( + StreamSlice( + partition={}, + cursor_slice={ + _LOWER_SLICE_BOUNDARY_FIELD: 0, + _UPPER_SLICE_BOUNDARY_FIELD: 99, + }, + ) + ) + + assert children == [ + StreamSlice( + partition={}, + cursor_slice={_LOWER_SLICE_BOUNDARY_FIELD: 0, _UPPER_SLICE_BOUNDARY_FIELD: 49}, + ), + StreamSlice( + partition={}, + cursor_slice={_LOWER_SLICE_BOUNDARY_FIELD: 50, _UPPER_SLICE_BOUNDARY_FIELD: 99}, + ), + ] + + def test_given_multiple_datetime_formats_when_reduce_then_round_trip_format_correctly( + self, + ) -> None: + cursor = self._cursor(datetime_format="%Y-%m-%d %H:%M:%S") + + children = cursor.split_request_window( + self._slice("2024-01-01 00:00:00", "2024-01-01 23:59:59") + ) + + assert children == [ + self._slice("2024-01-01 00:00:00", "2024-01-01 11:59:59"), + self._slice("2024-01-01 12:00:00", "2024-01-01 23:59:59"), + ] + + def test_given_non_utc_timezone_offset_when_reduce_then_preserve_offset_in_round_trip( + self, + ) -> None: + cursor = self._cursor(datetime_format="%Y-%m-%dT%H:%M:%S%z") + + children = cursor.split_request_window( + self._slice("2024-01-01T00:00:00+0530", "2024-01-01T23:59:59+0530") + ) + + assert children == [ + self._slice("2024-01-01T00:00:00+0530", "2024-01-01T11:59:59+0530"), + self._slice("2024-01-01T12:00:00+0530", "2024-01-01T23:59:59+0530"), + ] diff --git a/unit_tests/sources/streams/http/test_http_client.py b/unit_tests/sources/streams/http/test_http_client.py index 8724c48e49..ad3ee5acd1 100644 --- a/unit_tests/sources/streams/http/test_http_client.py +++ b/unit_tests/sources/streams/http/test_http_client.py @@ -34,6 +34,9 @@ from airbyte_cdk.sources.streams.http.page_size_reduction_exception import ( PageSizeReductionRequiredException, ) +from airbyte_cdk.sources.streams.http.request_window_split_exception import ( + RequestWindowSplitRequiredException, +) from airbyte_cdk.sources.streams.http.requests_native_auth import ( Oauth2Authenticator, TokenAuthenticator, @@ -1838,3 +1841,107 @@ def test_given_no_reduce_page_size_action_then_log_the_response_as_a_page(): logged = [json.loads(message.log.message) for message in message_repository.consume_queue()] assert [entry["http"].get("is_auxiliary") for entry in logged] == [None] + + +def test_send_raises_request_window_splitting_required_exception_with_split_request_window_response_action(): + mocked_session = MagicMock(spec=requests.Session) + http_client = HttpClient( + name="test", + logger=MagicMock(), + error_handler=HttpStatusErrorHandler( + logger=MagicMock(), + error_mapping={ + 400: ErrorResolution( + ResponseAction.SPLIT_REQUEST_WINDOW, + FailureType.config_error, + "test reduce request window message", + ) + }, + ), + session=mocked_session, + ) + mocked_response = requests.Response() + mocked_response.status_code = 400 + mocked_session.send.return_value = mocked_response + + # the retriever is responsible for reducing and re-reading the window, so the backoff handlers must not retry + with pytest.raises(RequestWindowSplitRequiredException) as exception: + http_client.send_request(http_method="get", url="https://airbyte.io", request_kwargs={}) + + assert http_client._session.send.call_count == 1 + assert "test reduce request window message" in exception.value.internal_message + assert exception.value.message == ( + "The API rejected the current request window of stream test and requires a smaller one. If this " + "message ends a sync, the stream is not set up to split its window: add `request_window_splitting` " + "to its retriever, or remove the `SPLIT_REQUEST_WINDOW` action from its error handler." + ) + # the classifying error's own failure_type is preserved rather than defaulting to config_error + assert exception.value.failure_type == FailureType.config_error + + +def test_given_no_classified_failure_type_when_split_request_window_then_default_to_config_error(): + mocked_session = MagicMock(spec=requests.Session) + http_client = HttpClient( + name="test", + logger=MagicMock(), + error_handler=HttpStatusErrorHandler( + logger=MagicMock(), + error_mapping={400: ErrorResolution(ResponseAction.SPLIT_REQUEST_WINDOW, None, None)}, + ), + session=mocked_session, + ) + mocked_response = requests.Response() + mocked_response.status_code = 400 + mocked_session.send.return_value = mocked_response + + with pytest.raises(RequestWindowSplitRequiredException) as exception: + http_client.send_request(http_method="get", url="https://airbyte.io", request_kwargs={}) + + assert exception.value.failure_type == FailureType.config_error + + +def test_given_split_request_window_action_then_log_the_response_as_an_auxiliary_request(): + """ + Mirrors test_given_reduce_page_size_action_then_log_the_response_as_an_auxiliary_request: a response + resolving to SPLIT_REQUEST_WINDOW never becomes a page either - the retriever splits the window and + re-reads it - so it must not count against the Connector Builder's per-slice page limit. + """ + message_repository = InMemoryMessageRepository(Level.DEBUG) + mocked_session = MagicMock(spec=requests.Session) + http_client = HttpClient( + name="test", + logger=MagicMock(), + error_handler=HttpStatusErrorHandler( + logger=MagicMock(), + error_mapping={ + 400: ErrorResolution( + ResponseAction.SPLIT_REQUEST_WINDOW, + FailureType.config_error, + "test reduce request window message", + ) + }, + ), + session=mocked_session, + message_repository=message_repository, + ) + mocked_response = requests.Response() + mocked_response.status_code = 400 + mocked_response.request = requests.Request(method="GET", url="https://airbyte.io").prepare() + mocked_session.send.return_value = mocked_response + + with pytest.raises(RequestWindowSplitRequiredException): + http_client.send_request( + http_method="get", + url="https://airbyte.io", + request_kwargs={}, + log_formatter=lambda response: format_http_message( + response, "a title", "a description", "test" + ), + ) + + logged = [json.loads(message.log.message) for message in message_repository.consume_queue()] + assert [entry["http"]["is_auxiliary"] for entry in logged] == [True] + assert logged[0]["http"]["title"] == ( + "Stream 'test' request window rejected, retrying with a smaller window" + ) + assert "no records" in logged[0]["http"]["description"]