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 aa546c73a..5da2eba4e 100644 --- a/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py +++ b/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py @@ -3910,6 +3910,15 @@ def _create_async_job_status_mapping( f"API status {status} is already set for CDK status {cdk_status}. Please ensure API statuses are only provided once" ) api_status_to_cdk_status[status] = self._get_async_job_status(cdk_status) + + if not any( + cdk_status in (AsyncJobStatus.COMPLETED, AsyncJobStatus.SKIPPED) + for cdk_status in api_status_to_cdk_status.values() + ): + raise ValueError( + "AsyncJobStatusMap must map at least one API status to `completed` or `skipped`. " + "Without a terminal success status, async jobs can never finish and the stream polls until `polling_job_timeout` expires." + ) return api_status_to_cdk_status def _get_async_job_status(self, status: str) -> AsyncJobStatus: diff --git a/airbyte_cdk/sources/declarative/requesters/http_job_repository.py b/airbyte_cdk/sources/declarative/requesters/http_job_repository.py index d837ed902..f62da2a41 100644 --- a/airbyte_cdk/sources/declarative/requesters/http_job_repository.py +++ b/airbyte_cdk/sources/declarative/requesters/http_job_repository.py @@ -102,14 +102,18 @@ def _get_validated_job_status(self, response: requests.Response) -> AsyncJobStat AsyncJobStatus: The validated job status. Raises: - ValueError: If the API status is unknown. + AirbyteTracedException: If the API status is not mapped to a CDK status. The failure type is + `config_error` so that the orchestrator stops the sync immediately instead of polling the + job until `polling_job_timeout` expires. """ api_status = next(iter(self.status_extractor.extract_records(response)), None) job_status = self.status_mapping.get(str(api_status), None) if job_status is None: - raise ValueError( - f"API status `{api_status}` is unknown. Contact the connector developer to make sure this status is supported." + raise AirbyteTracedException( + message=f'Async job status "{api_status}" is not supported by the connector.', + internal_message=f"Async job status `{api_status}` is missing from the connector's `status_mapping`, which declares: {sorted(self.status_mapping.keys())}. Contact the connector developer to add support for this status.", + failure_type=FailureType.config_error, ) return job_status 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 dce40e624..cec772f1e 100644 --- a/unit_tests/sources/declarative/async_job/test_job_orchestrator.py +++ b/unit_tests/sources/declarative/async_job/test_job_orchestrator.py @@ -198,6 +198,32 @@ def test_given_failure_when_create_and_get_completed_partitions_then_raise_excep == [call(_A_STREAM_SLICE)] * _MAX_NUMBER_OF_ATTEMPTS ) + @mock.patch(sleep_mock_target) + def test_given_traced_config_error_when_update_job_status_then_abort_running_jobs_immediately( + self, mock_sleep: MagicMock + ) -> None: + self._job_repository.start.return_value = self._job_for_a_slice + config_error = AirbyteTracedException( + "Async job status is not supported by the connector.", + failure_type=FailureType.config_error, + ) + + def update_status(_jobs: Set[AsyncJob]) -> None: + if self._job_repository.update_jobs_status.call_count == 1: + raise config_error + self._job_for_a_slice.update_status(AsyncJobStatus.COMPLETED) + + self._job_repository.update_jobs_status.side_effect = update_status + orchestrator = self._orchestrator([_A_STREAM_SLICE]) + + with pytest.raises(AirbyteTracedException) as exception_info: + list(orchestrator.create_and_get_completed_partitions()) + + assert exception_info.value is config_error + self._job_repository.update_jobs_status.assert_called_once_with({self._job_for_a_slice}) + self._job_repository.abort.assert_called_once_with(self._job_for_a_slice) + assert len(orchestrator._job_tracker._jobs) == 0 + def test_when_fetch_records_then_yield_records_from_each_job(self) -> None: self._job_repository.fetch_records.return_value = [_ANY_RECORD] orchestrator = self._orchestrator([_A_STREAM_SLICE]) 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 21c99adc7..b70e315ca 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 @@ -4861,6 +4861,52 @@ def test_create_async_retriever(): assert download_retriever_record_selector.schema_normalization._config.name == "NoTransform" +def test_create_async_retriever_requires_terminal_success_status(): + definition = { + "type": "AsyncRetriever", + "status_mapping": { + "running": ["pending"], + "completed": [], + "failed": ["failed"], + "timeout": ["timeout"], + }, + "record_selector": { + "type": "RecordSelector", + "extractor": {"type": "DpathExtractor", "field_path": []}, + }, + "status_extractor": {"type": "DpathExtractor", "field_path": ["status"]}, + "creation_requester": { + "type": "HttpRequester", + "path": "/jobs", + "url_base": "https://api.test.com", + "http_method": "POST", + }, + "polling_requester": { + "type": "HttpRequester", + "path": "/jobs/{{ creation_response['id'] }}", + "url_base": "https://api.test.com", + "http_method": "GET", + }, + "download_requester": { + "type": "HttpRequester", + "path": "{{ download_target }}", + "url_base": "", + "http_method": "GET", + }, + } + + with pytest.raises(ValueError, match="AsyncJobStatusMap must map"): + factory.create_component( + model_type=AsyncRetrieverModel, + component_definition=definition, + name="test_stream", + primary_key=None, + stream_slicer=None, + transformations=[], + config={}, + ) + + def test_api_budget(): manifest = { "type": "DeclarativeSource", diff --git a/unit_tests/sources/declarative/requesters/test_http_job_repository.py b/unit_tests/sources/declarative/requesters/test_http_job_repository.py index 473c3d99e..c56d43bba 100644 --- a/unit_tests/sources/declarative/requesters/test_http_job_repository.py +++ b/unit_tests/sources/declarative/requesters/test_http_job_repository.py @@ -8,6 +8,7 @@ import pytest +from airbyte_cdk.models import FailureType from airbyte_cdk.sources.declarative.async_job.status import AsyncJobStatus from airbyte_cdk.sources.declarative.decoders import NoopDecoder from airbyte_cdk.sources.declarative.decoders.json_decoder import JsonDecoder @@ -34,6 +35,7 @@ from airbyte_cdk.sources.types import StreamSlice from airbyte_cdk.sources.utils.transform import TransformConfig, TypeTransformer from airbyte_cdk.test.mock_http import HttpMocker, HttpRequest, HttpResponse +from airbyte_cdk.utils import AirbyteTracedException _ANY_CONFIG = {} _ANY_SLICE = StreamSlice(partition={}, cursor_slice={}) @@ -119,8 +121,22 @@ def test_given_unknown_status_when_update_jobs_status_then_raise_error(self) -> ) job = self._repository.start(_ANY_SLICE) - with pytest.raises(ValueError): + with pytest.raises(AirbyteTracedException) as exception_info: self._repository.update_jobs_status([job]) + assert exception_info.value.failure_type == FailureType.config_error + assert "invalid_status" in str(exception_info.value) + + def test_given_missing_status_when_update_jobs_status_then_raise_error(self) -> None: + self._mock_create_response(_A_JOB_ID) + self._http_mocker.get( + HttpRequest(url=f"{_EXPORT_URL}/{_A_JOB_ID}"), + HttpResponse(body=json.dumps({"id": _A_JOB_ID})), + ) + job = self._repository.start(_ANY_SLICE) + + with pytest.raises(AirbyteTracedException) as exception_info: + self._repository.update_jobs_status([job]) + assert exception_info.value.failure_type == FailureType.config_error def test_given_multiple_jobs_when_update_jobs_status_then_all_the_jobs_are_updated( self,