diff --git a/airbyte_cdk/sources/declarative/async_job/job.py b/airbyte_cdk/sources/declarative/async_job/job.py index ea83b7456c..597b0d1991 100644 --- a/airbyte_cdk/sources/declarative/async_job/job.py +++ b/airbyte_cdk/sources/declarative/async_job/job.py @@ -1,7 +1,7 @@ # Copyright (c) 2024 Airbyte, Inc., all rights reserved. -from datetime import timedelta +from datetime import datetime, timedelta, timezone from typing import Optional from airbyte_cdk.sources.declarative.async_job.timer import Timer @@ -19,11 +19,17 @@ class AsyncJob: """ def __init__( - self, api_job_id: str, job_parameters: StreamSlice, timeout: Optional[timedelta] = None + self, + api_job_id: str, + job_parameters: StreamSlice, + timeout: Optional[timedelta] = None, + is_creation_failure: bool = False, ) -> None: self._api_job_id = api_job_id self._job_parameters = job_parameters self._status = AsyncJobStatus.RUNNING + self._retry_after: Optional[datetime] = None + self._is_creation_failure = is_creation_failure timeout = timeout if timeout else timedelta(minutes=60) self._timer = Timer(timeout) @@ -54,5 +60,23 @@ def update_status(self, status: AsyncJobStatus) -> None: self._status = status + def is_creation_failure(self) -> bool: + """Return True if this job was never actually created on the API side.""" + return self._is_creation_failure + + def set_retry_after(self, retry_after: datetime) -> None: + """Set the earliest time this job can be retried.""" + self._retry_after = retry_after + + def retry_deferred(self) -> bool: + """Return True if a deferred retry has been scheduled.""" + return self._retry_after is not None + + def ready_to_retry(self) -> bool: + """Return True if the job has no deferred retry or the wait period has elapsed.""" + if self._retry_after is None: + return True + return datetime.now(tz=timezone.utc) >= self._retry_after + def __repr__(self) -> str: return f"AsyncJob(api_job_id={self.api_job_id()}, job_parameters={self.job_parameters()}, status={self.status()})" diff --git a/airbyte_cdk/sources/declarative/async_job/job_orchestrator.py b/airbyte_cdk/sources/declarative/async_job/job_orchestrator.py index 784abb6b54..3607ad0264 100644 --- a/airbyte_cdk/sources/declarative/async_job/job_orchestrator.py +++ b/airbyte_cdk/sources/declarative/async_job/job_orchestrator.py @@ -5,7 +5,7 @@ import time import traceback import uuid -from datetime import timedelta +from datetime import datetime, timedelta, timezone from typing import ( Any, Generator, @@ -168,6 +168,7 @@ def __init__( exceptions_to_break_on: Iterable[Type[Exception]] = tuple(), has_bulk_parent: bool = False, job_max_retry: Optional[int] = None, + failed_retry_wait_time_in_seconds: Optional[int] = None, ) -> None: """ If the stream slices provided as a parameters relies on a async job streams that relies on the same JobTracker, `has_bulk_parent` @@ -181,6 +182,9 @@ def __init__( "An AsyncJobStatus has been either removed or added which means the logic of this class needs to be reviewed. Once the logic has been updated, please update _KNOWN_JOB_STATUSES" ) + if failed_retry_wait_time_in_seconds is not None and failed_retry_wait_time_in_seconds <= 0: + raise ValueError("failed_retry_wait_time_in_seconds must be >= 1") + self._job_repository: AsyncJobRepository = job_repository self._slice_iterator = LookaheadIterator(slices) self._running_partitions: List[AsyncPartition] = [] @@ -189,6 +193,7 @@ def __init__( self._exceptions_to_break_on: Tuple[Type[Exception], ...] = tuple(exceptions_to_break_on) self._has_bulk_parent = has_bulk_parent self._job_max_retry = job_max_retry + self._failed_retry_wait_time_in_seconds = failed_retry_wait_time_in_seconds self._non_breaking_exceptions: List[Exception] = [] @@ -196,6 +201,29 @@ def _replace_failed_jobs(self, partition: AsyncPartition) -> None: failed_status_jobs = (AsyncJobStatus.FAILED, AsyncJobStatus.TIMED_OUT) jobs_to_replace = [job for job in partition.jobs if job.status() in failed_status_jobs] for job in jobs_to_replace: + if ( + self._failed_retry_wait_time_in_seconds is not None + and job.status() == AsyncJobStatus.FAILED + and not job.is_creation_failure() + ): + if not job.ready_to_retry(): + lazy_log( + LOGGER, + logging.DEBUG, + lambda: f"Job {job.api_job_id()} is not ready to retry yet (deferred). Skipping.", + ) + continue + if not job.retry_deferred(): + job.set_retry_after( + datetime.now(tz=timezone.utc) + + timedelta(seconds=self._failed_retry_wait_time_in_seconds) + ) + lazy_log( + LOGGER, + logging.INFO, + lambda: f"Job {job.api_job_id()} failed. Deferring retry for {self._failed_retry_wait_time_in_seconds} seconds.", + ) + continue new_job = self._start_job(job.job_parameters(), job.api_job_id()) partition.replace_job(job, [new_job]) @@ -281,7 +309,12 @@ def _keep_api_budget_with_failed_job( return job def _create_failed_job(self, stream_slice: StreamSlice) -> AsyncJob: - job = AsyncJob(f"{uuid.uuid4()} - Job that could not start", stream_slice, _NO_TIMEOUT) + job = AsyncJob( + f"{uuid.uuid4()} - Job that could not start", + stream_slice, + _NO_TIMEOUT, + is_creation_failure=True, + ) job.update_status(AsyncJobStatus.FAILED) return job diff --git a/airbyte_cdk/sources/declarative/declarative_component_schema.yaml b/airbyte_cdk/sources/declarative/declarative_component_schema.yaml index 49a005afba..6fbda8f006 100644 --- a/airbyte_cdk/sources/declarative/declarative_component_schema.yaml +++ b/airbyte_cdk/sources/declarative/declarative_component_schema.yaml @@ -4107,6 +4107,14 @@ definitions: - type: string interpolation_context: - config + failed_retry_wait_time_in_seconds: + description: Time in seconds to wait before retrying a failed async job. Only applies to jobs that ran on the API side and reported a FAILED status (e.g. report generation failed due to a cooldown). Creation failures (HTTP errors when starting a job, such as 429s) and TIMED_OUT jobs are retried immediately and are not affected by this setting. When set, the orchestrator defers retry of real failed jobs until the wait time has elapsed, without blocking other jobs. + anyOf: + - type: integer + minimum: 1 + - type: string + interpolation_context: + - config download_target_requester: description: Requester component that describes how to prepare HTTP requests to send to the source API to extract the url from polling response by the completed async job. anyOf: diff --git a/airbyte_cdk/sources/declarative/models/declarative_component_schema.py b/airbyte_cdk/sources/declarative/models/declarative_component_schema.py index 4415855277..931cef7f1b 100644 --- a/airbyte_cdk/sources/declarative/models/declarative_component_schema.py +++ b/airbyte_cdk/sources/declarative/models/declarative_component_schema.py @@ -3074,6 +3074,10 @@ class AsyncRetriever(BaseModel): None, description="The time in minutes after which the single Async Job should be considered as Timed Out.", ) + failed_retry_wait_time_in_seconds: Optional[Union[int, str]] = Field( + None, + description="Time in seconds to wait before retrying a failed async job. Only applies to jobs that ran on the API side and reported a FAILED status (e.g. report generation failed due to a cooldown). Creation failures (HTTP errors when starting a job, such as 429s) and TIMED_OUT jobs are retried immediately and are not affected by this setting. When set, the orchestrator defers retry of real failed jobs until the wait time has elapsed, without blocking other jobs.", + ) download_target_requester: Optional[Union[HttpRequester, CustomRequester]] = Field( None, description="Requester component that describes how to prepare HTTP requests to send to the source API to extract the url from polling response by the completed async job.", 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 8b806d2f59..bc9b4f82c0 100644 --- a/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py +++ b/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py @@ -3955,6 +3955,17 @@ def _get_job_timeout() -> datetime.timedelta: job_timeout=_get_job_timeout(), ) + failed_retry_wait_time_in_seconds: Optional[int] = ( + int( + InterpolatedString.create( + str(model.failed_retry_wait_time_in_seconds), + parameters={}, + ).eval(config) + ) + if model.failed_retry_wait_time_in_seconds + else None + ) + async_job_partition_router = AsyncJobPartitionRouter( job_orchestrator_factory=lambda stream_slices: AsyncJobOrchestrator( job_repository, @@ -3966,6 +3977,7 @@ def _get_job_timeout() -> datetime.timedelta: # set the `job_max_retry` to 1 for the `Connector Builder`` use-case. # `None` == default retry is set to 3 attempts, under the hood. job_max_retry=1 if self._emit_connector_builder_messages else None, + failed_retry_wait_time_in_seconds=failed_retry_wait_time_in_seconds, ), stream_slicer=stream_slicer, config=config, diff --git a/unit_tests/sources/declarative/async_job/test_job.py b/unit_tests/sources/declarative/async_job/test_job.py index 6399433e44..975297a2ee 100644 --- a/unit_tests/sources/declarative/async_job/test_job.py +++ b/unit_tests/sources/declarative/async_job/test_job.py @@ -1,8 +1,10 @@ # Copyright (c) 2024 Airbyte, Inc., all rights reserved. import time -from datetime import timedelta -from unittest import TestCase +from datetime import datetime, timedelta, timezone +from typing import Optional + +import pytest from airbyte_cdk.sources.declarative.async_job.job import AsyncJob from airbyte_cdk.sources.declarative.async_job.status import AsyncJobStatus @@ -14,19 +16,30 @@ _IMMEDIATELY_TIMED_OUT = timedelta(microseconds=1) -class AsyncJobTest(TestCase): - def test_given_timer_is_not_out_when_status_then_return_actual_status(self) -> None: - job = AsyncJob(_AN_API_JOB_ID, _ANY_STREAM_SLICE, _A_VERY_BIG_TIMEOUT) - assert job.status() == AsyncJobStatus.RUNNING +def test_given_timer_is_not_out_when_status_then_return_actual_status() -> None: + job = AsyncJob(_AN_API_JOB_ID, _ANY_STREAM_SLICE, _A_VERY_BIG_TIMEOUT) + assert job.status() == AsyncJobStatus.RUNNING + + +def test_given_timer_is_out_when_status_then_return_timed_out() -> None: + job = AsyncJob(_AN_API_JOB_ID, _ANY_STREAM_SLICE, _IMMEDIATELY_TIMED_OUT) + time.sleep(0.001) + assert job.status() == AsyncJobStatus.TIMED_OUT - def test_given_timer_is_out_when_status_then_return_timed_out(self) -> None: - job = AsyncJob(_AN_API_JOB_ID, _ANY_STREAM_SLICE, _IMMEDIATELY_TIMED_OUT) - time.sleep(0.001) - assert job.status() == AsyncJobStatus.TIMED_OUT - def test_given_status_is_terminal_when_update_status_then_stop_timer(self) -> None: - """ - This test will become important once we will print stats associated with jobs. As for now, we stop the timer but do not return any - metrics regarding the timer so it is not useful. - """ - pass +@pytest.mark.parametrize( + "retry_after_offset,expected_deferred,expected_ready", + [ + pytest.param(None, False, True, id="no_retry_after_set"), + pytest.param(timedelta(hours=1), True, False, id="retry_after_in_future"), + pytest.param(-timedelta(seconds=1), True, True, id="retry_after_in_past"), + ], +) +def test_retry_after( + retry_after_offset: Optional[timedelta], expected_deferred: bool, expected_ready: bool +) -> None: + job = AsyncJob(_AN_API_JOB_ID, _ANY_STREAM_SLICE, _A_VERY_BIG_TIMEOUT) + if retry_after_offset is not None: + job.set_retry_after(datetime.now(tz=timezone.utc) + retry_after_offset) + assert job.retry_deferred() == expected_deferred + assert job.ready_to_retry() == expected_ready diff --git a/unit_tests/sources/declarative/async_job/test_job_orchestrator.py b/unit_tests/sources/declarative/async_job/test_job_orchestrator.py index 1558e5987f..dce40e6242 100644 --- a/unit_tests/sources/declarative/async_job/test_job_orchestrator.py +++ b/unit_tests/sources/declarative/async_job/test_job_orchestrator.py @@ -4,6 +4,7 @@ import sys import threading import time +from datetime import datetime, timedelta, timezone from typing import Callable, List, Mapping, Optional, Set, Tuple from unittest import TestCase, mock from unittest.mock import MagicMock, Mock, call @@ -416,6 +417,189 @@ def test_given_skipped_does_not_retry(self, mock_sleep: MagicMock) -> None: # start is called only once — SKIPPED does not trigger a retry assert self._job_repository.start.call_count == 1 + @mock.patch(sleep_mock_target) + def test_given_failed_retry_wait_time_when_job_fails_then_defers_retry( + self, mock_sleep: MagicMock + ) -> None: + """When failed_retry_wait_time_in_seconds is set and a job fails, the retry should + be deferred: first call sets retry_after timestamp and skips, subsequent calls skip + until cooldown elapses, then the job is replaced.""" + job_tracker = JobTracker(_NO_JOB_LIMIT) + job = self._an_async_job("deferred-job", _A_STREAM_SLICE) + job_tracker._jobs.add("deferred-job") + partition = AsyncPartition([job], _A_STREAM_SLICE) + orchestrator = AsyncJobOrchestrator( + self._job_repository, + [], + job_tracker, + self._message_repository, + failed_retry_wait_time_in_seconds=1800, + ) + + job.update_status(AsyncJobStatus.FAILED) + + # First call: should set retry_after and NOT replace + orchestrator._replace_failed_jobs(partition) + assert job.retry_deferred() + assert not job.ready_to_retry() + self._job_repository.start.assert_not_called() + + # Second call while cooldown hasn't elapsed: should still skip + orchestrator._replace_failed_jobs(partition) + self._job_repository.start.assert_not_called() + + # Simulate cooldown elapsed by setting retry_after to the past + job.set_retry_after(datetime.now(tz=timezone.utc) - timedelta(seconds=1)) + replacement_job = self._an_async_job("replacement-job", _A_STREAM_SLICE) + self._job_repository.start.return_value = replacement_job + + # Third call after cooldown: should replace the job + orchestrator._replace_failed_jobs(partition) + self._job_repository.start.assert_called_once() + + @mock.patch(sleep_mock_target) + def test_given_no_failed_retry_wait_time_when_job_fails_then_replaces_immediately( + self, mock_sleep: MagicMock + ) -> None: + """Without failed_retry_wait_time_in_seconds, FAILED jobs are replaced immediately + (existing behavior).""" + job_tracker = JobTracker(_NO_JOB_LIMIT) + job = self._an_async_job("immediate-job", _A_STREAM_SLICE) + job_tracker._jobs.add("immediate-job") + partition = AsyncPartition([job], _A_STREAM_SLICE) + replacement_job = self._an_async_job("replacement-job", _A_STREAM_SLICE) + self._job_repository.start.return_value = replacement_job + orchestrator = AsyncJobOrchestrator( + self._job_repository, + [], + job_tracker, + self._message_repository, + ) + + job.update_status(AsyncJobStatus.FAILED) + orchestrator._replace_failed_jobs(partition) + + self._job_repository.start.assert_called_once() + + @mock.patch(sleep_mock_target) + def test_given_failed_retry_wait_time_when_timed_out_job_then_replaces_immediately( + self, mock_sleep: MagicMock + ) -> None: + """TIMED_OUT jobs should be replaced immediately even when + failed_retry_wait_time_in_seconds is set (only FAILED gets deferred).""" + job_tracker = JobTracker(_NO_JOB_LIMIT) + job = self._an_async_job("timed-out-job", _A_STREAM_SLICE) + job_tracker._jobs.add("timed-out-job") + partition = AsyncPartition([job], _A_STREAM_SLICE) + replacement_job = self._an_async_job("replacement-job", _A_STREAM_SLICE) + self._job_repository.start.return_value = replacement_job + orchestrator = AsyncJobOrchestrator( + self._job_repository, + [], + job_tracker, + self._message_repository, + failed_retry_wait_time_in_seconds=1800, + ) + + job.update_status(AsyncJobStatus.TIMED_OUT) + orchestrator._replace_failed_jobs(partition) + + self._job_repository.start.assert_called_once() + + @mock.patch(sleep_mock_target) + def test_given_creation_failure_job_when_cooldown_configured_then_replaces_immediately( + self, mock_sleep: MagicMock + ) -> None: + """Creation-failure jobs (created when API rejects with 429) should be retried + immediately even when failed_retry_wait_time_in_seconds is set. Only real FAILED + jobs (report created but got FATAL) should be deferred.""" + job_tracker = JobTracker(_NO_JOB_LIMIT) + creation_failure_job = AsyncJob( + "creation-failure-job", _A_STREAM_SLICE, is_creation_failure=True + ) + creation_failure_job.update_status(AsyncJobStatus.FAILED) + job_tracker._jobs.add("creation-failure-job") + partition = AsyncPartition([creation_failure_job], _A_STREAM_SLICE) + replacement_job = self._an_async_job("replacement-job", _A_STREAM_SLICE) + self._job_repository.start.return_value = replacement_job + orchestrator = AsyncJobOrchestrator( + self._job_repository, + [], + job_tracker, + self._message_repository, + failed_retry_wait_time_in_seconds=1800, + ) + + orchestrator._replace_failed_jobs(partition) + self._job_repository.start.assert_called_once() + assert not creation_failure_job.retry_deferred() + + @mock.patch(sleep_mock_target) + def test_given_real_failed_then_cooldown_elapses_then_start_returns_creation_failure_then_no_rearm( + self, mock_sleep: MagicMock + ) -> None: + """Regression test for the original bug: real FAILED -> arm cooldown -> cooldown + elapses -> _start_job returns a creation-failure job (API still rejects) -> the + replacement must NOT re-arm cooldown; it should be replaced immediately on the + next tick.""" + job_tracker = JobTracker(_NO_JOB_LIMIT) + real_job = self._an_async_job("real-job", _A_STREAM_SLICE) + job_tracker._jobs.add("real-job") + partition = AsyncPartition([real_job], _A_STREAM_SLICE) + orchestrator = AsyncJobOrchestrator( + self._job_repository, + [], + job_tracker, + self._message_repository, + failed_retry_wait_time_in_seconds=1800, + ) + + real_job.update_status(AsyncJobStatus.FAILED) + + # Tick 1: arms cooldown on the real FAILED job + orchestrator._replace_failed_jobs(partition) + assert real_job.retry_deferred() + self._job_repository.start.assert_not_called() + + # Simulate cooldown elapsed + real_job.set_retry_after(datetime.now(tz=timezone.utc) - timedelta(seconds=1)) + + # _start_job returns a creation-failure job (API still rejects with 429) + creation_failure_replacement = AsyncJob( + "creation-failure-replacement", _A_STREAM_SLICE, is_creation_failure=True + ) + creation_failure_replacement.update_status(AsyncJobStatus.FAILED) + self._job_repository.start.return_value = creation_failure_replacement + job_tracker._jobs.add("creation-failure-replacement") + + # Tick 2: cooldown elapsed, replaces with creation-failure job + orchestrator._replace_failed_jobs(partition) + self._job_repository.start.assert_called_once() + + # The replacement is a creation-failure -> should NOT have cooldown armed + replaced_job = list(partition.jobs)[0] + assert replaced_job.is_creation_failure() + assert not replaced_job.retry_deferred() + + # Tick 3: creation-failure job should be replaced immediately (no deferral) + second_replacement = self._an_async_job("second-replacement", _A_STREAM_SLICE) + self._job_repository.start.return_value = second_replacement + job_tracker._jobs.add("second-replacement") + orchestrator._replace_failed_jobs(partition) + assert self._job_repository.start.call_count == 2 + + def test_create_failed_job_tags_as_creation_failure(self) -> None: + """Verify _create_failed_job produces a job marked as is_creation_failure.""" + orchestrator = AsyncJobOrchestrator( + self._job_repository, + [], + JobTracker(_NO_JOB_LIMIT), + self._message_repository, + ) + job = orchestrator._create_failed_job(_A_STREAM_SLICE) + assert job.is_creation_failure() + assert job.status() == AsyncJobStatus.FAILED + def _mock_repository(self) -> None: self._job_repository = Mock(spec=AsyncJobRepository)