diff --git a/.github/actions/private-git-deps-auth/action.yml b/.github/actions/private-git-deps-auth/action.yml new file mode 100644 index 000000000..3d026979f --- /dev/null +++ b/.github/actions/private-git-deps-auth/action.yml @@ -0,0 +1,30 @@ +# TK-TODO: Temporary. `airbyte-server-models` is consumed from the private airbyte-ops-mcp repo +# until it is published to PyPI (airbytehq/airbyte-ops-mcp#1326). Remove this action and its +# call sites once pyproject.toml pins a released version. +name: Authenticate git for private dependency repos +description: Authenticate git for private dependency repos +inputs: + app-id: + description: GitHub App ID. + required: true + private-key: + description: GitHub App private key. + required: true +runs: + using: composite + steps: + - name: Create GitHub App token + id: get-app-token + uses: actions/create-github-app-token@67018539274d69449ef7c02e8e71183d1719ab42 # v2.1.4 + with: + owner: "airbytehq" + repositories: "airbyte-ops-mcp" + permission-contents: read + app-id: ${{ inputs.app-id }} + private-key: ${{ inputs.private-key }} + - name: Configure git credentials for airbyte-ops-mcp + shell: bash + env: + GH_APP_TOKEN: ${{ steps.get-app-token.outputs.token }} + run: | + git config --global url."https://x-access-token:${GH_APP_TOKEN}@github.com/airbytehq/airbyte-ops-mcp".insteadOf "https://github.com/airbytehq/airbyte-ops-mcp" diff --git a/.github/workflows/autofix.yml b/.github/workflows/autofix.yml index 2159659ee..f3c529c2c 100644 --- a/.github/workflows/autofix.yml +++ b/.github/workflows/autofix.yml @@ -23,6 +23,11 @@ jobs: version: "latest" python-version: '3.11' + - name: Authenticate git for private dependencies + uses: ./.github/actions/private-git-deps-auth + with: + app-id: ${{ secrets.OCTAVIA_BOT_APP_ID }} + private-key: ${{ secrets.OCTAVIA_BOT_PRIVATE_KEY }} - name: Install dev dependencies # We only need helper libraries run: uv sync --group dev diff --git a/.github/workflows/fix-pr-command.yml b/.github/workflows/fix-pr-command.yml index ff91cd6f7..210e1f07f 100644 --- a/.github/workflows/fix-pr-command.yml +++ b/.github/workflows/fix-pr-command.yml @@ -86,6 +86,11 @@ jobs: with: version: "latest" python-version: ${{ matrix.python-version }} + - name: Authenticate git for private dependencies + uses: ./.github/actions/private-git-deps-auth + with: + app-id: ${{ secrets.OCTAVIA_BOT_APP_ID }} + private-key: ${{ secrets.OCTAVIA_BOT_PRIVATE_KEY }} - name: Install dev dependencies # We only need helper libraries run: uv sync --group dev diff --git a/.github/workflows/pydoc_preview.yml b/.github/workflows/pydoc_preview.yml index cba9061c8..80218fe30 100644 --- a/.github/workflows/pydoc_preview.yml +++ b/.github/workflows/pydoc_preview.yml @@ -25,6 +25,11 @@ jobs: version: "latest" python-version: '3.11' + - name: Authenticate git for private dependencies + uses: ./.github/actions/private-git-deps-auth + with: + app-id: ${{ secrets.OCTAVIA_BOT_APP_ID }} + private-key: ${{ secrets.OCTAVIA_BOT_PRIVATE_KEY }} - name: Install dev dependencies # We only need helper libraries run: uv sync --group dev diff --git a/.github/workflows/pydoc_publish.yml b/.github/workflows/pydoc_publish.yml index db8198f4c..f033ed894 100644 --- a/.github/workflows/pydoc_publish.yml +++ b/.github/workflows/pydoc_publish.yml @@ -41,6 +41,11 @@ jobs: - name: Setup Pages uses: actions/configure-pages@983d7736d9b0ae728b81ab479565c72886d7745b # v5 + - name: Authenticate git for private dependencies + uses: ./.github/actions/private-git-deps-auth + with: + app-id: ${{ secrets.OCTAVIA_BOT_APP_ID }} + private-key: ${{ secrets.OCTAVIA_BOT_PRIVATE_KEY }} - name: Install dev dependencies # We only need helper libraries run: uv sync --group dev diff --git a/.github/workflows/python_lint.yml b/.github/workflows/python_lint.yml index 8cddeeebe..f62280b92 100644 --- a/.github/workflows/python_lint.yml +++ b/.github/workflows/python_lint.yml @@ -24,6 +24,11 @@ jobs: with: version: "latest" python-version: '3.11' + - name: Authenticate git for private dependencies + uses: ./.github/actions/private-git-deps-auth + with: + app-id: ${{ secrets.OCTAVIA_BOT_APP_ID }} + private-key: ${{ secrets.OCTAVIA_BOT_PRIVATE_KEY }} - name: Install dev dependencies # We only need helper libraries run: uv sync --group dev @@ -44,6 +49,11 @@ jobs: with: version: "latest" python-version: '3.11' + - name: Authenticate git for private dependencies + uses: ./.github/actions/private-git-deps-auth + with: + app-id: ${{ secrets.OCTAVIA_BOT_APP_ID }} + private-key: ${{ secrets.OCTAVIA_BOT_PRIVATE_KEY }} - name: Install dev dependencies # We only need helper libraries run: uv sync --group dev @@ -64,6 +74,11 @@ jobs: with: version: "latest" python-version: '3.11' + - name: Authenticate git for private dependencies + uses: ./.github/actions/private-git-deps-auth + with: + app-id: ${{ secrets.OCTAVIA_BOT_APP_ID }} + private-key: ${{ secrets.OCTAVIA_BOT_PRIVATE_KEY }} - name: Install dev dependencies # We only need helper libraries run: uv sync --group dev @@ -86,6 +101,11 @@ jobs: with: version: "latest" python-version: '3.11' + - name: Authenticate git for private dependencies + uses: ./.github/actions/private-git-deps-auth + with: + app-id: ${{ secrets.OCTAVIA_BOT_APP_ID }} + private-key: ${{ secrets.OCTAVIA_BOT_PRIVATE_KEY }} - name: Install dev dependencies # We only need helper libraries run: uv sync --group dev diff --git a/.github/workflows/python_pytest.yml b/.github/workflows/python_pytest.yml index c74862a48..2a5a1d6ad 100644 --- a/.github/workflows/python_pytest.yml +++ b/.github/workflows/python_pytest.yml @@ -42,6 +42,11 @@ jobs: with: version: "latest" python-version: '3.11' + - name: Authenticate git for private dependencies + uses: ./.github/actions/private-git-deps-auth + with: + app-id: ${{ secrets.OCTAVIA_BOT_APP_ID }} + private-key: ${{ secrets.OCTAVIA_BOT_PRIVATE_KEY }} - name: Install dev dependencies # We only need helper libraries run: uv sync --group dev @@ -130,6 +135,11 @@ jobs: with: version: "latest" python-version: ${{ matrix.python-version }} + - name: Authenticate git for private dependencies + uses: ./.github/actions/private-git-deps-auth + with: + app-id: ${{ secrets.OCTAVIA_BOT_APP_ID }} + private-key: ${{ secrets.OCTAVIA_BOT_PRIVATE_KEY }} - name: Install dev dependencies # We only need helper libraries run: uv sync --group dev @@ -233,6 +243,11 @@ jobs: with: version: "latest" python-version: '3.11' + - name: Authenticate git for private dependencies + uses: ./.github/actions/private-git-deps-auth + with: + app-id: ${{ secrets.OCTAVIA_BOT_APP_ID }} + private-key: ${{ secrets.OCTAVIA_BOT_PRIVATE_KEY }} - name: Install dev dependencies # We only need helper libraries run: uv sync --group dev @@ -311,6 +326,11 @@ jobs: with: version: "latest" python-version: '3.11' + - name: Authenticate git for private dependencies + uses: ./.github/actions/private-git-deps-auth + with: + app-id: ${{ secrets.OCTAVIA_BOT_APP_ID }} + private-key: ${{ secrets.OCTAVIA_BOT_PRIVATE_KEY }} - name: Install dev dependencies # We only need helper libraries run: uv sync --group dev @@ -357,6 +377,11 @@ jobs: with: version: "latest" python-version: '3.11' + - name: Authenticate git for private dependencies + uses: ./.github/actions/private-git-deps-auth + with: + app-id: ${{ secrets.OCTAVIA_BOT_APP_ID }} + private-key: ${{ secrets.OCTAVIA_BOT_PRIVATE_KEY }} - name: Install dev dependencies # We only need helper libraries run: uv sync --group dev diff --git a/.github/workflows/test-pr-command.yml b/.github/workflows/test-pr-command.yml index 4658cf3ad..060859824 100644 --- a/.github/workflows/test-pr-command.yml +++ b/.github/workflows/test-pr-command.yml @@ -113,6 +113,11 @@ jobs: with: version: "latest" python-version: ${{ matrix.python-version }} + - name: Authenticate git for private dependencies + uses: ./.github/actions/private-git-deps-auth + with: + app-id: ${{ secrets.OCTAVIA_BOT_APP_ID }} + private-key: ${{ secrets.OCTAVIA_BOT_PRIVATE_KEY }} - name: Install dev dependencies # We only need helper libraries run: uv sync --group dev diff --git a/.github/workflows/uv-lock-command.yml b/.github/workflows/uv-lock-command.yml index 0f09317d9..f41b2ee22 100644 --- a/.github/workflows/uv-lock-command.yml +++ b/.github/workflows/uv-lock-command.yml @@ -92,6 +92,12 @@ jobs: # Run `uv lock` + - name: Authenticate git for private dependencies + uses: ./.github/actions/private-git-deps-auth + with: + app-id: ${{ secrets.OCTAVIA_BOT_APP_ID }} + private-key: ${{ secrets.OCTAVIA_BOT_PRIVATE_KEY }} + - name: Run `uv lock` run: uv lock diff --git a/airbyte/_util/api_imports.py b/airbyte/_util/api_imports.py index df367256e..f92e3f587 100644 --- a/airbyte/_util/api_imports.py +++ b/airbyte/_util/api_imports.py @@ -6,9 +6,9 @@ for type hinting purposes. Design Guidelines: -- No modules except `api_util` and `api_imports` should import from `airbyte_api`. -- If a module needs to import from `airbyte_api`, it should import from `api_imports` (this module) - instead. +- No modules except `api_util` and `api_imports` should import from `airbyte_server_models`. +- If a module needs to import from `airbyte_server_models`, it should import from + `api_imports` (this module) instead. - This module is divided into two sections: internal-use classes and public-use classes. - Public-use classes should be carefully reviewed to ensure that they are necessary for public use and that we are willing to support them as part of PyAirbyte. @@ -21,7 +21,7 @@ # Internal-Use Classes # These classes are used internally to cache API responses. -from airbyte_api.models import ( +from airbyte_server_models.public_api import ( ConnectionResponse, DestinationResponse, JobResponse, @@ -32,7 +32,7 @@ # This class is used to represent the status of a job. It may be used in # type hints for public functions that return a job status. -from airbyte_api.models import JobStatusEnum # Alias not needed +from airbyte_server_models.public_api import JobStatusEnum # Alias not needed __all__: list[str] = [ diff --git a/airbyte/_util/api_util.py b/airbyte/_util/api_util.py index bcb5012f3..bc09a7501 100644 --- a/airbyte/_util/api_util.py +++ b/airbyte/_util/api_util.py @@ -6,22 +6,61 @@ PyAirbyte. Classes and functions from the Airbyte API external library should always be wrapped in PyAirbyte classes - unless there's a very compelling reason to surface these models intentionally. -Similarly, modules outside of this file should try to avoid interfacing with `airbyte_api` library -directly. This will ensure a single source of truth when mapping between the `airbyte` and -`airbyte_api` libraries. +Similarly, modules outside of this file should try to avoid interfacing with +`airbyte_server_models` directly. This will ensure a single source of truth when mapping +between the `airbyte` and `airbyte_server_models` libraries. """ from __future__ import annotations import base64 import json +import logging from http import HTTPStatus -from typing import TYPE_CHECKING, Any, Literal +from typing import TYPE_CHECKING, Any, Literal, TypeVar, cast, overload -import airbyte_api import requests -from airbyte_api import api, models -from airbyte_api.errors import SDKError +from airbyte_server_models import public_api as models + +# The Config API models live in an underscore-prefixed module on purpose: the Config API is +# Airbyte-internal and may change without notice, so the private name is a deliberate signal +# to consumers. PyAirbyte accepts that contract, hence the `PLC2701` (import-private-name) +# suppressions below. +from airbyte_server_models._config_api import ( + AirbyteCatalog, # noqa: PLC2701 + BuilderProjectForDefinitionRequestBody, # noqa: PLC2701 + BuilderProjectForDefinitionResponse, # noqa: PLC2701 + CheckConnectionRead, # noqa: PLC2701 + ConnectionIdRequestBody, # noqa: PLC2701 + ConnectionState, # noqa: PLC2701 + ConnectionStateCreateOrUpdate, # noqa: PLC2701 + ConnectorBuilderProjectIdWithWorkspaceId, # noqa: PLC2701 + ConnectorBuilderProjectRead, # noqa: PLC2701 + ConnectorBuilderProjectTestingValues, # noqa: PLC2701 + ConnectorBuilderProjectTestingValuesUpdate, # noqa: PLC2701 + DestinationIdRequestBody, # noqa: PLC2701 + ListOrganizationsByUserRequestBody, # noqa: PLC2701 + ListWorkspacesInOrganizationRequestBody, # noqa: PLC2701 + OrganizationIdRequestBody, # noqa: PLC2701 + OrganizationInfoRead, # noqa: PLC2701 + OrganizationRead, + OrganizationReadList, # noqa: PLC2701 + Pagination, # noqa: PLC2701 + PermissionRead, + PermissionReadList, # noqa: PLC2701 + SourceDefinitionSpecification, # noqa: PLC2701 + SourceIdRequestBody, # noqa: PLC2701 + UserAuthIdRequestBody, # noqa: PLC2701 + UserIdRequestBody, # noqa: PLC2701 + UserRead, # noqa: PLC2701 + WebBackendConnectionRead, # noqa: PLC2701 + WebBackendConnectionRequestBody, # noqa: PLC2701 + WebBackendConnectionUpdate, # noqa: PLC2701 + WorkspaceIdRequestBody, # noqa: PLC2701 + WorkspaceRead, + WorkspaceReadList, # noqa: PLC2701 +) +from pydantic import BaseModel, ValidationError from airbyte._util.meta import AIRBYTE_ANALYTIC_SOURCE_HEADER, get_cloud_api_analytic_source from airbyte.constants import CLOUD_API_ROOT, CLOUD_CONFIG_API_ROOT, CLOUD_CONFIG_API_ROOT_ENV_VAR @@ -43,15 +82,15 @@ if TYPE_CHECKING: from collections.abc import Callable - from airbyte_api.models import ( - DestinationConfiguration, - ) + +logger = logging.getLogger(__name__) JOB_WAIT_INTERVAL_SECS = 2.0 JOB_WAIT_TIMEOUT_SECS_DEFAULT = 60 * 60 # 1 hour PAGE_SIZE = 100 JWT_PART_COUNT = 3 +_T = TypeVar("_T", bound=BaseModel) FORBIDDEN_RESOURCE_GUIDANCE = ( "Airbyte Cloud returns 403 both for IDs that don't exist and for resources outside your " "access. Check the ID, and that it belongs to the workspace you passed." @@ -86,39 +125,34 @@ def _get_page_limit(remaining: int | None) -> int: return min(remaining, PAGE_SIZE) -def _get_sdk_error_context(error: SDKError) -> dict[str, Any]: - """Extract context information from an SDKError for debugging. +def _get_api_error_context(response: requests.Response) -> dict[str, Any]: + """Extract request and response details from a failed API call for debugging. - This helper extracts the actual request URL and other useful debugging - information from the Speakeasy SDK's SDKError exception. The SDK stores - the raw response object which contains the request details. + Produces the same context keys the Speakeasy SDK error wrapper used to surface. """ context: dict[str, Any] = { - "status_code": error.status_code, - "error_message": error.message, + "status_code": response.status_code, + "error_message": response.text, } - - if error.raw_response is not None: - request = error.raw_response.request - if request is not None: - context["request_url"] = str(request.url) - context["request_method"] = request.method - context["response_content_type"] = error.raw_response.headers.get("content-type") - + if response.request is not None: + context["request_url"] = str(response.request.url) + context["request_method"] = response.request.method + context["response_content_type"] = response.headers.get("content-type") return context -def _wrap_sdk_error( - error: SDKError, base_context: dict[str, Any] | None = None +def _wrap_api_error( + response: requests.Response, base_context: dict[str, Any] | None = None ) -> AirbyteCloudError: - """Wrap an SDKError with additional context for debugging. + """Wrap a failed API response with additional context for debugging. - This function converts a Speakeasy SDK error into an AirbyteCloudError with - full URL context, making it easier to debug API issues like 403 and 404 errors. + Mirrors the former SDK error-wrapping semantics: 403 and 404 become + `AirbyteMissingResourceError` (403 with extra guidance), everything else + becomes `AirbyteCloudError`. """ - sdk_context = _get_sdk_error_context(error) - merged_context = {**(base_context or {}), **sdk_context} - status_code = sdk_context.get("status_code") + api_context = _get_api_error_context(response) + merged_context = {**(base_context or {}), **api_context} + status_code = api_context.get("status_code") is_forbidden = status_code == HTTPStatus.FORBIDDEN error_type = ( AirbyteMissingResourceError @@ -130,7 +164,7 @@ def _wrap_sdk_error( "The requested resource was not found, or these credentials can't access it " "(HTTP 403)." if is_forbidden - else f"API error occurred: {error.message}" + else f"API error occurred: {response.text}" ), guidance=FORBIDDEN_RESOURCE_GUIDANCE if is_forbidden else None, context=merged_context, @@ -212,30 +246,19 @@ def get_web_url_root(api_root: str) -> str: return api_root -def get_airbyte_server_instance( +def _resolve_bearer_token( *, api_root: str, client_id: SecretString | None, client_secret: SecretString | None, bearer_token: SecretString | None, -) -> airbyte_api.AirbyteAPI: - """Get an Airbyte API instance. + timeout: tuple[float, float] | None = None, +) -> SecretString: + """Resolve the bearer token to use, minting one from client credentials if needed. Supports two authentication methods (mutually exclusive): 1. OAuth2 client credentials (client_id + client_secret) 2. Bearer token authentication - - Args: - api_root: The API root URL. - client_id: OAuth2 client ID (required if not using bearer_token). - client_secret: OAuth2 client secret (required if not using bearer_token). - bearer_token: Pre-generated bearer token (alternative to client credentials). - - Returns: - An authenticated AirbyteAPI instance. - - Raises: - AirbyteLibInputError: If authentication parameters are invalid. """ # Guard: must provide either bearer token OR both client credentials if bearer_token is None and (client_id is None or client_secret is None): @@ -251,33 +274,176 @@ def get_airbyte_server_instance( guidance="Provide either client_id and client_secret, or bearer_token, but not both.", ) - client = requests.Session() - client.headers[AIRBYTE_ANALYTIC_SOURCE_HEADER] = get_cloud_api_analytic_source() - - # Option 1: Bearer token authentication - if bearer_token is not None: - return airbyte_api.AirbyteAPI( - security=models.Security( - bearer_auth=bearer_token, - ), - server_url=api_root, - client=client, + if bearer_token is None: + # Client credentials flow (guaranteed non-None by first guard) + bearer_token = get_bearer_token( + client_id=cast(SecretString, client_id), + client_secret=cast(SecretString, client_secret), + api_root=api_root, + timeout=timeout, ) + return bearer_token - # Option 2: Client credentials flow (guaranteed non-None by first guard) - return airbyte_api.AirbyteAPI( - security=models.Security( - client_credentials=models.SchemeClientCredentials( - client_id=client_id, - client_secret=client_secret, - token_url=api_root + "/applications/token", - # e.g. https://api.airbyte.com/v1/applications/token - ), - ), - server_url=api_root, - client=client, +def _api_headers( + *, + api_root: str, + client_id: SecretString | None, + client_secret: SecretString | None, + bearer_token: SecretString | None, +) -> dict[str, str]: + """Build Public API headers, minting a bearer token from client credentials if needed.""" + resolved_token = _resolve_bearer_token( + api_root=api_root, + client_id=client_id, + client_secret=client_secret, + bearer_token=bearer_token, + ) + return { + "Content-Type": "application/json", + "Accept": "application/json", + "Authorization": f"Bearer {resolved_token}", + "User-Agent": "PyAirbyte Client", + AIRBYTE_ANALYTIC_SOURCE_HEADER: get_cloud_api_analytic_source(), + } + + +def _serialize_public_api_request(request: BaseModel) -> dict[str, Any]: + body = request.model_dump(mode="json", by_alias=True, exclude_none=True) + # Free-form connector configuration is sent verbatim (nulls included): + # exclude_none only applies to the request's own top-level fields. + configuration = getattr(request, "configuration", None) + if isinstance(configuration, BaseModel): + body["configuration"] = configuration.model_dump(mode="json", by_alias=True) + return body + + +@overload +def _make_public_api_request( # Mirrors the API surface. + *, + method: Literal["GET", "POST", "PATCH", "PUT", "DELETE"], + api_root: str, + path: str, + response_model: type[_T], + client_id: SecretString | None, + client_secret: SecretString | None, + bearer_token: SecretString | None, + request: BaseModel | None = None, + params: dict[str, Any] | None = None, + error_context: dict[str, Any] | None = None, +) -> _T | None: ... + + +@overload +def _make_public_api_request( # Mirrors the API surface. + *, + method: Literal["GET", "POST", "PATCH", "PUT", "DELETE"], + api_root: str, + path: str, + response_model: None, + client_id: SecretString | None, + client_secret: SecretString | None, + bearer_token: SecretString | None, + request: BaseModel | None = None, + params: dict[str, Any] | None = None, + error_context: dict[str, Any] | None = None, +) -> dict[str, Any] | None: ... + + +def _make_public_api_request( # noqa: PLR0913 # Mirrors the API surface. + *, + method: Literal["GET", "POST", "PATCH", "PUT", "DELETE"], + api_root: str, + path: str, + response_model: type[_T] | None, + client_id: SecretString | None, + client_secret: SecretString | None, + bearer_token: SecretString | None, + request: BaseModel | None = None, + params: dict[str, Any] | None = None, + error_context: dict[str, Any] | None = None, +) -> _T | None: + """Make a Public API request and decode the response into `response_model`. + + Raises `AirbyteMissingResourceError` on 403/404 and `AirbyteCloudError` on + other non-2xx responses, matching the former Speakeasy SDK error wrapping. + Returns `None` when the response has no content. + """ + headers = _api_headers( + api_root=api_root, + client_id=client_id, + client_secret=client_secret, + bearer_token=bearer_token, ) + full_url = api_root.rstrip("/") + path + response = requests.request( + method=method, + url=full_url, + headers=headers, + params=params, + json=(_serialize_public_api_request(request) if request is not None else None), + # No timeout, matching the former SDK behavior. + ) + if not status_ok(response.status_code): + raise _wrap_api_error(response, error_context) + + if response.status_code == HTTPStatus.NO_CONTENT or not response.content: + return None + + try: + body = response.json() + except requests.exceptions.JSONDecodeError as ex: + raise AirbyteCloudError( + message=f"Public API response for {path} did not match the expected schema.", + context={ + **(error_context or {}), + "full_url": full_url, + "path": path, + "response": response.text, + }, + ) from ex + + if response_model is None: + return body + + try: + return response_model.model_validate(body) + except ValidationError as ex: + raise AirbyteCloudError( + message=f"Public API response for {path} did not match the expected schema.", + context={ + **(error_context or {}), + "full_url": full_url, + "path": path, + "response": response.text, + }, + ) from ex + + +def _decode_list_items( + body: Any, # noqa: ANN401 # Raw JSON payload + *, + item_model: type[_T], + list_path: str, +) -> list[_T]: + """Decode each item in a list endpoint's `data` payload individually. + + A single malformed item is skipped with a warning instead of failing the + whole page, so one bad record can't take down a listing. + """ + items: list[_T] = [] + for item in (body or {}).get("data") or []: + try: + items.append(item_model.model_validate(item)) + except ValidationError: + item_id = item.get("id") if isinstance(item, dict) else getattr(item, "id", None) + logger.warning( + "Skipping %s item that failed schema validation on %s (id=%s).", + getattr(item_model, "__name__", str(item_model)), + list_path, + item_id, + ) + return items # Get workspace @@ -292,32 +458,24 @@ def get_workspace( bearer_token: SecretString | None, ) -> models.WorkspaceResponse: """Get a workspace object.""" - airbyte_instance = get_airbyte_server_instance( + base_context = {"workspace_id": workspace_id, "api_root": api_root} + workspace = _make_public_api_request( + method="GET", api_root=api_root, + path=f"/workspaces/{workspace_id}", + response_model=models.WorkspaceResponse, client_id=client_id, client_secret=client_secret, bearer_token=bearer_token, + error_context=base_context, ) - base_context = {"workspace_id": workspace_id, "api_root": api_root} - try: - response = airbyte_instance.workspaces.get_workspace( - api.GetWorkspaceRequest( - workspace_id=workspace_id, - ), - ) - except SDKError as e: - raise _wrap_sdk_error(e, base_context) from e - if status_ok(response.status_code) and response.workspace_response: - return response.workspace_response + if workspace: + return workspace raise AirbyteMissingResourceError( resource_type="workspace", - context={ - "workspace_id": workspace_id, - "request_url": response.raw_response.url, - "status_code": response.status_code, - }, + context=base_context, ) @@ -332,40 +490,34 @@ def create_workspace( region_id: str | None = None, ) -> models.WorkspaceResponse: """Create a workspace.""" - airbyte_instance = get_airbyte_server_instance( - api_root=api_root, - client_id=client_id, - client_secret=client_secret, - bearer_token=bearer_token, - ) base_context = { "name": name, "organization_id": organization_id, "region_id": region_id, "api_root": api_root, } - try: - response = airbyte_instance.workspaces.create_workspace( - request=models.WorkspaceCreateRequest( - name=name, - organization_id=organization_id, - region_id=region_id, - ) - ) - except SDKError as e: - raise _wrap_sdk_error(e, base_context) from e + workspace = _make_public_api_request( + method="POST", + api_root=api_root, + path="/workspaces", + request=models.WorkspaceCreateRequest( + name=name, + organization_id=organization_id, + region_id=region_id, + ), + response_model=models.WorkspaceResponse, + client_id=client_id, + client_secret=client_secret, + bearer_token=bearer_token, + error_context=base_context, + ) - if status_ok(response.status_code) and response.workspace_response: - return response.workspace_response + if workspace: + return workspace raise AirbyteCloudError( message="Could not create workspace.", - context={ - **base_context, - "request_url": response.raw_response.url, - "status_code": response.status_code, - }, - response=response, + context=base_context, ) @@ -379,38 +531,29 @@ def rename_workspace( bearer_token: SecretString | None, ) -> models.WorkspaceResponse: """Rename a workspace.""" - airbyte_instance = get_airbyte_server_instance( - api_root=api_root, - client_id=client_id, - client_secret=client_secret, - bearer_token=bearer_token, - ) base_context = { "workspace_id": workspace_id, "name": name, "api_root": api_root, } - try: - response = airbyte_instance.workspaces.update_workspace( - api.UpdateWorkspaceRequest( - workspace_id=workspace_id, - workspace_update_request=models.WorkspaceUpdateRequest(name=name), - ) - ) - except SDKError as e: - raise _wrap_sdk_error(e, base_context) from e + workspace = _make_public_api_request( + method="PATCH", + api_root=api_root, + path=f"/workspaces/{workspace_id}", + request=models.WorkspaceUpdateRequest(name=name), + response_model=models.WorkspaceResponse, + client_id=client_id, + client_secret=client_secret, + bearer_token=bearer_token, + error_context=base_context, + ) - if status_ok(response.status_code) and response.workspace_response: - return response.workspace_response + if workspace: + return workspace raise AirbyteCloudError( message="Could not rename workspace.", - context={ - **base_context, - "request_url": response.raw_response.url, - "status_code": response.status_code, - }, - response=response, + context=base_context, ) @@ -480,32 +623,19 @@ def permanently_delete_workspace( connection_ids=[connection.connection_id for connection in connections], ) - airbyte_instance = get_airbyte_server_instance( + _make_public_api_request( + method="DELETE", + api_root=api_root, + path=f"/workspaces/{workspace_id}", + response_model=None, client_id=client_id, client_secret=client_secret, bearer_token=bearer_token, - api_root=api_root, + error_context={ + "workspace_id": workspace_id, + "api_root": api_root, + }, ) - base_context = { - "workspace_id": workspace_id, - "api_root": api_root, - } - try: - response = airbyte_instance.workspaces.delete_workspace( - api.DeleteWorkspaceRequest(workspace_id=workspace_id), - ) - except SDKError as e: - raise _wrap_sdk_error(e, base_context) from e - - if not status_ok(response.status_code): - raise AirbyteCloudError( - context={ - **base_context, - "request_url": response.raw_response.url, - "status_code": response.status_code, - }, - response=response, - ) # List resources @@ -528,41 +658,32 @@ def list_connections( _validate_pagination_params(limit=limit) name_filter = (lambda n: n == name) if name is not None else name_filter or (lambda _: True) - airbyte_instance = get_airbyte_server_instance( - client_id=client_id, - client_secret=client_secret, - bearer_token=bearer_token, - api_root=api_root, - ) result: list[models.ConnectionResponse] = [] current_offset = 0 remaining = limit - base_context = {"workspace_id": workspace_id, "api_root": api_root} while remaining is None or remaining > 0: - try: - response = airbyte_instance.connections.list_connections( - api.ListConnectionsRequest( - workspace_ids=[workspace_id], - offset=current_offset, - limit=PAGE_SIZE, - ), - ) - 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, - }, - ) - assert response.connections_response is not None - page_data = response.connections_response.data - if not page_data: + body = _make_public_api_request( + method="GET", + api_root=api_root, + path="/connections", + params={ + "workspaceIds": [workspace_id], + "offset": current_offset, + "limit": PAGE_SIZE, + }, + response_model=None, + client_id=client_id, + client_secret=client_secret, + bearer_token=bearer_token, + error_context={"workspace_id": workspace_id, "api_root": api_root}, + ) + raw_items = (body or {}).get("data") or [] + if not raw_items: break + page_data = _decode_list_items( + body, item_model=models.ConnectionResponse, list_path="/connections" + ) matching_connections = [ connection for connection in page_data if name_filter(connection.name) ] @@ -573,10 +694,10 @@ def list_connections( if remaining is not None: remaining -= len(page_results) - if not response.connections_response.next: + if not (body or {}).get("next"): break - current_offset += len(page_data) + current_offset += len(raw_items) return result @@ -609,52 +730,42 @@ def list_workspaces( has_name_filter = name is not None or name_filter is not None name_filter = (lambda n: n == name) if name is not None else name_filter or (lambda _: True) - airbyte_instance: airbyte_api.AirbyteAPI = get_airbyte_server_instance( - client_id=client_id, - client_secret=client_secret, - bearer_token=bearer_token, - api_root=api_root, - ) result: list[models.WorkspaceResponse] = [] current_offset = 0 remaining = limit - base_context = {"workspace_id": workspace_id, "api_root": api_root} while remaining is None or remaining > 0: page_limit = PAGE_SIZE if has_name_filter else _get_page_limit(remaining) - try: - response: api.ListWorkspacesResponse = airbyte_instance.workspaces.list_workspaces( - api.ListWorkspacesRequest( - offset=current_offset, - limit=page_limit, - ), - ) - 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, - }, - ) - - assert response.workspaces_response is not None - page_data = response.workspaces_response.data - if not page_data: + body = _make_public_api_request( + method="GET", + api_root=api_root, + path="/workspaces", + params={ + "offset": current_offset, + "limit": page_limit, + }, + response_model=None, + client_id=client_id, + client_secret=client_secret, + bearer_token=bearer_token, + error_context={"workspace_id": workspace_id, "api_root": api_root}, + ) + raw_items = (body or {}).get("data") or [] + if not raw_items: break + page_data = _decode_list_items( + body, item_model=models.WorkspaceResponse, list_path="/workspaces" + ) matching_workspaces = [workspace for workspace in page_data if name_filter(workspace.name)] page_results = matching_workspaces if remaining is None else matching_workspaces[:remaining] result += page_results if remaining is not None: remaining -= len(page_results) - if not response.workspaces_response.next: + if not (body or {}).get("next"): break - current_offset += len(page_data) + current_offset += len(raw_items) return result @@ -676,52 +787,40 @@ def list_sources( _validate_pagination_params(limit=limit) name_filter = (lambda n: n == name) if name is not None else name_filter or (lambda _: True) - _ = workspace_id # Not used (yet) - airbyte_instance: airbyte_api.AirbyteAPI = get_airbyte_server_instance( - client_id=client_id, - client_secret=client_secret, - bearer_token=bearer_token, - api_root=api_root, - ) result: list[models.SourceResponse] = [] current_offset = 0 remaining = limit - base_context = {"workspace_id": workspace_id, "api_root": api_root} while remaining is None or remaining > 0: - try: - response: api.ListSourcesResponse = airbyte_instance.sources.list_sources( - api.ListSourcesRequest( - workspace_ids=[workspace_id], - offset=current_offset, - limit=PAGE_SIZE, - ), - ) - 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, - }, - ) - assert response.sources_response is not None - page_data = response.sources_response.data - if not page_data: + body = _make_public_api_request( + method="GET", + api_root=api_root, + path="/sources", + params={ + "workspaceIds": [workspace_id], + "offset": current_offset, + "limit": PAGE_SIZE, + }, + response_model=None, + client_id=client_id, + client_secret=client_secret, + bearer_token=bearer_token, + error_context={"workspace_id": workspace_id, "api_root": api_root}, + ) + raw_items = (body or {}).get("data") or [] + if not raw_items: break + page_data = _decode_list_items(body, item_model=models.SourceResponse, list_path="/sources") matching_sources = [source for source in page_data if name_filter(source.name)] page_results = matching_sources if remaining is None else matching_sources[:remaining] result += page_results if remaining is not None: remaining -= len(page_results) - if not response.sources_response.next: + if not (body or {}).get("next"): break - current_offset += len(page_data) + current_offset += len(raw_items) return result @@ -743,42 +842,32 @@ def list_destinations( _validate_pagination_params(limit=limit) name_filter = (lambda n: n == name) if name is not None else name_filter or (lambda _: True) - _ = workspace_id # Not used (yet) - airbyte_instance = get_airbyte_server_instance( - client_id=client_id, - client_secret=client_secret, - bearer_token=bearer_token, - api_root=api_root, - ) result: list[models.DestinationResponse] = [] current_offset = 0 remaining = limit - base_context = {"workspace_id": workspace_id, "api_root": api_root} while remaining is None or remaining > 0: - try: - response = airbyte_instance.destinations.list_destinations( - api.ListDestinationsRequest( - workspace_ids=[workspace_id], - offset=current_offset, - limit=PAGE_SIZE, - ), - ) - 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, - }, - ) - assert response.destinations_response is not None - page_data = response.destinations_response.data - if not page_data: + body = _make_public_api_request( + method="GET", + api_root=api_root, + path="/destinations", + params={ + "workspaceIds": [workspace_id], + "offset": current_offset, + "limit": PAGE_SIZE, + }, + response_model=None, + client_id=client_id, + client_secret=client_secret, + bearer_token=bearer_token, + error_context={"workspace_id": workspace_id, "api_root": api_root}, + ) + raw_items = (body or {}).get("data") or [] + if not raw_items: break + page_data = _decode_list_items( + body, item_model=models.DestinationResponse, list_path="/destinations" + ) matching_destinations = [ destination for destination in page_data if name_filter(destination.name) ] @@ -789,10 +878,10 @@ def list_destinations( if remaining is not None: remaining -= len(page_results) - if not response.destinations_response.next: + if not (body or {}).get("next"): break - current_offset += len(page_data) + current_offset += len(raw_items) return result @@ -811,37 +900,29 @@ def get_connection( ) -> models.ConnectionResponse: """Get a connection.""" _ = workspace_id # Not used (yet) - airbyte_instance = get_airbyte_server_instance( - client_id=client_id, - client_secret=client_secret, - bearer_token=bearer_token, - api_root=api_root, - ) base_context = { "workspace_id": workspace_id, "connection_id": connection_id, "api_root": api_root, } - try: - response = airbyte_instance.connections.get_connection( - api.GetConnectionRequest( - connection_id=connection_id, - ), - ) - except SDKError as e: - raise _wrap_sdk_error(e, base_context) from e + connection = _make_public_api_request( + method="GET", + api_root=api_root, + path=f"/connections/{connection_id}", + response_model=models.ConnectionResponse, + client_id=client_id, + client_secret=client_secret, + bearer_token=bearer_token, + error_context=base_context, + ) - if status_ok(response.status_code) and response.connection_response: - return response.connection_response + if connection: + return connection raise AirbyteMissingResourceError( resource_name_or_id=connection_id, resource_type="connection", - log_text=response.raw_response.text, - context={ - "request_url": response.raw_response.url, - "status_code": response.status_code, - }, + context=base_context, ) @@ -861,37 +942,36 @@ def run_connection( If raise_on_failure is True, this will raise an exception if the connection fails. """ _ = workspace_id # Not used (yet) - airbyte_instance = get_airbyte_server_instance( - client_id=client_id, - client_secret=client_secret, - bearer_token=bearer_token, - api_root=api_root, - ) try: - response = airbyte_instance.jobs.create_job( - models.JobCreateRequest( + job = _make_public_api_request( + method="POST", + api_root=api_root, + path="/jobs", + request=models.JobCreateRequest( connection_id=connection_id, job_type=models.JobTypeEnum.SYNC, ), + response_model=models.JobResponse, + client_id=client_id, + client_secret=client_secret, + bearer_token=bearer_token, + error_context={"workspace_id": workspace_id, "connection_id": connection_id}, ) - except SDKError as e: + except AirbyteCloudError as e: raise AirbyteConnectionSyncError( connection_id=connection_id, - message=f"API error occurred: {e.message}", - context={"workspace_id": workspace_id, **_get_sdk_error_context(e)}, + message=e.message, + context={"workspace_id": workspace_id, **(e.context or {})}, ) from e - if status_ok(response.status_code) and response.job_response: - return response.job_response + if job: + return job raise AirbyteConnectionSyncError( connection_id=connection_id, context={ "workspace_id": workspace_id, - "request_url": response.raw_response.url, - "status_code": response.status_code, }, - response=response, ) @@ -933,12 +1013,6 @@ def get_job_logs( # noqa: PLR0913 # Too many arguments - needed for auth flexi A list of JobResponse objects. """ _validate_pagination_params(limit=limit) - airbyte_instance = get_airbyte_server_instance( - client_id=client_id, - client_secret=client_secret, - bearer_token=bearer_token, - api_root=api_root, - ) result: list[models.JobResponse] = [] current_offset = offset or 0 remaining = limit @@ -961,61 +1035,36 @@ def get_job_logs( # noqa: PLR0913 # Too many arguments - needed for auth flexi while remaining is None or remaining > 0: page_limit = _get_page_limit(remaining) - try: - response: api.ListJobsResponse = airbyte_instance.jobs.list_jobs( - api.ListJobsRequest( - workspace_ids=[workspace_id], - connection_id=connection_id, - limit=page_limit, - offset=current_offset, - order_by=order_by, - job_type=job_type_value, - ), - ) - 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( - 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, - }, - ) + body = _make_public_api_request( + method="GET", + api_root=api_root, + path="/jobs", + params={ + "workspaceIds": [workspace_id], + "connectionId": connection_id, + "limit": page_limit, + "offset": current_offset, + "orderBy": order_by, + "jobType": job_type_value, + }, + response_model=None, + client_id=client_id, + client_secret=client_secret, + bearer_token=bearer_token, + error_context=base_context, + ) - page_data: list[models.JobResponse] = list(response.jobs_response.data) - if not page_data: + raw_items = (body or {}).get("data") or [] + if not raw_items: break + page_data = _decode_list_items(body, item_model=models.JobResponse, list_path="/jobs") result += page_data if remaining is not None: remaining -= len(page_data) - current_offset += len(page_data) + current_offset += len(raw_items) - if not response.jobs_response.next or len(page_data) < page_limit: + if not (body or {}).get("next") or len(raw_items) < page_limit: break return result if limit is None else result[:limit] @@ -1030,32 +1079,23 @@ def get_job_info( bearer_token: SecretString | None, ) -> models.JobResponse: """Get a job.""" - airbyte_instance = get_airbyte_server_instance( + job = _make_public_api_request( + method="GET", + api_root=api_root, + path=f"/jobs/{job_id}", + response_model=models.JobResponse, client_id=client_id, client_secret=client_secret, bearer_token=bearer_token, - api_root=api_root, + error_context={"job_id": job_id}, ) - try: - response = airbyte_instance.jobs.get_job( - api.GetJobRequest( - job_id=job_id, - ), - ) - except SDKError as e: - raise _wrap_sdk_error(e, {"job_id": job_id}) from e - if status_ok(response.status_code) and response.job_response: - return response.job_response + if job: + return job raise AirbyteMissingResourceError( resource_name_or_id=str(job_id), resource_type="job", - log_text=response.raw_response.text, - context={ - "request_url": response.raw_response.url, - "status_code": response.status_code, - }, ) @@ -1068,49 +1108,22 @@ def cancel_job( bearer_token: SecretString | None, ) -> models.JobResponse: """Cancel a running job.""" - airbyte_instance = get_airbyte_server_instance( + job = _make_public_api_request( + method="DELETE", + api_root=api_root, + path=f"/jobs/{job_id}", + response_model=models.JobResponse, client_id=client_id, client_secret=client_secret, bearer_token=bearer_token, - api_root=api_root, + error_context={"job_id": job_id}, ) - response = airbyte_instance.jobs.cancel_job( - api.CancelJobRequest( - job_id=job_id, - ), - ) - if status_ok(response.status_code): - if response.job_response: - return response.job_response - raise AirbyteCloudError( - message="Job cancellation response payload was empty.", - response=response, - context={ - "request_url": response.raw_response.url, - "status_code": response.status_code, - }, - ) - - if response.status_code == HTTPStatus.NOT_FOUND: - raise AirbyteMissingResourceError( - resource_name_or_id=str(job_id), - resource_type="job", - response=response, - log_text=response.raw_response.text, - context={ - "request_url": response.raw_response.url, - "status_code": response.status_code, - }, - ) + if job: + return job raise AirbyteCloudError( - message="Could not cancel job.", - response=response, - log_text=response.raw_response.text, - context={ - "request_url": response.raw_response.url, - "status_code": response.status_code, - }, + message="Job cancellation response payload was empty.", + context={"job_id": job_id}, ) @@ -1132,31 +1145,27 @@ def create_source( Either `definition_id` or `config[sourceType]` must be provided. """ - airbyte_instance = get_airbyte_server_instance( - client_id=client_id, - client_secret=client_secret, - bearer_token=bearer_token, + source = _make_public_api_request( + method="POST", api_root=api_root, - ) - response: api.CreateSourceResponse = airbyte_instance.sources.create_source( - models.SourceCreateRequest( + path="/sources", + request=models.SourceCreateRequest( name=name, workspace_id=workspace_id, - configuration=config, # Speakeasy API wants a dataclass, not a dict + configuration=config, definition_id=definition_id or None, # Only used for custom sources secret_id=None, # For OAuth, not yet supported ), + response_model=models.SourceResponse, + client_id=client_id, + client_secret=client_secret, + bearer_token=bearer_token, ) - if status_ok(response.status_code) and response.source_response: - return response.source_response + if source: + return source raise AirbyteCloudError( message="Could not create source.", - context={ - "request_url": response.raw_response.url, - "status_code": response.status_code, - }, - response=response, ) @@ -1172,37 +1181,27 @@ def get_source( Secrets in the returned configuration are redacted by the API. """ - airbyte_instance = get_airbyte_server_instance( + base_context = {"source_id": source_id, "api_root": api_root} + body = _make_public_api_request( + method="GET", + api_root=api_root, + path=f"/sources/{source_id}", + response_model=None, client_id=client_id, client_secret=client_secret, bearer_token=bearer_token, - api_root=api_root, + error_context=base_context, ) - base_context = {"source_id": source_id, "api_root": api_root} - try: - response = airbyte_instance.sources.get_source( - api.GetSourceRequest( - source_id=source_id, - ), - ) - except SDKError as e: - raise _wrap_sdk_error(e, base_context) from e - - if status_ok(response.status_code) and response.source_response: - raw_response: dict[str, Any] = json.loads(response.raw_response.text) - source = response.source_response - config = raw_response.get("configuration") or {} - source.configuration = config # pyrefly: ignore[bad-assignment] # Raw config. + + if body: + source = models.SourceResponse.model_validate(body) + source.configuration = body.get("configuration") or {} # pyrefly: ignore[bad-assignment] return source raise AirbyteMissingResourceError( resource_name_or_id=source_id, resource_type="source", - log_text=response.raw_response.text, - context={ - "request_url": response.raw_response.url, - "status_code": response.status_code, - }, + context=base_context, ) @@ -1263,25 +1262,16 @@ def delete_source( }, ) - airbyte_instance = get_airbyte_server_instance( + _make_public_api_request( + method="DELETE", + api_root=api_root, + path=f"/sources/{source_id}", + response_model=None, client_id=client_id, client_secret=client_secret, bearer_token=bearer_token, - api_root=api_root, - ) - response = airbyte_instance.sources.delete_source( - api.DeleteSourceRequest( - source_id=source_id, - ), + error_context={"source_id": source_id}, ) - if not status_ok(response.status_code): - raise AirbyteCloudError( - context={ - "source_id": source_id, - "request_url": response.raw_response.url, - "status_code": response.status_code, - }, - ) def patch_source( @@ -1311,32 +1301,28 @@ def patch_source( Returns: Updated SourceResponse object """ - airbyte_instance = get_airbyte_server_instance( + source = _make_public_api_request( + method="PATCH", + api_root=api_root, + path=f"/sources/{source_id}", + request=models.SourcePatchRequest( + name=name, + configuration=config, + ), + response_model=models.SourceResponse, client_id=client_id, client_secret=client_secret, bearer_token=bearer_token, - api_root=api_root, + error_context={"source_id": source_id}, ) - response = airbyte_instance.sources.patch_source( - api.PatchSourceRequest( - source_id=source_id, - source_patch_request=models.SourcePatchRequest( - name=name, - configuration=config, - ), - ), - ) - if status_ok(response.status_code) and response.source_response: - return response.source_response + if source: + return source raise AirbyteCloudError( message="Could not update source.", context={ "source_id": source_id, - "request_url": response.raw_response.url, - "status_code": response.status_code, }, - response=response, ) @@ -1344,12 +1330,9 @@ def patch_source( def _get_destination_type_str( - destination: DestinationConfiguration | dict[str, Any], + destination: dict[str, Any], ) -> str: - if isinstance(destination, dict): - destination_type = destination.get("destinationType") - else: - destination_type = getattr(destination, "DESTINATION_TYPE", None) + destination_type = destination.get("destinationType") if not destination_type or not isinstance(destination_type, str): raise AirbyteLibInputError( @@ -1369,42 +1352,38 @@ def create_destination( name: str, *, workspace_id: str, - config: DestinationConfiguration | dict[str, Any], + config: dict[str, Any], api_root: str, client_id: SecretString | None, client_secret: SecretString | None, bearer_token: SecretString | None, ) -> models.DestinationResponse: """Get a connection.""" - airbyte_instance = get_airbyte_server_instance( - client_id=client_id, - client_secret=client_secret, - bearer_token=bearer_token, - api_root=api_root, - ) definition_id_override: str | None = None if _get_destination_type_str(config) == "dev-null": # TODO: We have to hard-code the definition ID for dev-null destination. # https://github.com/airbytehq/PyAirbyte/issues/743 definition_id_override = "a7bcc9d8-13b3-4e49-b80d-d020b90045e3" - response: api.CreateDestinationResponse = airbyte_instance.destinations.create_destination( - models.DestinationCreateRequest( + destination = _make_public_api_request( + method="POST", + api_root=api_root, + path="/destinations", + request=models.DestinationCreateRequest( definition_id=definition_id_override, name=name, workspace_id=workspace_id, - configuration=config, # Speakeasy API wants a dataclass, not a dict + configuration=config, ), + response_model=models.DestinationResponse, + client_id=client_id, + client_secret=client_secret, + bearer_token=bearer_token, ) - if status_ok(response.status_code) and response.destination_response: - return response.destination_response + if destination: + return destination raise AirbyteCloudError( message="Could not create destination.", - context={ - "request_url": response.raw_response.url, - "status_code": response.status_code, - }, - response=response, ) @@ -1420,37 +1399,29 @@ def get_destination( Secrets in the returned configuration are redacted by the API. """ - airbyte_instance = get_airbyte_server_instance( + base_context = {"destination_id": destination_id, "api_root": api_root} + body = _make_public_api_request( + method="GET", + api_root=api_root, + path=f"/destinations/{destination_id}", + response_model=None, client_id=client_id, client_secret=client_secret, bearer_token=bearer_token, - api_root=api_root, + error_context=base_context, ) - base_context = {"destination_id": destination_id, "api_root": api_root} - try: - response = airbyte_instance.destinations.get_destination( - api.GetDestinationRequest( - destination_id=destination_id, - ), - ) - except SDKError as e: - raise _wrap_sdk_error(e, base_context) from e - if status_ok(response.status_code) and response.destination_response: - raw_response: dict[str, Any] = json.loads(response.raw_response.text) - response.destination_response.configuration = ( - raw_response.get("configuration") or {} # pyrefly: ignore[bad-assignment] + if body: + destination = models.DestinationResponse.model_validate(body) + destination.configuration = ( + body.get("configuration") or {} # pyrefly: ignore[bad-assignment] ) - return response.destination_response + return destination raise AirbyteMissingResourceError( resource_name_or_id=destination_id, resource_type="destination", - log_text=response.raw_response.text, - context={ - "request_url": response.raw_response.url, - "status_code": response.status_code, - }, + context=base_context, ) @@ -1532,29 +1503,22 @@ def get_source_definition( bearer_token: SecretString | None, ) -> models.DefinitionResponse: """Get a source connector definition, including its `docker_repository` name.""" - airbyte_instance = get_airbyte_server_instance( + definition = _make_public_api_request( + method="GET", + api_root=api_root, + path=f"/workspaces/{workspace_id}/definitions/sources/{definition_id}", + response_model=models.DefinitionResponse, client_id=client_id, client_secret=client_secret, bearer_token=bearer_token, - api_root=api_root, + error_context={"definition_id": definition_id, "workspace_id": workspace_id}, ) - response = airbyte_instance.source_definitions.get_source_definition( - api.GetSourceDefinitionRequest( - definition_id=definition_id, - workspace_id=workspace_id, - ), - ) - if status_ok(response.status_code) and response.definition_response: - return response.definition_response + if definition: + return definition raise AirbyteMissingResourceError( resource_name_or_id=definition_id, resource_type="source definition", - log_text=response.raw_response.text, - context={ - "request_url": response.raw_response.url, - "status_code": response.status_code, - }, ) @@ -1568,29 +1532,22 @@ def get_destination_definition( bearer_token: SecretString | None, ) -> models.DefinitionResponse: """Get a destination connector definition, including its `docker_repository` name.""" - airbyte_instance = get_airbyte_server_instance( + definition = _make_public_api_request( + method="GET", + api_root=api_root, + path=f"/workspaces/{workspace_id}/definitions/destinations/{definition_id}", + response_model=models.DefinitionResponse, client_id=client_id, client_secret=client_secret, bearer_token=bearer_token, - api_root=api_root, - ) - response = airbyte_instance.destination_definitions.get_destination_definition( - api.GetDestinationDefinitionRequest( - definition_id=definition_id, - workspace_id=workspace_id, - ), + error_context={"definition_id": definition_id, "workspace_id": workspace_id}, ) - if status_ok(response.status_code) and response.definition_response: - return response.definition_response + if definition: + return definition raise AirbyteMissingResourceError( resource_name_or_id=definition_id, resource_type="destination definition", - log_text=response.raw_response.text, - context={ - "request_url": response.raw_response.url, - "status_code": response.status_code, - }, ) @@ -1652,25 +1609,16 @@ def delete_destination( }, ) - airbyte_instance = get_airbyte_server_instance( + _make_public_api_request( + method="DELETE", + api_root=api_root, + path=f"/destinations/{destination_id}", + response_model=None, client_id=client_id, client_secret=client_secret, bearer_token=bearer_token, - api_root=api_root, - ) - response = airbyte_instance.destinations.delete_destination( - api.DeleteDestinationRequest( - destination_id=destination_id, - ), + error_context={"destination_id": destination_id}, ) - if not status_ok(response.status_code): - raise AirbyteCloudError( - context={ - "destination_id": destination_id, - "request_url": response.raw_response.url, - "status_code": response.status_code, - }, - ) def patch_destination( @@ -1681,7 +1629,7 @@ def patch_destination( client_secret: SecretString | None, bearer_token: SecretString | None, name: str | None = None, - config: DestinationConfiguration | dict[str, Any] | None = None, + config: dict[str, Any] | None = None, ) -> models.DestinationResponse: """Update/patch a destination configuration. @@ -1700,32 +1648,28 @@ def patch_destination( Returns: Updated DestinationResponse object """ - airbyte_instance = get_airbyte_server_instance( + destination = _make_public_api_request( + method="PATCH", + api_root=api_root, + path=f"/destinations/{destination_id}", + request=models.DestinationPatchRequest( + name=name, + configuration=config, + ), + response_model=models.DestinationResponse, client_id=client_id, client_secret=client_secret, bearer_token=bearer_token, - api_root=api_root, - ) - response = airbyte_instance.destinations.patch_destination( - api.PatchDestinationRequest( - destination_id=destination_id, - destination_patch_request=models.DestinationPatchRequest( - name=name, - configuration=config, - ), - ), + error_context={"destination_id": destination_id}, ) - if status_ok(response.status_code) and response.destination_response: - return response.destination_response + if destination: + return destination raise AirbyteCloudError( message="Could not update destination.", context={ "destination_id": destination_id, - "request_url": response.raw_response.url, - "status_code": response.status_code, }, - response=response, ) @@ -1755,9 +1699,9 @@ def build_stream_configurations( def build_connection_schedule( schedule_type: str, cron_expression: str | None = None, -) -> models.AirbyteAPIConnectionSchedule: +) -> models.AirbyteApiConnectionSchedule: """Build a connection schedule object.""" - return models.AirbyteAPIConnectionSchedule( + return models.AirbyteApiConnectionSchedule( schedule_type=models.ScheduleTypeEnum(schedule_type), cron_expression=cron_expression, ) @@ -1777,33 +1721,33 @@ def create_connection( # noqa: PLR0913 # Too many arguments selected_stream_names: list[str], ) -> models.ConnectionResponse: _ = workspace_id # Not used (yet) - airbyte_instance = get_airbyte_server_instance( - client_id=client_id, - client_secret=client_secret, - bearer_token=bearer_token, - api_root=api_root, - ) stream_configurations_obj = build_stream_configurations(selected_stream_names) - response = airbyte_instance.connections.create_connection( - models.ConnectionCreateRequest( + connection = _make_public_api_request( + method="POST", + api_root=api_root, + path="/connections", + request=models.ConnectionCreateRequest( name=name, source_id=source_id, destination_id=destination_id, configurations=stream_configurations_obj, prefix=prefix, ), + response_model=models.ConnectionResponse, + client_id=client_id, + client_secret=client_secret, + bearer_token=bearer_token, + error_context={"source_id": source_id, "destination_id": destination_id}, ) - if not status_ok(response.status_code) or response.connection_response is None: + if connection is None: raise AirbyteCloudError( context={ "source_id": source_id, "destination_id": destination_id, - "request_url": response.raw_response.url, - "status_code": response.status_code, }, ) - return response.connection_response + return connection def get_connection_by_name( @@ -1916,25 +1860,16 @@ def delete_connection( ) _ = workspace_id # Not used (yet) - airbyte_instance = get_airbyte_server_instance( + _make_public_api_request( + method="DELETE", + api_root=api_root, + path=f"/connections/{connection_id}", + response_model=None, client_id=client_id, client_secret=client_secret, bearer_token=bearer_token, - api_root=api_root, + error_context={"connection_id": connection_id}, ) - response = airbyte_instance.connections.delete_connection( - api.DeleteConnectionRequest( - connection_id=connection_id, - ), - ) - if not status_ok(response.status_code): - raise AirbyteCloudError( - context={ - "connection_id": connection_id, - "request_url": response.raw_response.url, - "status_code": response.status_code, - }, - ) def patch_connection( # noqa: PLR0913 # Too many arguments @@ -1945,8 +1880,8 @@ def patch_connection( # noqa: PLR0913 # Too many arguments client_secret: SecretString | None, bearer_token: SecretString | None, name: str | None = None, - configurations: models.StreamConfigurationsInput | None = None, - schedule: models.AirbyteAPIConnectionSchedule | None = None, + configurations: models.StreamConfigurations | None = None, + schedule: models.AirbyteApiConnectionSchedule | None = None, prefix: str | None = None, status: str | models.ConnectionStatusEnum | None = None, ) -> models.ConnectionResponse: @@ -1970,53 +1905,45 @@ def patch_connection( # noqa: PLR0913 # Too many arguments Returns: Updated ConnectionResponse object """ - airbyte_instance = get_airbyte_server_instance( - client_id=client_id, - client_secret=client_secret, - bearer_token=bearer_token, - api_root=api_root, - ) if isinstance(status, str): - try: - status_value = models.ConnectionStatusEnum(status) - except ValueError: - valid_statuses = ", ".join( - connection_status.value for connection_status in models.ConnectionStatusEnum - ) + valid_statuses = ", ".join( + connection_status.value for connection_status in models.ConnectionStatusEnum + ) + if status not in {m.value for m in models.ConnectionStatusEnum}: raise AirbyteLibInputError( message=f"`status` must be one of: {valid_statuses}.", input_value=status, - ) from None + ) + status_value = models.ConnectionStatusEnum(status) else: status_value = status - try: - response = airbyte_instance.connections.patch_connection( - api.PatchConnectionRequest( - connection_id=connection_id, - connection_patch_request=models.ConnectionPatchRequest( - name=name, - configurations=configurations, - schedule=schedule, - prefix=prefix, - status=status_value, - ), - ), - ) - except SDKError as e: - raise _wrap_sdk_error(e, {"connection_id": connection_id}) from e + connection = _make_public_api_request( + method="PATCH", + api_root=api_root, + path=f"/connections/{connection_id}", + request=models.ConnectionPatchRequest( + name=name, + configurations=configurations, + schedule=schedule, + prefix=prefix, + status=status_value, + ), + response_model=models.ConnectionResponse, + client_id=client_id, + client_secret=client_secret, + bearer_token=bearer_token, + error_context={"connection_id": connection_id}, + ) - if status_ok(response.status_code) and response.connection_response: - return response.connection_response + if connection: + return connection raise AirbyteCloudError( message="Could not update connection.", context={ "connection_id": connection_id, - "request_url": response.raw_response.url, - "status_code": response.status_code, }, - response=response, ) @@ -2054,17 +1981,48 @@ def get_bearer_token( return SecretString(response.json()["access_token"]) -def _make_config_api_request( +@overload +def _make_config_api_request( # Mirrors the API surface. *, api_root: str, path: str, - json: dict[str, Any], + request: BaseModel | dict[str, Any], + response_model: type[_T], client_id: SecretString | None, client_secret: SecretString | None, bearer_token: SecretString | None, config_api_root: str | None = None, timeout: tuple[float, float] | None = None, -) -> dict[str, Any]: +) -> _T: ... + + +@overload +def _make_config_api_request( # Mirrors the API surface. + *, + api_root: str, + path: str, + request: BaseModel | dict[str, Any], + response_model: None = None, + client_id: SecretString | None, + client_secret: SecretString | None, + bearer_token: SecretString | None, + config_api_root: str | None = None, + timeout: tuple[float, float] | None = None, +) -> dict[str, Any]: ... + + +def _make_config_api_request( # noqa: PLR0913 # Mirrors the API surface. + *, + api_root: str, + path: str, + request: BaseModel | dict[str, Any], + response_model: type[_T] | None = None, + client_id: SecretString | None, + client_secret: SecretString | None, + bearer_token: SecretString | None, + config_api_root: str | None = None, + timeout: tuple[float, float] | None = None, +) -> _T | dict[str, Any]: config_api_root = get_config_api_root(api_root, config_api_root=config_api_root) headers = _config_api_headers( api_root=api_root, @@ -2078,7 +2036,11 @@ def _make_config_api_request( method="POST", url=full_url, headers=headers, - json=json, + json=( + request.model_dump(mode="json", exclude_none=True) + if isinstance(request, BaseModel) + else request + ), timeout=timeout, ) if not status_ok(response.status_code): @@ -2112,7 +2074,34 @@ def _make_config_api_request( if response.status_code == HTTPStatus.NO_CONTENT: return {} - return response.json() + try: + body = response.json() + except requests.exceptions.JSONDecodeError as ex: + if response_model is None: + raise + raise AirbyteCloudError( + message=f"Config API response for {path} did not match the expected schema.", + context={ + "full_url": full_url, + "path": path, + "response": response.text, + }, + ) from ex + + if response_model is None: + return body + + try: + return response_model.model_validate(body) + except ValidationError as ex: + raise AirbyteCloudError( + message=f"Config API response for {path} did not match the expected schema.", + context={ + "full_url": full_url, + "path": path, + "response": response.text, + }, + ) from ex def _config_api_headers( @@ -2241,12 +2230,17 @@ def check_connector( """ _ = workspace_id # Not used (yet) + request: BaseModel + if connector_type == "source": + request = SourceIdRequestBody(sourceId=actor_id) + else: + request = DestinationIdRequestBody(destinationId=actor_id) + try: json_result = _make_config_api_request( path=f"/{connector_type}s/check_connection", - json={ - f"{connector_type}Id": actor_id, - }, + request=request, + response_model=CheckConnectionRead, api_root=api_root, config_api_root=config_api_root, client_id=client_id, @@ -2263,7 +2257,10 @@ def check_connector( "finish setup in Airbyte Cloud.", ) raise - result, message = json_result.get("status"), json_result.get("message") + result, message = ( + json_result.status.value if json_result.status is not None else None, + json_result.message, + ) if result == "succeeded": return True, None @@ -2275,7 +2272,7 @@ def check_connector( context={ "actor_id": actor_id, "connector_type": str(connector_type), - "response": json_result, + "response": json_result.model_dump(mode="json", by_alias=True, exclude_none=True), }, ) @@ -2330,30 +2327,26 @@ def create_custom_yaml_source_definition( bearer_token: SecretString | None, ) -> models.DeclarativeSourceDefinitionResponse: """Create a custom YAML source definition.""" - airbyte_instance = get_airbyte_server_instance( + definition = _make_public_api_request( + method="POST", api_root=api_root, + path=f"/workspaces/{workspace_id}/definitions/declarative_sources", + request=models.CreateDeclarativeSourceDefinitionRequest( + name=name, + manifest=manifest, + ), + response_model=models.DeclarativeSourceDefinitionResponse, client_id=client_id, client_secret=client_secret, bearer_token=bearer_token, + error_context={"name": name, "workspace_id": workspace_id}, ) - - request_body = models.CreateDeclarativeSourceDefinitionRequest( - name=name, - manifest=manifest, - ) - request = api.CreateDeclarativeSourceDefinitionRequest( - workspace_id=workspace_id, - create_declarative_source_definition_request=request_body, - ) - response = airbyte_instance.declarative_source_definitions.create_declarative_source_definition( - request - ) - if response.declarative_source_definition_response is None: + if definition is None: raise AirbyteCloudError( message="Failed to create custom YAML source definition", context={"name": name, "workspace_id": workspace_id}, ) - return response.declarative_source_definition_response + return definition def list_custom_yaml_source_definitions( @@ -2365,32 +2358,27 @@ def list_custom_yaml_source_definitions( bearer_token: SecretString | None, ) -> list[models.DeclarativeSourceDefinitionResponse]: """List all custom YAML source definitions in a workspace.""" - airbyte_instance = get_airbyte_server_instance( + path = f"/workspaces/{workspace_id}/definitions/declarative_sources" + body = _make_public_api_request( + method="GET", api_root=api_root, + path=path, + response_model=None, client_id=client_id, client_secret=client_secret, bearer_token=bearer_token, + error_context={"workspace_id": workspace_id}, ) - - request = api.ListDeclarativeSourceDefinitionsRequest( - workspace_id=workspace_id, - ) - response = airbyte_instance.declarative_source_definitions.list_declarative_source_definitions( - request - ) - if ( - not status_ok(response.status_code) - or response.declarative_source_definitions_response is None - ): + if body is None: raise AirbyteCloudError( message="Failed to list custom YAML source definitions", context={ "workspace_id": workspace_id, - "request_url": response.raw_response.url, - "status_code": response.status_code, }, ) - return response.declarative_source_definitions_response.data + return _decode_list_items( + body, item_model=models.DeclarativeSourceDefinitionResponse, list_path=path + ) def get_custom_yaml_source_definition( @@ -2403,34 +2391,28 @@ def get_custom_yaml_source_definition( bearer_token: SecretString | None, ) -> models.DeclarativeSourceDefinitionResponse: """Get a specific custom YAML source definition.""" - airbyte_instance = get_airbyte_server_instance( + definition = _make_public_api_request( + method="GET", api_root=api_root, + path=f"/workspaces/{workspace_id}/definitions/declarative_sources/{definition_id}", + response_model=models.DeclarativeSourceDefinitionResponse, client_id=client_id, client_secret=client_secret, bearer_token=bearer_token, + error_context={ + "workspace_id": workspace_id, + "definition_id": definition_id, + }, ) - - request = api.GetDeclarativeSourceDefinitionRequest( - workspace_id=workspace_id, - definition_id=definition_id, - ) - response = airbyte_instance.declarative_source_definitions.get_declarative_source_definition( - request - ) - if ( - not status_ok(response.status_code) - or response.declarative_source_definition_response is None - ): + if definition is None: raise AirbyteCloudError( message="Failed to get custom YAML source definition", context={ "workspace_id": workspace_id, "definition_id": definition_id, - "request_url": response.raw_response.url, - "status_code": response.status_code, }, ) - return response.declarative_source_definition_response + return definition def update_custom_yaml_source_definition( @@ -2444,38 +2426,31 @@ def update_custom_yaml_source_definition( bearer_token: SecretString | None, ) -> models.DeclarativeSourceDefinitionResponse: """Update a custom YAML source definition.""" - airbyte_instance = get_airbyte_server_instance( + definition = _make_public_api_request( + method="PUT", api_root=api_root, + path=f"/workspaces/{workspace_id}/definitions/declarative_sources/{definition_id}", + request=models.UpdateDeclarativeSourceDefinitionRequest( + manifest=manifest, + ), + response_model=models.DeclarativeSourceDefinitionResponse, client_id=client_id, client_secret=client_secret, bearer_token=bearer_token, + error_context={ + "workspace_id": workspace_id, + "definition_id": definition_id, + }, ) - - request_body = models.UpdateDeclarativeSourceDefinitionRequest( - manifest=manifest, - ) - request = api.UpdateDeclarativeSourceDefinitionRequest( - workspace_id=workspace_id, - definition_id=definition_id, - update_declarative_source_definition_request=request_body, - ) - response = airbyte_instance.declarative_source_definitions.update_declarative_source_definition( - request - ) - if ( - not status_ok(response.status_code) - or response.declarative_source_definition_response is None - ): + if definition is None: raise AirbyteCloudError( message="Failed to update custom YAML source definition", context={ "workspace_id": workspace_id, "definition_id": definition_id, - "request_url": response.raw_response.url, - "status_code": response.status_code, }, ) - return response.declarative_source_definition_response + return definition def delete_custom_yaml_source_definition( @@ -2534,29 +2509,20 @@ def delete_custom_yaml_source_definition( # Else proceed with deletion - airbyte_instance = get_airbyte_server_instance( + _make_public_api_request( + method="DELETE", api_root=api_root, + path=f"/workspaces/{workspace_id}/definitions/declarative_sources/{definition_id}", + response_model=None, client_id=client_id, client_secret=client_secret, bearer_token=bearer_token, + error_context={ + "workspace_id": workspace_id, + "definition_id": definition_id, + }, ) - request = api.DeleteDeclarativeSourceDefinitionRequest( - workspace_id=workspace_id, - definition_id=definition_id, - ) - response = airbyte_instance.declarative_source_definitions.delete_declarative_source_definition( - request - ) - if not status_ok(response.status_code): - raise AirbyteCloudError( - context={ - "workspace_id": workspace_id, - "definition_id": definition_id, - "response": response, - }, - ) - def get_connector_builder_project_for_definition_id( *, @@ -2567,7 +2533,7 @@ def get_connector_builder_project_for_definition_id( client_secret: SecretString | None, bearer_token: SecretString | None, config_api_root: str | None = None, -) -> dict[str, Any]: +) -> BuilderProjectForDefinitionResponse: """Get the connector builder project info for a declarative source definition. Uses the Config API endpoint: @@ -2585,15 +2551,16 @@ def get_connector_builder_project_for_definition_id( config_api_root: Optional explicit Config API root URL. Returns: - A dict containing 'builderProjectId' and 'workspaceId' (the workspace that + A response containing 'builderProjectId' and 'workspaceId' (the workspace that owns the builder project, which may differ from the caller's workspace). """ return _make_config_api_request( path="/connector_builder_projects/get_for_definition_id", - json={ - "actorDefinitionId": definition_id, - "workspaceId": workspace_id, - }, + request=BuilderProjectForDefinitionRequestBody( + actorDefinitionId=definition_id, + workspaceId=workspace_id, + ), + response_model=BuilderProjectForDefinitionResponse, api_root=api_root, config_api_root=config_api_root, client_id=client_id, @@ -2617,7 +2584,7 @@ def list_connector_builder_projects( """ response = _make_config_api_request( path="/connector_builder_projects/list", - json={"workspaceId": workspace_id}, + request={"workspaceId": workspace_id}, api_root=api_root, config_api_root=config_api_root, client_id=client_id, @@ -2636,7 +2603,7 @@ def get_connector_builder_project( client_secret: SecretString | None, bearer_token: SecretString | None, config_api_root: str | None = None, -) -> dict[str, Any]: +) -> ConnectorBuilderProjectRead: """Get a connector builder project, including the draft manifest if one exists. Uses the Config API endpoint: @@ -2654,17 +2621,18 @@ def get_connector_builder_project( config_api_root: Optional explicit Config API root URL. Returns: - A dictionary containing the builder project details. Key fields include: + A generated response containing the builder project details. Key fields include: - builderProject: The project metadata (name, hasDraft, etc.) - declarativeManifest: The draft manifest data (if hasDraft is True), which contains a 'manifest' field with the actual YAML manifest dict. """ return _make_config_api_request( path="/connector_builder_projects/get_with_manifest", - json={ - "workspaceId": workspace_id, - "builderProjectId": builder_project_id, - }, + request=ConnectorBuilderProjectIdWithWorkspaceId( + workspaceId=workspace_id, + builderProjectId=builder_project_id, + ), + response_model=ConnectorBuilderProjectRead, api_root=api_root, config_api_root=config_api_root, client_id=client_id, @@ -2700,7 +2668,7 @@ def update_connector_builder_project( # noqa: PLR0913 _make_config_api_request( path="/connector_builder_projects/update", - json={ + request={ "workspaceId": workspace_id, "builderProjectId": builder_project_id, "builderProject": builder_project, @@ -2724,7 +2692,7 @@ def update_connector_builder_project_testing_values( # noqa: PLR0913 client_secret: SecretString | None, bearer_token: SecretString | None, config_api_root: str | None = None, -) -> dict[str, Any]: +) -> ConnectorBuilderProjectTestingValues: """Update the testing values for a connector builder project. This call replaces the entire testing values object stored for the project. @@ -2750,12 +2718,13 @@ def update_connector_builder_project_testing_values( # noqa: PLR0913 """ return _make_config_api_request( path="/connector_builder_projects/update_testing_values", - json={ - "workspaceId": workspace_id, - "builderProjectId": builder_project_id, - "testingValues": testing_values, - "spec": spec, - }, + request=ConnectorBuilderProjectTestingValuesUpdate( + workspaceId=workspace_id, + builderProjectId=builder_project_id, + testingValues=ConnectorBuilderProjectTestingValues.model_validate(testing_values), + spec=SourceDefinitionSpecification.model_validate(spec), + ), + response_model=ConnectorBuilderProjectTestingValues, api_root=api_root, config_api_root=config_api_root, client_id=client_id, @@ -2787,23 +2756,23 @@ def list_organizations_for_user( Returns: List of OrganizationResponse objects containing organization_id, organization_name, email """ - airbyte_instance = get_airbyte_server_instance( + body = _make_public_api_request( + method="GET", api_root=api_root, + path="/organizations", + response_model=None, client_id=client_id, client_secret=client_secret, bearer_token=bearer_token, ) - response = airbyte_instance.organizations.list_organizations_for_user() - if status_ok(response.status_code) and response.organizations_response: - return response.organizations_response.data + if body is not None: + return _decode_list_items( + body, item_model=models.OrganizationResponse, list_path="/organizations" + ) raise AirbyteCloudError( message="Failed to list organizations for user.", - context={ - "request_url": response.raw_response.url, - "status_code": response.status_code, - }, ) @@ -2817,7 +2786,7 @@ def list_organizations_for_user_id( config_api_root: str | None = None, name_contains: str | None = None, limit: int | None = None, -) -> list[dict[str, Any]]: +) -> list[OrganizationRead]: """List organizations the given user is a member of. Uses the Config API endpoint: POST /v1/organizations/list_by_user_id @@ -2836,29 +2805,23 @@ def list_organizations_for_user_id( limit: Optional maximum number of organizations to return Returns: - List of organization dictionaries containing organizationId, organizationName, email, etc. + List of organization models containing organizationId, organizationName, email, etc. """ _validate_pagination_params(limit=limit) - result: list[dict[str, Any]] = [] + result: list[OrganizationRead] = [] page_size = PAGE_SIZE - payload: dict[str, Any] = { - "userId": user_id, - "pagination": { - "pageSize": page_size, - "rowOffset": 0, - }, - } - if name_contains is not None: - payload["nameContains"] = name_contains + row_offset = 0 while True: json_result = _make_config_api_request( path="/organizations/list_by_user_id", - json={ - **payload, - "pagination": payload["pagination"].copy(), - }, + request=ListOrganizationsByUserRequestBody( + userId=user_id, + pagination=Pagination(pageSize=page_size, rowOffset=row_offset), + nameContains=name_contains, + ), + response_model=OrganizationReadList, api_root=api_root, config_api_root=config_api_root, client_id=client_id, @@ -2866,17 +2829,7 @@ def list_organizations_for_user_id( bearer_token=bearer_token, ) - if not isinstance(json_result, dict): - raise AirbyteCloudError( - message="The organizations API returned an unexpected response.", - context={"response": json_result}, - ) - organizations = json_result.get("organizations", []) - if not isinstance(organizations, list): - raise AirbyteCloudError( - message="The organizations API returned an unexpected response.", - context={"response": json_result}, - ) + organizations = json_result.organizations if not organizations: break @@ -2889,7 +2842,7 @@ def list_organizations_for_user_id( if len(organizations) < page_size: break - payload["pagination"]["rowOffset"] += page_size + row_offset += page_size return result @@ -2904,7 +2857,7 @@ def list_workspaces_in_organization( config_api_root: str | None = None, name_contains: str | None = None, limit: int | None = None, -) -> list[dict[str, Any]]: +) -> list[WorkspaceRead]: """List workspaces within a specific organization. Uses the Config API endpoint: POST /v1/workspaces/list_by_organization_id @@ -2920,28 +2873,24 @@ def list_workspaces_in_organization( limit: Optional maximum number of workspaces to return Returns: - List of workspace dictionaries containing workspaceId, organizationId, name, slug, etc. + List of workspace models containing workspaceId, organizationId, name, slug, etc. """ _validate_pagination_params(limit=limit) - result: list[dict[str, Any]] = [] + result: list[WorkspaceRead] = [] page_size = 100 - # Build base payload - payload: dict[str, Any] = { - "organizationId": organization_id, - "pagination": { - "pageSize": page_size, - "rowOffset": 0, - }, - } - if name_contains: - payload["nameContains"] = name_contains + row_offset = 0 # Fetch pages until we have all results or reach the limit while True: json_result = _make_config_api_request( path="/workspaces/list_by_organization_id", - json=payload, + request=ListWorkspacesInOrganizationRequestBody( + organizationId=organization_id, + pagination=Pagination(pageSize=page_size, rowOffset=row_offset), + nameContains=name_contains, + ), + response_model=WorkspaceReadList, api_root=api_root, config_api_root=config_api_root, client_id=client_id, @@ -2949,7 +2898,7 @@ def list_workspaces_in_organization( bearer_token=bearer_token, ) - workspaces = json_result.get("workspaces", []) + workspaces = json_result.workspaces # If no results returned, we've exhausted all pages if not workspaces: @@ -2966,7 +2915,7 @@ def list_workspaces_in_organization( break # Bump offset for next iteration - payload["pagination"]["rowOffset"] += page_size + row_offset += page_size return result @@ -3019,7 +2968,7 @@ def list_workspaces_by_user( # noqa: PLR0913 # Mirrors list_workspaces_in_orga while True: json_result = _make_config_api_request( path="/workspaces/list_by_user_id", - json={ + request={ **payload, "pagination": payload["pagination"].copy(), }, @@ -3062,7 +3011,7 @@ def get_workspace_organization_info( bearer_token: SecretString | None, config_api_root: str | None = None, timeout: tuple[float, float] | None = None, -) -> dict[str, Any]: +) -> OrganizationInfoRead: """Get organization info for a workspace. Uses the Config API endpoint: POST /v1/workspaces/get_organization_info @@ -3080,15 +3029,16 @@ def get_workspace_organization_info( timeout: Optional connect and read timeout for the request. Returns: - Dictionary containing organization info: + Generated organization info response: - organizationId: The organization ID - organizationName: The organization name - sso: Whether SSO is enabled - billing: Billing information (optional) """ - result = _make_config_api_request( + return _make_config_api_request( path="/workspaces/get_organization_info", - json={"workspaceId": workspace_id}, + request=WorkspaceIdRequestBody(workspaceId=workspace_id), + response_model=OrganizationInfoRead, api_root=api_root, config_api_root=config_api_root, client_id=client_id, @@ -3096,12 +3046,6 @@ def get_workspace_organization_info( bearer_token=bearer_token, timeout=timeout, ) - if isinstance(result, dict): - return result - raise AirbyteCloudError( - message="The workspace API returned an unexpected response.", - context={"response": result}, - ) def get_connection_state( @@ -3128,15 +3072,17 @@ def get_connection_state( Returns: Dictionary containing the connection state. """ - return _make_config_api_request( + response = _make_config_api_request( path="/state/get", - json={"connectionId": connection_id}, + request=ConnectionIdRequestBody(connectionId=connection_id), + response_model=ConnectionState, api_root=api_root, config_api_root=config_api_root, client_id=client_id, client_secret=client_secret, bearer_token=bearer_token, ) + return response.model_dump(mode="json", by_alias=True, exclude_none=True) def replace_connection_state( @@ -3181,21 +3127,25 @@ def replace_connection_state( Dictionary containing the updated ConnectionState object. """ try: - return _make_config_api_request( + response = _make_config_api_request( path="/state/create_or_update_safe", - json={ - "connectionId": connection_id, - "connectionState": { - **connection_state_dict, - "connectionId": connection_id, - }, - }, + request=ConnectionStateCreateOrUpdate( + connectionId=connection_id, + connectionState=ConnectionState.model_validate( + { + **connection_state_dict, + "connectionId": connection_id, + } + ), + ), + response_model=ConnectionState, api_root=api_root, config_api_root=config_api_root, client_id=client_id, client_secret=client_secret, bearer_token=bearer_token, ) + return response.model_dump(mode="json", by_alias=True, exclude_none=True) except AirbyteCloudError as ex: if ex.context and ex.context.get("status_code") == HTTPStatus.LOCKED: raise AirbyteConnectionSyncActiveError( @@ -3233,15 +3183,20 @@ def get_connection_catalog( Returns: Dictionary containing the connection info with syncCatalog. """ - return _make_config_api_request( + response = _make_config_api_request( path="/web_backend/connections/get", - json={"connectionId": connection_id, "withRefreshedCatalog": False}, + request=WebBackendConnectionRequestBody( + connectionId=connection_id, + withRefreshedCatalog=False, + ), + response_model=WebBackendConnectionRead, api_root=api_root, config_api_root=config_api_root, client_id=client_id, client_secret=client_secret, bearer_token=bearer_token, ) + return response.model_dump(mode="json", by_alias=True, exclude_none=True) def replace_connection_catalog( @@ -3273,21 +3228,23 @@ def replace_connection_catalog( Returns: Dictionary containing the updated WebBackendConnectionRead response. """ - return _make_config_api_request( + response = _make_config_api_request( path="/web_backend/connections/update", - json={ - "connectionId": connection_id, - "syncCatalog": configured_catalog_dict, + request=WebBackendConnectionUpdate( + connectionId=connection_id, + syncCatalog=AirbyteCatalog.model_validate(configured_catalog_dict), # Resets are destructive and cause customer-side data outage. # If a reset is desired, caller will need to decide & manage. - "skipReset": True, - }, + skipReset=True, + ), + response_model=WebBackendConnectionRead, api_root=api_root, config_api_root=config_api_root, client_id=client_id, client_secret=client_secret, bearer_token=bearer_token, ) + return response.model_dump(mode="json", by_alias=True, exclude_none=True) def get_organization_info( @@ -3298,7 +3255,7 @@ def get_organization_info( client_id: SecretString | None, client_secret: SecretString | None, bearer_token: SecretString | None, -) -> dict[str, Any]: +) -> OrganizationInfoRead: """Get organization info including billing status. Uses the Config API endpoint: POST /v1/organizations/get_organization_info @@ -3312,27 +3269,22 @@ def get_organization_info( config_api_root: Optional explicit Config API root URL. Returns: - Dictionary containing organization info: + Generated organization info response: - organizationId: The organization ID - organizationName: The organization name - sso: Whether SSO is enabled - billing: Billing information (optional, contains paymentStatus, subscriptionStatus, etc.) """ - result = _make_config_api_request( + return _make_config_api_request( path="/organizations/get_organization_info", - json={"organizationId": organization_id}, + request=OrganizationIdRequestBody(organizationId=organization_id), + response_model=OrganizationInfoRead, api_root=api_root, config_api_root=config_api_root, client_id=client_id, client_secret=client_secret, bearer_token=bearer_token, ) - if isinstance(result, dict): - return result - raise AirbyteCloudError( - message="The organization API returned an unexpected response.", - context={"response": result}, - ) def get_user_id_from_bearer_token(bearer_token: SecretString) -> str: @@ -3376,14 +3328,12 @@ def get_user_by_auth_id( client_secret: SecretString | None, bearer_token: SecretString | None, timeout: tuple[float, float] | None = None, -) -> dict[str, Any]: +) -> UserRead: """Get an Airbyte user by the authentication provider user ID.""" - result = _make_config_api_request( + return _make_config_api_request( path="/users/get_by_auth_id", - json={ - "authUserId": auth_user_id, - "authProvider": "keycloak", - }, + request=UserAuthIdRequestBody(authUserId=auth_user_id), + response_model=UserRead, api_root=api_root, config_api_root=config_api_root, client_id=client_id, @@ -3391,13 +3341,6 @@ def get_user_by_auth_id( bearer_token=bearer_token, timeout=timeout, ) - if isinstance(result, dict): - return result - - raise AirbyteCloudError( - message="The user API returned an unexpected response.", - context={"response": result}, - ) def update_user_default_workspace( @@ -3416,7 +3359,7 @@ def update_user_default_workspace( """ result = _make_config_api_request( path="/users/update", - json={ + request={ "userId": user_id, "defaultWorkspaceId": workspace_id, }, @@ -3453,7 +3396,7 @@ def get_workspace_config_api( """ result = _make_config_api_request( path="/workspaces/get", - json={ + request={ "workspaceId": workspace_id, "includeTombstone": True, }, @@ -3480,28 +3423,19 @@ def list_permissions_for_user( client_id: SecretString | None, client_secret: SecretString | None, bearer_token: SecretString | None, -) -> list[dict[str, Any]]: +) -> list[PermissionRead]: """List permissions granted to an Airbyte user.""" result = _make_config_api_request( path="/permissions/list_by_user", - json={"userId": user_id}, + request=UserIdRequestBody(userId=user_id), + response_model=PermissionReadList, api_root=api_root, config_api_root=config_api_root, client_id=client_id, client_secret=client_secret, bearer_token=bearer_token, ) - if isinstance(result, list): - return result - - permissions = result.get("permissions") if isinstance(result, dict) else None - if isinstance(permissions, list): - return permissions - - raise AirbyteCloudError( - message="The permissions API returned an unexpected response.", - context={"response": result}, - ) + return result.permissions # Billing status constants (using tuples for safe `in` checks with unhashable types) diff --git a/airbyte/_util/telemetry.py b/airbyte/_util/telemetry.py index 8091964a6..6af5b0269 100644 --- a/airbyte/_util/telemetry.py +++ b/airbyte/_util/telemetry.py @@ -35,7 +35,7 @@ import os import sys from contextlib import suppress -from enum import Enum +from enum import StrEnum from functools import lru_cache from pathlib import Path from typing import Any, cast @@ -175,14 +175,14 @@ def _get_analytics_id() -> str | None: _ANALYTICS_ID = _get_analytics_id() -class EventState(str, Enum): +class EventState(StrEnum): STARTED = "started" FAILED = "failed" SUCCEEDED = "succeeded" CANCELED = "canceled" -class EventType(str, Enum): +class EventType(StrEnum): INSTALL = "install" SYNC = "sync" VALIDATE = "validate" diff --git a/airbyte/caches/_utils/_cache_to_dest.py b/airbyte/caches/_utils/_cache_to_dest.py index 518d75f62..27b944c45 100644 --- a/airbyte/caches/_utils/_cache_to_dest.py +++ b/airbyte/caches/_utils/_cache_to_dest.py @@ -6,16 +6,6 @@ from pathlib import Path from typing import TYPE_CHECKING, Any -from airbyte_api.models import ( - BatchedStandardInserts, - DatasetLocation, - DestinationBigquery, - DestinationDuckdb, - DestinationPostgres, - DestinationSnowflake, - UsernameAndPassword, -) - from airbyte.exceptions import AirbyteLibInputError from airbyte.secrets.base import SecretString @@ -23,7 +13,6 @@ if TYPE_CHECKING: from collections.abc import Callable - from airbyte._util import api_util from airbyte.caches.base import CacheBase from airbyte.caches.bigquery import BigQueryCache from airbyte.caches.duckdb import DuckDBCache @@ -37,9 +26,9 @@ def cache_to_destination_configuration( cache: CacheBase, -) -> api_util.DestinationConfiguration: +) -> dict[str, Any]: """Get the destination configuration from the cache.""" - conversion_fn_map: dict[str, Callable[[Any], api_util.DestinationConfiguration]] = { + conversion_fn_map: dict[str, Callable[[Any], dict[str, Any]]] = { "BigQueryCache": bigquery_cache_to_destination_configuration, "bigquery": bigquery_cache_to_destination_configuration, "DuckDBCache": duckdb_cache_to_destination_configuration, @@ -67,59 +56,71 @@ def cache_to_destination_configuration( def duckdb_cache_to_destination_configuration( cache: DuckDBCache, -) -> DestinationDuckdb: +) -> dict[str, Any]: """Get the destination configuration from the DuckDB cache.""" - return DestinationDuckdb( - destination_path=str(cache.db_path), - schema=cache.schema_name, - ) + return { + "destination_path": str(cache.db_path), + "destinationType": "duckdb", + "schema": cache.schema_name, + } def motherduck_cache_to_destination_configuration( cache: MotherDuckCache, -) -> DestinationDuckdb: +) -> dict[str, Any]: """Get the destination configuration from the DuckDB cache.""" - return DestinationDuckdb( - destination_path=cache.db_path, - schema=cache.schema_name, - motherduck_api_key=cache.api_key, - ) + return { + "destination_path": cache.db_path, + "destinationType": "duckdb", + "motherduck_api_key": cache.api_key, + "schema": cache.schema_name, + } def postgres_cache_to_destination_configuration( cache: PostgresCache, -) -> DestinationPostgres: +) -> dict[str, Any]: """Get the destination configuration from the Postgres cache.""" - return DestinationPostgres( - database=cache.database, - host=cache.host, - password=cache.password, - port=cache.port, - schema=cache.schema_name, - username=cache.username, - ) + return { + "database": cache.database, + "host": cache.host, + "username": cache.username, + "destinationType": "postgres", + "disable_type_dedupe": False, + "drop_cascade": False, + "password": cache.password, + "port": cache.port, + "schema": cache.schema_name, + "ssl": False, + "unconstrained_number": False, + } def snowflake_cache_to_destination_configuration( cache: SnowflakeCache, -) -> DestinationSnowflake: +) -> dict[str, Any]: """Get the destination configuration from the Snowflake cache.""" - return DestinationSnowflake( - host=f"{cache.account}.snowflakecomputing.com", - database=cache.get_database_name().upper(), - schema=cache.schema_name.upper(), - warehouse=cache.warehouse, - role=cache.role, - username=cache.username, - credentials=UsernameAndPassword( - password=cache.password, # pyrefly: ignore[bad-argument-type] - ), - ) + return { + "host": f"{cache.account}.snowflakecomputing.com", + "database": cache.get_database_name().upper(), + "schema": cache.schema_name.upper(), + "warehouse": cache.warehouse, + "role": cache.role, + "username": cache.username, + "credentials": { + "password": cache.password, + "auth_type": "Username and Password", + }, + "destinationType": "snowflake", + "disable_type_dedupe": False, + "retention_period_days": 1, + "use_merge_for_upsert": False, + } def bigquery_cache_to_destination_configuration( cache: BigQueryCache, -) -> DestinationBigquery: +) -> dict[str, Any]: """Get the destination configuration from the BigQuery cache.""" credentials_json: str | None = ( SecretString(Path(cache.credentials_path).read_text(encoding="utf-8")) @@ -127,10 +128,13 @@ def bigquery_cache_to_destination_configuration( else None ) - return DestinationBigquery( - project_id=cache.project_name, - dataset_id=cache.dataset_name, - dataset_location=DatasetLocation(cache.dataset_location), - credentials_json=credentials_json, - loading_method=BatchedStandardInserts(), - ) + return { + "project_id": cache.project_name, + "dataset_id": cache.dataset_name, + "dataset_location": cache.dataset_location, + "cdc_deletion_mode": "Hard delete", + "credentials_json": credentials_json, + "destinationType": "bigquery", + "disable_type_dedupe": False, + "loading_method": {"method": "Standard"}, + } diff --git a/airbyte/caches/_utils/_dest_to_cache.py b/airbyte/caches/_utils/_dest_to_cache.py index ecc0a15a9..a9c408c15 100644 --- a/airbyte/caches/_utils/_dest_to_cache.py +++ b/airbyte/caches/_utils/_dest_to_cache.py @@ -8,12 +8,7 @@ from pathlib import Path from typing import TYPE_CHECKING, Any -from airbyte_api.models import ( - DestinationBigquery, - DestinationDuckdb, - DestinationPostgres, - DestinationSnowflake, -) +from pydantic import BaseModel from airbyte.caches.base import CacheBase from airbyte.caches.bigquery import BigQueryCache @@ -29,7 +24,6 @@ if TYPE_CHECKING: from collections.abc import Callable - from airbyte._util import api_util from airbyte.caches.base import CacheBase @@ -50,40 +44,60 @@ def get_supported_destination_types() -> set[str]: return _SUPPORTED_DESTINATION_TYPES +def _required_key( + destination_configuration: dict[str, Any], + key: str, +) -> Any: # noqa: ANN401 + """Read a required key from a destination configuration dict.""" + value = destination_configuration.get(key) + if value is None: + raise AirbyteLibInputError( + message=f"Destination configuration is missing required key '{key}'.", + context={"destination_configuration": destination_configuration}, + ) + return value + + def destination_to_cache( - destination_configuration: api_util.DestinationConfiguration | dict[str, Any], + destination_configuration: dict[str, Any] | BaseModel, *, schema_name: str | None = None, ) -> CacheBase: """Get the destination configuration from the cache.""" - conversion_fn_map: dict[str, Callable[[Any], CacheBase]] = { + conversion_fn_map: dict[str, Callable[[dict[str, Any]], CacheBase]] = { "bigquery": bigquery_destination_to_cache, "duckdb": duckdb_destination_to_cache, "motherduck": motherduck_destination_to_cache, "postgres": postgres_destination_to_cache, "snowflake": snowflake_destination_to_cache, } + if isinstance(destination_configuration, BaseModel): + destination_configuration = destination_configuration.model_dump(mode="json", by_alias=True) + if isinstance(destination_configuration, dict): - try: - destination_type = ( - destination_configuration.get("DESTINATION_TYPE") - or destination_configuration["destinationType"] - ) - if hasattr(destination_type, "value"): - destination_type = destination_type.value - elif hasattr(destination_type, "_value_"): - destination_type = destination_type._value_ - else: - destination_type = str(destination_type) - except KeyError as ex: + destination_type = destination_configuration.get( + "DESTINATION_TYPE" + ) or destination_configuration.get("destinationType") + if destination_type is None: raise AirbyteLibInputError( message=( "Missing 'destinationType' in keys " f"{list(destination_configuration.keys())}." ), - ) from ex + ) from None + if hasattr(destination_type, "value"): + destination_type = destination_type.value + elif hasattr(destination_type, "_value_"): + destination_type = destination_type._value_ + else: + destination_type = str(destination_type) else: - destination_type = destination_configuration.DESTINATION_TYPE.value + raise AirbyteLibInputError( + message=( + "Destination configuration must be a dict or a model, " + f"got {type(destination_configuration).__name__}." + ), + ) if destination_type not in conversion_fn_map: raise AirbyteLibInputError( @@ -106,7 +120,7 @@ def destination_to_cache( def bigquery_destination_to_cache( - destination_configuration: DestinationBigquery | dict[str, Any], + destination_configuration: dict[str, Any], ) -> BigQueryCache: """Create a new BigQuery cache from the destination configuration. @@ -118,19 +132,7 @@ def bigquery_destination_to_cache( temporary file and used directly. Otherwise we fall back to the `BIGQUERY_CREDENTIALS_PATH` secret/env-var (the cloud API path). """ - # Extract credentials_json before converting to the Pydantic model, - # because DestinationBigquery may strip or obfuscate the field. - raw_credentials_json: str | None = None - if isinstance(destination_configuration, dict): - raw_credentials_json = destination_configuration.get("credentials_json") - filtered = { - k: v - for k, v in destination_configuration.items() - if k not in {"destinationType", "DESTINATION_TYPE"} - } - destination_configuration = DestinationBigquery(**filtered) - elif hasattr(destination_configuration, "credentials_json"): - raw_credentials_json = destination_configuration.credentials_json + raw_credentials_json: str | None = destination_configuration.get("credentials_json") if raw_credentials_json and "****" not in raw_credentials_json: # Plaintext credentials from a local destination config — write to @@ -146,26 +148,18 @@ def bigquery_destination_to_cache( credentials_path = get_secret("BIGQUERY_CREDENTIALS_PATH") return BigQueryCache( - project_name=destination_configuration.project_id, - dataset_name=destination_configuration.dataset_id, + project_name=_required_key(destination_configuration, "project_id"), + dataset_name=_required_key(destination_configuration, "dataset_id"), credentials_path=credentials_path, - dataset_location=destination_configuration.dataset_location, + dataset_location=destination_configuration.get("dataset_location"), ) def duckdb_destination_to_cache( - destination_configuration: DestinationDuckdb | dict[str, Any], + destination_configuration: dict[str, Any], ) -> DuckDBCache: """Create a new DuckDB cache from the destination configuration.""" - if isinstance(destination_configuration, dict): - filtered = { - k: v - for k, v in destination_configuration.items() - if k not in {"destinationType", "DESTINATION_TYPE"} - } - destination_configuration = DestinationDuckdb(**filtered) - - db_path = destination_configuration.destination_path + db_path = _required_key(destination_configuration, "destination_path") # The DuckDB destination Docker container mounts a host directory to # `/local` inside the container. Paths written as `/local/foo.duckdb` @@ -179,61 +173,44 @@ def duckdb_destination_to_cache( return DuckDBCache( db_path=db_path, - schema_name=destination_configuration.schema or "main", + schema_name=destination_configuration.get("schema") or "main", ) def motherduck_destination_to_cache( - destination_configuration: DestinationDuckdb | dict[str, Any], + destination_configuration: dict[str, Any], ) -> MotherDuckCache: """Create a new MotherDuck cache from the destination configuration.""" - if isinstance(destination_configuration, dict): - filtered = { - k: v - for k, v in destination_configuration.items() - if k not in {"destinationType", "DESTINATION_TYPE"} - } - destination_configuration = DestinationDuckdb(**filtered) - - if not destination_configuration.motherduck_api_key: + if not destination_configuration.get("motherduck_api_key"): raise AirbyteLibInputError(message="MotherDuck API key is required for MotherDuck cache.") return MotherDuckCache( - database=destination_configuration.destination_path, - schema_name=destination_configuration.schema or "main", - api_key=SecretString(destination_configuration.motherduck_api_key), + database=_required_key(destination_configuration, "destination_path"), + schema_name=destination_configuration.get("schema") or "main", + api_key=SecretString(destination_configuration["motherduck_api_key"]), ) def postgres_destination_to_cache( - destination_configuration: DestinationPostgres | dict[str, Any], + destination_configuration: dict[str, Any], ) -> PostgresCache: """Create a new Postgres cache from the destination configuration.""" - if isinstance(destination_configuration, dict): - # Strip dispatch keys before constructing the model object. - filtered = { - k: v - for k, v in destination_configuration.items() - if k not in {"destinationType", "DESTINATION_TYPE"} - } - destination_configuration = DestinationPostgres(**filtered) - - port: int = int(destination_configuration.port) if destination_configuration.port else 5432 - if not destination_configuration.password: + port: int = int(destination_configuration.get("port") or 5432) + if not destination_configuration.get("password"): raise AirbyteLibInputError(message="Password is required for Postgres cache.") return PostgresCache( - database=destination_configuration.database, - host=destination_configuration.host, - password=destination_configuration.password, + database=_required_key(destination_configuration, "database"), + host=_required_key(destination_configuration, "host"), + password=destination_configuration["password"], port=port, - schema_name=destination_configuration.schema or "public", - username=destination_configuration.username, + schema_name=destination_configuration.get("schema") or "public", + username=_required_key(destination_configuration, "username"), ) def snowflake_destination_to_cache( - destination_configuration: DestinationSnowflake | dict[str, Any], + destination_configuration: dict[str, Any], password_secret_name: str = SNOWFLAKE_PASSWORD_SECRET_NAME, ) -> SnowflakeCache: """Create a new Snowflake cache from the destination configuration. @@ -241,21 +218,10 @@ def snowflake_destination_to_cache( We may have to inject credentials, because they are obfuscated when config is returned from the REST API. """ - if isinstance(destination_configuration, dict): - filtered = { - k: v - for k, v in destination_configuration.items() - if k not in {"destinationType", "DESTINATION_TYPE"} - } - destination_configuration = DestinationSnowflake(**filtered) - snowflake_password: str | None = None - if ( - destination_configuration.credentials - and hasattr(destination_configuration.credentials, "password") - and isinstance(destination_configuration.credentials.password, str) - ): - destination_password = str(destination_configuration.credentials.password) + credentials = destination_configuration.get("credentials") or {} + if isinstance(credentials, dict) and isinstance(credentials.get("password"), str): + destination_password = str(credentials["password"]) if "****" in destination_password: try: snowflake_password = get_secret(password_secret_name) @@ -272,11 +238,11 @@ def snowflake_destination_to_cache( snowflake_password = get_secret(password_secret_name) return SnowflakeCache( - account=destination_configuration.host.split(".snowflakecomputing")[0], - database=destination_configuration.database, - schema_name=destination_configuration.schema, - warehouse=destination_configuration.warehouse, - role=destination_configuration.role, - username=destination_configuration.username, + account=_required_key(destination_configuration, "host").split(".snowflakecomputing")[0], + database=_required_key(destination_configuration, "database"), + schema_name=_required_key(destination_configuration, "schema"), + warehouse=_required_key(destination_configuration, "warehouse"), + role=_required_key(destination_configuration, "role"), + username=_required_key(destination_configuration, "username"), password=snowflake_password, ) diff --git a/airbyte/caches/bigquery.py b/airbyte/caches/bigquery.py index fe5596f9f..b8dc859b7 100644 --- a/airbyte/caches/bigquery.py +++ b/airbyte/caches/bigquery.py @@ -17,9 +17,7 @@ from __future__ import annotations -from typing import TYPE_CHECKING, ClassVar, NoReturn - -from airbyte_api.models import DestinationBigquery +from typing import TYPE_CHECKING, Any, ClassVar, NoReturn from airbyte._processors.sql.bigquery import BigQueryConfig, BigQuerySqlProcessor from airbyte.caches._utils._cache_to_dest import ( @@ -41,10 +39,9 @@ class BigQueryCache(BigQueryConfig, CacheBase): _sql_processor_class: ClassVar[type[SqlProcessorBase]] = BigQuerySqlProcessor paired_destination_name: ClassVar[str | None] = "destination-bigquery" - paired_destination_config_class: ClassVar[type | None] = DestinationBigquery @property - def paired_destination_config(self) -> DestinationBigquery: + def paired_destination_config(self) -> dict[str, Any]: """Return a dictionary of destination configuration values.""" return bigquery_cache_to_destination_configuration(cache=self) diff --git a/airbyte/caches/duckdb.py b/airbyte/caches/duckdb.py index c27a661fa..a8d1fe2e4 100644 --- a/airbyte/caches/duckdb.py +++ b/airbyte/caches/duckdb.py @@ -17,9 +17,8 @@ from __future__ import annotations import warnings -from typing import TYPE_CHECKING, ClassVar +from typing import TYPE_CHECKING, Any, ClassVar -from airbyte_api.models import DestinationDuckdb from duckdb_engine import DuckDBEngineWarning from airbyte._processors.sql.duckdb import DuckDBConfig, DuckDBSqlProcessor @@ -46,10 +45,9 @@ class DuckDBCache(DuckDBConfig, CacheBase): _sql_processor_class: ClassVar[type[SqlProcessorBase]] = DuckDBSqlProcessor paired_destination_name: ClassVar[str | None] = "destination-duckdb" - paired_destination_config_class: ClassVar[type | None] = DestinationDuckdb @property - def paired_destination_config(self) -> DestinationDuckdb: + def paired_destination_config(self) -> dict[str, Any]: """Return a dictionary of destination configuration values.""" return duckdb_cache_to_destination_configuration(cache=self) diff --git a/airbyte/caches/motherduck.py b/airbyte/caches/motherduck.py index 0cb6e70bb..17d327c14 100644 --- a/airbyte/caches/motherduck.py +++ b/airbyte/caches/motherduck.py @@ -17,9 +17,8 @@ from __future__ import annotations import warnings -from typing import TYPE_CHECKING, ClassVar +from typing import TYPE_CHECKING, Any, ClassVar -from airbyte_api.models import DestinationDuckdb from duckdb_engine import DuckDBEngineWarning from overrides import overrides from pydantic import Field @@ -74,10 +73,9 @@ class MotherDuckCache(MotherDuckConfig, DuckDBCache): _sql_processor_class: ClassVar[type[SqlProcessorBase]] = MotherDuckSqlProcessor paired_destination_name: ClassVar[str | None] = "destination-motherduck" - paired_destination_config_class: ClassVar[type | None] = DestinationDuckdb @property - def paired_destination_config(self) -> DestinationDuckdb: + def paired_destination_config(self) -> dict[str, Any]: """Return a dictionary of destination configuration values.""" return motherduck_cache_to_destination_configuration(cache=self) diff --git a/airbyte/caches/postgres.py b/airbyte/caches/postgres.py index 678baad8a..6500dc8f5 100644 --- a/airbyte/caches/postgres.py +++ b/airbyte/caches/postgres.py @@ -19,9 +19,7 @@ from __future__ import annotations -from typing import TYPE_CHECKING, ClassVar - -from airbyte_api.models import DestinationPostgres +from typing import TYPE_CHECKING, Any, ClassVar from airbyte._processors.sql.postgres import PostgresConfig, PostgresSqlProcessor from airbyte.caches._utils._cache_to_dest import ( @@ -43,22 +41,22 @@ class PostgresCache(PostgresConfig, CacheBase): _sql_processor_class: ClassVar[type[SqlProcessorBase]] = PostgresSqlProcessor paired_destination_name: ClassVar[str | None] = "destination-postgres" - paired_destination_config_class: ClassVar[type | None] = DestinationPostgres @property - def paired_destination_config(self) -> DestinationPostgres: + def paired_destination_config(self) -> dict[str, Any]: """Return a dictionary of destination configuration values.""" return postgres_cache_to_destination_configuration(cache=self) - def clone_as_cloud_destination_config(self) -> DestinationPostgres: - """Return a DestinationPostgres instance with the same configuration.""" - return DestinationPostgres( - host=self.host, - port=self.port, - username=self.username, - password=self.password, - database=self.database, - ) + def clone_as_cloud_destination_config(self) -> dict[str, Any]: + """Return a cloud Postgres destination configuration with the same configuration.""" + return { + "host": self.host, + "port": self.port, + "username": self.username, + "password": self.password, + "database": self.database, + "destinationType": "postgres", + } # Expose the Cache class and also the Config class. diff --git a/airbyte/caches/snowflake.py b/airbyte/caches/snowflake.py index e878c8b3b..56267089d 100644 --- a/airbyte/caches/snowflake.py +++ b/airbyte/caches/snowflake.py @@ -59,9 +59,7 @@ from __future__ import annotations -from typing import ClassVar - -from airbyte_api.models import DestinationSnowflake +from typing import Any, ClassVar from airbyte._processors.sql.snowflake import SnowflakeConfig, SnowflakeSqlProcessor from airbyte.caches._utils._cache_to_dest import ( @@ -79,10 +77,9 @@ class SnowflakeCache(SnowflakeConfig, CacheBase): _sql_processor_class: ClassVar[type[SqlProcessorBase]] = SnowflakeSqlProcessor paired_destination_name: ClassVar[str | None] = "destination-snowflake" - paired_destination_config_class: ClassVar[type | None] = DestinationSnowflake @property - def paired_destination_config(self) -> DestinationSnowflake: + def paired_destination_config(self) -> dict[str, Any]: """Return a dictionary of destination configuration values.""" return snowflake_cache_to_destination_configuration(cache=self) diff --git a/airbyte/cloud/client.py b/airbyte/cloud/client.py index 9f83eaf6c..7e3ac4570 100644 --- a/airbyte/cloud/client.py +++ b/airbyte/cloud/client.py @@ -96,7 +96,9 @@ if TYPE_CHECKING: - from collections.abc import Callable, Mapping + from collections.abc import Callable + + from airbyte_server_models._config_api import OrganizationRead from airbyte.secrets.base import SecretString @@ -113,6 +115,13 @@ def _organization_has_feature( return organization.is_feature_enabled(feature) +def _api_record_to_dict(record: Any) -> dict[str, Any]: # noqa: ANN401 + """Dump a typed API read model to its JSON-mapped dict, passing dicts through.""" + if isinstance(record, dict): + return record + return record.model_dump(mode="json", by_alias=True, exclude_none=True) + + @dataclass(init=False, kw_only=True) class CloudClient: """Authenticated client for Airbyte Cloud and self-managed Airbyte APIs.""" @@ -509,7 +518,10 @@ def matches_name(workspace_name: str) -> bool: limit=limit, page_size=page_size, ) - return [CloudWorkspaceInfo.from_mapping(workspace) for workspace in workspaces] + return [ + CloudWorkspaceInfo.from_mapping(_api_record_to_dict(workspace)) + for workspace in workspaces + ] def _list_workspaces_in_organizations( self, @@ -537,7 +549,8 @@ def _list_workspaces_in_organizations( limit=None if name is not None or name_filter is not None else remaining_limit, ) organization_workspaces = [ - CloudWorkspaceInfo.from_mapping(workspace) for workspace in workspaces + CloudWorkspaceInfo.from_mapping(_api_record_to_dict(workspace)) + for workspace in workspaces ] if name is not None: organization_workspaces = [ @@ -621,9 +634,10 @@ def _get_workspace_parent_organization_id(self, workspace_id: str) -> str: client_secret=self.client_secret, bearer_token=self._get_config_api_bearer_token(), ) - resolved_organization_id = organization.get("organizationId") - if isinstance(resolved_organization_id, str) and resolved_organization_id: - return resolved_organization_id + organization_record = _api_record_to_dict(organization) + organization_id_value = organization_record.get("organizationId") + if organization_id_value: + return str(organization_id_value) raise exc.AirbyteCloudError( message="The workspace response did not include an organization ID.", context={"workspace_id": workspace_id, "response": organization}, @@ -648,7 +662,7 @@ def _get_authenticated_user_info(self) -> dict[str, Any]: guidance="Provide either client credentials or a bearer token.", ) auth_user_id = api_util.get_user_id_from_bearer_token(bearer_token) - self._authenticated_user_info = api_util.get_user_by_auth_id( + user_info = api_util.get_user_by_auth_id( auth_user_id, api_root=self.public_api_root, config_api_root=self.config_api_root, @@ -656,6 +670,7 @@ def _get_authenticated_user_info(self) -> dict[str, Any]: client_secret=self.client_secret, bearer_token=bearer_token, ) + self._authenticated_user_info = _api_record_to_dict(user_info) return self._authenticated_user_info def _get_authenticated_user_id(self) -> str: @@ -723,7 +738,7 @@ def _get_user_permissions(self) -> tuple[dict[str, Any], ...]: """Get and cache permissions for the authenticated user.""" if self._user_permissions is None: self._user_permissions = tuple( - permission + _api_record_to_dict(permission) for permission in api_util.list_permissions_for_user( self._get_authenticated_user_id(), api_root=self.public_api_root, @@ -732,7 +747,7 @@ def _get_user_permissions(self) -> tuple[dict[str, Any], ...]: client_secret=self.client_secret, bearer_token=self._get_config_api_bearer_token(), ) - if isinstance(permission, dict) + if isinstance(permission, dict) or hasattr(permission, "model_dump") ) return self._user_permissions @@ -745,11 +760,7 @@ def _get_membership_organization_ids(self) -> tuple[str, ...]: organization_ids: list[str] = [] for permission in permissions: permission_organization_id = permission.get("organizationId") - if ( - isinstance(permission_organization_id, str) - and permission_organization_id - and permission_organization_id not in organization_ids - ): + if permission_organization_id and permission_organization_id not in organization_ids: organization_ids.append(permission_organization_id) self._membership_organization_ids = tuple(organization_ids) return self._membership_organization_ids @@ -802,15 +813,16 @@ def _get_workspace_organization(self, workspace_id: str) -> CloudOrganizationInf # organization unknown. self._workspace_organizations[workspace_id] = None return None - organization_id = organization.get("organizationId") + organization_record = _api_record_to_dict(organization) + organization_id = organization_record.get("organizationId") if not isinstance(organization_id, str) or not organization_id: self._workspace_organizations[workspace_id] = None return None organization_info = CloudOrganizationInfo( # pyrefly: ignore[missing-argument] organization_id=organization_id, organization_name=( - organization.get("organizationName") - if isinstance(organization.get("organizationName"), str) + organization_record.get("organizationName") + if isinstance(organization_record.get("organizationName"), str) else None ), ) @@ -1128,9 +1140,7 @@ def _get_organization_candidates( except AirbyteCloudError: pass else: - candidate_name = organization_info.get("organizationName") - if isinstance(candidate_name, str): - organization_name = candidate_name + organization_name = _api_record_to_dict(organization_info).get("organizationName") candidates.append( { "organization_id": organization_id, @@ -1230,7 +1240,7 @@ def _list_organizations_by_user_id( """List organizations via the Config API, with server-side search and paging.""" user_id = self._get_authenticated_user_id() return [ - self._organization_from_mapping(organization) + self._organization_from_model(organization) for organization in api_util.list_organizations_for_user_id( user_id=user_id, api_root=self.public_api_root, @@ -1243,15 +1253,15 @@ def _list_organizations_by_user_id( ) ] - def _organization_from_mapping( + def _organization_from_model( self, - organization: Mapping[str, Any], + organization: OrganizationRead, ) -> CloudOrganization: - """Build a `CloudOrganization` from a Config API organization mapping.""" + """Build a `CloudOrganization` from a Config API organization model.""" return CloudOrganization( - organization_id=organization["organizationId"], - organization_name=organization.get("organizationName"), - email=organization.get("email"), + organization_id=str(organization.organizationId), + organization_name=organization.organizationName, + email=organization.email, client_id=self.client_id, client_secret=self.client_secret, bearer_token=self.bearer_token, @@ -1263,7 +1273,7 @@ def _fetch_organizations(self) -> list[CloudOrganization]: """Fetch all organizations available to this client.""" return [ CloudOrganization( - organization_id=organization.organization_id, + organization_id=str(organization.organization_id), organization_name=organization.organization_name, email=organization.email, client_id=self.client_id, @@ -1297,9 +1307,15 @@ def _get_organization_by_id(self, organization_id: str) -> CloudOrganization | N ) except AirbyteCloudError: return None - if not isinstance(organization_info.get("organizationId"), str): - return None - return self._organization_from_mapping(organization_info) + return CloudOrganization( + organization_id=str(organization_info.organizationId), + organization_name=organization_info.organizationName, + client_id=self.client_id, + client_secret=self.client_secret, + bearer_token=self.bearer_token, + public_api_root=self.public_api_root, + config_api_root=self.config_api_root, + ) def _search_organizations_by_name( self, diff --git a/airbyte/cloud/connectors.py b/airbyte/cloud/connectors.py index b2fcf54b4..3e713d33f 100644 --- a/airbyte/cloud/connectors.py +++ b/airbyte/cloud/connectors.py @@ -57,7 +57,7 @@ from http import HTTPStatus from pathlib import Path -from typing import TYPE_CHECKING, Any, ClassVar, Literal, Protocol +from typing import TYPE_CHECKING, Any, ClassVar, Literal, Protocol, cast import requests import yaml @@ -237,9 +237,9 @@ def connector_type(self) -> ConnectorType: bearer_token=self.workspace.bearer_token, ) self._connector_info = ( - CloudSourceInfo.from_api_response(response) + CloudSourceInfo.from_api_response(cast(Any, response)) if self._connector_type is ConnectorType.SOURCE - else CloudDestinationInfo.from_api_response(response) + else CloudDestinationInfo.from_api_response(cast(Any, response)) ) return self._connector_type @@ -1344,11 +1344,17 @@ def connector_builder_project_id(self) -> str | None: bearer_token=self.workspace.bearer_token, config_api_root=self.workspace.config_api_root, ) - self._connector_builder_project_id = result.get("builderProjectId") + builder_project_id = ( + result.get("builderProjectId") if isinstance(result, dict) else result.builderProjectId + ) + workspace_id = result.get("workspaceId") if isinstance(result, dict) else result.workspaceId + self._connector_builder_project_id = ( + str(builder_project_id) if builder_project_id is not None else None + ) self._connector_builder_project_id_fetched = True # The builder project may live in a different workspace than the caller's. # We must use the project's owning workspace ID when fetching its data. - self._builder_project_workspace_id = result.get("workspaceId") + self._builder_project_workspace_id = str(workspace_id) if workspace_id is not None else None return self._connector_builder_project_id @@ -1412,7 +1418,7 @@ def get_builder_project_data( }, ) - self._builder_project_data = api_util.get_connector_builder_project( + project = api_util.get_connector_builder_project( workspace_id=self._builder_project_workspace_id or self.workspace.workspace_id, builder_project_id=builder_project_id, api_root=self.workspace.api_root, @@ -1421,6 +1427,11 @@ def get_builder_project_data( bearer_token=self.workspace.bearer_token, config_api_root=self.workspace.config_api_root, ) + self._builder_project_data = ( + project + if isinstance(project, dict) + else project.model_dump(mode="json", by_alias=True, exclude_none=True) + ) return self._builder_project_data @property diff --git a/airbyte/cloud/models.py b/airbyte/cloud/models.py index 0b3019941..454425b1b 100644 --- a/airbyte/cloud/models.py +++ b/airbyte/cloud/models.py @@ -6,7 +6,7 @@ from collections.abc import Mapping from dataclasses import asdict, dataclass, is_dataclass from enum import Enum -from typing import Any, Literal, Protocol +from typing import Any, Literal, Protocol, cast from pydantic import BaseModel, ConfigDict, Field @@ -28,7 +28,7 @@ class _WorkspaceResponseLike(Protocol): workspace_id: str name: str data_residency: str - notifications: object + notifications: Any class _ScheduleResponseLike(Protocol): @@ -45,26 +45,44 @@ class _ConnectionResponseLike(Protocol): name: str configurations: Any prefix: str | None - namespace_definition: object | None + namespace_definition: Any namespace_format: str | None - schedule: _ScheduleResponseLike | None - status: object + schedule: Any + status: Any class _JobResponseLike(Protocol): job_id: int - status: object + status: Any bytes_synced: int | None rows_synced: int | None start_time: str +def _missing_str_enum_member(cls: type, value: object) -> object | None: + """`Enum._missing_` factory: accept unknown string values as pseudo-members. + + API-fed enums must tolerate values newer than this library's build, so an + unrecognized status decodes into a pseudo-member holding the raw value + instead of raising `ValueError`. + """ + if not isinstance(value, str): + return None + member = cast(Any, str.__new__(cls, value)) + member._name_ = value.upper() + member._value_ = value + return member + + class ConnectionStatus(StrEnum): """Mirrors the API's connection status.""" ACTIVE = "active" INACTIVE = "inactive" DEPRECATED = "deprecated" + LOCKED = "locked" + + _missing_ = classmethod(_missing_str_enum_member) class ConnectorFeature(StrEnum): @@ -130,8 +148,8 @@ class _DestinationResponseLike(Protocol): class _DeclarativeSourceDefinitionResponseLike(Protocol): id: str name: str - manifest: dict[str, Any] | None - version: object + manifest: Any + version: Any @dataclass @@ -172,9 +190,12 @@ class JobStatusEnum(str, Enum): FAILED = "failed" SUCCEEDED = "succeeded" CANCELLED = "cancelled" + QUEUED = "queued" + + _missing_ = classmethod(_missing_str_enum_member) -class JobTypeEnum(str, Enum): +class JobTypeEnum(StrEnum): """Job type values for Airbyte Cloud jobs.""" SYNC = "sync" @@ -182,6 +203,8 @@ class JobTypeEnum(str, Enum): REFRESH = "refresh" CLEAR = "clear" + _missing_ = classmethod(_missing_str_enum_member) + class WorkspacePrivilegeScope(str, Enum): """How broadly `list_workspaces` searches for workspaces.""" @@ -557,7 +580,11 @@ def from_api_response( return cls( definition_id=definition.id, name=definition.name, - manifest=definition.manifest, + manifest=( + definition.manifest.model_dump(mode="json", by_alias=True) + if isinstance(definition.manifest, BaseModel) + else definition.manifest + ), version=str(definition.version), ) @@ -570,6 +597,13 @@ def _configuration_dict(configuration: object) -> dict[str, Any] | None: return {str(key): value for key, value in asdict(configuration).items()} if isinstance(configuration, Mapping): return {str(key): value for key, value in configuration.items()} + if isinstance(configuration, BaseModel): + root_value = getattr(configuration, "root", configuration) + if isinstance(root_value, Mapping): + return {str(key): value for key, value in root_value.items()} + dumped = configuration.model_dump(mode="json") + if isinstance(dumped, Mapping): + return {str(key): value for key, value in dumped.items()} return None diff --git a/airbyte/cloud/organizations.py b/airbyte/cloud/organizations.py index 1b4efc437..30851d4a0 100644 --- a/airbyte/cloud/organizations.py +++ b/airbyte/cloud/organizations.py @@ -5,7 +5,7 @@ import logging from functools import cached_property -from typing import Any +from typing import TYPE_CHECKING import requests @@ -16,6 +16,10 @@ from airbyte.secrets.base import SecretString +if TYPE_CHECKING: + from airbyte_server_models._config_api import OrganizationInfoRead + + logger = logging.getLogger(__name__) @@ -56,16 +60,20 @@ def __init__( config_api_root=config_api_root, organization_id=organization_id, ) - self._organization_info: dict[str, Any] | None = None + self._organization_info: OrganizationInfoRead | None = None self._organization_info_fetch_failed: bool = False - def _fetch_organization_info(self, *, force_refresh: bool = False) -> dict[str, Any]: + def _fetch_organization_info( + self, + *, + force_refresh: bool = False, + ) -> OrganizationInfoRead | None: """Fetch and cache organization info including billing status.""" if force_refresh: self._organization_info_fetch_failed = False if self._organization_info_fetch_failed and self._organization_info is None: - return {} + return None if not force_refresh and self._organization_info is not None: return self._organization_info @@ -83,7 +91,7 @@ def _fetch_organization_info(self, *, force_refresh: bool = False) -> dict[str, logger.debug("Failed to fetch organization info.", exc_info=ex) if self._organization_info is None: self._organization_info_fetch_failed = True - return self._organization_info or {} + return self._organization_info else: return self._organization_info @@ -93,15 +101,12 @@ def organization_name(self) -> str | None: if self._organization_name is not None: return self._organization_name info = self._fetch_organization_info() - return info.get("organizationName") + return info.organizationName if info is not None else None @property def email(self) -> str | None: """Email associated with the organization.""" - if self._email is not None: - return self._email - info = self._fetch_organization_info() - return info.get("email") + return self._email def get_billing_status(self) -> CloudOrganizationBillingInfo: """Fetch billing status for the organization or raise on failure.""" @@ -119,7 +124,12 @@ def get_billing_status(self) -> CloudOrganizationBillingInfo: message="Failed to retrieve organization billing information.", context={"organization_id": self.organization_id}, ) from ex - billing = info.get("billing") + info_record = ( + info + if isinstance(info, dict) + else info.model_dump(mode="json", by_alias=True, exclude_none=True) + ) + billing = info_record.get("billing") if not isinstance(billing, dict): raise AirbyteCloudError( message="Organization info did not include billing details.", @@ -139,13 +149,23 @@ def get_billing_status(self) -> CloudOrganizationBillingInfo: def payment_status(self) -> str | None: """Payment status of the organization.""" info = self._fetch_organization_info() - return (info.get("billing") or {}).get("paymentStatus") + billing = info.billing if info is not None else None + return ( + billing.paymentStatus.value + if billing is not None and billing.paymentStatus is not None + else None + ) @property def subscription_status(self) -> str | None: """Subscription status of the organization.""" info = self._fetch_organization_info() - return (info.get("billing") or {}).get("subscriptionStatus") + billing = info.billing if info is not None else None + return ( + billing.subscriptionStatus.value + if billing is not None and billing.subscriptionStatus is not None + else None + ) @property def is_account_locked(self) -> bool: diff --git a/airbyte/cloud/sync_results.py b/airbyte/cloud/sync_results.py index 3a15cb94f..893d5d5aa 100644 --- a/airbyte/cloud/sync_results.py +++ b/airbyte/cloud/sync_results.py @@ -105,6 +105,12 @@ from dataclasses import asdict, dataclass from typing import TYPE_CHECKING, Any, final +# The Config API models live in an underscore-prefixed module on purpose: the Config API is +# Airbyte-internal and may change without notice, so the private name is a deliberate signal +# to consumers. PyAirbyte accepts that contract, hence the `PLC2701` (import-private-name) +# suppressions below. +from airbyte_server_models._config_api import JobIdRequestBody, JobInfoRead # noqa: PLC2701 + from airbyte_cdk.utils.datetime_helpers import ab_datetime_parse from airbyte._util import api_util @@ -327,12 +333,13 @@ def start_time(self) -> datetime: api_root=self.workspace.api_root, config_api_root=self.workspace.config_api_root, path="/jobs/get", - json={"id": self.job_id}, + request=JobIdRequestBody(id=self.job_id), + response_model=JobInfoRead, client_id=self.workspace.client_id, client_secret=self.workspace.client_secret, bearer_token=self.workspace.bearer_token, ) - raw_start_time = job_info_raw.get("startTime") + raw_start_time = job_info_raw.job.startedAt if job_info_raw.job else None if raw_start_time: return ab_datetime_parse(raw_start_time) raise @@ -346,13 +353,12 @@ def _fetch_job_with_attempts(self) -> dict[str, Any]: api_root=self.workspace.api_root, config_api_root=self.workspace.config_api_root, path="/jobs/get", - json={ - "id": self.job_id, - }, + request=JobIdRequestBody(id=self.job_id), + response_model=JobInfoRead, client_id=self.workspace.client_id, client_secret=self.workspace.client_secret, bearer_token=self.workspace.bearer_token, - ) + ).model_dump(mode="json", by_alias=True, exclude_none=True) return self._job_with_attempts_info def get_attempts(self) -> list[SyncAttempt]: diff --git a/airbyte/cloud/workspaces.py b/airbyte/cloud/workspaces.py index 202e9ec12..ebff3009b 100644 --- a/airbyte/cloud/workspaces.py +++ b/airbyte/cloud/workspaces.py @@ -75,6 +75,8 @@ if TYPE_CHECKING: from collections.abc import Callable + from airbyte_server_models._config_api import OrganizationInfoRead + from airbyte.cloud.organizations import CloudOrganization from airbyte.sources.base import Source @@ -298,7 +300,7 @@ def workspace_url(self) -> str | None: return f"{get_web_url_root(self.api_root)}/workspaces/{self.workspace_id}" @cached_property - def _organization_info(self) -> dict[str, Any]: + def _organization_info(self) -> OrganizationInfoRead: """Fetch and cache organization info for this workspace. Uses the Config API endpoint for an efficient O(1) lookup. @@ -359,8 +361,8 @@ def get_organization( raise return None - organization_id = info.get("organizationId") - organization_name = info.get("organizationName") + organization_id = str(info.organizationId) + organization_name = info.organizationName # Validate that both organization_id and organization_name are non-null and non-empty if not organization_id or not organization_name: diff --git a/airbyte/mcp/_sso_auth.py b/airbyte/mcp/_sso_auth.py index 615694a2f..55193b864 100644 --- a/airbyte/mcp/_sso_auth.py +++ b/airbyte/mcp/_sso_auth.py @@ -169,7 +169,7 @@ def _token_has_airbyte_user(access_token: str) -> bool: return False raise - user_id = user.get("userId") + user_id = user.get("userId") if isinstance(user, dict) else str(user.userId or "") return isinstance(user_id, str) and bool(user_id) diff --git a/airbyte/mcp/_user_identity.py b/airbyte/mcp/_user_identity.py index a6c13040e..4ffb035a4 100644 --- a/airbyte/mcp/_user_identity.py +++ b/airbyte/mcp/_user_identity.py @@ -155,6 +155,18 @@ def airbyte_user_properties() -> dict[str, Any]: return {"airbyte_user_id": current_airbyte_user_id()} +def _record_string(record: Any, key: str) -> str | None: # noqa: ANN401 + """Read a string field from a Config API response, whether a model or a raw dict. + + Model fields may be non-`str` types (e.g. `UUID`); coerce those to `str` so + callers keep the string semantics the raw-JSON path had. + """ + value = record.get(key) if isinstance(record, dict) else getattr(record, key, None) + if isinstance(value, str): + return value + return str(value) if value is not None else None + + def _lookup_airbyte_user( auth_user_id: str, *, @@ -171,10 +183,10 @@ def _lookup_airbyte_user( bearer_token=bearer_token, timeout=(USER_ID_LOOKUP_TIMEOUT_SECONDS, USER_ID_LOOKUP_TIMEOUT_SECONDS), ) - user_id = user.get("userId") + user_id = _record_string(user, "userId") if not isinstance(user_id, str) or not user_id: return None - default_workspace_id = user.get("defaultWorkspaceId") + default_workspace_id = _record_string(user, "defaultWorkspaceId") return AirbyteUser( user_id=user_id, default_workspace_id=( @@ -276,7 +288,7 @@ def _lookup_workspace_organization_id( bearer_token=bearer_token, timeout=(USER_ID_LOOKUP_TIMEOUT_SECONDS, USER_ID_LOOKUP_TIMEOUT_SECONDS), ) - organization_id = organization.get("organizationId") + organization_id = _record_string(organization, "organizationId") return organization_id if isinstance(organization_id, str) and organization_id else None diff --git a/airbyte/mcp/interactive/_shared_models.py b/airbyte/mcp/interactive/_shared_models.py index 011c86872..b559f4d37 100644 --- a/airbyte/mcp/interactive/_shared_models.py +++ b/airbyte/mcp/interactive/_shared_models.py @@ -3,7 +3,7 @@ from __future__ import annotations -from enum import Enum +from enum import StrEnum from pydantic import BaseModel, Field @@ -11,7 +11,7 @@ from airbyte.registry import ConnectorType -class SupportLevel(str, Enum): +class SupportLevel(StrEnum): """Connector support levels ordered by precedence.""" ARCHIVED = "archived" diff --git a/airbyte/mcp/interactive/_workspace_sync_status_ui.py b/airbyte/mcp/interactive/_workspace_sync_status_ui.py index 4bff26fff..9f2448342 100644 --- a/airbyte/mcp/interactive/_workspace_sync_status_ui.py +++ b/airbyte/mcp/interactive/_workspace_sync_status_ui.py @@ -646,14 +646,14 @@ def _status_pie_section( with Div(style=_status_pie_chart_style(status_pie_rows)): PieChart( # pyrefly: ignore[missing-argument] data=status_pie_rows, - data_key="connections", - name_key="status", + dataKey="connections", + nameKey="status", height=360, - inner_radius=87, + innerRadius=87, padding_angle=2, - show_label=True, - show_legend=True, - show_tooltip=True, + showLabel=True, + showLegend=True, + showTooltip=True, ) with Row(gap=3, css_class="flex-wrap"): for status_row in status_pie_rows: diff --git a/airbyte/registry.py b/airbyte/registry.py index c715cd4c9..24e29623a 100644 --- a/airbyte/registry.py +++ b/airbyte/registry.py @@ -9,7 +9,7 @@ import os import warnings from copy import copy -from enum import Enum +from enum import StrEnum from pathlib import Path from typing import Any, Self, cast @@ -19,7 +19,6 @@ from airbyte import exceptions as exc from airbyte._registry_utils import fetch_registry_version_date, parse_changelog_html -from airbyte._util.compat import StrEnum from airbyte._util.meta import is_docker_installed from airbyte.constants import AIRBYTE_OFFLINE_MODE from airbyte.logs import _warn_once @@ -46,7 +45,7 @@ ) -class InstallType(str, Enum): +class InstallType(StrEnum): """The type of installation for a connector.""" YAML = "yaml" @@ -67,7 +66,7 @@ class InstallType(str, Enum): """All connectors in the registry (environment-independent).""" -class Language(str, Enum): +class Language(StrEnum): """The language of a connector.""" PYTHON = InstallType.PYTHON.value diff --git a/airbyte/secrets/base.py b/airbyte/secrets/base.py index a4e0b12d4..20bdbf52d 100644 --- a/airbyte/secrets/base.py +++ b/airbyte/secrets/base.py @@ -5,7 +5,7 @@ import json from abc import ABC, abstractmethod -from enum import Enum +from enum import StrEnum from typing import TYPE_CHECKING, Any, cast from pydantic_core import CoreSchema, core_schema @@ -20,7 +20,7 @@ from pydantic.json_schema import JsonSchemaValue -class SecretSourceEnum(str, Enum): +class SecretSourceEnum(StrEnum): """Enumeration of secret sources supported by PyAirbyte.""" ENV = "env" diff --git a/airbyte/strategies.py b/airbyte/strategies.py index 03bcd65f0..4fbd5a17e 100644 --- a/airbyte/strategies.py +++ b/airbyte/strategies.py @@ -4,7 +4,7 @@ from __future__ import annotations -from enum import Enum +from enum import StrEnum from airbyte_protocol.models import DestinationSyncMode @@ -17,7 +17,7 @@ _AUTO = "auto" -class WriteStrategy(str, Enum): +class WriteStrategy(StrEnum): """Read strategies for PyAirbyte. Read strategies set a preferred method for writing data to a destination. The actual method used @@ -53,7 +53,7 @@ class WriteStrategy(str, Enum): """ -class WriteMethod(str, Enum): +class WriteMethod(StrEnum): """Write methods for PyAirbyte. Unlike write strategies, write methods are expected to be fully resolved and do not require any diff --git a/pyproject.toml b/pyproject.toml index a7cab8894..a2914f1b7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -8,9 +8,9 @@ authors = [ requires-python = ">=3.11,<3.13" dynamic = ["version"] dependencies = [ - "airbyte-api>=0.53.0,<1.0", "airbyte-cdk>=7.28.3,<8.0.0", "airbyte-protocol-models-pdv2>=0.13.0,<1.0", + "airbyte-server-models", "cryptography>=44.0.0,<47.0.0,!=45.0.0,!=45.0.1", # 45.0.0/45.0.1 broke load_pem_private_key() "cyclopts>=4.0,<5.0", # TODO: Restore broader version constraints after verifying compatibility @@ -114,6 +114,13 @@ fallback-version = "0.0.0" [tool.hatch.metadata] allow-direct-references = true +[tool.uv.sources] +# TK-TODO: `airbyte-server-models` is consumed from a git branch while the package is +# unreleased. Pin this to a published PyPI version (and remove this source entry) +# before the next PyAirbyte release. +# CI authenticates to the private repo via .github/actions/private-git-deps-auth (Octavia bot app). +airbyte-server-models = { git = "https://github.com/airbytehq/airbyte-ops-mcp.git", rev = "b314c99aaa0ddc0e53ce79f94c1d82d26f94a5f0", subdirectory = "airbyte-server-models" } + [tool.hatch.build.targets.wheel] packages = ["airbyte"] diff --git a/tests/integration_tests/cloud/conftest.py b/tests/integration_tests/cloud/conftest.py index a85a33b1f..11d369a0a 100644 --- a/tests/integration_tests/cloud/conftest.py +++ b/tests/integration_tests/cloud/conftest.py @@ -19,12 +19,6 @@ from airbyte.secrets.google_gsm import GoogleGSMSecretManager from airbyte.sources.base import Source from airbyte.sources.util import get_source -from airbyte_api.models import ( - DestinationBigquery, - DestinationDuckdb, - DestinationPostgres, - DestinationSnowflake, -) AIRBYTE_CLOUD_WORKSPACE_ID = "19d7a891-8e0e-40ac-8a8c-5faf8d11e47c" @@ -135,9 +129,7 @@ def deployable_dummy_destination( @pytest.fixture(scope="function") def new_deployable_destination( request, -) -> ( - DestinationDuckdb | DestinationPostgres | DestinationBigquery | DestinationSnowflake -): +) -> dict: """This is a placeholder fixture that will be overridden by pytest_generate_tests().""" return request.getfixturevalue(request.param) diff --git a/tests/integration_tests/cloud/test_cloud_api_util.py b/tests/integration_tests/cloud/test_cloud_api_util.py index 4ffaef657..e700cc8f7 100644 --- a/tests/integration_tests/cloud/test_cloud_api_util.py +++ b/tests/integration_tests/cloud/test_cloud_api_util.py @@ -21,10 +21,8 @@ ) from airbyte.cloud.connectors import ConnectorType from airbyte.secrets.base import SecretString -from airbyte_api.models import ( - DestinationDuckdb, +from airbyte_server_models.public_api.models import ( DestinationResponse, - SourceFaker, SourceResponse, WorkspaceResponse, ) @@ -112,7 +110,7 @@ def test_create_and_delete_source( airbyte_cloud_client_secret: SecretString, ) -> None: new_resource_name = "deleteme-source-faker" + text_util.generate_random_suffix() - source_config = SourceFaker() + source_config = {"sourceType": "faker"} source = api_util.create_source( name=new_resource_name, workspace_id=workspace_id, @@ -146,10 +144,11 @@ def test_create_and_delete_destination( new_resource_name = ( "deleteme-destination-faker" + text_util.generate_random_suffix() ) - destination_config = DestinationDuckdb( - destination_path="temp_db", - motherduck_api_key=motherduck_api_key, - ) + destination_config = { + "destinationType": "duckdb", + "destination_path": "temp_db", + "motherduck_api_key": motherduck_api_key, + } destination = api_util.create_destination( name=new_resource_name, @@ -192,7 +191,7 @@ def test_create_and_delete_connection( name=new_source_name, api_root=airbyte_cloud_api_root, workspace_id=workspace_id, - config=SourceFaker(), + config={"sourceType": "faker"}, client_id=airbyte_cloud_client_id, client_secret=airbyte_cloud_client_secret, bearer_token=None, @@ -205,10 +204,11 @@ def test_create_and_delete_connection( name=new_destination_name, api_root=airbyte_cloud_api_root, workspace_id=workspace_id, - config=DestinationDuckdb( - destination_path="temp_db", - motherduck_api_key=motherduck_api_key, - ), + config={ + "destinationType": "duckdb", + "destination_path": "temp_db", + "motherduck_api_key": motherduck_api_key, + }, client_id=airbyte_cloud_client_id, client_secret=airbyte_cloud_client_secret, bearer_token=None, diff --git a/tests/unit_tests/test_cache_to_dest.py b/tests/unit_tests/test_cache_to_dest.py new file mode 100644 index 000000000..76cbfeea6 --- /dev/null +++ b/tests/unit_tests/test_cache_to_dest.py @@ -0,0 +1,121 @@ +# Copyright (c) 2026 Airbyte, Inc., all rights reserved. +"""Golden tests for cache-to-cloud-destination configuration dicts. + +These freeze the exact wire JSON the removed Speakeasy `airbyte_api` destination +config models serialized to, so the dict-based replacements stay byte-identical. +""" + +from __future__ import annotations + +import pytest +from airbyte.caches.bigquery import BigQueryCache +from airbyte.caches.duckdb import DuckDBCache +from airbyte.caches.motherduck import MotherDuckCache +from airbyte.caches.postgres import PostgresCache +from airbyte.caches.snowflake import SnowflakeCache + +from airbyte.caches._utils import _cache_to_dest + + +@pytest.mark.parametrize( + ("cache", "expected"), + [ + pytest.param( + DuckDBCache.model_construct(db_path="/tmp/db.duckdb", schema_name="main"), + { + "destination_path": "/tmp/db.duckdb", + "destinationType": "duckdb", + "schema": "main", + }, + id="duckdb", + ), + pytest.param( + MotherDuckCache.model_construct( + db_path="md:db", schema_name="main", api_key="token" + ), + { + "destination_path": "md:db", + "destinationType": "duckdb", + "motherduck_api_key": "token", + "schema": "main", + }, + id="motherduck", + ), + pytest.param( + PostgresCache.model_construct( + database="db", + host="host", + username="user", + password="pass", + port=5432, + schema_name="public", + ), + { + "database": "db", + "host": "host", + "username": "user", + "destinationType": "postgres", + "disable_type_dedupe": False, + "drop_cascade": False, + "password": "pass", + "port": 5432, + "schema": "public", + "ssl": False, + "unconstrained_number": False, + }, + id="postgres", + ), + pytest.param( + SnowflakeCache.model_construct( + account="acct", + database="db", + warehouse="wh", + role="role", + username="user", + password="pass", + schema_name="public", + ), + { + "host": "acct.snowflakecomputing.com", + "database": "DB", + "schema": "PUBLIC", + "warehouse": "wh", + "role": "role", + "username": "user", + "credentials": { + "password": "pass", + "auth_type": "Username and Password", + }, + "destinationType": "snowflake", + "disable_type_dedupe": False, + "retention_period_days": 1, + "use_merge_for_upsert": False, + }, + id="snowflake", + ), + pytest.param( + BigQueryCache.model_construct( + project_name="proj", + dataset_name="ds", + dataset_location="US", + credentials_path=None, + ), + { + "project_id": "proj", + "dataset_id": "ds", + "dataset_location": "US", + "cdc_deletion_mode": "Hard delete", + "credentials_json": None, + "destinationType": "bigquery", + "disable_type_dedupe": False, + "loading_method": {"method": "Standard"}, + }, + id="bigquery", + ), + ], +) +def test_cache_to_destination_configuration_matches_speakeasy_wire_json( + cache: object, + expected: dict, +) -> None: + assert _cache_to_dest.cache_to_destination_configuration(cache) == expected diff --git a/tests/unit_tests/test_cloud_api_util.py b/tests/unit_tests/test_cloud_api_util.py index aa1e90a5b..8dbfaa0d4 100644 --- a/tests/unit_tests/test_cloud_api_util.py +++ b/tests/unit_tests/test_cloud_api_util.py @@ -3,9 +3,10 @@ from __future__ import annotations -from collections.abc import Callable +import json from types import SimpleNamespace from unittest.mock import Mock +from uuid import UUID import pytest import requests @@ -17,10 +18,31 @@ AirbyteWorkspaceNotEmptyError, AirbyteLibInputError, ) +from airbyte.cloud.models import CloudConnectionInfo, CloudJobInfo, JobStatusEnum from airbyte.registry import ConnectorType from airbyte.secrets.base import SecretString -from airbyte_api import api, models -from airbyte_api.errors import SDKError +from airbyte_server_models.public_api import models +from airbyte_server_models._config_api import ( + CheckConnectionRead, + JobConfigType, + OrganizationInfoRead, + OrganizationRead, + OrganizationReadList, + PermissionRead, + PermissionReadList, + PermissionType, + Status4, + SynchronousJobRead, + UserRead, +) + +USER_ID = UUID("00000000-0000-0000-0000-000000000001") +ORGANIZATION_ID = UUID("00000000-0000-0000-0000-000000000002") +PERMISSION_ID = UUID("00000000-0000-0000-0000-000000000003") + + +def _id_for_index(index: int) -> UUID: + return UUID(int=index + 10) def _job_response(job_id: int) -> models.JobResponse: @@ -34,28 +56,14 @@ def _job_response(job_id: int) -> models.JobResponse: ) -def _list_jobs_response( - data: list[models.JobResponse], - *, - next_page: str | None, -) -> api.ListJobsResponse: - """Create a paginated jobs API response.""" - raw_response = requests.Response() - raw_response.url = "https://api.airbyte.com/v1/jobs" - return api.ListJobsResponse( - content_type="application/json", - status_code=200, - raw_response=raw_response, - jobs_response=models.JobsResponse( - data=data, - next=next_page, - ), - ) +def _job_wire(job_id: int) -> dict[str, object]: + """Create a minimal job wire dict for pagination tests.""" + return _job_response(job_id).model_dump(mode="json", by_alias=True) def _connection_response(name: str, index: int) -> models.ConnectionResponse: """Create a minimal connection response for pagination tests.""" - return models.ConnectionResponse( + return models.ConnectionResponse.model_construct( configurations={}, connection_id=f"connection-{index}", created_at=index, @@ -69,6 +77,23 @@ def _connection_response(name: str, index: int) -> models.ConnectionResponse: ) +def _connection_wire(name: str, index: int) -> dict[str, object]: + """Create a minimal connection wire dict for pagination tests.""" + return { + "connectionId": f"connection-{index}", + "name": name, + "sourceId": f"source-{index}", + "destinationId": f"destination-{index}", + "workspaceId": "workspace-id", + "status": "active", + "schedule": {"scheduleType": "manual"}, + "nonBreakingSchemaUpdatesBehavior": "ignore", + "configurations": {"streams": []}, + "createdAt": index, + "tags": [], + } + + def _workspace_response(name: str, index: int) -> models.WorkspaceResponse: """Create a minimal workspace response for pagination tests.""" return models.WorkspaceResponse( @@ -79,42 +104,46 @@ def _workspace_response(name: str, index: int) -> models.WorkspaceResponse: ) -def _list_connections_response( - data: list[models.ConnectionResponse], - *, - next_page: str | None, -) -> api.ListConnectionsResponse: - """Create a paginated connections API response.""" - raw_response = requests.Response() - raw_response.url = "https://api.airbyte.com/v1/connections" - return api.ListConnectionsResponse( - content_type="application/json", - status_code=200, - raw_response=raw_response, - connections_response=models.ConnectionsResponse( - data=data, - next=next_page, - ), - ) +def _workspace_wire(name: str, index: int) -> dict[str, object]: + """Create a minimal workspace wire dict for pagination tests.""" + return _workspace_response(name, index).model_dump(mode="json", by_alias=True) -def _list_workspaces_response( - data: list[models.WorkspaceResponse], - *, - next_page: str | None, -) -> api.ListWorkspacesResponse: - """Create a paginated workspaces API response.""" - raw_response = requests.Response() - raw_response.url = "https://api.airbyte.com/v1/workspaces" - return api.ListWorkspacesResponse( - content_type="application/json", - status_code=200, - raw_response=raw_response, - workspaces_response=models.WorkspacesResponse( - data=data, - next=next_page, - ), - ) +def _page_body(items: list[dict[str, object]], next_page: str | None = None) -> dict: + """Create a paged list-endpoint wire body.""" + return {"data": items, "next": next_page} + + +def _public_api_spy( + monkeypatch: pytest.MonkeyPatch, + results: list[object], +) -> list[dict[str, object]]: + """Patch `_make_public_api_request`, capturing call kwargs and returning queued results.""" + captured: list[dict[str, object]] = [] + pages = list(results) + + def fake_request(**kwargs: object) -> object: + captured.append(kwargs) + result = pages.pop(0) + response_model = kwargs.get("response_model") + if response_model is not None and isinstance(result, dict): + return response_model.model_validate(result) + return result + + monkeypatch.setattr(api_util, "_make_public_api_request", fake_request) + return captured + + +def _http_error_response( + status_code: int, *, url: str, text: str = "" +) -> requests.Response: + """Build a `requests.Response` like a failed Public API call.""" + response = requests.Response() + response.status_code = status_code + response.url = url + response._content = text.encode() + response.request = requests.Request("GET", url).prepare() + return response @pytest.mark.parametrize( @@ -144,20 +173,19 @@ def _list_workspaces_response( ), ], ) -def test_wrap_sdk_error_classifies_missing_or_forbidden( +def test_wrap_api_error_classifies_missing_or_forbidden( status_code: int, expected_error_type: type[AirbyteCloudError], expected_message: str, expected_guidance: str | None, ) -> None: - raw_response = requests.Response() - raw_response.status_code = status_code - raw_response.url = "https://api.airbyte.com/v1/workspaces/workspace-id" - error = SDKError( - "Workspace lookup failed.", status_code, "response body", raw_response + raw_response = _http_error_response( + status_code, + url="https://api.airbyte.com/v1/workspaces/workspace-id", + text="Workspace lookup failed.", ) - wrapped = api_util._wrap_sdk_error(error, {"workspace_id": "workspace-id"}) + wrapped = api_util._wrap_api_error(raw_response, {"workspace_id": "workspace-id"}) assert type(wrapped) is expected_error_type assert wrapped.get_message() == expected_message @@ -212,89 +240,19 @@ def test_get_user_id_from_bearer_token_rejects_invalid_tokens( api_util.get_user_id_from_bearer_token(SecretString(token)) -@pytest.mark.parametrize( - ( - "helper", - "kwargs", - "expected_path", - "expected_json", - "fake_response", - "expected_result", - ), - [ - pytest.param( - api_util.get_user_by_auth_id, - {"auth_user_id": "auth-user-id"}, - "/users/get_by_auth_id", - {"authUserId": "auth-user-id", "authProvider": "keycloak"}, - {"userId": "user-id"}, - {"userId": "user-id"}, - id="user-by-auth-id", - ), - pytest.param( - api_util.list_permissions_for_user, - {"user_id": "user-id"}, - "/permissions/list_by_user", - {"userId": "user-id"}, - [{"permissionType": "organization_member", "organizationId": "org-id"}], - [{"permissionType": "organization_member", "organizationId": "org-id"}], - id="permissions-list", - ), - pytest.param( - api_util.list_permissions_for_user, - {"user_id": "user-id"}, - "/permissions/list_by_user", - {"userId": "user-id"}, - { - "permissions": [ - { - "permissionType": "organization_member", - "organizationId": "org-id", - } - ] - }, - [{"permissionType": "organization_member", "organizationId": "org-id"}], - id="permissions-envelope", - ), - pytest.param( - api_util.update_user_default_workspace, - {"user_id": "user-id", "workspace_id": "workspace-id"}, - "/users/update", - {"userId": "user-id", "defaultWorkspaceId": "workspace-id"}, - {"userId": "user-id", "defaultWorkspaceId": "workspace-id"}, - {"userId": "user-id", "defaultWorkspaceId": "workspace-id"}, - id="update-user-default-workspace", - ), - pytest.param( - api_util.get_workspace_config_api, - {"workspace_id": "workspace-id"}, - "/workspaces/get", - {"workspaceId": "workspace-id", "includeTombstone": True}, - {"workspaceId": "workspace-id", "organizationId": "org-id"}, - {"workspaceId": "workspace-id", "organizationId": "org-id"}, - id="workspace-get", - ), - ], -) -def test_config_api_helpers_forward_requests( +def test_get_user_by_auth_id_forwards_typed_request( monkeypatch: pytest.MonkeyPatch, - helper: Callable[..., object], - kwargs: dict[str, str], - expected_path: str, - expected_json: dict[str, str], - fake_response: object, - expected_result: object, ) -> None: captured: dict[str, object] = {} def fake_config_request(**request_kwargs: object) -> object: captured.update(request_kwargs) - return fake_response + return UserRead(userId=USER_ID, email="user@example.com", metadata={}) monkeypatch.setattr(api_util, "_make_config_api_request", fake_config_request) - result = helper( - **kwargs, + result = api_util.get_user_by_auth_id( + "auth-user-id", api_root="https://api.example", config_api_root="https://config.example", client_id=None, @@ -302,124 +260,140 @@ def fake_config_request(**request_kwargs: object) -> object: bearer_token=SecretString("token"), ) - assert result == expected_result - assert captured["path"] == expected_path - assert captured["json"] == expected_json + assert result.userId == USER_ID + assert captured["path"] == "/users/get_by_auth_id" + assert captured["request"].model_dump(mode="json", exclude_none=True) == { + "authUserId": "auth-user-id" + } assert captured["config_api_root"] == "https://config.example" -@pytest.mark.parametrize( - ("helper", "kwargs", "response", "expected_message"), - [ - pytest.param( - api_util.get_user_by_auth_id, - {"auth_user_id": "auth-user-id"}, - [], - "user API returned an unexpected response", - id="user-list", - ), - pytest.param( - api_util.get_user_by_auth_id, - {"auth_user_id": "auth-user-id"}, - "unexpected", - "user API returned an unexpected response", - id="user-string", - ), - pytest.param( - api_util.list_permissions_for_user, - {"user_id": "user-id"}, - "unexpected", - "permissions API returned an unexpected response", - id="permissions-string", - ), - pytest.param( - api_util.list_permissions_for_user, - {"user_id": "user-id"}, - {"permissions": {}}, - "permissions API returned an unexpected response", - id="permissions-non-list-envelope", - ), - pytest.param( - api_util.list_organizations_for_user_id, - {"user_id": "user-id"}, - [], - "organizations API returned an unexpected response", - id="organizations-list", - ), - pytest.param( - api_util.list_organizations_for_user_id, - {"user_id": "user-id"}, - {"organizations": {}}, - "organizations API returned an unexpected response", - id="organizations-non-list-envelope", - ), - pytest.param( - api_util.get_organization_info, - {"organization_id": "organization-id"}, - [], - "organization API returned an unexpected response", - id="organization-list", - ), - pytest.param( - api_util.get_organization_info, - {"organization_id": "organization-id"}, - "unexpected", - "organization API returned an unexpected response", - id="organization-string", - ), - pytest.param( - api_util.get_workspace_organization_info, - {"workspace_id": "workspace-id"}, - [], - "workspace API returned an unexpected response", - id="workspace-list", - ), - pytest.param( - api_util.get_workspace_organization_info, - {"workspace_id": "workspace-id"}, - "unexpected", - "workspace API returned an unexpected response", - id="workspace-string", - ), - pytest.param( - api_util.update_user_default_workspace, - {"user_id": "user-id", "workspace_id": "workspace-id"}, - "unexpected", - "user API returned an unexpected response", - id="update-user-string", - ), - pytest.param( - api_util.get_workspace_config_api, - {"workspace_id": "workspace-id"}, - [], - "workspace API returned an unexpected response", - id="workspace-config-list", - ), - ], -) -def test_config_api_helpers_reject_unexpected_response( +def test_list_permissions_for_user_forwards_typed_request( monkeypatch: pytest.MonkeyPatch, - helper: Callable[..., object], - kwargs: dict[str, str], - response: object, - expected_message: str, ) -> None: - monkeypatch.setattr( - api_util, - "_make_config_api_request", - lambda **_: response, + captured: dict[str, object] = {} + + permission = PermissionRead( + permissionId=PERMISSION_ID, + permissionType=PermissionType.organization_member, + userId=USER_ID, + organizationId=ORGANIZATION_ID, ) - with pytest.raises(AirbyteCloudError, match=expected_message) as exc_info: - helper( - **kwargs, + def fake_config_request(**request_kwargs: object) -> object: + captured.update(request_kwargs) + return PermissionReadList(permissions=[permission]) + + monkeypatch.setattr(api_util, "_make_config_api_request", fake_config_request) + + result = api_util.list_permissions_for_user( + str(USER_ID), + api_root="https://api.example", + client_id=None, + client_secret=None, + bearer_token=SecretString("token"), + ) + + assert result == [permission] + assert captured["request"].model_dump(mode="json", exclude_none=True) == { + "userId": str(USER_ID) + } + + +def test_make_config_api_request_wraps_response_validation_error( + monkeypatch: pytest.MonkeyPatch, +) -> None: + response = requests.Response() + response.status_code = 200 + response._content = b'{"invalid": true}' + response.url = "https://config.example/users/get_by_auth_id" + monkeypatch.setattr(requests, "request", lambda **_: response) + + with pytest.raises( + AirbyteCloudError, match="did not match the expected schema" + ) as exc_info: + api_util._make_config_api_request( + api_root="https://api.example", + path="/users/get_by_auth_id", + request=api_util.UserAuthIdRequestBody(authUserId="auth-user-id"), + response_model=UserRead, + client_id=None, + client_secret=None, + bearer_token=SecretString("token"), + config_api_root="https://config.example", + ) + + assert exc_info.value.context == { + "full_url": "https://config.example/users/get_by_auth_id", + "path": "/users/get_by_auth_id", + "response": '{"invalid": true}', + } + + +def test_make_config_api_request_wraps_invalid_json_response( + monkeypatch: pytest.MonkeyPatch, +) -> None: + response = requests.Response() + response.status_code = 200 + response._content = b"not json" + response.url = "https://config.example/users/get_by_auth_id" + + def raise_json_error() -> object: + raise requests.exceptions.JSONDecodeError("x", "doc", 0) + + monkeypatch.setattr(response, "json", raise_json_error) + monkeypatch.setattr(requests, "request", lambda **_: response) + + with pytest.raises( + AirbyteCloudError, match="did not match the expected schema" + ) as exc_info: + api_util._make_config_api_request( api_root="https://api.example", + path="/users/get_by_auth_id", + request=api_util.UserAuthIdRequestBody(authUserId="auth-user-id"), + response_model=UserRead, client_id=None, client_secret=None, bearer_token=SecretString("token"), + config_api_root="https://config.example", ) - assert exc_info.value.context == {"response": response} + assert exc_info.value.context == { + "full_url": "https://config.example/users/get_by_auth_id", + "path": "/users/get_by_auth_id", + "response": "not json", + } + + +def test_get_organization_info_forwards_typed_request( + monkeypatch: pytest.MonkeyPatch, +) -> None: + captured: dict[str, object] = {} + + monkeypatch.setattr( + api_util, + "_make_config_api_request", + lambda **kwargs: captured.update(kwargs) + or OrganizationInfoRead( + organizationId=ORGANIZATION_ID, + organizationName="Organization", + sso=False, + scim=False, + ), + ) + + result = api_util.get_organization_info( + str(ORGANIZATION_ID), + api_root="https://api.example", + client_id=None, + client_secret=None, + bearer_token=SecretString("token"), + ) + + assert result.organizationId == ORGANIZATION_ID + assert captured["request"].model_dump(mode="json", exclude_none=True) == { + "organizationId": str(ORGANIZATION_ID) + } @pytest.mark.parametrize( @@ -453,29 +427,38 @@ def test_list_organizations_for_user_id_paginates_and_forwards_filters( ) -> None: requests: list[dict[str, object]] = [] pages = [ - { - "organizations": [ - {"organizationId": f"organization-{index}"} for index in range(100) + OrganizationReadList( + organizations=[ + OrganizationRead( + organizationId=_id_for_index(index), + organizationName=f"Organization {index}", + email=f"organization-{index}@example.com", + ) + for index in range(100) ] - }, - { - "organizations": [ - {"organizationId": f"organization-{index}"} for index in range(100, 200) + ), + OrganizationReadList( + organizations=[ + OrganizationRead( + organizationId=_id_for_index(index), + organizationName=f"Organization {index}", + email=f"organization-{index}@example.com", + ) + for index in range(100, 200) ] - }, - {"organizations": []}, + ), + OrganizationReadList(organizations=[]), ] - def fake_config_request(**kwargs: object) -> dict[str, object]: - request = kwargs["json"] - assert isinstance(request, dict) - requests.append(request) + def fake_config_request(**kwargs: object) -> OrganizationReadList: + request = kwargs["request"] + requests.append(request.model_dump(mode="json", exclude_none=True)) return pages.pop(0) monkeypatch.setattr(api_util, "_make_config_api_request", fake_config_request) result = api_util.list_organizations_for_user_id( - "user-id", + str(USER_ID), api_root="https://api.airbyte.com/v1/", client_id=SecretString("client-id"), client_secret=SecretString("client-secret"), @@ -484,12 +467,14 @@ def fake_config_request(**kwargs: object) -> dict[str, object]: limit=limit, ) - result_ids = [organization["organizationId"] for organization in result] + result_ids = [str(organization.organizationId) for organization in result] assert len(result_ids) == expected_count - assert result_ids[0] == "organization-0" - assert result_ids[-1] == expected_last + assert result_ids[0] == str(_id_for_index(0)) + assert result_ids[-1] == str( + _id_for_index(int(expected_last.removeprefix("organization-"))) + ) assert requests[0] == { - "userId": "user-id", + "userId": str(USER_ID), "pagination": {"pageSize": 100, "rowOffset": 0}, **({"nameContains": name_contains} if name_contains is not None else {}), } @@ -497,6 +482,50 @@ def fake_config_request(**kwargs: object) -> dict[str, object]: assert len(requests) == request_count +@pytest.mark.parametrize( + ("connector_type", "request_field"), + [("source", "sourceId"), ("destination", "destinationId")], +) +def test_check_connector_uses_connector_specific_request_body( + monkeypatch: pytest.MonkeyPatch, + connector_type: str, + request_field: str, +) -> None: + captured: dict[str, object] = {} + actor_id = str(_id_for_index(50)) + + def fake_config_request(**kwargs: object) -> CheckConnectionRead: + captured.update(kwargs) + return CheckConnectionRead( + status=Status4.succeeded, + jobInfo=SynchronousJobRead( + id=_id_for_index(51), + configType=( + JobConfigType.check_connection_source + if connector_type == "source" + else JobConfigType.check_connection_destination + ), + createdAt=1, + endedAt=2, + succeeded=True, + ), + ) + + monkeypatch.setattr(api_util, "_make_config_api_request", fake_config_request) + + assert api_util.check_connector( + actor_id=actor_id, + connector_type=connector_type, # type: ignore[arg-type] + api_root="https://api.example", + client_id=None, + client_secret=None, + bearer_token=SecretString("token"), + ) == (True, None) + assert captured["request"].model_dump(mode="json", exclude_none=True) == { + request_field: actor_id + } + + @pytest.mark.parametrize( ( "page_lengths", @@ -556,9 +585,9 @@ def test_list_workspaces_by_user_paginates_and_respects_limit( next_workspace_id += page_length def fake_config_request(**kwargs: object) -> dict[str, object]: - json_request = kwargs["json"] + json_request = kwargs["request"] assert isinstance(json_request, dict) - captured_requests.append({"path": kwargs["path"], "json": json_request}) + captured_requests.append({"path": kwargs["path"], "request": json_request}) return pages.pop(0) monkeypatch.setattr(api_util, "_make_config_api_request", fake_config_request) @@ -580,7 +609,7 @@ def fake_config_request(**kwargs: object) -> dict[str, object]: assert [request["path"] for request in captured_requests] == [ "/workspaces/list_by_user_id" ] * len(expected_offsets) - assert [request["json"] for request in captured_requests] == [ + assert [request["request"] for request in captured_requests] == [ { "userId": "user-id", "pagination": {"pageSize": 100, "rowOffset": offset}, @@ -628,9 +657,9 @@ def test_list_workspaces_by_user_filters_each_page_before_limit( ] def fake_config_request(**kwargs: object) -> dict[str, object]: - json_request = kwargs["json"] + json_request = kwargs["request"] assert isinstance(json_request, dict) - captured_requests.append({"path": kwargs["path"], "json": json_request}) + captured_requests.append({"path": kwargs["path"], "request": json_request}) return pages.pop(0) monkeypatch.setattr(api_util, "_make_config_api_request", fake_config_request) @@ -648,7 +677,7 @@ def fake_config_request(**kwargs: object) -> dict[str, object]: assert [workspace["name"] for workspace in result] == expected_names assert [ - request["json"]["pagination"]["rowOffset"] for request in captured_requests + request["request"]["pagination"]["rowOffset"] for request in captured_requests ] == (expected_offsets) assert [request["path"] for request in captured_requests] == [ "/workspaces/list_by_user_id" @@ -661,7 +690,7 @@ def test_list_workspaces_by_user_honors_page_size( captured_payloads: list[dict[str, object]] = [] def fake_config_request(**kwargs: object) -> dict[str, object]: - json_request = kwargs["json"] + json_request = kwargs["request"] assert isinstance(json_request, dict) captured_payloads.append(json_request) return { @@ -693,36 +722,14 @@ def fake_config_request(**kwargs: object) -> dict[str, object]: def test_create_workspace_forwards_request( monkeypatch: pytest.MonkeyPatch, ) -> None: - captured_request = None - - def create_workspace( - *, - request: models.WorkspaceCreateRequest, - ) -> api.CreateWorkspaceResponse: - nonlocal captured_request - captured_request = request - raw_response = requests.Response() - raw_response.url = "https://api.airbyte.com/v1/workspaces" - return api.CreateWorkspaceResponse( - content_type="application/json", - status_code=200, - raw_response=raw_response, - workspace_response=_workspace_response("New workspace", 1), - ) - - airbyte_instance = SimpleNamespace( - workspaces=SimpleNamespace(create_workspace=create_workspace) - ) - monkeypatch.setattr( - api_util, - "get_airbyte_server_instance", - lambda **_: airbyte_instance, - ) + captured = _public_api_spy(monkeypatch, [_workspace_response("New workspace", 1)]) + organization_id = str(_id_for_index(0)) + region_id = str(_id_for_index(1)) workspace = api_util.create_workspace( name="New workspace", - organization_id="organization-id", - region_id="us-east", + organization_id=organization_id, + region_id=region_id, api_root="https://api.airbyte.com/v1", client_id=None, client_secret=None, @@ -730,38 +737,19 @@ def create_workspace( ) assert workspace.workspace_id == "workspace-1" - assert captured_request is not None + assert captured[0]["method"] == "POST" + assert captured[0]["path"] == "/workspaces" + captured_request = captured[0]["request"] assert captured_request.name == "New workspace" - assert captured_request.organization_id == "organization-id" - assert captured_request.region_id == "us-east" + assert str(captured_request.organization_id) == organization_id + assert str(captured_request.region_id) == region_id def test_rename_workspace_forwards_request( monkeypatch: pytest.MonkeyPatch, ) -> None: - captured_request = None - - def update_workspace( - request: api.UpdateWorkspaceRequest, - ) -> api.UpdateWorkspaceResponse: - nonlocal captured_request - captured_request = request - raw_response = requests.Response() - raw_response.url = "https://api.airbyte.com/v1/workspaces/workspace-1" - return api.UpdateWorkspaceResponse( - content_type="application/json", - status_code=200, - raw_response=raw_response, - workspace_response=_workspace_response("Renamed workspace", 1), - ) - - airbyte_instance = SimpleNamespace( - workspaces=SimpleNamespace(update_workspace=update_workspace) - ) - monkeypatch.setattr( - api_util, - "get_airbyte_server_instance", - lambda **_: airbyte_instance, + captured = _public_api_spy( + monkeypatch, [_workspace_response("Renamed workspace", 1)] ) workspace = api_util.rename_workspace( @@ -774,40 +762,75 @@ def update_workspace( ) assert workspace.name == "Renamed workspace" - assert captured_request is not None - assert captured_request.workspace_id == "workspace-1" - assert captured_request.workspace_update_request.name == "Renamed workspace" + assert captured[0]["method"] == "PATCH" + assert captured[0]["path"] == "/workspaces/workspace-1" + assert captured[0]["request"].name == "Renamed workspace" -def test_patch_connection_normalizes_status_string( +def test_public_api_request_sends_configuration_verbatim( monkeypatch: pytest.MonkeyPatch, ) -> None: - """Verify string connection statuses are normalized for the generated SDK.""" - captured_request = None - - def patch_connection( - request: api.PatchConnectionRequest, - ) -> api.PatchConnectionResponse: - nonlocal captured_request - captured_request = request - raw_response = requests.Response() - raw_response.url = "https://api.airbyte.com/v1/connections/connection-1" - return api.PatchConnectionResponse( - content_type="application/json", + """Connector configuration nulls are kept; top-level None fields are dropped.""" + captured_bodies: list[dict[str, object]] = [] + + def fake_request(**kwargs: object) -> SimpleNamespace: + captured_bodies.append(kwargs["json"]) + kind = "destination" if "destinations" in kwargs["url"] else "source" + response_body = { + f"{kind}Id": f"{kind}-1", + "name": "n", + f"{kind}Type": "faker", + "definitionId": "def-1", + "workspaceId": str(_id_for_index(0)), + "configuration": {}, + "createdAt": 1, + } + return SimpleNamespace( status_code=200, - raw_response=raw_response, - connection_response=_connection_response("Connection", 1), + content=b"x", + json=lambda: response_body, + text=json.dumps(response_body), ) - airbyte_instance = SimpleNamespace( - connections=SimpleNamespace(patch_connection=patch_connection) + monkeypatch.setattr(api_util.requests, "request", fake_request) + kwargs = { + "api_root": "https://api.airbyte.com/v1", + "client_id": None, + "client_secret": None, + "bearer_token": SecretString("token"), + } + + api_util.create_source( + name="n", + workspace_id=str(_id_for_index(0)), + config={"seed": None, "sourceType": "faker", "count": 5}, + **kwargs, ) - monkeypatch.setattr( - api_util, - "get_airbyte_server_instance", - lambda **_: airbyte_instance, + api_util.patch_destination( + destination_id="destination-1", + config={"seed": None, "destinationType": "motherduck"}, + **kwargs, ) + assert captured_bodies[0]["configuration"] == { + "seed": None, + "sourceType": "faker", + "count": 5, + } + assert "definitionId" not in captured_bodies[0] + assert captured_bodies[1]["configuration"] == { + "seed": None, + "destinationType": "motherduck", + } + assert "name" not in captured_bodies[1] + + +def test_patch_connection_normalizes_status_string( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Verify string connection statuses are normalized for the API request model.""" + captured = _public_api_spy(monkeypatch, [_connection_response("Connection", 1)]) + api_util.patch_connection( connection_id="connection-1", api_root="https://api.airbyte.com/v1", @@ -817,22 +840,17 @@ def patch_connection( status="inactive", ) - assert captured_request is not None - assert ( - captured_request.connection_patch_request.status - == models.ConnectionStatusEnum.INACTIVE - ) + assert captured[0]["request"].status == models.ConnectionStatusEnum.INACTIVE def test_patch_connection_rejects_invalid_status( monkeypatch: pytest.MonkeyPatch, ) -> None: """Verify invalid string connection statuses produce PyAirbyte errors.""" - airbyte_instance = SimpleNamespace(connections=SimpleNamespace()) monkeypatch.setattr( api_util, - "get_airbyte_server_instance", - lambda **_: airbyte_instance, + "_make_public_api_request", + Mock(side_effect=AssertionError("API should not be called")), ) with pytest.raises(AirbyteLibInputError, match="`status` must be one of"): @@ -859,7 +877,7 @@ def test_permanently_delete_workspace_requires_safe_name( workspace_name: str, should_delete: bool, ) -> None: - delete_calls = 0 + delete_calls = [] def get_workspace(**_: object) -> models.WorkspaceResponse: return models.WorkspaceResponse( @@ -869,29 +887,13 @@ def get_workspace(**_: object) -> models.WorkspaceResponse: workspace_id="workspace-1", ) - def delete_workspace( - request: api.DeleteWorkspaceRequest, - ) -> api.DeleteWorkspaceResponse: - nonlocal delete_calls - delete_calls += 1 - assert request.workspace_id == "workspace-1" - raw_response = requests.Response() - raw_response.url = "https://api.airbyte.com/v1/workspaces/workspace-1" - return api.DeleteWorkspaceResponse( - content_type="", - status_code=204, - raw_response=raw_response, - ) + def fake_request(**kwargs: object) -> None: + if kwargs["method"] == "DELETE": + delete_calls.append(kwargs["path"]) + return None - airbyte_instance = SimpleNamespace( - workspaces=SimpleNamespace(delete_workspace=delete_workspace) - ) monkeypatch.setattr(api_util, "get_workspace", get_workspace) - monkeypatch.setattr( - api_util, - "get_airbyte_server_instance", - lambda **_: airbyte_instance, - ) + monkeypatch.setattr(api_util, "_make_public_api_request", fake_request) monkeypatch.setattr(api_util, "list_connections", lambda **_: []) if should_delete: @@ -902,7 +904,7 @@ def delete_workspace( client_secret=None, bearer_token=SecretString("token"), ) - assert delete_calls == 1 + assert delete_calls == ["/workspaces/workspace-1"] else: with pytest.raises(AirbyteLibInputError): api_util.permanently_delete_workspace( @@ -912,34 +914,17 @@ def delete_workspace( client_secret=None, bearer_token=SecretString("token"), ) - assert delete_calls == 0 + assert delete_calls == [] def test_permanently_delete_workspace_requires_empty_workspace( monkeypatch: pytest.MonkeyPatch, ) -> None: - delete_calls = 0 - - def delete_workspace( - request: api.DeleteWorkspaceRequest, - ) -> api.DeleteWorkspaceResponse: - nonlocal delete_calls - delete_calls += 1 - raw_response = requests.Response() - raw_response.url = "https://api.airbyte.com/v1/workspaces/workspace-1" - return api.DeleteWorkspaceResponse( - content_type="", - status_code=204, - raw_response=raw_response, - ) - - airbyte_instance = SimpleNamespace( - workspaces=SimpleNamespace(delete_workspace=delete_workspace) - ) + delete_calls = [] monkeypatch.setattr( api_util, - "get_airbyte_server_instance", - lambda **_: airbyte_instance, + "_make_public_api_request", + lambda **kwargs: delete_calls.append(kwargs) or None, ) monkeypatch.setattr( api_util, @@ -959,7 +944,7 @@ def delete_workspace( assert exc_info.value.workspace_id == "workspace-id" assert exc_info.value.connection_ids == ["connection-1"] - assert delete_calls == 0 + assert delete_calls == [] @pytest.mark.parametrize( @@ -968,12 +953,11 @@ def delete_workspace( pytest.param( {"limit": 1, "name_filter": lambda name: name == "target"}, [ - _list_connections_response( + _page_body( [ - _connection_response("miss", 1), - _connection_response("target", 2), + _connection_wire("miss", 1), + _connection_wire("target", 2), ], - next_page=None, ), ], ["target"], @@ -983,16 +967,15 @@ def delete_workspace( pytest.param( {"limit": 2, "name_filter": lambda name: name == "target"}, [ - _list_connections_response( - [_connection_response("target", 1)], + _page_body( + [_connection_wire("target", 1)], next_page="next", ), - _list_connections_response( + _page_body( [ - _connection_response("target", 2), - _connection_response("extra", 3), + _connection_wire("target", 2), + _connection_wire("extra", 3), ], - next_page=None, ), ], ["target", "target"], @@ -1002,12 +985,11 @@ def delete_workspace( pytest.param( {"name": ""}, [ - _list_connections_response( + _page_body( [ - _connection_response("", 1), - _connection_response("non-empty", 2), + _connection_wire("", 1), + _connection_wire("non-empty", 2), ], - next_page=None, ), ], [""], @@ -1017,13 +999,12 @@ def delete_workspace( pytest.param( {}, [ - _list_connections_response( - [_connection_response("first", 1)], + _page_body( + [_connection_wire("first", 1)], next_page="next", ), - _list_connections_response( - [_connection_response("second", 2)], - next_page=None, + _page_body( + [_connection_wire("second", 2)], ), ], ["first", "second"], @@ -1035,28 +1016,12 @@ def delete_workspace( def test_list_connections_paginates_resources( monkeypatch: pytest.MonkeyPatch, kwargs: dict, - pages: list[api.ListConnectionsResponse], + pages: list[dict], expected_names: list[str], expected_requests: list[tuple[int | None, int | None]], ) -> None: """Verify resource list pagination, filtering, and request sizing.""" - captured_requests: list[api.ListConnectionsRequest] = [] - - def list_connections( - request: api.ListConnectionsRequest, - ) -> api.ListConnectionsResponse: - """Capture connection list requests and return queued pages.""" - captured_requests.append(request) - return pages.pop(0) - - airbyte_instance = SimpleNamespace( - connections=SimpleNamespace(list_connections=list_connections), - ) - monkeypatch.setattr( - api_util, - "get_airbyte_server_instance", - lambda **_: airbyte_instance, - ) + captured = _public_api_spy(monkeypatch, pages) result = api_util.list_connections( workspace_id="workspace-id", @@ -1069,7 +1034,7 @@ def list_connections( assert [connection.name for connection in result] == expected_names assert [ - (request.limit, request.offset) for request in captured_requests + (call["params"]["limit"], call["params"]["offset"]) for call in captured ] == expected_requests @@ -1077,32 +1042,12 @@ def test_list_workspaces_does_not_filter_by_workspace_id( monkeypatch: pytest.MonkeyPatch, ) -> None: """Verify workspace listing fetches accessible workspaces across pages.""" - captured_requests: list[api.ListWorkspacesRequest] = [] - pages = [ - _list_workspaces_response( - [_workspace_response("first", 1)], - next_page="next", - ), - _list_workspaces_response( - [_workspace_response("second", 2)], - next_page=None, - ), - ] - - def list_workspaces( - request: api.ListWorkspacesRequest, - ) -> api.ListWorkspacesResponse: - """Capture workspace list requests and return queued pages.""" - captured_requests.append(request) - return pages.pop(0) - - airbyte_instance = SimpleNamespace( - workspaces=SimpleNamespace(list_workspaces=list_workspaces), - ) - monkeypatch.setattr( - api_util, - "get_airbyte_server_instance", - lambda **_: airbyte_instance, + captured = _public_api_spy( + monkeypatch, + [ + _page_body([_workspace_wire("first", 1)], next_page="next"), + _page_body([_workspace_wire("second", 2)]), + ], ) result = api_util.list_workspaces( @@ -1118,40 +1063,21 @@ def list_workspaces( "workspace-2", ] assert [ - (request.workspace_ids, request.limit, request.offset) - for request in captured_requests + (call["params"]["limit"], call["params"]["offset"]) for call in captured ] == [ - (None, 100, 0), - (None, 100, 1), + (100, 0), + (100, 1), ] + assert all("workspaceIds" not in call["params"] for call in captured) def test_list_workspaces_caps_unfiltered_api_page_size( monkeypatch: pytest.MonkeyPatch, ) -> None: """Verify unfiltered workspace listing uses requested limit as API page size.""" - captured_requests: list[api.ListWorkspacesRequest] = [] - pages = [ - _list_workspaces_response( - [_workspace_response("first", 1)], - next_page="next", - ), - ] - - def list_workspaces( - request: api.ListWorkspacesRequest, - ) -> api.ListWorkspacesResponse: - """Capture workspace list requests and return queued pages.""" - captured_requests.append(request) - return pages.pop(0) - - airbyte_instance = SimpleNamespace( - workspaces=SimpleNamespace(list_workspaces=list_workspaces), - ) - monkeypatch.setattr( - api_util, - "get_airbyte_server_instance", - lambda **_: airbyte_instance, + captured = _public_api_spy( + monkeypatch, + [_page_body([_workspace_wire("first", 1)], next_page="next")], ) result = api_util.list_workspaces( @@ -1165,9 +1091,8 @@ def list_workspaces( assert [workspace.workspace_id for workspace in result] == ["workspace-1"] assert [ - (request.workspace_ids, request.limit, request.offset) - for request in captured_requests - ] == [(None, 1, 0)] + (call["params"]["limit"], call["params"]["offset"]) for call in captured + ] == [(1, 0)] @pytest.mark.parametrize("limit", [0, -1]) @@ -1186,28 +1111,12 @@ def test_list_connections_rejects_invalid_limits(limit: int) -> None: def test_get_job_logs_paginates_until_limit(monkeypatch: pytest.MonkeyPatch) -> None: """Verify job log pagination stops after collecting the requested limit.""" - captured_requests: list[api.ListJobsRequest] = [] - pages = [ - _list_jobs_response( - [_job_response(job_id) for job_id in range(100)], - next_page="next", - ), - _list_jobs_response( - [_job_response(job_id) for job_id in range(100, 150)], - next_page=None, - ), - ] - - def list_jobs(request: api.ListJobsRequest) -> api.ListJobsResponse: - """Capture job list requests and return queued pages.""" - captured_requests.append(request) - return pages.pop(0) - - airbyte_instance = SimpleNamespace(jobs=SimpleNamespace(list_jobs=list_jobs)) - monkeypatch.setattr( - api_util, - "get_airbyte_server_instance", - lambda **_: airbyte_instance, + captured = _public_api_spy( + monkeypatch, + [ + _page_body([_job_wire(job_id) for job_id in range(100)], next_page="next"), + _page_body([_job_wire(job_id) for job_id in range(100, 150)]), + ], ) result = api_util.get_job_logs( @@ -1221,7 +1130,9 @@ def list_jobs(request: api.ListJobsRequest) -> api.ListJobsResponse: ) assert [job.job_id for job in result] == list(range(150)) - assert [(request.limit, request.offset) for request in captured_requests] == [ + assert [ + (call["params"]["limit"], call["params"]["offset"]) for call in captured + ] == [ (100, 0), (50, 100), ] @@ -1231,28 +1142,12 @@ def test_get_job_logs_uses_offset_and_allows_unbounded_limit( monkeypatch: pytest.MonkeyPatch, ) -> None: """Verify job log pagination preserves offset and treats `None` as unbounded.""" - captured_requests: list[api.ListJobsRequest] = [] - pages = [ - _list_jobs_response( - [_job_response(job_id) for job_id in range(100)], - next_page="next", - ), - _list_jobs_response( - [_job_response(job_id) for job_id in range(100, 125)], - next_page=None, - ), - ] - - def list_jobs(request: api.ListJobsRequest) -> api.ListJobsResponse: - """Capture job list requests and return queued pages.""" - captured_requests.append(request) - return pages.pop(0) - - airbyte_instance = SimpleNamespace(jobs=SimpleNamespace(list_jobs=list_jobs)) - monkeypatch.setattr( - api_util, - "get_airbyte_server_instance", - lambda **_: airbyte_instance, + captured = _public_api_spy( + monkeypatch, + [ + _page_body([_job_wire(job_id) for job_id in range(100)], next_page="next"), + _page_body([_job_wire(job_id) for job_id in range(100, 125)]), + ], ) result = api_util.get_job_logs( @@ -1267,7 +1162,9 @@ def list_jobs(request: api.ListJobsRequest) -> api.ListJobsResponse: ) assert [job.job_id for job in result] == list(range(125)) - assert [(request.limit, request.offset) for request in captured_requests] == [ + assert [ + (call["params"]["limit"], call["params"]["offset"]) for call in captured + ] == [ (100, 10), (100, 110), ] @@ -1277,28 +1174,8 @@ def test_cancel_job_forwards_request_and_returns_job_response( monkeypatch: pytest.MonkeyPatch, ) -> None: """Verify cancelling a job forwards its ID and returns the API job response.""" - captured_request: api.CancelJobRequest | None = None job_response = _job_response(42) - raw_response = requests.Response() - raw_response.url = "https://api.airbyte.com/v1/jobs/42" - - def cancel_job(request: api.CancelJobRequest) -> api.CancelJobResponse: - """Capture the cancellation request.""" - nonlocal captured_request - captured_request = request - return api.CancelJobResponse( - content_type="application/json", - status_code=200, - raw_response=raw_response, - job_response=job_response, - ) - - airbyte_instance = SimpleNamespace(jobs=SimpleNamespace(cancel_job=cancel_job)) - monkeypatch.setattr( - api_util, - "get_airbyte_server_instance", - lambda **_: airbyte_instance, - ) + captured = _public_api_spy(monkeypatch, [job_response]) result = api_util.cancel_job( job_id=42, @@ -1309,32 +1186,19 @@ def cancel_job(request: api.CancelJobRequest) -> api.CancelJobResponse: ) assert result is job_response - assert captured_request is not None - assert captured_request.job_id == 42 + assert captured[0]["method"] == "DELETE" + assert captured[0]["path"] == "/jobs/42" def test_cancel_job_raises_for_non_ok_response( monkeypatch: pytest.MonkeyPatch, ) -> None: """Verify cancelling a job raises when the API response is not successful.""" - raw_response = requests.Response() - raw_response.status_code = 404 - raw_response.url = "https://api.airbyte.com/v1/jobs/42" - - def cancel_job(request: api.CancelJobRequest) -> api.CancelJobResponse: - """Return a not-found cancellation response.""" - _ = request - return api.CancelJobResponse( - content_type="application/json", - status_code=404, - raw_response=raw_response, - ) - - airbyte_instance = SimpleNamespace(jobs=SimpleNamespace(cancel_job=cancel_job)) + error_response = _http_error_response(404, url="https://api.airbyte.com/v1/jobs/42") monkeypatch.setattr( api_util, - "get_airbyte_server_instance", - lambda **_: airbyte_instance, + "_make_public_api_request", + Mock(side_effect=api_util._wrap_api_error(error_response)), ) with pytest.raises(AirbyteMissingResourceError): @@ -1351,24 +1215,11 @@ def test_cancel_job_raises_airbyte_error_for_non_not_found_response( monkeypatch: pytest.MonkeyPatch, ) -> None: """Verify non-not-found cancellation failures use a general API error.""" - raw_response = requests.Response() - raw_response.status_code = 409 - raw_response.url = "https://api.airbyte.com/v1/jobs/42" - - def cancel_job(request: api.CancelJobRequest) -> api.CancelJobResponse: - """Return a conflict cancellation response.""" - _ = request - return api.CancelJobResponse( - content_type="application/json", - status_code=409, - raw_response=raw_response, - ) - - airbyte_instance = SimpleNamespace(jobs=SimpleNamespace(cancel_job=cancel_job)) + error_response = _http_error_response(409, url="https://api.airbyte.com/v1/jobs/42") monkeypatch.setattr( api_util, - "get_airbyte_server_instance", - lambda **_: airbyte_instance, + "_make_public_api_request", + Mock(side_effect=api_util._wrap_api_error(error_response)), ) with pytest.raises(AirbyteCloudError) as error: @@ -1466,7 +1317,7 @@ def fake_request(**kwargs: object) -> SimpleNamespace: api_util._make_config_api_request( path="/workspaces/get", - json={}, + request={}, api_root="https://api.airbyte.com/v1", client_id=None, client_secret=None, @@ -1489,7 +1340,7 @@ def test_config_api_request_handles_no_content_response( result = api_util._make_config_api_request( path="/connector_builder_projects/update", - json={}, + request={}, api_root="https://api.airbyte.com/v1", config_api_root="https://config.airbyte.com/v1", client_id=None, @@ -1522,7 +1373,7 @@ def test_list_connector_builder_projects(monkeypatch: pytest.MonkeyPatch) -> Non assert result == projects config_api_request.assert_called_once_with( path="/connector_builder_projects/list", - json={"workspaceId": "workspace-id"}, + request={"workspaceId": "workspace-id"}, api_root="https://api.airbyte.com/v1", config_api_root="https://config.airbyte.com/v1", client_id=None, @@ -1577,7 +1428,7 @@ def test_update_connector_builder_project_payload( assert result is None config_api_request.assert_called_once_with( path="/connector_builder_projects/update", - json={ + request={ "workspaceId": "owner-workspace", "builderProjectId": "builder-project-id", "builderProject": expected_builder_project, @@ -1627,7 +1478,7 @@ def test_config_api_request_maps_forbidden_as_missing_resource( with pytest.raises(expected_error_type) as exc_info: api_util._make_config_api_request( path="/workspaces/get", - json={"workspaceId": "workspace-id"}, + request={"workspaceId": "workspace-id"}, api_root="https://api.airbyte.com/v1", config_api_root="https://config.airbyte.com/v1", client_id=None, @@ -1654,21 +1505,42 @@ def test_config_api_request_maps_forbidden_as_missing_resource( ) -def test_public_api_client_sends_analytic_source_header( +def test_public_api_request_sends_bearer_and_analytic_source_headers( monkeypatch: pytest.MonkeyPatch, ) -> None: + """The Public API path sends bearer auth plus the analytic source header.""" monkeypatch.setattr(meta, "_MCP_MODE_ENABLED", True) monkeypatch.setattr(constants, "_HOSTED_MCP_MODE_ENABLED", False) + captured: dict[str, object] = {} + + def fake_request(**kwargs: object) -> SimpleNamespace: + captured.update(kwargs) + return SimpleNamespace( + status_code=200, + content=b"{}", + json=lambda: {}, + text="{}", + url="https://api.airbyte.com/v1/x", + headers={}, + request=SimpleNamespace(url="https://api.airbyte.com/v1/x", method="GET"), + ) + + monkeypatch.setattr(api_util.requests, "request", fake_request) - airbyte_instance = api_util.get_airbyte_server_instance( + api_util._make_public_api_request( + method="GET", api_root="https://api.airbyte.com/v1", + path="/x", + response_model=None, client_id=None, client_secret=None, bearer_token=SecretString("token"), ) - session = airbyte_instance.sdk_configuration.client - assert session.headers[meta.AIRBYTE_ANALYTIC_SOURCE_HEADER] == "pyairbyte-mcp-local" + headers = captured["headers"] + assert isinstance(headers, dict) + assert headers["Authorization"] == "Bearer token" + assert headers[meta.AIRBYTE_ANALYTIC_SOURCE_HEADER] == "pyairbyte-mcp-local" def test_get_bearer_token_sends_analytic_source_header( @@ -1694,29 +1566,25 @@ def fake_post(**kwargs: object) -> SimpleNamespace: assert headers[meta.AIRBYTE_ANALYTIC_SOURCE_HEADER] == "pyairbyte" -def _sdk_404_error(resource_type: str) -> SDKError: - """Create an SDKError like the Speakeasy SDK raises on a 404.""" - raw_response = requests.Response() - raw_response.status_code = 404 - raw_response.url = "https://api.airbyte.com/v1/connectors/connector-id" - return SDKError( - "API error occurred: Status 404", - 404, - f'{{"resourceType":"{resource_type}"}}', - raw_response, +def _api_404_error() -> AirbyteCloudError: + """Wrap a failed Public API response like a 404 connector lookup.""" + return api_util._wrap_api_error( + _http_error_response( + 404, + url="https://api.airbyte.com/v1/connectors/connector-id", + text="not found", + ) ) -def _sdk_status_error(status_code: int) -> SDKError: - """Create an SDKError like the Speakeasy SDK raises on a non-2xx status.""" - raw_response = requests.Response() - raw_response.status_code = status_code - raw_response.url = "https://api.airbyte.com/v1/connectors/connector-id" - return SDKError( - f"API error occurred: Status {status_code}", - status_code, - '{"message":"Caller does not have the required permissions"}', - raw_response, +def _api_status_error(status_code: int) -> AirbyteCloudError: + """Wrap a failed Public API response for a non-2xx status.""" + return api_util._wrap_api_error( + _http_error_response( + status_code, + url="https://api.airbyte.com/v1/connectors/connector-id", + text='{"message":"Caller does not have the required permissions"}', + ) ) @@ -1727,23 +1595,28 @@ def _sdk_status_error(status_code: int) -> SDKError: pytest.param("patch-connection", 400, id="patch-connection"), ], ) -def test_api_util_calls_wrap_sdk_errors_with_status_context( +def test_api_util_calls_wrap_api_errors_with_status_context( monkeypatch: pytest.MonkeyPatch, operation: str, status_code: int, ) -> None: - """SDK errors from job and connection calls retain status and request context.""" - sdk_error = _sdk_status_error(status_code) - get_job = Mock(side_effect=sdk_error) - patch_connection = Mock(side_effect=sdk_error) - airbyte_instance = SimpleNamespace( - jobs=SimpleNamespace(get_job=get_job), - connections=SimpleNamespace(patch_connection=patch_connection), - ) + """Errors from job and connection calls retain status and request context.""" + api_error = _api_status_error(status_code) + + def fake_request(**kwargs: object) -> object: + context = dict(api_error.context or {}) + context.update(kwargs.get("error_context") or {}) + error_cls = ( + api_util.AirbyteMissingResourceError + if status_code == 403 + else api_util.AirbyteCloudError + ) + raise error_cls(message=api_error.get_message(), context=context) + monkeypatch.setattr( api_util, - "get_airbyte_server_instance", - lambda **_: airbyte_instance, + "_make_public_api_request", + fake_request, ) with pytest.raises(AirbyteCloudError) as exc_info: @@ -1766,17 +1639,12 @@ def test_api_util_calls_wrap_sdk_errors_with_status_context( assert exc_info.value.context is not None assert exc_info.value.context["status_code"] == status_code - assert exc_info.value.__cause__ is sdk_error if operation == "get-job": assert isinstance(exc_info.value, AirbyteMissingResourceError) assert exc_info.value.context["job_id"] == 42 - get_job.assert_called_once_with(api.GetJobRequest(job_id=42)) - patch_connection.assert_not_called() else: assert not isinstance(exc_info.value, AirbyteMissingResourceError) assert exc_info.value.context["connection_id"] == "connection-1" - patch_connection.assert_called_once() - get_job.assert_not_called() @pytest.mark.parametrize( @@ -1792,42 +1660,31 @@ def test_get_connector_falls_back_to_destination( source_status: int, ) -> None: """A bare destination ID must still resolve when the source lookup 404s or 403s.""" - raw_response = requests.Response() - raw_response.status_code = 200 - raw_response.url = "https://api.airbyte.com/v1/destinations/dest-id" - raw_response._content = ( - b'{"destinationId":"dest-id","name":"dest",' - b'"destinationType":"duckdb","workspaceId":"ws"}' - ) - destination_response = models.DestinationResponse( - configuration=models.DestinationDuckdb(destination_path="/tmp/test.duckdb"), - created_at=0, - definition_id="definition-id", - destination_id="dest-id", - destination_type="duckdb", - name="dest", - workspace_id="ws", - ) - airbyte_instance = SimpleNamespace( - sources=SimpleNamespace( - get_source=Mock(side_effect=_sdk_status_error(source_status)), - ), - destinations=SimpleNamespace( - get_destination=Mock( - return_value=api.GetDestinationResponse( - content_type="application/json", - status_code=200, - raw_response=raw_response, - destination_response=destination_response, + destination_wire = { + "destinationId": "dest-id", + "name": "dest", + "destinationType": "duckdb", + "definitionId": "definition-id", + "workspaceId": "ws", + "configuration": {"destination_path": "/tmp/test.duckdb"}, + "createdAt": 0, + } + calls: list[dict[str, object]] = [] + + def fake_request(**kwargs: object) -> object: + calls.append(kwargs) + if kwargs["path"].startswith("/sources/"): + raise api_util._wrap_api_error( + _http_error_response( + source_status, + url=f"https://api.airbyte.com/v1{kwargs['path']}", + text="boom", ), - ), - ), - ) - monkeypatch.setattr( - api_util, - "get_airbyte_server_instance", - lambda **_: airbyte_instance, - ) + kwargs.get("error_context"), + ) + return destination_wire + + monkeypatch.setattr(api_util, "_make_public_api_request", fake_request) connector_type, response = api_util.get_connector( "dest-id", @@ -1838,27 +1695,17 @@ def test_get_connector_falls_back_to_destination( ) assert connector_type is ConnectorType.DESTINATION - assert response is destination_response + assert response.destination_id == "dest-id" def test_get_connector_raises_missing_resource_when_neither_found( monkeypatch: pytest.MonkeyPatch, ) -> None: - airbyte_instance = SimpleNamespace( - sources=SimpleNamespace( - get_source=Mock(side_effect=_sdk_404_error("SOURCE_CONNECTION")), - ), - destinations=SimpleNamespace( - get_destination=Mock( - side_effect=_sdk_404_error("DESTINATION_CONNECTION"), - ), - ), - ) - monkeypatch.setattr( - api_util, - "get_airbyte_server_instance", - lambda **_: airbyte_instance, - ) + def fake_request(**kwargs: object) -> object: + _ = kwargs + raise _api_404_error() + + monkeypatch.setattr(api_util, "_make_public_api_request", fake_request) with pytest.raises(AirbyteMissingResourceError): api_util.get_connector( @@ -1897,18 +1744,24 @@ def test_get_connector_error_fallback( expected_status: int, expect_destination_call: bool, ) -> None: - get_destination = Mock( - side_effect=_sdk_status_error(destination_status or 200), - ) - airbyte_instance = SimpleNamespace( - sources=SimpleNamespace( - get_source=Mock(side_effect=_sdk_status_error(source_status)), - ), - destinations=SimpleNamespace(get_destination=get_destination), - ) - monkeypatch.setattr( - api_util, "get_airbyte_server_instance", lambda **_: airbyte_instance - ) + calls: list[str] = [] + + def fake_request(**kwargs: object) -> object: + path = kwargs["path"] + calls.append(path) + status = ( + source_status + if path.startswith("/sources/") + else (destination_status or 200) + ) + raise api_util._wrap_api_error( + _http_error_response( + status, url=f"https://api.airbyte.com/v1{path}", text="boom" + ), + kwargs.get("error_context"), + ) + + monkeypatch.setattr(api_util, "_make_public_api_request", fake_request) with pytest.raises(AirbyteCloudError) as exc_info: api_util.get_connector( @@ -1922,31 +1775,25 @@ def test_get_connector_error_fallback( context = exc_info.value.context or {} assert context["status_code"] == expected_status assert ("source_id" in context) is (expected_status == source_status) - assert get_destination.called is expect_destination_call + assert ( + any(path.startswith("/destinations/") for path in calls) + is expect_destination_call + ) -def test_get_connector_does_not_fall_back_on_non_raised_5xx_response( +def test_get_connector_does_not_fall_back_on_5xx_response( monkeypatch: pytest.MonkeyPatch, ) -> None: - """A 5xx the SDK returns instead of raising is not read as "maybe a destination".""" - get_destination = Mock() - airbyte_instance = SimpleNamespace( - sources=SimpleNamespace( - get_source=Mock( - return_value=SimpleNamespace( - status_code=500, - source_response=None, - raw_response=SimpleNamespace(text="boom", url="https://api"), - ), - ), - ), - destinations=SimpleNamespace(get_destination=get_destination), - ) - monkeypatch.setattr( - api_util, "get_airbyte_server_instance", lambda **_: airbyte_instance - ) + """A 5xx source lookup is not read as "maybe a destination".""" + calls: list[str] = [] - with pytest.raises(AirbyteMissingResourceError) as exc_info: + def fake_request(**kwargs: object) -> object: + calls.append(kwargs["path"]) + raise _api_status_error(500) + + monkeypatch.setattr(api_util, "_make_public_api_request", fake_request) + + with pytest.raises(AirbyteCloudError) as exc_info: api_util.get_connector( "connector-id", api_root="https://api.airbyte.com/v1", @@ -1956,25 +1803,16 @@ def test_get_connector_does_not_fall_back_on_non_raised_5xx_response( ) assert (exc_info.value.context or {})["status_code"] == 500 - get_destination.assert_not_called() + assert not any(path.startswith("/destinations/") for path in calls) -def test_get_source_reraises_non_404_sdk_error_as_airbyte_error( +def test_get_source_reraises_non_404_api_error_as_airbyte_error( monkeypatch: pytest.MonkeyPatch, ) -> None: - raw_response = requests.Response() - raw_response.status_code = 500 - raw_response.url = "https://api.airbyte.com/v1/sources/source-id" - error = SDKError( - "API error occurred: Status 500", 500, "response body", raw_response - ) - airbyte_instance = SimpleNamespace( - sources=SimpleNamespace(get_source=Mock(side_effect=error)), - ) monkeypatch.setattr( api_util, - "get_airbyte_server_instance", - lambda **_: airbyte_instance, + "_make_public_api_request", + Mock(side_effect=_api_status_error(500)), ) with pytest.raises(AirbyteCloudError) as exc_info: @@ -1987,3 +1825,203 @@ def test_get_source_reraises_non_404_sdk_error_as_airbyte_error( ) assert type(exc_info.value) is AirbyteCloudError + + +_FIELD_FILTERING_MAPPER = { + "type": "field-filtering", + "mapperConfiguration": {"targetField": "foo"}, +} + + +def _connection_wire_with_mapper(mapper: dict) -> dict[str, object]: + """Create a connection wire dict carrying a single stream mapper.""" + return { + "connectionId": "connection-1", + "name": "Connection", + "sourceId": "source-1", + "destinationId": "destination-1", + "workspaceId": "workspace-id", + "status": "active", + "schedule": {"scheduleType": "manual"}, + "nonBreakingSchemaUpdatesBehavior": "ignore", + "configurations": {"streams": [{"name": "s", "mappers": [mapper]}]}, + "createdAt": 0, + "tags": [], + } + + +@pytest.mark.parametrize( + "mapper", + [ + pytest.param(_FIELD_FILTERING_MAPPER, id="field-filtering"), + pytest.param( + {"type": "future-mapper", "mapperConfiguration": {"someKey": 1}}, + id="unknown-mapper-type", + ), + ], +) +def test_get_connection_decodes_stream_mappers( + monkeypatch: pytest.MonkeyPatch, + mapper: dict, +) -> None: + """Connections with field-filtering or unknown mappers decode without failing.""" + _public_api_spy(monkeypatch, [_connection_wire_with_mapper(mapper)]) + + connection = api_util.get_connection( + workspace_id="workspace-id", + connection_id="connection-1", + api_root="https://api.airbyte.com/v1", + client_id=None, + client_secret=None, + bearer_token=SecretString("token"), + ) + + stream_mapper = connection.configurations.streams[0].mappers[0] + assert stream_mapper.type == mapper["type"] + if mapper["type"] == "field-filtering": + assert stream_mapper.mapper_configuration.target_field == "foo" + else: + assert stream_mapper.mapper_configuration == {"someKey": 1} + + info = CloudConnectionInfo.from_api_response(connection) + assert info.connection_id == "connection-1" + + +@pytest.mark.parametrize( + ("status", "expected"), + [ + pytest.param("queued", JobStatusEnum.QUEUED, id="queued"), + pytest.param("some_future_status", "some_future_status", id="unknown-status"), + ], +) +def test_get_job_info_decodes_queued_and_unknown_statuses( + monkeypatch: pytest.MonkeyPatch, + status: str, + expected: object, +) -> None: + """Queued and future job statuses decode without raising.""" + _public_api_spy( + monkeypatch, + [ + { + "jobId": 42, + "status": status, + "jobType": "sync", + "startTime": "2026-01-01T00:00:00Z", + "connectionId": "connection-id", + } + ], + ) + + job = api_util.get_job_info( + job_id=42, + api_root="https://api.airbyte.com/v1", + client_id=None, + client_secret=None, + bearer_token=SecretString("token"), + ) + + assert job.status == expected + job_info = CloudJobInfo.from_api_response(job) + assert job_info.status == expected + + +def test_list_connections_skips_invalid_items_and_logs( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A malformed item in a paged list is skipped with a warning. + + Pagination must advance by the RAW item count, not the decoded count: + a skipped item must not shift the next page's offset (which would yield + duplicates), and a page of only invalid items must not stop pagination + while `next` is still set. + """ + invalid_item = {"connectionId": "connection-bad"} + invalid_only = {"connectionId": "connection-bad-2"} + captured = _public_api_spy( + monkeypatch, + [ + # Page 1: one invalid item among three raw items. + _page_body( + [ + _connection_wire("good", 1), + invalid_item, + _connection_wire("also-good", 2), + ], + next_page="next", + ), + # Page 2: every item fails to decode, but `next` is still set. + _page_body([invalid_only], next_page="next"), + _page_body([_connection_wire("third", 3)]), + ], + ) + + warnings: list[str] = [] + + def record_warning(message: str, *args: object, **kwargs: object) -> None: + warnings.append(message % args if args else message) + + monkeypatch.setattr(api_util.logger, "warning", record_warning) + + result = api_util.list_connections( + workspace_id="workspace-id", + api_root="https://api.airbyte.com/v1", + client_id=None, + client_secret=None, + bearer_token=SecretString("token"), + ) + + # No duplicates: offset advanced by the raw page lengths (3, then 1). + assert [connection.name for connection in result] == ["good", "also-good", "third"] + assert [call["params"]["offset"] for call in captured] == [0, 3, 4] + assert any("failed schema validation" in message for message in warnings) + + +def test_cloud_connection_info_configurations_keep_snake_case_keys( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """`CloudConnectionInfo.configurations` serializes with snake_case keys (MCP output).""" + _public_api_spy( + monkeypatch, + [ + { + "connectionId": "connection-1", + "name": "Connection", + "sourceId": "source-1", + "destinationId": "destination-1", + "workspaceId": "workspace-id", + "status": "active", + "schedule": {"scheduleType": "manual"}, + "nonBreakingSchemaUpdatesBehavior": "ignore", + "configurations": { + "streams": [ + { + "name": "users", + "syncMode": "incremental_append", + "cursorField": ["updated_at"], + "primaryKey": [["id"]], + } + ] + }, + "createdAt": 0, + "tags": [], + } + ], + ) + + connection = api_util.get_connection( + workspace_id="workspace-id", + connection_id="connection-1", + api_root="https://api.airbyte.com/v1", + client_id=None, + client_secret=None, + bearer_token=SecretString("token"), + ) + + info = CloudConnectionInfo.from_api_response(connection) + dumped = info.model_dump(mode="json") + stream = dumped["configurations"]["streams"][0] + assert stream["sync_mode"] == "incremental_append" + assert stream["cursor_field"] == ["updated_at"] + assert stream["primary_key"] == [["id"]] + assert "syncMode" not in stream diff --git a/tests/unit_tests/test_cloud_client.py b/tests/unit_tests/test_cloud_client.py index 485ca0f95..2f7922890 100644 --- a/tests/unit_tests/test_cloud_client.py +++ b/tests/unit_tests/test_cloud_client.py @@ -8,7 +8,7 @@ from airbyte import exceptions as exc from airbyte.cloud.client import CloudClient from airbyte.cloud.models import WorkspacePrivilegeScope -from airbyte_api import models +from airbyte_server_models.public_api import models def _api_patches( @@ -123,7 +123,7 @@ def test_resolve_default_workspace_id_uses_exactly_one_direct_grant( patches[4], patch( "airbyte._util.api_util.get_workspace", - return_value=models.WorkspaceResponse( + return_value=models.WorkspaceResponse.model_construct( data_residency="auto", name="Workspace", notifications=models.NotificationsConfig(), @@ -153,7 +153,7 @@ def get_workspace( **kwargs: object, ) -> models.WorkspaceResponse: if workspace_id == "live-workspace": - return models.WorkspaceResponse( + return models.WorkspaceResponse.model_construct( data_residency="auto", name="Live Workspace", notifications=models.NotificationsConfig(), @@ -215,7 +215,7 @@ def test_direct_workspace_validation_is_capped() -> None: patch( "airbyte._util.api_util.get_workspace", side_effect=[ - models.WorkspaceResponse( + models.WorkspaceResponse.model_construct( data_residency="auto", name=f"Workspace {index}", notifications=models.NotificationsConfig(), @@ -252,6 +252,22 @@ def test_direct_workspace_validation_is_capped() -> None: assert get_workspace_organization_info.call_count == 25 +def test_workspace_parent_organization_id_missing_returns_none() -> None: + """An organization record without `organizationId` resolves to None, not 'None'.""" + patches = _api_patches(user={"userId": "user-id"}) + with ( + patches[0], + patches[1], + patches[2], + patches[3] as get_workspace_organization_info, + patches[4], + ): + get_workspace_organization_info.return_value = {"organizationName": "Org"} + client = CloudClient(bearer_token="token") + + assert client.get_workspace_parent_organization_id("ws-1") is None + + def test_workspace_organization_failure_is_cached() -> None: patches = _api_patches(user={"userId": "user-id"}) with ( @@ -283,7 +299,7 @@ def test_default_context_resolves_workspace_when_organization_lookup_fails() -> patches[4], patch( "airbyte._util.api_util.get_workspace", - return_value=models.WorkspaceResponse( + return_value=models.WorkspaceResponse.model_construct( data_residency="auto", name="Workspace 1", notifications=models.NotificationsConfig(), @@ -322,7 +338,7 @@ def test_default_context_enriches_configured_workspace() -> None: patches[4], patch( "airbyte._util.api_util.get_workspace", - return_value=models.WorkspaceResponse( + return_value=models.WorkspaceResponse.model_construct( data_residency="auto", name="Configured", notifications=models.NotificationsConfig(), @@ -464,7 +480,7 @@ def test_list_workspaces_skips_stale_grant_before_valid_grant_with_limit() -> No resource_type="workspace", resource_name_or_id="stale-workspace", ) - valid_workspace = models.WorkspaceResponse( + valid_workspace = models.WorkspaceResponse.model_construct( data_residency="auto", name="Valid workspace", notifications=models.NotificationsConfig(), @@ -528,7 +544,7 @@ def test_list_workspaces_defaults_to_direct_memberships() -> None: patches[4], patch( "airbyte._util.api_util.get_workspace", - return_value=models.WorkspaceResponse( + return_value=models.WorkspaceResponse.model_construct( data_residency="auto", name="Workspace 1", notifications=models.NotificationsConfig(), @@ -953,13 +969,13 @@ def test_list_workspaces_uses_direct_grants_when_memberships_are_ambiguous() -> patch( "airbyte._util.api_util.get_workspace", side_effect=[ - models.WorkspaceResponse( + models.WorkspaceResponse.model_construct( data_residency="auto", name="Workspace 1", notifications=models.NotificationsConfig(), workspace_id="workspace-1", ), - models.WorkspaceResponse( + models.WorkspaceResponse.model_construct( data_residency="auto", name="Workspace 2", notifications=models.NotificationsConfig(), @@ -1012,7 +1028,7 @@ def test_get_default_context_for_user_is_bounded_to_permission_derived_scope() - patch( "airbyte._util.api_util.get_workspace", side_effect=[ - models.WorkspaceResponse( + models.WorkspaceResponse.model_construct( data_residency="auto", name=f"Workspace {index}", notifications=models.NotificationsConfig(), @@ -1085,7 +1101,7 @@ def test_get_default_context_for_user_truncates_workspace_memberships() -> None: patch( "airbyte._util.api_util.get_workspace", side_effect=[ - models.WorkspaceResponse( + models.WorkspaceResponse.model_construct( data_residency="auto", name=f"Workspace {index}", notifications=models.NotificationsConfig(), @@ -1139,7 +1155,7 @@ def test_get_default_context_for_user_hints_setter_when_stored_default_missing() patches[4], patch( "airbyte._util.api_util.get_workspace", - return_value=models.WorkspaceResponse( + return_value=models.WorkspaceResponse.model_construct( data_residency="auto", name="Workspace 1", notifications=models.NotificationsConfig(), @@ -1173,7 +1189,7 @@ def test_get_default_context_for_user_omits_setter_hint_when_default_stored() -> patches[4], patch( "airbyte._util.api_util.get_workspace", - return_value=models.WorkspaceResponse( + return_value=models.WorkspaceResponse.model_construct( data_residency="auto", name="Workspace 1", notifications=models.NotificationsConfig(), diff --git a/tests/unit_tests/test_cloud_connections.py b/tests/unit_tests/test_cloud_connections.py index b27ec66cb..40506a846 100644 --- a/tests/unit_tests/test_cloud_connections.py +++ b/tests/unit_tests/test_cloud_connections.py @@ -24,8 +24,7 @@ AirbyteMissingResourceError, AirbyteLibInputError, ) -from airbyte_api import models -from airbyte_api.errors import SDKError +from airbyte_server_models.public_api import models def _job_response( @@ -35,7 +34,7 @@ def _job_response( connection_id: str = "connection-id", ) -> models.JobResponse: """Create a minimal job response.""" - return models.JobResponse( + return models.JobResponse.model_construct( connection_id=connection_id, job_id=job_id, job_type=models.JobTypeEnum.SYNC, @@ -570,7 +569,7 @@ def test_set_schedule_accepts_quartz_cron( ) -> None: """Verify Quartz cron expressions are passed through to the API.""" connection = _connection() - captured: list[models.AirbyteAPIConnectionSchedule] = [] + captured: list[models.AirbyteApiConnectionSchedule] = [] def patch_connection( *, @@ -579,11 +578,11 @@ def patch_connection( client_id: object, client_secret: object, bearer_token: object, - schedule: models.AirbyteAPIConnectionSchedule, + schedule: models.AirbyteApiConnectionSchedule, ) -> models.ConnectionResponse: _ = (connection_id, api_root, client_id, client_secret, bearer_token) captured.append(schedule) - return models.ConnectionResponse( + return models.ConnectionResponse.model_construct( connection_id="connection-id", created_at=0, destination_id="destination-id", @@ -715,21 +714,19 @@ def run_connection(**kwargs: object) -> None: fetch_mock.assert_not_called() -def test_run_connection_wraps_sdk_error( +def test_run_connection_wraps_api_error( monkeypatch: pytest.MonkeyPatch, ) -> None: - """Verify `run_connection` wraps a create_job SDKError in AirbyteConnectionSyncError.""" - airbyte_instance = MagicMock() - airbyte_instance.jobs.create_job.side_effect = SDKError( - message="Status 409", - status_code=409, - body="...", - raw_response=MagicMock(), - ) + """Verify `run_connection` wraps a create_job API error in AirbyteConnectionSyncError.""" monkeypatch.setattr( api_util, - "get_airbyte_server_instance", - MagicMock(return_value=airbyte_instance), + "_make_public_api_request", + MagicMock( + side_effect=AirbyteCloudError( + message="API error occurred: Status 409", + context={"status_code": 409}, + ) + ), ) with pytest.raises(AirbyteConnectionSyncError) as exc_info: diff --git a/tests/unit_tests/test_cloud_credentials.py b/tests/unit_tests/test_cloud_credentials.py index d949039f4..a63ff6e20 100644 --- a/tests/unit_tests/test_cloud_credentials.py +++ b/tests/unit_tests/test_cloud_credentials.py @@ -9,7 +9,7 @@ import pytest import requests -from airbyte_api import models +from airbyte_server_models.public_api import models from airbyte import constants from airbyte._util import api_util @@ -45,6 +45,16 @@ ) from airbyte.mcp import cloud as mcp_cloud from airbyte.secrets.base import SecretString +from airbyte_server_models._config_api import ( + Notification, + NotificationType, + OrganizationInfoRead, + OrganizationRead, + PermissionRead, + PermissionType, + UserRead, + WorkspaceRead, +) def _raise(error: Exception) -> Callable[..., NoReturn]: @@ -54,6 +64,70 @@ def _raiser(*_args: object, **_kwargs: object) -> NoReturn: return _raiser +def _user_response(user_id: str = "user-id") -> UserRead: + return UserRead.model_construct( + userId=user_id, + email="user@example.com", + metadata={}, + ) + + +def _organization_info_response( + *, + organization_id: str = "organization-id", + organization_name: str = "Organization", +) -> OrganizationInfoRead: + return OrganizationInfoRead.model_construct( + organizationId=organization_id, + organizationName=organization_name, + sso=False, + scim=False, + ) + + +def _permission_response( + *, + organization_id: str | None = "organization-id", +) -> PermissionRead: + return PermissionRead.model_construct( + permissionId="permission-id", + permissionType=PermissionType.organization_member, + userId="user-id", + organizationId=organization_id, + ) + + +def _organization_response( + *, + organization_id: str = "organization-id", + organization_name: str = "Organization", + email: str = "test@example.com", +) -> OrganizationRead: + return OrganizationRead.model_construct( + organizationId=organization_id, + organizationName=organization_name, + email=email, + ) + + +def _workspace_config_response( + *, + workspace_id: str, + name: str, + organization_id: str = "organization-id", + notifications: list[Notification] | None = None, +) -> WorkspaceRead: + return WorkspaceRead.model_construct( + workspaceId=workspace_id, + customerId="customer-id", + name=name, + slug=name.casefold().replace(" ", "-"), + initialSetupComplete=False, + organizationId=organization_id, + notifications=notifications, + ) + + def _stub_organization_features( monkeypatch: pytest.MonkeyPatch, *, enabled: bool = False ) -> None: @@ -78,20 +152,21 @@ def _patch_workspace_discovery( lambda _: "auth-user-id", ) monkeypatch.setattr( - api_util, - "get_user_by_auth_id", - lambda *_, **__: {"userId": "user-id"}, + api_util, "get_user_by_auth_id", lambda *_, **__: _user_response() ) monkeypatch.setattr( api_util, "list_permissions_for_user", - lambda *_, **__: permissions or [], + lambda *_, **__: [ + {"permissionId": "permission-id", "userId": "user-id", **permission} + for permission in (permissions or []) + ], ) - def fake_get_workspace_organization_info(**_: object) -> dict[str, object]: + def fake_get_workspace_organization_info(**_: object) -> OrganizationInfoRead: if parent_organization_id is None: pytest.fail("workspace parent lookup should not be called") - return {"organizationId": parent_organization_id} + return _organization_info_response(organization_id=parent_organization_id) monkeypatch.setattr( api_util, @@ -106,9 +181,15 @@ def fake_get_workspace_organization_info(**_: object) -> dict[str, object]: def fake_list_workspaces_in_organization( **kwargs: object, - ) -> list[dict[str, object]]: + ) -> list[WorkspaceRead]: captured.update(kwargs) - return org_scoped_result if org_scoped_result is not None else [] + return [ + _workspace_config_response( + workspace_id=str(workspace["workspaceId"]), + name=str(workspace["name"]), + ) + for workspace in (org_scoped_result or []) + ] monkeypatch.setattr( api_util, @@ -395,13 +476,17 @@ def fake_list_workspaces_in_organization( *, limit: int | None = None, **_: object, - ) -> list[dict[str, object]]: + ) -> list[WorkspaceRead]: nonlocal captured_limit captured_limit = limit return [ - {"name": "miss", "workspaceId": "workspace-miss"}, - {"name": "target-one", "workspaceId": "workspace-target-one"}, - {"name": "target-two", "workspaceId": "workspace-target-two"}, + _workspace_config_response(workspace_id="workspace-miss", name="miss"), + _workspace_config_response( + workspace_id="workspace-target-one", name="target-one" + ), + _workspace_config_response( + workspace_id="workspace-target-two", name="target-two" + ), ] monkeypatch.setattr( @@ -432,11 +517,13 @@ def test_cloud_client_list_workspaces_matches_exact_name_after_server_filter( def fake_list_workspaces_in_organization( **kwargs: object, - ) -> list[dict[str, object]]: + ) -> list[WorkspaceRead]: captured.update(kwargs) return [ - {"name": "Production-old", "workspaceId": "workspace-production-old"}, - {"name": "Prod", "workspaceId": "workspace-prod"}, + _workspace_config_response( + workspace_id="workspace-production-old", name="Production-old" + ), + _workspace_config_response(workspace_id="workspace-prod", name="Prod"), ] monkeypatch.setattr( @@ -471,7 +558,7 @@ def fake_create_workspace( ) -> models.WorkspaceResponse: nonlocal captured_organization_id captured_organization_id = organization_id - return models.WorkspaceResponse( + return models.WorkspaceResponse.model_construct( data_residency="auto", name="New workspace", notifications=models.NotificationsConfig(), @@ -495,13 +582,19 @@ def test_cloud_client_list_workspaces_accepts_api_notification_list( ) -> None: def fake_list_workspaces_in_organization( **_: object, - ) -> list[dict[str, object]]: + ) -> list[WorkspaceRead]: return [ - { - "workspaceId": "workspace-id", - "name": "Workspace", - "notifications": [{"sendOnSuccess": True}], - } + _workspace_config_response( + workspace_id="workspace-id", + name="Workspace", + notifications=[ + Notification.model_construct( + notificationType=NotificationType.slack, + sendOnSuccess=True, + sendOnFailure=False, + ) + ], + ) ] monkeypatch.setattr( @@ -516,7 +609,7 @@ def fake_list_workspaces_in_organization( ).list_workspaces(organization_id="organization-id") assert len(workspaces) == 1 - assert workspaces[0].notifications == [{"sendOnSuccess": True}] + assert workspaces[0].notifications[0]["sendOnSuccess"] is True def test_cloud_workspace_info_accepts_api_notification_mapping() -> None: @@ -536,7 +629,7 @@ def test_cloud_client_rename_workspace_forwards_inputs( def fake_rename_workspace(**kwargs: object) -> models.WorkspaceResponse: captured_kwargs.update(kwargs) - return models.WorkspaceResponse( + return models.WorkspaceResponse.model_construct( data_residency="auto", name="Renamed workspace", notifications=models.NotificationsConfig(), @@ -808,7 +901,7 @@ def test_cloud_workspace_rename_forwards_inputs( def fake_rename_workspace(**kwargs: object) -> models.WorkspaceResponse: captured_kwargs.update(kwargs) - return models.WorkspaceResponse( + return models.WorkspaceResponse.model_construct( data_residency="auto", name="Renamed workspace", notifications=models.NotificationsConfig(), @@ -907,7 +1000,7 @@ def test_cloud_client_get_organization_uses_default_organization_id( api_util, "list_organizations_for_user", lambda **_: [ - models.OrganizationResponse( + models.OrganizationResponse.model_construct( organization_id="default-org", organization_name="Default Org", email="test@example.com", @@ -917,11 +1010,10 @@ def test_cloud_client_get_organization_uses_default_organization_id( monkeypatch.setattr( api_util, "get_organization_info", - lambda **_: { - "organizationId": "default-org", - "organizationName": "Default Org", - "email": "test@example.com", - }, + lambda **_: _organization_info_response( + organization_id="default-org", + organization_name="Default Org", + ), ) organization = CloudClient( @@ -969,27 +1061,29 @@ def test_cloud_client_get_organization_resolves_default_context( api_util, "get_user_id_from_bearer_token", lambda _: "auth-user-id" ) monkeypatch.setattr( - api_util, "get_user_by_auth_id", lambda *_, **__: {"userId": "user-id"} + api_util, "get_user_by_auth_id", lambda *_, **__: _user_response() ) monkeypatch.setattr( api_util, "list_permissions_for_user", - lambda *_, **__: permissions or [], + lambda *_, **__: [ + {"permissionId": "permission-id", "userId": "user-id", **permission} + for permission in (permissions or []) + ], ) monkeypatch.setattr( api_util, "get_workspace_organization_info", - lambda **_: {"organizationId": parent_organization_id} + lambda **_: _organization_info_response(organization_id=parent_organization_id) if parent_organization_id is not None else pytest.fail("workspace parent lookup should not be called"), ) monkeypatch.setattr( api_util, "get_organization_info", - lambda **kwargs: { - "organizationId": kwargs["organization_id"], - "organizationName": "Organization", - }, + lambda **kwargs: _organization_info_response( + organization_id=str(kwargs["organization_id"]), + ), ) organization = CloudClient(bearer_token="token", **client_kwargs).get_organization() @@ -1009,9 +1103,9 @@ def test_cloud_client_get_organization_rejects_ambiguous_default_context( monkeypatch.setattr( api_util, "get_organization_info", - lambda **kwargs: { - "organizationName": f"Organization {kwargs['organization_id']}" - }, + lambda **kwargs: _organization_info_response( + organization_name=f"Organization {kwargs['organization_id']}", + ), ) with pytest.raises(AirbyteLibInputError) as exc_info: @@ -1078,13 +1172,9 @@ def test_cloud_client_get_organization_uses_single_config_lookup( ) -> None: calls: list[dict[str, object]] = [] - def fake_get_organization_info(**kwargs: object) -> dict[str, object]: + def fake_get_organization_info(**kwargs: object) -> OrganizationInfoRead: calls.append(kwargs) - return { - "organizationId": "organization-id", - "organizationName": "Organization", - "email": "test@example.com", - } + return _organization_info_response() monkeypatch.setattr(api_util, "get_organization_info", fake_get_organization_info) monkeypatch.setattr( @@ -1107,7 +1197,7 @@ def fake_get_organization_info(**kwargs: object) -> dict[str, object]: pytest.param([], id="empty"), pytest.param( [ - models.OrganizationResponse( + models.OrganizationResponse.model_construct( organization_id="organization-id", organization_name="Organization", email="test@example.com", @@ -1117,12 +1207,12 @@ def fake_get_organization_info(**kwargs: object) -> dict[str, object]: ), pytest.param( [ - models.OrganizationResponse( + models.OrganizationResponse.model_construct( organization_id="organization-id-1", organization_name="Organization 1", email="one@example.com", ), - models.OrganizationResponse( + models.OrganizationResponse.model_construct( organization_id="organization-id-2", organization_name="Organization 2", email="two@example.com", @@ -1164,17 +1254,17 @@ def test_cloud_client_list_organizations_uses_config_api_for_filter_or_limit( ) -> None: captured: dict[str, object] = {} organizations = [ - models.OrganizationResponse( + models.OrganizationResponse.model_construct( organization_id="organization-id-1", organization_name="Development", email="one@example.com", ), - models.OrganizationResponse( + models.OrganizationResponse.model_construct( organization_id="organization-id-2", organization_name="development-copy", email="two@example.com", ), - models.OrganizationResponse( + models.OrganizationResponse.model_construct( organization_id="organization-id-3", organization_name="Production", email="three@example.com", @@ -1183,7 +1273,7 @@ def test_cloud_client_list_organizations_uses_config_api_for_filter_or_limit( def fake_list_organizations_for_user_id( **kwargs: object, - ) -> list[dict[str, object]]: + ) -> list[OrganizationRead]: captured.update(kwargs) filtered = organizations if kwargs["name_contains"] is not None: @@ -1195,11 +1285,11 @@ def fake_list_organizations_for_user_id( ] limit = kwargs["limit"] return [ - { - "organizationId": organization.organization_id, - "organizationName": organization.organization_name, - "email": organization.email, - } + _organization_response( + organization_id=organization.organization_id, + organization_name=organization.organization_name, + email=organization.email, + ) for organization in filtered[: limit if isinstance(limit, int) else None] ] @@ -1212,7 +1302,7 @@ def fake_list_organizations_for_user_id( api_util, "get_user_id_from_bearer_token", lambda _: "auth-user-id" ) monkeypatch.setattr( - api_util, "get_user_by_auth_id", lambda *_, **__: {"userId": "user-id"} + api_util, "get_user_by_auth_id", lambda *_, **__: _user_response() ) result = CloudClient(bearer_token="token").list_organizations( @@ -1229,7 +1319,7 @@ def test_cloud_client_list_organizations_falls_back_to_public_listing( monkeypatch: pytest.MonkeyPatch, ) -> None: organizations = [ - models.OrganizationResponse( + models.OrganizationResponse.model_construct( organization_id="organization-id", organization_name="Development", email="test@example.com", @@ -1243,7 +1333,7 @@ def test_cloud_client_list_organizations_falls_back_to_public_listing( monkeypatch.setattr( api_util, "get_user_by_auth_id", - lambda *_, **__: {"userId": "user-id"}, + lambda *_, **__: _user_response(), ) monkeypatch.setattr( api_util, @@ -1280,10 +1370,10 @@ def fake_get_bearer_token(**_: object) -> SecretString: def fake_list_organizations_for_user_id( **kwargs: object, - ) -> list[dict[str, object]]: + ) -> list[OrganizationRead]: nonlocal captured_bearer_token captured_bearer_token = kwargs["bearer_token"] - return [{"organizationId": "organization-id"}] + return [_organization_response(organization_id="organization-id")] monkeypatch.setattr(api_util, "get_bearer_token", fake_get_bearer_token) monkeypatch.setattr( @@ -1294,7 +1384,7 @@ def fake_list_organizations_for_user_id( monkeypatch.setattr( api_util, "get_user_by_auth_id", - lambda *_, **__: {"userId": "user-id"}, + lambda *_, **__: _user_response(), ) monkeypatch.setattr( api_util, @@ -1326,12 +1416,9 @@ def fake_get_bearer_token(**_: object) -> SecretString: token_calls += 1 return issued_token - def fake_get_organization_info(**kwargs: object) -> dict[str, object]: + def fake_get_organization_info(**kwargs: object) -> OrganizationInfoRead: captured_bearer_tokens.append(kwargs["bearer_token"]) - return { - "organizationId": "organization-id", - "organizationName": "Organization", - } + return _organization_info_response() monkeypatch.setattr(api_util, "get_bearer_token", fake_get_bearer_token) monkeypatch.setattr(api_util, "get_organization_info", fake_get_organization_info) @@ -1360,7 +1447,7 @@ def test_cloud_client_list_organizations_has_no_default_limit( monkeypatch: pytest.MonkeyPatch, ) -> None: organizations = [ - models.OrganizationResponse( + models.OrganizationResponse.model_construct( organization_id=f"organization-id-{index}", organization_name=f"Organization {index}", email=f"test-{index}@example.com", @@ -1397,7 +1484,7 @@ def test_cloud_client_list_organizations_reports_ambiguity_candidates( monkeypatch.setattr( api_util, "get_user_by_auth_id", - lambda *_, **__: {"userId": "user-id"}, + lambda *_, **__: _user_response(), ) monkeypatch.setattr( api_util, @@ -1570,15 +1657,18 @@ def test_cloud_client_list_workspaces_resolves_single_membership_and_caches_it( captured = _patch_workspace_discovery(monkeypatch, permissions=permissions) calls = {"user": 0, "permissions": 0} - def fake_get_user_by_auth_id(*_: object, **__: object) -> dict[str, object]: + def fake_get_user_by_auth_id(*_: object, **__: object) -> UserRead: calls["user"] += 1 - return {"userId": "user-id"} + return _user_response() def fake_list_permissions_for_user( *_: object, **__: object - ) -> list[dict[str, object]]: + ) -> list[PermissionRead]: calls["permissions"] += 1 - return permissions + return [ + {"permissionId": "permission-id", "userId": "user-id", **permission} + for permission in permissions + ] monkeypatch.setattr(api_util, "get_user_by_auth_id", fake_get_user_by_auth_id) monkeypatch.setattr( @@ -3119,6 +3209,24 @@ def test_cloud_organization_get_billing_status_requires_billing( organization.get_billing_status() +def test_cloud_organization_payment_status_none_when_billing_status_missing( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A missing paymentStatus yields None instead of an AttributeError.""" + monkeypatch.setattr( + api_util, + "get_organization_info", + lambda **_: SimpleNamespace( + billing=SimpleNamespace( + paymentStatus=None, + subscriptionStatus="subscribed", + ) + ), + ) + organization = CloudOrganization(organization_id="organization-id") + assert organization.payment_status is None + + def test_cloud_organization_get_billing_status_wraps_transport_error( monkeypatch: pytest.MonkeyPatch, ) -> None: diff --git a/tests/unit_tests/test_cloud_models.py b/tests/unit_tests/test_cloud_models.py index 9825a9acd..1bfaab60f 100644 --- a/tests/unit_tests/test_cloud_models.py +++ b/tests/unit_tests/test_cloud_models.py @@ -15,10 +15,8 @@ ConnectionSchedule, ConnectionStatus, ) -from airbyte_api.models import ( - DestinationDuckdb, +from airbyte_server_models.public_api.models import ( DestinationResponse, - SourceFaker, SourceResponse, ) @@ -28,7 +26,7 @@ [ pytest.param( SourceResponse( - configuration=SourceFaker(), + configuration={}, created_at=1, definition_id="source-faker-definition", name="Test source", @@ -42,7 +40,7 @@ id="source", ), pytest.param( - SourceResponse( + SourceResponse.model_construct( configuration=None, created_at=1, definition_id="source-empty-definition", @@ -58,7 +56,7 @@ ), pytest.param( DestinationResponse( - configuration=DestinationDuckdb(destination_path="/tmp/test.duckdb"), + configuration={"destination_path": "/tmp/test.duckdb"}, created_at=1, definition_id="destination-duckdb-definition", destination_id="destination-id", @@ -193,7 +191,7 @@ def test_connection_schedule_friendly_description( pytest.param( CloudSource._from_source_response, SourceResponse( - configuration=SourceFaker(), + configuration={}, created_at=1, definition_id="source-faker-definition", name="Test source", @@ -207,7 +205,7 @@ def test_connection_schedule_friendly_description( pytest.param( CloudDestination._from_destination_response, DestinationResponse( - configuration=DestinationDuckdb(destination_path="/tmp/test.duckdb"), + configuration={"destination_path": "/tmp/test.duckdb"}, created_at=1, definition_id="destination-duckdb-definition", destination_id="destination-id", diff --git a/tests/unit_tests/test_deferred_cloud_credentials.py b/tests/unit_tests/test_deferred_cloud_credentials.py index fc38e9c31..eecc4051e 100644 --- a/tests/unit_tests/test_deferred_cloud_credentials.py +++ b/tests/unit_tests/test_deferred_cloud_credentials.py @@ -522,12 +522,30 @@ def test_get_connector_metadata_by_definition_id( ) +_CHECK_CONNECTION_JOB_INFO = { + "id": "00000000-0000-0000-0000-000000000001", + "configType": "check_connection_source", + "status": "succeeded", + "createdAt": 0, + "endedAt": 1, + "succeeded": True, +} + + @responses.activate @pytest.mark.parametrize( ("status_code", "body", "expected"), [ - (200, {"status": "succeeded"}, True), - (200, {"status": "failed", "message": "Invalid credential abc123"}, False), + (200, {"status": "succeeded", "jobInfo": _CHECK_CONNECTION_JOB_INFO}, True), + ( + 200, + { + "status": "failed", + "message": "Invalid credential abc123", + "jobInfo": _CHECK_CONNECTION_JOB_INFO, + }, + False, + ), (422, {"message": "Missing required settings abc123"}, False), (500, {"message": "Unexpected failure abc123"}, None), ], diff --git a/tests/unit_tests/test_mcp_airbyte_user_id.py b/tests/unit_tests/test_mcp_airbyte_user_id.py index e1444d983..430bc6fc5 100644 --- a/tests/unit_tests/test_mcp_airbyte_user_id.py +++ b/tests/unit_tests/test_mcp_airbyte_user_id.py @@ -268,7 +268,11 @@ def test_lookup_request_has_a_finite_http_timeout( """A hung Config API can't pin the lookup's worker thread indefinitely.""" request = MagicMock() request.return_value.status_code = 200 - request.return_value.json.return_value = {"userId": USER_A} + request.return_value.json.return_value = { + "userId": USER_A, + "email": "user-a@example.com", + "metadata": {}, + } monkeypatch.setattr(api_util.requests, "request", request) _call("keycloak-a") diff --git a/tests/unit_tests/test_public_api_imports.py b/tests/unit_tests/test_public_api_imports.py index 89af07d32..945f248b0 100644 --- a/tests/unit_tests/test_public_api_imports.py +++ b/tests/unit_tests/test_public_api_imports.py @@ -11,9 +11,18 @@ REPO_ROOT = Path(__file__).parents[2] PUBLIC_MODULE_RESTRICTED_IMPORTS = { - Path("airbyte/cli"): ("airbyte_api", "airbyte._util.api_imports"), - Path("airbyte/mcp"): ("airbyte_api", "airbyte._util.api_imports"), - Path("airbyte/cloud"): ("airbyte_api", "airbyte._util.api_imports"), + Path("airbyte/cli"): ( + "airbyte_server_models.public_api", + "airbyte._util.api_imports", + ), + Path("airbyte/mcp"): ( + "airbyte_server_models.public_api", + "airbyte._util.api_imports", + ), + Path("airbyte/cloud"): ( + "airbyte_server_models.public_api", + "airbyte._util.api_imports", + ), } PUBLIC_MODULE_RESTRICTED_REFERENCES = { Path("airbyte/cli"): ("api_util.models",), diff --git a/uv.lock b/uv.lock index 250907800..54180d528 100644 --- a/uv.lock +++ b/uv.lock @@ -107,9 +107,9 @@ wheels = [ name = "airbyte" source = { editable = "." } dependencies = [ - { name = "airbyte-api" }, { name = "airbyte-cdk" }, { name = "airbyte-protocol-models-pdv2" }, + { name = "airbyte-server-models" }, { name = "cryptography" }, { name = "cyclopts" }, { name = "duckdb" }, @@ -186,9 +186,9 @@ dev = [ [package.metadata] requires-dist = [ - { name = "airbyte-api", specifier = ">=0.53.0,<1.0" }, { name = "airbyte-cdk", specifier = ">=7.28.3,<8.0.0" }, { name = "airbyte-protocol-models-pdv2", specifier = ">=0.13.0,<1.0" }, + { name = "airbyte-server-models", git = "https://github.com/airbytehq/airbyte-ops-mcp.git?subdirectory=airbyte-server-models&rev=b314c99aaa0ddc0e53ce79f94c1d82d26f94a5f0" }, { name = "cryptography", specifier = ">=44.0.0,!=45.0.0,!=45.0.1,<47.0.0" }, { name = "cyclopts", specifier = ">=4.0,<5.0" }, { name = "ddtrace", marker = "extra == 'datadog'", specifier = ">=4.15.2,<5.0" }, @@ -260,31 +260,6 @@ dev = [ { name = "viztracer", specifier = ">=0.16.3,<1.1.0" }, ] -[[package]] -name = "airbyte-api" -version = "0.53.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "certifi" }, - { name = "charset-normalizer" }, - { name = "dataclasses-json" }, - { name = "idna" }, - { name = "jsonpath-python" }, - { name = "marshmallow" }, - { name = "mypy-extensions" }, - { name = "packaging" }, - { name = "python-dateutil" }, - { name = "requests" }, - { name = "six" }, - { name = "typing-extensions" }, - { name = "typing-inspect" }, - { name = "urllib3" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/ad/23/8debf9c4ca8652c9ceb60b8074e654191ba2053001616151993b39d6cdcb/airbyte-api-0.53.0.tar.gz", hash = "sha256:f054ed170f9a691c3304e93a5212670fd2a38e5debb667c72d5ef8eb89cf7e9d", size = 330090, upload-time = "2025-10-02T19:55:43.378Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/7a/f5/ca44b6f3919f049a4441755958c042fcea09a44e876352a6005f5e5a9e46/airbyte_api-0.53.0-py3-none-any.whl", hash = "sha256:0bd86d5122789a7c97115a4b52ce492b013d0ea8eeab597b6495b0b310332f28", size = 773825, upload-time = "2025-10-02T19:55:42.043Z" }, -] - [[package]] name = "airbyte-cdk" version = "7.28.3" @@ -358,6 +333,14 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/61/24/414fe44348ca7c7ae71030012e77302b3e41d4ee217cc5931dabfbf90496/airbyte_protocol_models_pdv2-0.18.0-py3-none-any.whl", hash = "sha256:62b08659b3da0f03316c60d84736f12dbfd79340a49947155a1687c13314db8b", size = 12805, upload-time = "2025-06-27T21:01:35.637Z" }, ] +[[package]] +name = "airbyte-server-models" +version = "0.0.0.post872.dev0+b314c99" +source = { git = "https://github.com/airbytehq/airbyte-ops-mcp.git?subdirectory=airbyte-server-models&rev=b314c99aaa0ddc0e53ce79f94c1d82d26f94a5f0#b314c99aaa0ddc0e53ce79f94c1d82d26f94a5f0" } +dependencies = [ + { name = "pydantic", extra = ["email"] }, +] + [[package]] name = "annotated-types" version = "0.7.0" @@ -792,19 +775,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/12/a3/2e00fececc34a99ae3a5d5702a5dd29c5371e4ed016647301a2b9bcc1976/cyclopts-4.5.0-py3-none-any.whl", hash = "sha256:305b9aa90a9cd0916f0a450b43e50ad5df9c252680731a0719edfb9b20381bf5", size = 199772, upload-time = "2026-01-16T02:07:14.707Z" }, ] -[[package]] -name = "dataclasses-json" -version = "0.6.7" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "marshmallow" }, - { name = "typing-inspect" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/64/a4/f71d9cf3a5ac257c993b5ca3f93df5f7fb395c725e7f1e6479d2514173c3/dataclasses_json-0.6.7.tar.gz", hash = "sha256:b6b3e528266ea45b9535223bc53ca645f5208833c29229e847b3f26a1cc55fc0", size = 32227, upload-time = "2024-06-09T16:20:19.103Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/c3/be/d0d44e092656fe7a06b55e6103cbce807cdbdee17884a5367c68c9860853/dataclasses_json-0.6.7-py3-none-any.whl", hash = "sha256:0dbf33f26c8d5305befd61b39d2b3414e8a407bedc2834dea9b8d642666fb40a", size = 28686, upload-time = "2024-06-09T16:20:16.715Z" }, -] - [[package]] name = "dateparser" version = "1.2.2" @@ -1886,15 +1856,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/13/6a/71937d4760bf6f3beabc640cf18578683b5247845db5ee4dbd8c66898def/joserfc-1.5.0-py3-none-any.whl", hash = "sha256:eaaded4f4c6717a761baa41b4067307d0c246b9d5e38acd44e80a332f5ddaf24", size = 68388, upload-time = "2025-11-30T06:01:50.43Z" }, ] -[[package]] -name = "jsonpath-python" -version = "1.1.4" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/b8/bf/626a72f2d093c5eb4f4de55b443714afa7231beeae40d4a1c69b5c5aa4d1/jsonpath_python-1.1.4.tar.gz", hash = "sha256:bb3e13854e4807c078a1503ae2d87c211b8bff4d9b40b6455ed583b3b50a7fdd", size = 84766, upload-time = "2025-11-25T12:08:39.521Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/ee/bc/52e5bf0d9839e082b976c19afcab7561d0d719c7627483bf5dc251d27eed/jsonpath_python-1.1.4-py3-none-any.whl", hash = "sha256:8700cb8610c44da6e5e9bff50232779c44bf7dc5bc62662d49319ee746898442", size = 12687, upload-time = "2025-11-25T12:08:38.453Z" }, -] - [[package]] name = "jsonref" version = "1.1.0" @@ -2047,18 +2008,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/e5/f1/216fc1bbfd74011693a4fd837e7026152e89c4bcf3e77b6692fba9923123/markupsafe-3.0.3-cp312-cp312-win_arm64.whl", hash = "sha256:35add3b638a5d900e807944a078b51922212fb3dedb01633a8defc4b01a3c85f", size = 13906, upload-time = "2025-09-27T18:36:40.689Z" }, ] -[[package]] -name = "marshmallow" -version = "3.26.2" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "packaging" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/55/79/de6c16cc902f4fc372236926b0ce2ab7845268dcc30fb2fbb7f71b418631/marshmallow-3.26.2.tar.gz", hash = "sha256:bbe2adb5a03e6e3571b573f42527c6fe926e17467833660bebd11593ab8dfd57", size = 222095, upload-time = "2025-12-22T06:53:53.309Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/be/2f/5108cb3ee4ba6501748c4908b908e55f42a5b66245b4cfe0c99326e1ef6e/marshmallow-3.26.2-py3-none-any.whl", hash = "sha256:013fa8a3c4c276c24d26d84ce934dc964e2aa794345a0f8c7e5a7191482c8a73", size = 50964, upload-time = "2025-12-22T06:53:51.801Z" }, -] - [[package]] name = "mcp" version = "2.2.0" @@ -2202,15 +2151,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/b7/da/7d22601b625e241d4f23ef1ebff8acfc60da633c9e7e7922e24d10f592b3/multidict-6.7.0-py3-none-any.whl", hash = "sha256:394fc5c42a333c9ffc3e421a4c85e08580d990e08b99f6bf35b4132114c5dcb3", size = 12317, upload-time = "2025-10-06T14:52:29.272Z" }, ] -[[package]] -name = "mypy-extensions" -version = "1.1.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/a2/6e/371856a3fb9d31ca8dac321cda606860fa4548858c0cc45d9d1d4ca2628b/mypy_extensions-1.1.0.tar.gz", hash = "sha256:52e68efc3284861e772bbcd66823fde5ae21fd2fdb51c62a211403730b916558", size = 6343, upload-time = "2025-04-22T14:54:24.164Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/79/7b/2c79738432f5c924bef5071f933bcc9efd0473bac3b4aa584a6f7c1c8df8/mypy_extensions-1.1.0-py3-none-any.whl", hash = "sha256:1be4cccdb0f2482337c4743e60421de3a356cd97508abadd57d47403e94f5505", size = 4963, upload-time = "2025-04-22T14:54:22.983Z" }, -] - [[package]] name = "nexus-rpc" version = "1.2.0" @@ -4047,19 +3987,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/18/67/36e9267722cc04a6b9f15c7f3441c2363321a3ea07da7ae0c0707beb2a9c/typing_extensions-4.15.0-py3-none-any.whl", hash = "sha256:f0fa19c6845758ab08074a0cfa8b7aecb71c999ca73d62883bc25cc018c4e548", size = 44614, upload-time = "2025-08-25T13:49:24.86Z" }, ] -[[package]] -name = "typing-inspect" -version = "0.9.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "mypy-extensions" }, - { name = "typing-extensions" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/dc/74/1789779d91f1961fa9438e9a8710cdae6bd138c80d7303996933d117264a/typing_inspect-0.9.0.tar.gz", hash = "sha256:b23fc42ff6f6ef6954e4852c1fb512cdd18dbea03134f91f856a95ccc9461f78", size = 13825, upload-time = "2023-05-24T20:25:47.612Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/65/f3/107a22063bf27bdccf2024833d3445f4eea42b2e598abfbd46f6a63b6cb0/typing_inspect-0.9.0-py3-none-any.whl", hash = "sha256:9ee6fc59062311ef8547596ab6b955e1b8aa46242d854bfc78f4f6b0eff35f9f", size = 8827, upload-time = "2023-05-24T20:25:45.287Z" }, -] - [[package]] name = "typing-inspection" version = "0.4.2"