diff --git a/airbyte/_util/api_util.py b/airbyte/_util/api_util.py index bb6ebc4da..193a3730c 100644 --- a/airbyte/_util/api_util.py +++ b/airbyte/_util/api_util.py @@ -3658,6 +3658,124 @@ def get_connection_catalog( ) +def run_source_discover_command( + source_id: str, + discover_request_id: str, + *, + api_root: str, + client_id: SecretString | None, + client_secret: SecretString | None, + bearer_token: SecretString | None, + config_api_root: str | None = None, +) -> dict[str, Any]: + """Start a source discover command. + + Uses the Config API endpoint: POST /v1/commands/run/discover. + """ + return _make_config_api_request( + path="/commands/run/discover", + json={"id": discover_request_id, "actor_id": source_id}, + api_root=api_root, + config_api_root=config_api_root, + client_id=client_id, + client_secret=client_secret, + bearer_token=bearer_token, + ) + + +def get_command_status( + command_id: str, + *, + api_root: str, + client_id: SecretString | None, + client_secret: SecretString | None, + bearer_token: SecretString | None, + config_api_root: str | None = None, +) -> dict[str, Any]: + """Get the status of a Config API command.""" + return _make_config_api_request( + path="/commands/status", + json={"id": command_id}, + api_root=api_root, + config_api_root=config_api_root, + client_id=client_id, + client_secret=client_secret, + bearer_token=bearer_token, + ) + + +def get_discover_command_output( + command_id: str, + *, + api_root: str, + client_id: SecretString | None, + client_secret: SecretString | None, + bearer_token: SecretString | None, + config_api_root: str | None = None, +) -> dict[str, Any]: + """Get the output of a completed source discover command.""" + return _make_config_api_request( + path="/commands/output/discover", + json={"id": command_id}, + api_root=api_root, + config_api_root=config_api_root, + client_id=client_id, + client_secret=client_secret, + bearer_token=bearer_token, + ) + + +def get_most_recent_source_catalog( + source_id: str, + *, + api_root: str, + client_id: SecretString | None, + client_secret: SecretString | None, + bearer_token: SecretString | None, + config_api_root: str | None = None, +) -> dict[str, Any]: + """Get the most recent saved source catalog.""" + return _make_config_api_request( + path="/sources/most_recent_source_actor_catalog", + json={"sourceId": source_id}, + api_root=api_root, + config_api_root=config_api_root, + client_id=client_id, + client_secret=client_secret, + bearer_token=bearer_token, + ) + + +def get_catalog_diff( + current_catalog_id: str, + new_catalog_id: str, + *, + api_root: str, + client_id: SecretString | None, + client_secret: SecretString | None, + bearer_token: SecretString | None, + connection_id: str | None = None, + config_api_root: str | None = None, +) -> dict[str, Any]: + """Compare two source catalogs without persisting changes to a connection.""" + request_body: dict[str, Any] = { + "current_catalog_id": current_catalog_id, + "new_catalog_id": new_catalog_id, + "persist_changes": False, + } + if connection_id is not None: + request_body["connection_id"] = connection_id + return _make_config_api_request( + path="/catalogs/diff", + json=request_body, + api_root=api_root, + config_api_root=config_api_root, + client_id=client_id, + client_secret=client_secret, + bearer_token=bearer_token, + ) + + def replace_connection_catalog( connection_id: str, configured_catalog_dict: dict[str, Any], diff --git a/airbyte/cloud/__init__.py b/airbyte/cloud/__init__.py index 2a2043f35..3eb25f5c4 100644 --- a/airbyte/cloud/__init__.py +++ b/airbyte/cloud/__init__.py @@ -94,6 +94,8 @@ from airbyte.cloud.connections import CloudConnection from airbyte.cloud.models import ( CloudDefaultContextInfo, + CloudSourceCatalogDiffStreamSummary, + CloudSourceCatalogDiscoverResult, CloudWorkspaceInfo, ConnectionStatus, JobStatusEnum, @@ -133,6 +135,8 @@ "CloudOrganization", "CloudWorkspace", "CloudConnection", + "CloudSourceCatalogDiffStreamSummary", + "CloudSourceCatalogDiscoverResult", "CloudClientConfig", "CloudDefaultContextInfo", "CloudWorkspaceInfo", diff --git a/airbyte/cloud/connections.py b/airbyte/cloud/connections.py index c296108cd..67435df9c 100644 --- a/airbyte/cloud/connections.py +++ b/airbyte/cloud/connections.py @@ -902,6 +902,7 @@ def import_raw_catalog(self, catalog: dict[str, Any]) -> None: bearer_token=self.workspace.bearer_token, config_api_root=self.workspace.config_api_root, ) + self._connection_info = None def rename(self, name: str) -> CloudConnection: """Rename the connection. diff --git a/airbyte/cloud/connectors.py b/airbyte/cloud/connectors.py index 0e08038c2..2035efb02 100644 --- a/airbyte/cloud/connectors.py +++ b/airbyte/cloud/connectors.py @@ -55,9 +55,11 @@ from __future__ import annotations +from datetime import UTC, datetime from http import HTTPStatus from pathlib import Path from typing import TYPE_CHECKING, Any, ClassVar, Literal, Protocol +from uuid import uuid4 import jsonschema import requests @@ -92,6 +94,8 @@ CheckResult, CloudCustomSourceDefinitionInfo, CloudDestinationInfo, + CloudSourceCatalogDiffStreamSummary, + CloudSourceCatalogDiscoverResult, CloudSourceInfo, ConnectorFeature, ConnectorType, @@ -129,6 +133,92 @@ """409 messages that carry no Cloud detail: the platform default and PyAirbyte's fallback.""" +def _get_catalog_discover_result( + workspace: CloudWorkspace, + discover_request_id: str, +) -> CloudSourceCatalogDiscoverResult: + """Get a source-discover result using a workspace's Config API credentials.""" + if not isinstance(discover_request_id, str) or not discover_request_id.strip(): + raise exc.AirbyteLibInputError( + message="`discover_request_id` must be a non-empty string.", + context={"discover_request_id": discover_request_id}, + ) + + command_status = api_util.get_command_status( + discover_request_id, + api_root=workspace.api_root, + config_api_root=workspace.config_api_root, + client_id=workspace.client_id, + client_secret=workspace.client_secret, + bearer_token=workspace.bearer_token, + ) + request_status = command_status["status"] + if request_status != "completed": + return CloudSourceCatalogDiscoverResult( + discover_request_id=discover_request_id, + request_status=request_status, + ) + + output = api_util.get_discover_command_output( + discover_request_id, + api_root=workspace.api_root, + config_api_root=workspace.config_api_root, + client_id=workspace.client_id, + client_secret=workspace.client_secret, + bearer_token=workspace.bearer_token, + ) + failure = output.get("failureReason") + failure_reason = failure.get("externalMessage") if isinstance(failure, dict) else None + catalog = output.get("catalog") + stream_count = len(catalog["streams"]) if catalog is not None else None + return CloudSourceCatalogDiscoverResult( + discover_request_id=discover_request_id, + request_status=request_status, + outcome=output["status"], + source_catalog_id=output.get("catalogId"), + stream_count=stream_count, + failure_reason=failure_reason, + ) + + +def _build_catalog_diff_result( + source_catalog_id: str, + response: dict[str, Any], +) -> CloudSourceCatalogDiscoverResult: + """Build a bounded summary of the Config API catalog-diff response.""" + transforms = response["catalog_diff"]["transforms"] + summaries = [] + for transform in transforms: + descriptor = transform["streamDescriptor"] + update_stream = transform.get("updateStream") or {} + field_transforms = update_stream.get("fieldTransforms", []) + attribute_transforms = update_stream.get("streamAttributeTransforms", []) + breaking_flags = [ + change["breaking"] + for change in [*field_transforms, *attribute_transforms] + if "breaking" in change + ] + summaries.append( + CloudSourceCatalogDiffStreamSummary( + stream_name=descriptor["name"], + stream_namespace=descriptor.get("namespace"), + transform_type=transform["transformType"], + breaking_flags=breaking_flags, + field_additions_count=sum( + change["transformType"] == "add_field" for change in field_transforms + ), + field_removals_count=sum( + change["transformType"] == "remove_field" for change in field_transforms + ), + ) + ) + return CloudSourceCatalogDiscoverResult( + source_catalog_id=source_catalog_id, + catalog_diff_summary=summaries, + schema_change=response["schema_change"], + ) + + class _ConnectorDefinitionLike(Protocol): """The connector definition fields PyAirbyte reads from the Cloud public API.""" @@ -1085,6 +1175,68 @@ def _fetch_connector_info(self) -> CloudSourceInfo: ) ) + def _validate_source_id(self) -> None: + if not self.source_id.strip(): + raise exc.AirbyteLibInputError( + message="`source_id` must be a non-empty string.", + context={"source_id": self.source_id}, + ) + + def start_catalog_discover(self) -> CloudSourceCatalogDiscoverResult: + """Start an asynchronous source catalog discover command.""" + self._validate_source_id() + discover_request_id = str(uuid4()) + api_util.run_source_discover_command( + source_id=self.source_id, + discover_request_id=discover_request_id, + api_root=self.workspace.api_root, + config_api_root=self.workspace.config_api_root, + client_id=self.workspace.client_id, + client_secret=self.workspace.client_secret, + bearer_token=self.workspace.bearer_token, + ) + return CloudSourceCatalogDiscoverResult( + source_id=self.source_id, + discover_request_id=discover_request_id, + request_status="pending", + ) + + def get_catalog_discover_result( + self, + discover_request_id: str, + ) -> CloudSourceCatalogDiscoverResult: + """Get the status and completed output of a source catalog discover.""" + result = _get_catalog_discover_result( + workspace=self.workspace, + discover_request_id=discover_request_id, + ) + return result.model_copy(update={"source_id": self.source_id}) + + def get_latest_saved_catalog(self) -> CloudSourceCatalogDiscoverResult: + """Get the timestamp and stream count of the latest saved catalog.""" + self._validate_source_id() + response = api_util.get_most_recent_source_catalog( + source_id=self.source_id, + api_root=self.workspace.api_root, + config_api_root=self.workspace.config_api_root, + client_id=self.workspace.client_id, + client_secret=self.workspace.client_secret, + bearer_token=self.workspace.bearer_token, + ) + catalog = response.get("catalog") + streams = catalog.get("streams", []) if isinstance(catalog, dict) else [] + if not streams: + return CloudSourceCatalogDiscoverResult(source_id=self.source_id) + + updated_at = response.get("updatedAt") + return CloudSourceCatalogDiscoverResult( + source_id=self.source_id, + stream_count=len(streams), + updated_at=( + datetime.fromtimestamp(updated_at, tz=UTC) if updated_at is not None else None + ), + ) + def rename(self, name: str) -> CloudSource: """Rename the source. diff --git a/airbyte/cloud/models.py b/airbyte/cloud/models.py index 59caccd83..9e1355b11 100644 --- a/airbyte/cloud/models.py +++ b/airbyte/cloud/models.py @@ -5,6 +5,7 @@ from collections.abc import Mapping from dataclasses import asdict, dataclass, is_dataclass +from datetime import datetime # noqa: TC003 # Needed at runtime for Pydantic field types. from enum import Enum from typing import Any, Literal, Protocol @@ -509,6 +510,32 @@ def from_api_response(cls, source: _SourceResponseLike) -> CloudSourceInfo: ) +class CloudSourceCatalogDiffStreamSummary(BaseModel): + """Summary of one stream transformation in a catalog diff.""" + + stream_name: str + stream_namespace: str | None = None + transform_type: str + breaking_flags: list[bool] = Field(default_factory=list) + field_additions_count: int = 0 + field_removals_count: int = 0 + + +class CloudSourceCatalogDiscoverResult(BaseModel): + """Status or outcome of discovering a Cloud source catalog.""" + + source_id: str | None = None + discover_request_id: str | None = None + request_status: Literal["pending", "running", "completed", "cancelled"] | None = None + outcome: Literal["succeeded", "failed"] | None = None + source_catalog_id: str | None = None + stream_count: int | None = None + updated_at: datetime | None = None + failure_reason: str | None = None + catalog_diff_summary: list[CloudSourceCatalogDiffStreamSummary] | None = None + schema_change: Literal["no_change", "non_breaking", "breaking"] | None = None + + class CloudDestinationInfo(BaseModel): """Information about an Airbyte Cloud destination.""" diff --git a/airbyte/cloud/workspaces.py b/airbyte/cloud/workspaces.py index 101e7d1eb..420433114 100644 --- a/airbyte/cloud/workspaces.py +++ b/airbyte/cloud/workspaces.py @@ -515,6 +515,83 @@ def get_source( connector_id=source_id, ) + def get_catalog_discover_result( + self, + discover_request_id: str, + ) -> cloud_connectors.CloudSourceCatalogDiscoverResult: + """Get the status and completed output for a discover request ID.""" + return cloud_connectors._get_catalog_discover_result( # noqa: SLF001 + workspace=self, + discover_request_id=discover_request_id, + ) + + def compare_source_catalogs( + self, + source_catalog_id: str, + *, + compare_to_source_catalog_id: str | None = None, + compare_to_connection_id: str | None = None, + ) -> cloud_connectors.CloudSourceCatalogDiscoverResult: + """Compare a new source catalog with another catalog or a connection.""" + if not source_catalog_id.strip(): + raise exc.AirbyteLibInputError( + message="`source_catalog_id` must be a non-empty string.", + ) + if compare_to_source_catalog_id is not None and compare_to_connection_id is not None: + raise exc.AirbyteLibInputError( + message=( + "Provide `compare_to_source_catalog_id` or " + "`compare_to_connection_id`, not both." + ), + ) + if compare_to_source_catalog_id is not None: + if not compare_to_source_catalog_id.strip(): + raise exc.AirbyteLibInputError( + message="`compare_to_source_catalog_id` must be a non-empty string.", + ) + current_catalog_id = compare_to_source_catalog_id + elif compare_to_connection_id is not None: + if not compare_to_connection_id.strip(): + raise exc.AirbyteLibInputError( + message="`compare_to_connection_id` must be a non-empty string.", + ) + connection = api_util.get_connection_catalog( + connection_id=compare_to_connection_id, + api_root=self.api_root, + config_api_root=self.config_api_root, + client_id=self.client_id, + client_secret=self.client_secret, + bearer_token=self.bearer_token, + ) + current_catalog_id = connection.get("catalogId") + if not current_catalog_id: + raise exc.AirbyteLibInputError( + message=("The connection has no saved source catalog to compare against."), + context={"connection_id": compare_to_connection_id}, + ) + else: + raise exc.AirbyteLibInputError( + message=( + "Looking up a source catalog by ID without comparing it is not supported yet." + ), + context={"source_catalog_id": source_catalog_id}, + ) + + diff = api_util.get_catalog_diff( + current_catalog_id=current_catalog_id, + new_catalog_id=source_catalog_id, + connection_id=compare_to_connection_id, + api_root=self.api_root, + config_api_root=self.config_api_root, + client_id=self.client_id, + client_secret=self.client_secret, + bearer_token=self.bearer_token, + ) + return cloud_connectors._build_catalog_diff_result( # noqa: SLF001 + source_catalog_id=source_catalog_id, + response=diff, + ) + def get_destination( self, destination_id: str, diff --git a/airbyte/mcp/_otel.py b/airbyte/mcp/_otel.py index 051402741..78b950c6b 100644 --- a/airbyte/mcp/_otel.py +++ b/airbyte/mcp/_otel.py @@ -52,6 +52,8 @@ r"workspaces/(?:list_by_organization_id|list_by_user_id|get_organization_info|get)|" r"state/(?:get|create_or_update_safe)|web_backend/connections/(?:get|update)|" r"users/(?:get_by_auth_id|update)|permissions/list_by_user|jobs/get|" + r"commands/(?:run/discover|status|output/discover)|catalogs/diff|" + r"sources/most_recent_source_actor_catalog|" rf"(?:sources|destinations)/{_UUID_PATTERN}/(?:execute|search|search-status|enablement)|" rf"workspaces/{_UUID_PATTERN}/skills/docs)" ) diff --git a/airbyte/mcp/cloud.py b/airbyte/mcp/cloud.py index 7e0b908a8..fa42f63ff 100644 --- a/airbyte/mcp/cloud.py +++ b/airbyte/mcp/cloud.py @@ -45,6 +45,7 @@ CloudDefaultContextInfo, CloudDefaultWorkspaceUpdateInfo, CloudOrganizationInfo, + CloudSourceCatalogDiscoverResult, ConnectionSchedule, ConnectionStatus, ConnectorFeature, @@ -4131,6 +4132,139 @@ def update_cloud_connection( ) +@mcp_tool(destructive=False, pipeline_change=False) +def discover_cloud_source_catalog( + ctx: Context, + *, + source_id: Annotated[ + str | None, + Field( + description=( + "The source ID. Starts a new asynchronous catalog discover by default; " + "set `no_refresh=True` to use the latest saved catalog instead." + ), + default=None, + ), + TraceArg.FINGERPRINT, + ] = None, + discover_request_id: Annotated[ + str | None, + Field( + description=( + "The request ID returned when a discover was started. Poll with this " + "value to retrieve its status and completed outcome." + ), + default=None, + ), + TraceArg.FINGERPRINT, + ] = None, + no_refresh: Annotated[ + bool, + Field( + description=( + "With `source_id`, return the latest saved catalog without starting a " + "discover. The latest-saved-catalog mode cannot be applied to a connection." + ), + default=False, + ), + ] = False, + source_catalog_id: Annotated[ + str | None, + Field( + description=( + "The newly discovered source catalog ID to compare. Comparing requires " + "exactly one `compare_to_*` argument; summary by ID alone is not supported yet." + ), + default=None, + ), + TraceArg.FINGERPRINT, + ] = None, + compare_to_source_catalog_id: Annotated[ + str | None, + Field( + description="Compare `source_catalog_id` against this saved source catalog ID.", + default=None, + ), + TraceArg.FINGERPRINT, + ] = None, + compare_to_connection_id: Annotated[ + str | None, + Field( + description="Compare `source_catalog_id` against this connection's current catalog.", + default=None, + ), + TraceArg.FINGERPRINT, + ] = None, + workspace_id: Annotated[ + str | None, + Field( + description=WORKSPACE_ID_TIP_TEXT, + default=None, + ), + TraceArg.OMIT, + ] = None, +) -> CloudSourceCatalogDiscoverResult: + """Start, check, or compare asynchronous Cloud source catalog discovers. + + Starting a discover can be slow. Poll the returned `discover_request_id` to get + its status and outcome. Use `no_refresh=True` with `source_id` to read the latest + saved catalog; this mode cannot be applied to a connection. Compare catalogs by + passing `source_catalog_id` and exactly one comparison target. A catalog summary + by ID alone is not supported yet. + """ + if compare_to_source_catalog_id is not None and compare_to_connection_id is not None: + raise AirbyteLibInputError( + message=( + "Provide `compare_to_source_catalog_id` or `compare_to_connection_id`, " "not both." + ), + ) + has_comparison_target = ( + compare_to_source_catalog_id is not None or compare_to_connection_id is not None + ) + if has_comparison_target and source_catalog_id is None: + raise AirbyteLibInputError( + message="Comparison targets require `source_catalog_id`.", + ) + if source_catalog_id is not None: + if source_id is not None or discover_request_id is not None or no_refresh: + raise AirbyteLibInputError( + message=( + "`source_catalog_id` cannot be combined with `source_id`, " + "`discover_request_id`, or `no_refresh`." + ), + ) + workspace: CloudWorkspace = _get_cloud_workspace(ctx, workspace_id) + return workspace.compare_source_catalogs( + source_catalog_id=source_catalog_id, + compare_to_source_catalog_id=compare_to_source_catalog_id, + compare_to_connection_id=compare_to_connection_id, + ) + + if discover_request_id is not None: + if source_id is not None: + raise AirbyteLibInputError( + message="Provide `source_id` or `discover_request_id`, not both.", + ) + if no_refresh: + raise AirbyteLibInputError( + message="`no_refresh=True` cannot be used with `discover_request_id`.", + ) + elif source_id is None: + raise AirbyteLibInputError( + message="Provide `source_id` or `discover_request_id`.", + ) + + workspace: CloudWorkspace = _get_cloud_workspace(ctx, workspace_id) + if discover_request_id is not None: + return workspace.get_catalog_discover_result(discover_request_id) + + assert source_id is not None + source = workspace.get_source(source_id=source_id) + if no_refresh: + return source.get_latest_saved_catalog() + return source.start_catalog_discover() + + @mcp_tool( read_only=True, idempotent=True, diff --git a/tests/unit_tests/fixtures/arg_trace_plan.json b/tests/unit_tests/fixtures/arg_trace_plan.json index 23b989e52..eb7415bee 100644 --- a/tests/unit_tests/fixtures/arg_trace_plan.json +++ b/tests/unit_tests/fixtures/arg_trace_plan.json @@ -66,6 +66,15 @@ "scenarios": "fingerprint list", "skip_preflight": "value allowed=[false,true]" }, + "discover_cloud_source_catalog": { + "compare_to_connection_id": "fingerprint", + "compare_to_source_catalog_id": "fingerprint", + "discover_request_id": "fingerprint", + "no_refresh": "value allowed=[false,true]", + "source_catalog_id": "fingerprint", + "source_id": "fingerprint", + "workspace_id": "omit" + }, "execute_external_api_query": { "action": "value allowed=[\"get\",\"list\",\"search\"]", "api_args": "hash", diff --git a/tests/unit_tests/test_cloud_api_util.py b/tests/unit_tests/test_cloud_api_util.py index f4907260b..7bb430a23 100644 --- a/tests/unit_tests/test_cloud_api_util.py +++ b/tests/unit_tests/test_cloud_api_util.py @@ -368,6 +368,91 @@ def test_get_user_id_from_bearer_token_rejects_invalid_tokens( "expected_result", ), [ + pytest.param( + api_util.get_connection_catalog, + {"connection_id": "connection-id"}, + "/web_backend/connections/get", + { + "connectionId": "connection-id", + "withRefreshedCatalog": False, + }, + {"catalogId": "saved-catalog-id"}, + {"catalogId": "saved-catalog-id"}, + id="connection-catalog", + ), + pytest.param( + api_util.run_source_discover_command, + { + "source_id": "source-id", + "discover_request_id": "request-id", + }, + "/commands/run/discover", + {"id": "request-id", "actor_id": "source-id"}, + {"id": "request-id"}, + {"id": "request-id"}, + id="run-source-discover", + ), + pytest.param( + api_util.get_command_status, + {"command_id": "request-id"}, + "/commands/status", + {"id": "request-id"}, + {"id": "request-id", "status": "pending"}, + {"id": "request-id", "status": "pending"}, + id="command-status", + ), + pytest.param( + api_util.get_discover_command_output, + {"command_id": "request-id"}, + "/commands/output/discover", + {"id": "request-id"}, + {"id": "request-id", "status": "succeeded"}, + {"id": "request-id", "status": "succeeded"}, + id="discover-command-output", + ), + pytest.param( + api_util.get_most_recent_source_catalog, + {"source_id": "source-id"}, + "/sources/most_recent_source_actor_catalog", + {"sourceId": "source-id"}, + {"catalog": {"streams": []}}, + {"catalog": {"streams": []}}, + id="most-recent-source-catalog", + ), + pytest.param( + api_util.get_catalog_diff, + { + "current_catalog_id": "old-catalog-id", + "new_catalog_id": "new-catalog-id", + }, + "/catalogs/diff", + { + "current_catalog_id": "old-catalog-id", + "new_catalog_id": "new-catalog-id", + "persist_changes": False, + }, + {"catalog_diff": {"transforms": []}, "schema_change": "no_change"}, + {"catalog_diff": {"transforms": []}, "schema_change": "no_change"}, + id="catalog-diff-no-connection", + ), + pytest.param( + api_util.get_catalog_diff, + { + "current_catalog_id": "old-catalog-id", + "new_catalog_id": "new-catalog-id", + "connection_id": "connection-id", + }, + "/catalogs/diff", + { + "current_catalog_id": "old-catalog-id", + "new_catalog_id": "new-catalog-id", + "connection_id": "connection-id", + "persist_changes": False, + }, + {"catalog_diff": {"transforms": []}, "schema_change": "no_change"}, + {"catalog_diff": {"transforms": []}, "schema_change": "no_change"}, + id="catalog-diff-with-connection", + ), pytest.param( api_util.get_user_by_auth_id, {"auth_user_id": "auth-user-id"}, diff --git a/tests/unit_tests/test_cloud_connection_catalog_edits.py b/tests/unit_tests/test_cloud_connection_catalog_edits.py new file mode 100644 index 000000000..ae6204631 --- /dev/null +++ b/tests/unit_tests/test_cloud_connection_catalog_edits.py @@ -0,0 +1,77 @@ +# Copyright (c) 2026 Airbyte, Inc., all rights reserved. +"""Unit tests for editing Cloud connection catalogs.""" + +from __future__ import annotations + +from typing import Any +from unittest.mock import MagicMock + +import pytest + +from airbyte._util import api_util +from airbyte.cloud.connections import CloudConnection +from airbyte.cloud.workspaces import CloudWorkspace + + +_CONNECTION_ID = "connection-id" +_API_ROOT = "https://api.example/v1" +_CONFIG_API_ROOT = "https://config.example/v1" + + +def _workspace() -> CloudWorkspace: + return CloudWorkspace( + workspace_id="workspace-id", + bearer_token="token", + api_root=_API_ROOT, + config_api_root=_CONFIG_API_ROOT, + ) + + +def _connection() -> CloudConnection: + return CloudConnection(workspace=_workspace(), connection_id=_CONNECTION_ID) + + +def _stream_entry( + name: str, + *, + namespace: str | None = None, +) -> dict[str, Any]: + stream: dict[str, Any] = { + "name": name, + "jsonSchema": {"type": "object"}, + "supportedSyncModes": ["full_refresh", "incremental"], + } + if namespace is not None: + stream["namespace"] = namespace + return { + "stream": stream, + "config": { + "syncMode": "full_refresh", + "destinationSyncMode": "append", + "selectedFields": [], + "fieldSelectionEnabled": False, + "primaryKey": [["id"]], + "selected": True, + }, + } + + +def _catalog(*entries: dict[str, Any]) -> dict[str, Any]: + return {"streams": list(entries)} + + +def test_import_raw_catalog_clears_cached_connection_info( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Importing a catalog invalidates previously cached connection information.""" + connection = _connection() + stale_connection_info = MagicMock() + stale_connection_info.configurations.streams = [MagicMock(name="old-stream")] + connection._connection_info = stale_connection_info + replace_catalog = MagicMock(return_value={}) + monkeypatch.setattr(api_util, "replace_connection_catalog", replace_catalog) + + connection.import_raw_catalog(_catalog(_stream_entry("new-stream"))) + + replace_catalog.assert_called_once() + assert connection._connection_info is None diff --git a/tests/unit_tests/test_cloud_source_catalog_discovery.py b/tests/unit_tests/test_cloud_source_catalog_discovery.py new file mode 100644 index 000000000..6fd60f1a2 --- /dev/null +++ b/tests/unit_tests/test_cloud_source_catalog_discovery.py @@ -0,0 +1,551 @@ +# Copyright (c) 2026 Airbyte, Inc., all rights reserved. +"""Tests for asynchronous Cloud source catalog discovery.""" + +from __future__ import annotations + +import asyncio +from datetime import datetime, timezone +from typing import Any +from unittest.mock import MagicMock +from uuid import UUID + +import pytest +from fastmcp import Client, FastMCP +from fastmcp.exceptions import ToolError + +from airbyte._util import api_util +from airbyte.cloud import connectors +from airbyte.cloud.connectors import CloudSource +from airbyte.cloud.models import ( + CloudSourceCatalogDiffStreamSummary, + CloudSourceCatalogDiscoverResult, +) +from airbyte.cloud.workspaces import CloudWorkspace +from airbyte.exceptions import AirbyteCloudError, AirbyteLibInputError +from airbyte.mcp import cloud as cloud_mcp + +_SOURCE_ID = "source-id" +_REQUEST_ID = "request-id" +_API_ROOT = "https://api.example/v1" +_CONFIG_API_ROOT = "https://config.example/v1" + + +def _workspace() -> CloudWorkspace: + return CloudWorkspace( + workspace_id="workspace-id", + bearer_token="token", + api_root=_API_ROOT, + config_api_root=_CONFIG_API_ROOT, + ) + + +def _source(source_id: str = _SOURCE_ID) -> CloudSource: + return CloudSource(workspace=_workspace(), connector_id=source_id) + + +def _server() -> FastMCP: + server = FastMCP("source-catalog-discovery-tests") + cloud_mcp.register_cloud_tools(server) + return server + + +async def _call_tool(server: FastMCP, arguments: dict[str, Any]) -> Any: + async with Client(server) as client: + return await client.call_tool("discover_cloud_source_catalog", arguments) + + +def test_start_catalog_discover_uses_a_generated_request_id( + monkeypatch: pytest.MonkeyPatch, +) -> None: + discover_request_id = UUID("12345678-1234-1234-1234-1234567890ab") + monkeypatch.setattr(connectors, "uuid4", lambda: discover_request_id) + run_discover = MagicMock(return_value={"id": str(discover_request_id)}) + monkeypatch.setattr(api_util, "run_source_discover_command", run_discover) + + result = _source().start_catalog_discover() + + assert result == CloudSourceCatalogDiscoverResult( + source_id=_SOURCE_ID, + discover_request_id=str(discover_request_id), + request_status="pending", + ) + run_discover.assert_called_once_with( + source_id=_SOURCE_ID, + discover_request_id=str(discover_request_id), + api_root=_API_ROOT, + config_api_root=_CONFIG_API_ROOT, + client_id=None, + client_secret=None, + bearer_token="token", + ) + + +@pytest.mark.parametrize("request_status", ["pending", "running", "cancelled"]) +def test_get_catalog_discover_result_returns_noncompleted_status( + monkeypatch: pytest.MonkeyPatch, + request_status: str, +) -> None: + status = MagicMock(return_value={"status": request_status}) + output = MagicMock() + monkeypatch.setattr(api_util, "get_command_status", status) + monkeypatch.setattr(api_util, "get_discover_command_output", output) + + result = _source().get_catalog_discover_result(_REQUEST_ID) + + assert result.source_id == _SOURCE_ID + assert result.discover_request_id == _REQUEST_ID + assert result.request_status == request_status + assert result.outcome is None + output.assert_not_called() + + +def test_get_catalog_discover_result_returns_success_summary( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr( + api_util, + "get_command_status", + MagicMock(return_value={"status": "completed"}), + ) + monkeypatch.setattr( + api_util, + "get_discover_command_output", + MagicMock( + return_value={ + "status": "succeeded", + "catalogId": "catalog-id", + "catalog": {"streams": [{}, {}, {}]}, + } + ), + ) + + result = _source().get_catalog_discover_result(_REQUEST_ID) + + assert result.request_status == "completed" + assert result.outcome == "succeeded" + assert result.source_catalog_id == "catalog-id" + assert result.stream_count == 3 + assert result.failure_reason is None + assert "catalog" not in result.model_dump() + + +def test_get_catalog_discover_result_exposes_external_failure_message( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr( + api_util, + "get_command_status", + MagicMock(return_value={"status": "completed"}), + ) + monkeypatch.setattr( + api_util, + "get_discover_command_output", + MagicMock( + return_value={ + "status": "failed", + "failureReason": { + "externalMessage": "The source could not be read.", + "internalMessage": "internal diagnostic", + }, + } + ), + ) + + result = _source().get_catalog_discover_result(_REQUEST_ID) + + assert result.outcome == "failed" + assert result.failure_reason == "The source could not be read." + assert "internal" not in result.model_dump_json() + + +def test_get_catalog_discover_result_propagates_unknown_request_error( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr( + api_util, + "get_command_status", + MagicMock(side_effect=AirbyteCloudError(message="Unknown request ID.")), + ) + + with pytest.raises(AirbyteCloudError, match="Unknown request ID"): + _source().get_catalog_discover_result(_REQUEST_ID) + + +def test_get_catalog_discover_result_rejects_empty_request_id() -> None: + with pytest.raises(AirbyteLibInputError, match="non-empty string"): + _source().get_catalog_discover_result(" ") + + +def test_get_latest_saved_catalog_returns_timestamp_and_stream_count( + monkeypatch: pytest.MonkeyPatch, +) -> None: + get_catalog = MagicMock( + return_value={ + "updatedAt": 1_700_000_000, + "catalog": {"streams": [{}, {}]}, + } + ) + monkeypatch.setattr(api_util, "get_most_recent_source_catalog", get_catalog) + + result = _source().get_latest_saved_catalog() + + assert result.source_id == _SOURCE_ID + assert result.updated_at == datetime(2023, 11, 14, 22, 13, 20, tzinfo=timezone.utc) + assert result.stream_count == 2 + get_catalog.assert_called_once_with( + source_id=_SOURCE_ID, + api_root=_API_ROOT, + config_api_root=_CONFIG_API_ROOT, + client_id=None, + client_secret=None, + bearer_token="token", + ) + + +def test_get_latest_saved_catalog_returns_none_when_catalog_is_empty( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr( + api_util, + "get_most_recent_source_catalog", + MagicMock(return_value={"updatedAt": 1_700_000_000, "catalog": {}}), + ) + + result = _source().get_latest_saved_catalog() + + assert result.updated_at is None + assert result.stream_count is None + + +def test_compare_source_catalogs_summarizes_diff_without_exposing_catalogs( + monkeypatch: pytest.MonkeyPatch, +) -> None: + response = { + "catalog_diff": { + "transforms": [ + { + "transformType": "update_stream", + "streamDescriptor": {"name": "users", "namespace": "public"}, + "updateStream": { + "fieldTransforms": [ + { + "transformType": "add_field", + "fieldName": ["email"], + "breaking": False, + }, + { + "transformType": "remove_field", + "fieldName": ["name"], + "breaking": True, + }, + ], + "streamAttributeTransforms": [ + {"transformType": "update_sync_mode", "breaking": True} + ], + }, + } + ] + }, + "schema_change": "breaking", + "merged_catalog": {"streams": [{"json_schema": {"secret": "not returned"}}]}, + } + get_diff = MagicMock(return_value=response) + monkeypatch.setattr(api_util, "get_catalog_diff", get_diff) + + result = _workspace().compare_source_catalogs( + "new-catalog-id", + compare_to_source_catalog_id="old-catalog-id", + ) + + assert result.source_catalog_id == "new-catalog-id" + assert result.schema_change == "breaking" + assert result.catalog_diff_summary == [ + CloudSourceCatalogDiffStreamSummary( + stream_name="users", + stream_namespace="public", + transform_type="update_stream", + breaking_flags=[False, True, True], + field_additions_count=1, + field_removals_count=1, + ) + ] + assert "merged_catalog" not in result.model_dump() + assert "json_schema" not in result.model_dump_json() + get_diff.assert_called_once_with( + current_catalog_id="old-catalog-id", + new_catalog_id="new-catalog-id", + connection_id=None, + api_root=_API_ROOT, + config_api_root=_CONFIG_API_ROOT, + client_id=None, + client_secret=None, + bearer_token="token", + ) + + +def test_compare_source_catalogs_uses_connection_catalog_id( + monkeypatch: pytest.MonkeyPatch, +) -> None: + get_connection = MagicMock(return_value={"catalogId": "connection-catalog-id"}) + get_diff = MagicMock( + return_value={ + "catalog_diff": {"transforms": []}, + "schema_change": "non_breaking", + "merged_catalog": {"streams": []}, + } + ) + monkeypatch.setattr(api_util, "get_connection_catalog", get_connection) + monkeypatch.setattr(api_util, "get_catalog_diff", get_diff) + + result = _workspace().compare_source_catalogs( + "new-catalog-id", + compare_to_connection_id="connection-id", + ) + + assert result.schema_change == "non_breaking" + assert result.catalog_diff_summary == [] + get_connection.assert_called_once_with( + connection_id="connection-id", + api_root=_API_ROOT, + config_api_root=_CONFIG_API_ROOT, + client_id=None, + client_secret=None, + bearer_token="token", + ) + get_diff.assert_called_once_with( + current_catalog_id="connection-catalog-id", + new_catalog_id="new-catalog-id", + connection_id="connection-id", + api_root=_API_ROOT, + config_api_root=_CONFIG_API_ROOT, + client_id=None, + client_secret=None, + bearer_token="token", + ) + + +def test_compare_source_catalogs_requires_connection_catalog_id( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr( + api_util, + "get_connection_catalog", + MagicMock(return_value={"catalogId": None}), + ) + get_diff = MagicMock() + monkeypatch.setattr(api_util, "get_catalog_diff", get_diff) + + with pytest.raises(AirbyteLibInputError, match="no saved source catalog"): + _workspace().compare_source_catalogs( + "new-catalog-id", + compare_to_connection_id="connection-id", + ) + + get_diff.assert_not_called() + + +def test_source_catalog_summary_without_comparison_is_not_supported() -> None: + with pytest.raises(AirbyteLibInputError, match="not supported yet"): + _workspace().compare_source_catalogs("catalog-id") + + +def test_discover_tool_starts_new_discovery_by_default( + monkeypatch: pytest.MonkeyPatch, +) -> None: + source = MagicMock() + source.start_catalog_discover.return_value = CloudSourceCatalogDiscoverResult( + source_id=_SOURCE_ID, + discover_request_id=_REQUEST_ID, + request_status="pending", + ) + workspace = MagicMock() + workspace.get_source.return_value = source + lookup = MagicMock(return_value=workspace) + monkeypatch.setattr(cloud_mcp, "_get_cloud_workspace", lookup) + + result = asyncio.run(_call_tool(_server(), {"source_id": _SOURCE_ID})) + + assert result.structured_content["request_status"] == "pending" + assert "catalog" not in result.structured_content + workspace.get_source.assert_called_once_with(source_id=_SOURCE_ID) + source.start_catalog_discover.assert_called_once_with() + source.get_latest_saved_catalog.assert_not_called() + + +def test_discover_tool_returns_latest_saved_catalog_when_no_refresh_is_true( + monkeypatch: pytest.MonkeyPatch, +) -> None: + source = MagicMock() + source.get_latest_saved_catalog.return_value = CloudSourceCatalogDiscoverResult( + source_id=_SOURCE_ID, + stream_count=2, + ) + workspace = MagicMock() + workspace.get_source.return_value = source + monkeypatch.setattr( + cloud_mcp, "_get_cloud_workspace", MagicMock(return_value=workspace) + ) + + result = asyncio.run( + _call_tool(_server(), {"source_id": _SOURCE_ID, "no_refresh": True}) + ) + + assert result.structured_content["stream_count"] == 2 + assert "catalog" not in result.structured_content + source.get_latest_saved_catalog.assert_called_once_with() + source.start_catalog_discover.assert_not_called() + + +def test_discover_tool_retrieves_result_by_request_id( + monkeypatch: pytest.MonkeyPatch, +) -> None: + workspace = MagicMock() + workspace.get_catalog_discover_result.return_value = ( + CloudSourceCatalogDiscoverResult( + discover_request_id=_REQUEST_ID, + request_status="running", + ) + ) + lookup = MagicMock(return_value=workspace) + monkeypatch.setattr(cloud_mcp, "_get_cloud_workspace", lookup) + + result = asyncio.run(_call_tool(_server(), {"discover_request_id": _REQUEST_ID})) + + assert result.structured_content["request_status"] == "running" + workspace.get_catalog_discover_result.assert_called_once_with(_REQUEST_ID) + workspace.get_source.assert_not_called() + + +@pytest.mark.parametrize( + ("arguments", "expected_call"), + [ + pytest.param( + { + "source_catalog_id": "new-catalog-id", + "compare_to_source_catalog_id": "old-catalog-id", + }, + { + "source_catalog_id": "new-catalog-id", + "compare_to_source_catalog_id": "old-catalog-id", + "compare_to_connection_id": None, + }, + id="compare-by-source-catalog-id", + ), + pytest.param( + { + "source_catalog_id": "new-catalog-id", + "compare_to_connection_id": "connection-id", + }, + { + "source_catalog_id": "new-catalog-id", + "compare_to_source_catalog_id": None, + "compare_to_connection_id": "connection-id", + }, + id="compare-by-connection-id", + ), + ], +) +def test_discover_tool_compares_catalogs( + monkeypatch: pytest.MonkeyPatch, + arguments: dict[str, Any], + expected_call: dict[str, Any], +) -> None: + workspace = MagicMock() + workspace.compare_source_catalogs.return_value = CloudSourceCatalogDiscoverResult( + source_catalog_id="new-catalog-id", + catalog_diff_summary=[], + schema_change="no_change", + ) + monkeypatch.setattr( + cloud_mcp, "_get_cloud_workspace", MagicMock(return_value=workspace) + ) + + result = asyncio.run(_call_tool(_server(), arguments)) + + assert result.structured_content["source_catalog_id"] == "new-catalog-id" + assert result.structured_content["catalog_diff_summary"] == [] + assert result.structured_content["schema_change"] == "no_change" + workspace.compare_source_catalogs.assert_called_once_with(**expected_call) + + +@pytest.mark.parametrize( + ("arguments", "error_message"), + [ + pytest.param( + {"source_id": _SOURCE_ID, "discover_request_id": _REQUEST_ID}, + "not both", + id="source-and-request", + ), + pytest.param( + {"discover_request_id": _REQUEST_ID, "no_refresh": True}, + "cannot be used", + id="no-refresh-and-request", + ), + pytest.param( + {"source_catalog_id": "catalog-id"}, + "not supported yet", + id="source-catalog-without-comparison", + ), + pytest.param( + {"source_catalog_id": "catalog-id", "source_id": _SOURCE_ID}, + "cannot be combined", + id="catalog-and-source", + ), + pytest.param( + {"source_catalog_id": "catalog-id", "discover_request_id": _REQUEST_ID}, + "cannot be combined", + id="catalog-and-request", + ), + pytest.param( + {"source_catalog_id": "catalog-id", "no_refresh": True}, + "cannot be combined", + id="catalog-and-no-refresh", + ), + pytest.param( + {"compare_to_source_catalog_id": "old-catalog-id"}, + "require `source_catalog_id`", + id="compare-without-source-catalog", + ), + pytest.param( + { + "source_catalog_id": "catalog-id", + "compare_to_source_catalog_id": "old-catalog-id", + "compare_to_connection_id": "connection-id", + }, + "not both", + id="both-comparison-targets", + ), + pytest.param({}, "Provide `source_id` or `discover_request_id`", id="neither"), + ], +) +def test_discover_tool_rejects_invalid_argument_combinations( + monkeypatch: pytest.MonkeyPatch, + arguments: dict[str, Any], + error_message: str, +) -> None: + workspace = MagicMock() + workspace.compare_source_catalogs.side_effect = AirbyteLibInputError( + message="not supported yet" + ) + lookup = MagicMock(return_value=workspace) + monkeypatch.setattr(cloud_mcp, "_get_cloud_workspace", lookup) + + with pytest.raises(ToolError, match=error_message): + asyncio.run(_call_tool(_server(), arguments)) + + if error_message == "not supported yet": + lookup.assert_called_once() + workspace.compare_source_catalogs.assert_called_once_with( + source_catalog_id="catalog-id", + compare_to_source_catalog_id=None, + compare_to_connection_id=None, + ) + else: + lookup.assert_not_called() + + +def test_discover_tool_is_advertised_as_nondestructive() -> None: + server = _server() + tool = asyncio.run(server.get_tool("discover_cloud_source_catalog")) + + assert tool.annotations.destructive_hint is False diff --git a/tests/unit_tests/test_mcp_otel.py b/tests/unit_tests/test_mcp_otel.py index fde65c00c..95b3d23e0 100644 --- a/tests/unit_tests/test_mcp_otel.py +++ b/tests/unit_tests/test_mcp_otel.py @@ -285,6 +285,30 @@ def echo() -> str: with capture_tool_spans() as spans: async with Client(app) as client: await client.call_tool("echo") + safe_routes = { + "safe-discover-command": ( + f"{observability.CLOUD_CONFIG_API_ROOT}/commands/run/discover" + ), + "safe-command-status": ( + f"{observability.CLOUD_CONFIG_API_ROOT}/commands/status" + ), + "safe-discover-output": ( + f"{observability.CLOUD_CONFIG_API_ROOT}/commands/output/discover" + ), + "safe-most-recent-catalog": ( + f"{observability.CLOUD_CONFIG_API_ROOT}/sources/" + "most_recent_source_actor_catalog" + ), + "safe-catalog-diff": ( + f"{observability.CLOUD_CONFIG_API_ROOT}/catalogs/diff" + ), + } + for span_name, route in safe_routes.items(): + with tracer.start_as_current_span( + span_name, kind=SpanKind.CLIENT + ) as span: + span.set_attribute("http.url", route) + span.set_attribute("http.request.method", "POST") with tracer.start_as_current_span( "safe-request", kind=SpanKind.CLIENT ) as span: @@ -306,12 +330,17 @@ def echo() -> str: safe = next(span for span in spans if span.name == "safe-request") unknown = next(span for span in spans if span.name == "unknown-request") + safe_route_spans = { + span.name: span for span in spans if span.name in safe_routes + } safe_attributes = safe.attributes or {} assert safe_attributes["http.url"] == f"{CLOUD_API_ROOT}/applications/token" assert safe_attributes["http.request.method"] == "POST" assert "url.query" not in safe_attributes assert "http.user_agent" not in safe_attributes assert "server.address" not in safe_attributes + for span_name, route in safe_routes.items(): + assert (safe_route_spans[span_name].attributes or {})["http.url"] == route assert (unknown.attributes or {})[ "url.full" ] == observability.REDACTED_PLACEHOLDER