diff --git a/airbyte/_util/api_util.py b/airbyte/_util/api_util.py index 2c4075d83..00c38a2c1 100644 --- a/airbyte/_util/api_util.py +++ b/airbyte/_util/api_util.py @@ -14,9 +14,12 @@ from __future__ import annotations import base64 +import copy import json +import logging +from dataclasses import replace from http import HTTPStatus -from typing import TYPE_CHECKING, Any, Literal +from typing import TYPE_CHECKING, Any, Literal, TypeVar import airbyte_api import requests @@ -42,6 +45,7 @@ if TYPE_CHECKING: from collections.abc import Callable + from enum import Enum from airbyte_api.models import ( DestinationConfiguration, @@ -63,6 +67,11 @@ DEFERRED_CREATE_TIMEOUT_SECS: tuple[float, float] = (5.0, 120.0) """Connect and read timeouts for a deferred-credential create on the Config API.""" +PUBLIC_API_FALLBACK_TIMEOUT_SECS: tuple[float, float] = (5.0, 60.0) +"""Connect and read timeouts for the raw Public API fallback GET.""" + +logger = logging.getLogger(__name__) +_T = TypeVar("_T") def _register_queued_job_status() -> None: @@ -571,17 +580,43 @@ def list_connections( ) except SDKError as e: raise _wrap_sdk_error(e, base_context) from e - - if not status_ok(response.status_code): - raise AirbyteCloudError( - context={ - "workspace_id": workspace_id, - "request_url": response.raw_response.url, - "status_code": response.status_code, + except ValueError as e: + logger.warning( + "Could not decode connections for workspace %s at offset %s: %s. " + "Retrying with lenient decoding.", + workspace_id, + current_offset, + e, + ) + raw_response = _get_public_api_json( + api_root=api_root, + path="/connections", + params={ + "workspaceIds": workspace_id, + "offset": current_offset, + "limit": PAGE_SIZE, }, + client_id=client_id, + client_secret=client_secret, + bearer_token=bearer_token, ) - assert response.connections_response is not None - page_data = response.connections_response.data + page_data = [ + _lenient_decode_connection(connection) for connection in raw_response["data"] + ] + next_page = raw_response.get("next") + else: + if not status_ok(response.status_code): + raise AirbyteCloudError( + context={ + "workspace_id": workspace_id, + "request_url": response.raw_response.url, + "status_code": response.status_code, + }, + ) + assert response.connections_response is not None + page_data = response.connections_response.data + next_page = response.connections_response.next + if not page_data: break @@ -595,7 +630,7 @@ def list_connections( if remaining is not None: remaining -= len(page_results) - if not response.connections_response.next: + if not next_page: break current_offset += len(page_data) @@ -852,6 +887,22 @@ def get_connection( ) except SDKError as e: raise _wrap_sdk_error(e, base_context) from e + except ValueError as e: + logger.warning( + "Could not decode connection %s: %s. Retrying with lenient decoding.", + connection_id, + e, + ) + return _lenient_decode_connection( + _get_public_api_json( + api_root=api_root, + path=f"/connections/{connection_id}", + params=None, + client_id=client_id, + client_secret=client_secret, + bearer_token=bearer_token, + ) + ) if status_ok(response.status_code) and response.connection_response: return response.connection_response @@ -996,39 +1047,69 @@ def get_job_logs( # noqa: PLR0913 # Too many arguments - needed for auth flexi ) except SDKError as e: raise _wrap_sdk_error(e, base_context) from e - - if not status_ok(response.status_code): - if response.status_code == HTTPStatus.NOT_FOUND: - raise AirbyteMissingResourceError( + except ValueError as e: + logger.warning( + "Could not decode jobs for connection %s at offset %s: %s. " + "Retrying with lenient decoding.", + connection_id, + current_offset, + e, + ) + params: dict[str, Any] = { + "workspaceIds": workspace_id, + "connectionId": connection_id, + "limit": page_limit, + "offset": current_offset, + } + if order_by is not None: + params["orderBy"] = order_by + if job_type_value is not None: + params["jobType"] = job_type_value.value + raw_response = _get_public_api_json( + api_root=api_root, + path="/jobs", + params=params, + client_id=client_id, + client_secret=client_secret, + bearer_token=bearer_token, + ) + page_data = [_lenient_decode_job(job) for job in raw_response["data"]] + next_page = raw_response.get("next") + else: + if not status_ok(response.status_code): + if response.status_code == HTTPStatus.NOT_FOUND: + raise AirbyteMissingResourceError( + response=response, + resource_type="job", + context={ + **base_context, + "request_url": response.raw_response.url, + "status_code": response.status_code, + }, + ) + raise AirbyteCloudError( + message="Failed to list jobs.", + response=response, + context={ + **base_context, + "request_url": response.raw_response.url, + "status_code": response.status_code, + }, + ) + if not response.jobs_response: + raise AirbyteCloudError( + message="Jobs response payload was empty.", response=response, - resource_type="job", context={ **base_context, "request_url": response.raw_response.url, "status_code": response.status_code, }, ) - raise AirbyteCloudError( - message="Failed to list jobs.", - response=response, - context={ - **base_context, - "request_url": response.raw_response.url, - "status_code": response.status_code, - }, - ) - if not response.jobs_response: - raise AirbyteCloudError( - message="Jobs response payload was empty.", - response=response, - context={ - **base_context, - "request_url": response.raw_response.url, - "status_code": response.status_code, - }, - ) - page_data: list[models.JobResponse] = list(response.jobs_response.data) + page_data = list(response.jobs_response.data) + next_page = response.jobs_response.next + if not page_data: break @@ -1037,7 +1118,7 @@ def get_job_logs( # noqa: PLR0913 # Too many arguments - needed for auth flexi remaining -= len(page_data) current_offset += len(page_data) - if not response.jobs_response.next or len(page_data) < page_limit: + if not next_page or len(page_data) < page_limit: break return result if limit is None else result[:limit] @@ -1066,6 +1147,22 @@ def get_job_info( ) except SDKError as e: raise _wrap_sdk_error(e, {"job_id": job_id}) from e + except ValueError as e: + logger.warning( + "Could not decode job %s: %s. Retrying with lenient decoding.", + job_id, + e, + ) + return _lenient_decode_job( + _get_public_api_json( + api_root=api_root, + path=f"/jobs/{job_id}", + params=None, + client_id=client_id, + client_secret=client_secret, + bearer_token=bearer_token, + ) + ) if status_ok(response.status_code) and response.job_response: return response.job_response @@ -2213,6 +2310,116 @@ def _config_api_headers( } +def _get_public_api_json( + *, + api_root: str, + path: str, + params: dict[str, Any] | None, + client_id: SecretString | None, + client_secret: SecretString | None, + bearer_token: SecretString | None, +) -> dict[str, Any]: + """Get raw JSON from the Public API.""" + full_url = api_root.rstrip("/") + path + response = requests.get( + full_url, + headers=_config_api_headers( + api_root=api_root, + client_id=client_id, + client_secret=client_secret, + bearer_token=bearer_token, + timeout=PUBLIC_API_FALLBACK_TIMEOUT_SECS, + ), + params=params, + timeout=PUBLIC_API_FALLBACK_TIMEOUT_SECS, + ) + if not status_ok(response.status_code): + try: + response.raise_for_status() + except requests.HTTPError as ex: + error_message = f"API request failed with status {response.status_code}" + error_context = { + "full_url": full_url, + "path": path, + "status_code": response.status_code, + "url": response.request.url, + "body": response.request.body, + "response": response.__dict__, + } + if response.status_code == HTTPStatus.NOT_FOUND: + raise AirbyteMissingResourceError( + message=error_message, + context=error_context, + ) from ex + if response.status_code == HTTPStatus.FORBIDDEN: + raise AirbyteMissingResourceError( + message=( + "The requested resource was not found, or these credentials can't " + "access it (HTTP 403)." + ), + guidance=FORBIDDEN_RESOURCE_GUIDANCE, + context=error_context, + ) from ex + raise AirbyteCloudError( + message=error_message, + context=error_context, + ) from ex + + return response.json() + + +def _decode_with_unknown_status( + raw: dict[str, Any], + model_cls: type[_T], + *, + status_enum: type[Enum], + placeholder: Enum, +) -> _T: + """Decode an SDK model, preserving a `status` value the SDK enum doesn't define.""" + raw_status = raw.get("status") + if isinstance(raw_status, str) and raw_status not in {status.value for status in status_enum}: + cleaned = {**raw, "status": placeholder.value} + decoded = airbyte_api.utils.unmarshal_json( + json.dumps(cleaned), + model_cls, + ) + return replace(decoded, status=raw_status) + + return airbyte_api.utils.unmarshal_json( + json.dumps(raw), + model_cls, + ) + + +def _lenient_decode_connection( + raw: dict[str, Any], +) -> models.ConnectionResponse: + """Decode a connection after dropping stream mappers and tolerating unknown status values.""" + cleaned = copy.deepcopy(raw) + configurations = cleaned.get("configurations") + if configurations is not None: + for stream in configurations.get("streams") or []: + if isinstance(stream, dict): + stream.pop("mappers", None) + + return _decode_with_unknown_status( + cleaned, + models.ConnectionResponse, + status_enum=models.ConnectionStatusEnum, + placeholder=models.ConnectionStatusEnum.INACTIVE, + ) + + +def _lenient_decode_job(raw: dict[str, Any]) -> models.JobResponse: + """Decode a job while tolerating status values unknown to the SDK.""" + return _decode_with_unknown_status( + raw, + models.JobResponse, + status_enum=models.JobStatusEnum, + placeholder=models.JobStatusEnum.RUNNING, + ) + + def create_connector_deferred( # noqa: PLR0913 # Mirrors the API surface. *, connector_type: Literal["source", "destination"], diff --git a/airbyte/cloud/models.py b/airbyte/cloud/models.py index 81b5079c9..995e38d45 100644 --- a/airbyte/cloud/models.py +++ b/airbyte/cloud/models.py @@ -64,6 +64,7 @@ class ConnectionStatus(StrEnum): ACTIVE = "active" INACTIVE = "inactive" + LOCKED = "locked" DEPRECATED = "deprecated" diff --git a/tests/unit_tests/test_cloud_api_util.py b/tests/unit_tests/test_cloud_api_util.py index 24000c510..2f9c562c7 100644 --- a/tests/unit_tests/test_cloud_api_util.py +++ b/tests/unit_tests/test_cloud_api_util.py @@ -6,6 +6,7 @@ import json from collections.abc import Callable from types import SimpleNamespace +from typing import Any from unittest.mock import Mock import pytest @@ -13,7 +14,12 @@ import responses from airbyte import constants from airbyte._util import api_util, meta -from airbyte.cloud.models import CloudJobInfo, JobStatusEnum +from airbyte.cloud.models import ( + CloudConnectionInfo, + CloudJobInfo, + ConnectionStatus, + JobStatusEnum, +) from airbyte.exceptions import ( AirbyteCloudError, AirbyteMissingResourceError, @@ -130,6 +136,16 @@ def _job_response(job_id: int) -> models.JobResponse: ) +def _raw_job_response(job_id: int, status: str) -> dict[str, Any]: + return { + "connectionId": "connection-id", + "jobId": job_id, + "jobType": "sync", + "startTime": "2026-01-01T00:00:00Z", + "status": status, + } + + def _list_jobs_response( data: list[models.JobResponse], *, @@ -165,6 +181,40 @@ def _connection_response(name: str, index: int) -> models.ConnectionResponse: ) +def _connection_with_field_filtering_mapper( + connection_id: str, +) -> dict[str, Any]: + return { + "connectionId": connection_id, + "name": "test connection", + "sourceId": "source-id", + "destinationId": "destination-id", + "workspaceId": "workspace-id", + "status": "active", + "schedule": {"scheduleType": "manual"}, + "dataResidency": "auto", + "nonBreakingSchemaUpdatesBehavior": "ignore", + "namespaceDefinition": "destination", + "createdAt": 1, + "tags": [], + "configurations": { + "streams": [ + { + "name": "leads", + "syncMode": "full_refresh_overwrite", + "mappers": [ + { + "id": "00000000-0000-0000-0000-000000000000", + "type": "field-filtering", + "mapperConfiguration": {"targetField": "foo"}, + } + ], + } + ] + }, + } + + def _workspace_response(name: str, index: int) -> models.WorkspaceResponse: """Create a minimal workspace response for pagination tests.""" return models.WorkspaceResponse( @@ -1169,6 +1219,196 @@ def list_connections( ] == expected_requests +@responses.activate +def test_get_connection_retries_without_undecodable_stream_mappers() -> None: + """A mapper the SDK cannot decode is omitted from the returned connection.""" + api_root = "https://api.airbyte.test/api/public/v1" + connection_id = "connection-id" + url = f"{api_root}/connections/{connection_id}" + responses.get(url, json=_connection_with_field_filtering_mapper(connection_id)) + + connection = api_util.get_connection( + workspace_id="workspace-id", + connection_id=connection_id, + api_root=api_root, + client_id=None, + client_secret=None, + bearer_token=SecretString("bearer-token"), + ) + + assert isinstance(connection, models.ConnectionResponse) + assert connection.configurations.streams[0].name == "leads" + assert len(responses.calls) == 2 + assert [call.request.url for call in responses.calls] == [url, url] + assert ( + responses.calls[1].request.req_kwargs["timeout"] + == api_util.PUBLIC_API_FALLBACK_TIMEOUT_SECS + ) + + +@responses.activate +def test_get_connection_fallback_404_raises_missing_resource_error() -> None: + """A missing resource during raw fallback raises the missing-resource error.""" + api_root = "https://api.airbyte.test/api/public/v1" + connection_id = "connection-id" + url = f"{api_root}/connections/{connection_id}" + responses.get(url, json=_connection_with_field_filtering_mapper(connection_id)) + responses.get(url, status=404, json={"message": "Not found"}) + + with pytest.raises(AirbyteMissingResourceError) as exc_info: + api_util.get_connection( + workspace_id="workspace-id", + connection_id=connection_id, + api_root=api_root, + client_id=None, + client_secret=None, + bearer_token=SecretString("bearer-token"), + ) + + assert len(responses.calls) == 2 + assert exc_info.value.context["status_code"] == 404 + assert exc_info.value.guidance is None + + +@responses.activate +def test_list_connections_retries_page_without_undecodable_stream_mappers() -> None: + """Connection listing falls back to raw JSON when a page has unknown mappers.""" + api_root = "https://api.airbyte.test/api/public/v1" + responses.get( + f"{api_root}/connections", + json={ + "data": [_connection_with_field_filtering_mapper("connection-id")], + "next": None, + }, + ) + + connections = api_util.list_connections( + workspace_id="workspace-id", + api_root=api_root, + client_id=None, + client_secret=None, + bearer_token=SecretString("bearer-token"), + ) + + assert [ + connection.configurations.streams[0].name for connection in connections + ] == ["leads"] + assert len(responses.calls) == 2 + + +@responses.activate +def test_get_connection_without_mappers_uses_sdk_response_once() -> None: + """Connections without mappers keep using the typed SDK response.""" + api_root = "https://api.airbyte.test/api/public/v1" + connection_id = "connection-id" + url = f"{api_root}/connections/{connection_id}" + body = _connection_with_field_filtering_mapper(connection_id) + body["configurations"]["streams"][0].pop("mappers") + responses.get(url, json=body) + + connection = api_util.get_connection( + workspace_id="workspace-id", + connection_id=connection_id, + api_root=api_root, + client_id=None, + client_secret=None, + bearer_token=SecretString("bearer-token"), + ) + + assert connection.configurations.streams[0].name == "leads" + assert len(responses.calls) == 1 + + +@responses.activate +def test_get_job_info_preserves_unknown_status() -> None: + """Job lookup retries when the SDK cannot decode a status value.""" + api_root = "https://api.airbyte.test/api/public/v1" + job_id = 42 + url = f"{api_root}/jobs/{job_id}" + responses.get(url, json=_raw_job_response(job_id, "future-status")) + + job = api_util.get_job_info( + job_id=job_id, + api_root=api_root, + client_id=None, + client_secret=None, + bearer_token=SecretString("bearer-token"), + ) + + assert job.status == "future-status" + assert len(responses.calls) == 2 + + +@responses.activate +def test_get_job_logs_preserves_unknown_statuses_and_stops_without_next_page() -> None: + """Job listing retries the page and stops when the raw response has no next page.""" + api_root = "https://api.airbyte.test/api/public/v1" + responses.get( + f"{api_root}/jobs", + json={ + "data": [ + _raw_job_response(42, "future-status"), + _raw_job_response(43, "succeeded"), + ] + }, + ) + + jobs = api_util.get_job_logs( + workspace_id="workspace-id", + connection_id="connection-id", + limit=2, + api_root=api_root, + client_id=None, + client_secret=None, + bearer_token=SecretString("bearer-token"), + ) + + assert len(jobs) == 2 + assert jobs[0].status == "future-status" + assert jobs[1].status is models.JobStatusEnum.SUCCEEDED + assert len(responses.calls) == 2 + + +@responses.activate +def test_get_connection_preserves_unknown_status() -> None: + """Connection lookup retries when the SDK cannot decode a status value.""" + api_root = "https://api.airbyte.test/api/public/v1" + connection_id = "connection-id" + url = f"{api_root}/connections/{connection_id}" + raw = _connection_with_field_filtering_mapper(connection_id) + raw["configurations"]["streams"][0].pop("mappers") + raw["status"] = "locked" + responses.get(url, json=raw) + + connection = api_util.get_connection( + workspace_id="workspace-id", + connection_id=connection_id, + api_root=api_root, + client_id=None, + client_secret=None, + bearer_token=SecretString("bearer-token"), + ) + + assert connection.status == "locked" + assert ( + CloudConnectionInfo.from_api_response(connection).status + is ConnectionStatus.LOCKED + ) + assert len(responses.calls) == 2 + + +def test_decode_with_unknown_status_keeps_known_enum() -> None: + """Known SDK status values are decoded to their enum members.""" + job = api_util._decode_with_unknown_status( + _raw_job_response(42, "succeeded"), + models.JobResponse, + status_enum=models.JobStatusEnum, + placeholder=models.JobStatusEnum.RUNNING, + ) + + assert job.status is models.JobStatusEnum.SUCCEEDED + + def test_list_workspaces_does_not_filter_by_workspace_id( monkeypatch: pytest.MonkeyPatch, ) -> None: