From a2ba515a491f12d14d3631afbb29e1a96c4c345f Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 16 Oct 2026 02:35:30 +0000 Subject: [PATCH 1/7] feat(mcp): add tools to refresh catalog and set stream sync mode on Cloud connections Co-Authored-By: AJ Steers --- airbyte/_util/api_util.py | 36 ++ airbyte/cloud/connections.py | 159 ++++++++ airbyte/mcp/cloud.py | 120 ++++++ .../test_cloud_connection_catalog_edits.py | 382 ++++++++++++++++++ 4 files changed, 697 insertions(+) create mode 100644 tests/unit_tests/test_cloud_connection_catalog_edits.py diff --git a/airbyte/_util/api_util.py b/airbyte/_util/api_util.py index 2be3b717f..d8d403eca 100644 --- a/airbyte/_util/api_util.py +++ b/airbyte/_util/api_util.py @@ -3576,6 +3576,42 @@ def replace_connection_state( raise +def get_refreshed_connection_catalog( + connection_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 a refreshed catalog for a connection. + + Uses the Config API endpoint: POST /v1/web_backend/connections/get + with `withRefreshedCatalog` enabled. + + Args: + connection_id: The connection ID to get a refreshed catalog for. + api_root: The API root URL. + client_id: OAuth client ID. + client_secret: OAuth client secret. + bearer_token: Bearer token for authentication (alternative to client credentials). + config_api_root: Optional explicit Config API root URL. + + Returns: + Dictionary containing the connection info with refreshed syncCatalog. + """ + return _make_config_api_request( + path="/web_backend/connections/get", + json={"connectionId": connection_id, "withRefreshedCatalog": True}, + api_root=api_root, + config_api_root=config_api_root, + client_id=client_id, + client_secret=client_secret, + bearer_token=bearer_token, + ) + + def get_connection_catalog( connection_id: str, *, diff --git a/airbyte/cloud/connections.py b/airbyte/cloud/connections.py index c296108cd..5eee92133 100644 --- a/airbyte/cloud/connections.py +++ b/airbyte/cloud/connections.py @@ -7,6 +7,8 @@ from http import HTTPStatus from typing import TYPE_CHECKING, Any, Literal, overload +from airbyte_protocol.models import DestinationSyncMode, SyncMode + from airbyte._util import api_util from airbyte.cloud._connection_catalog import ( _denormalize_catalog_to_api, @@ -903,6 +905,163 @@ def import_raw_catalog(self, catalog: dict[str, Any]) -> None: config_api_root=self.workspace.config_api_root, ) + def refresh_catalog(self) -> dict[str, Any]: + """Refresh the connection's catalog by re-discovering the source schema. + + Triggers a discover operation on the connection's source connector, then + replaces the connection's catalog with the refreshed result. This is + equivalent to clicking "Refresh source schema" in the Airbyte UI. + + Returns: + The refreshed raw syncCatalog dict. + + Raises: + AirbyteLibInputError: If the refreshed catalog is empty or missing. + """ + refreshed_response = api_util.get_refreshed_connection_catalog( + connection_id=self.connection_id, + api_root=self.workspace.api_root, + client_id=self.workspace.client_id, + client_secret=self.workspace.client_secret, + bearer_token=self.workspace.bearer_token, + config_api_root=self.workspace.config_api_root, + ) + refreshed_catalog = refreshed_response.get("syncCatalog") + if not isinstance(refreshed_catalog, dict) or not refreshed_catalog.get("streams"): + raise AirbyteLibInputError( + message="Refreshed catalog is empty or missing.", + context={"connection_id": self.connection_id}, + ) + + self.import_raw_catalog(refreshed_catalog) + return refreshed_catalog + + def set_stream_sync_mode( + self, + stream_name: str, + *, + sync_mode: SyncMode, + stream_namespace: str | None = None, + destination_sync_mode: DestinationSyncMode | None = None, + cursor_field: str | None = None, + ) -> None: + """Set the sync mode for one stream on this connection. + + Only the requested stream's config is changed. A `cursor_field` can also + be set with `full_refresh` to pre-configure it for a later incremental sync. + + Args: + stream_name: The name of the stream to modify. + sync_mode: The source sync mode to set. + stream_namespace: The namespace to select when stream names are ambiguous. + destination_sync_mode: The destination sync mode to set, if any. + cursor_field: The cursor field to set. Required for incremental mode if the + stream has no existing cursor, default cursor, or source-defined cursor. + + Raises: + AirbyteLibInputError: If the catalog or stream is missing, the stream name is + ambiguous, the sync mode is unsupported, or incremental mode has no cursor. + """ + catalog = self.dump_raw_catalog(normalize=False) + if ( + not isinstance(catalog, dict) + or not isinstance(catalog.get("streams"), list) + or not catalog["streams"] + ): + raise AirbyteLibInputError( + message="Connection catalog is empty or missing.", + context={"connection_id": self.connection_id}, + ) + + matching_entries: list[dict[str, Any]] = [] + for entry in catalog["streams"]: + stream_def = entry.get("stream", {}) + if stream_def.get("name") != stream_name: + continue + if stream_namespace is not None and stream_def.get("namespace") != stream_namespace: + continue + matching_entries.append(entry) + + if not matching_entries: + available_streams = [] + for entry in catalog["streams"]: + stream_def = entry.get("stream", {}) + name = stream_def.get("name") + namespace = stream_def.get("namespace") + available_streams.append(f"{namespace}.{name}" if namespace else name) + raise AirbyteLibInputError( + message=f"Stream '{stream_name}' not found in connection catalog.", + guidance="Use a stream name from `available_streams`.", + context={ + "connection_id": self.connection_id, + "stream_name": stream_name, + "stream_namespace": stream_namespace, + "available_streams": available_streams, + }, + ) + + if len(matching_entries) > 1: + matching_namespaces = [ + entry.get("stream", {}).get("namespace") for entry in matching_entries + ] + raise AirbyteLibInputError( + message=f"Stream name '{stream_name}' is ambiguous.", + guidance="Specify `stream_namespace` to select one of the matching streams.", + context={ + "connection_id": self.connection_id, + "stream_name": stream_name, + "matching_namespaces": matching_namespaces, + }, + ) + + target_entry = matching_entries[0] + stream_def = target_entry.get("stream", {}) + supported_sync_modes: list[str] = stream_def.get("supportedSyncModes", []) + if sync_mode.value not in supported_sync_modes: + raise AirbyteLibInputError( + message=( + f"Sync mode '{sync_mode.value}' is not supported by stream '{stream_name}'." + ), + guidance="Choose a mode listed in `supported_sync_modes`.", + context={ + "connection_id": self.connection_id, + "stream_name": stream_name, + "stream_namespace": stream_namespace, + "requested_sync_mode": sync_mode.value, + "supported_sync_modes": supported_sync_modes, + }, + ) + + config = target_entry.get("config") or {} + if sync_mode is SyncMode.incremental: + existing_cursor = config.get("cursorField") or stream_def.get("defaultCursorField") + source_defined_cursor = bool(stream_def.get("sourceDefinedCursor")) + if cursor_field is None and not existing_cursor and not source_defined_cursor: + raise AirbyteLibInputError( + message=( + f"Stream '{stream_name}' needs a cursor field before switching " + "to incremental sync mode." + ), + guidance=( + "Provide `cursor_field` or use a stream with a default or " + "source-defined cursor." + ), + context={ + "connection_id": self.connection_id, + "stream_name": stream_name, + "stream_namespace": stream_namespace, + }, + ) + + config["syncMode"] = sync_mode.value + if destination_sync_mode is not None: + config["destinationSyncMode"] = destination_sync_mode.value + if cursor_field is not None: + config["cursorField"] = [cursor_field] + target_entry["config"] = config + + self.import_raw_catalog(catalog) + def rename(self, name: str) -> CloudConnection: """Rename the connection. diff --git a/airbyte/mcp/cloud.py b/airbyte/mcp/cloud.py index 26013217b..0925464e4 100644 --- a/airbyte/mcp/cloud.py +++ b/airbyte/mcp/cloud.py @@ -19,6 +19,8 @@ from fastmcp_extensions import TraceArg, get_mcp_config, register_mcp_tools from pydantic import BaseModel, ConfigDict, Field +from airbyte_protocol.models import DestinationSyncMode, SyncMode + from airbyte import Destination, Source, get_destination, get_source from airbyte._direct_connectors import connector_docs from airbyte._direct_connectors.models import ( @@ -3971,6 +3973,124 @@ def set_cloud_connection_selected_streams( ) +@mcp_tool( + destructive=True, +) +def refresh_cloud_connection_catalog( + ctx: Context, + connection_id: Annotated[ + str, + Field(description="The ID of the connection to refresh the catalog for."), + ], + *, + workspace_id: Annotated[ + str | None, + Field( + description=WORKSPACE_ID_TIP_TEXT, + default=None, + ), + TraceArg.OMIT, + ], +) -> str: + """Refresh a connection's catalog by re-discovering its source schema.""" + check_guid_created_in_session(connection_id) + workspace: CloudWorkspace = _get_cloud_workspace(ctx, workspace_id) + connection = workspace.get_connection(connection_id=connection_id) + refreshed_catalog = connection.refresh_catalog() + stream_count = len(refreshed_catalog["streams"]) + + return ( + f"Successfully refreshed catalog for connection '{connection_id}'. " + f"Catalog now contains {stream_count} stream(s). " + f"URL: {connection.connection_url}" + ) + + +@mcp_tool( + destructive=True, +) +def set_cloud_connection_stream_sync_mode( + ctx: Context, + connection_id: Annotated[ + str, + Field(description="The ID of the connection to update."), + ], + stream_name: Annotated[ + str, + Field(description="The name of the stream to update."), + TraceArg.FINGERPRINT, + ], + sync_mode: Annotated[ + SyncMode, + Field(description="The source sync mode to set."), + ], + *, + stream_namespace: Annotated[ + str | None, + Field( + description="The namespace of the stream, required when its name is ambiguous.", + default=None, + ), + ], + destination_sync_mode: Annotated[ + DestinationSyncMode | None, + Field( + description=( + "The destination sync mode to set. If omitted, the current mode is preserved." + ), + default=None, + ), + ], + cursor_field: Annotated[ + str | None, + Field( + description=( + "The cursor field to set. Required for incremental mode when no existing, default, " + "or source-defined cursor is available. It can also be set during full-refresh " + "syncs to pre-configure a later incremental sync." + ), + default=None, + ), + TraceArg.FINGERPRINT, + ], + workspace_id: Annotated[ + str | None, + Field( + description=WORKSPACE_ID_TIP_TEXT, + default=None, + ), + TraceArg.OMIT, + ], +) -> str: + """Set the sync and destination modes for one stream in a connection.""" + check_guid_created_in_session(connection_id) + workspace: CloudWorkspace = _get_cloud_workspace(ctx, workspace_id) + connection = workspace.get_connection(connection_id=connection_id) + connection.set_stream_sync_mode( + stream_name=stream_name, + sync_mode=sync_mode, + stream_namespace=stream_namespace, + destination_sync_mode=destination_sync_mode, + cursor_field=cursor_field, + ) + + stream_identifier = ( + f"{stream_namespace}.{stream_name}" if stream_namespace is not None else stream_name + ) + updates = [] + if destination_sync_mode is not None: + updates.append(f"destination sync mode '{destination_sync_mode.value}'") + if cursor_field is not None: + updates.append(f"cursor field '{cursor_field}'") + update_details = f" with {', '.join(updates)}" if updates else "" + + return ( + f"Successfully set sync mode for stream '{stream_identifier}' " + f"on connection '{connection_id}' to '{sync_mode.value}'{update_details}. " + f"URL: {connection.connection_url}" + ) + + @mcp_tool( destructive=True, ) 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..369359040 --- /dev/null +++ b/tests/unit_tests/test_cloud_connection_catalog_edits.py @@ -0,0 +1,382 @@ +# Copyright (c) 2026 Airbyte, Inc., all rights reserved. +"""Unit tests for editing and refreshing Cloud connection catalogs.""" + +from __future__ import annotations + +from copy import deepcopy +from typing import Any +from unittest.mock import MagicMock + +import pytest +from airbyte_protocol.models import DestinationSyncMode, SyncMode + +from airbyte._util import api_util +from airbyte.cloud.connections import CloudConnection +from airbyte.cloud.workspaces import CloudWorkspace +from airbyte.exceptions import AirbyteLibInputError + + +_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, + sync_modes: list[str] | None = None, + source_defined_cursor: bool = False, + default_cursor_field: list[str] | None = None, + config: dict[str, Any] | None = None, +) -> dict[str, Any]: + stream: dict[str, Any] = { + "name": name, + "jsonSchema": {"type": "object"}, + "supportedSyncModes": sync_modes or ["full_refresh", "incremental"], + } + if namespace is not None: + stream["namespace"] = namespace + if source_defined_cursor: + stream["sourceDefinedCursor"] = True + if default_cursor_field is not None: + stream["defaultCursorField"] = default_cursor_field + return { + "stream": stream, + "config": config + or { + "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 _mock_catalog_api( + monkeypatch: pytest.MonkeyPatch, + catalog: dict[str, Any], +) -> tuple[MagicMock, MagicMock]: + get_catalog = MagicMock(return_value={"syncCatalog": catalog}) + replace_catalog = MagicMock(return_value={}) + monkeypatch.setattr(api_util, "get_connection_catalog", get_catalog) + monkeypatch.setattr(api_util, "replace_connection_catalog", replace_catalog) + return get_catalog, replace_catalog + + +def test_set_stream_sync_mode_updates_only_target_and_saves_raw_catalog( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Update the target stream without normalizing or changing other catalog data.""" + target = _stream_entry("users", namespace="public") + other = _stream_entry( + "orders", + namespace="sales", + config={ + "syncMode": "incremental", + "destinationSyncMode": "append_dedup", + "cursorField": ["changed_at"], + "selectedFields": [{"fieldPath": ["id"]}], + "fieldSelectionEnabled": True, + "primaryKey": [["order_id"]], + "selected": False, + }, + ) + original_other = deepcopy(other) + catalog = _catalog(target, other) + get_catalog, replace_catalog = _mock_catalog_api(monkeypatch, catalog) + + _connection().set_stream_sync_mode( + stream_name="users", + stream_namespace="public", + sync_mode=SyncMode.incremental, + destination_sync_mode=DestinationSyncMode.append_dedup, + cursor_field="updated_at", + ) + + get_catalog.assert_called_once_with( + connection_id=_CONNECTION_ID, + api_root=_API_ROOT, + client_id=None, + client_secret=None, + bearer_token="token", + config_api_root=_CONFIG_API_ROOT, + ) + assert target["config"] == { + "syncMode": "incremental", + "destinationSyncMode": "append_dedup", + "cursorField": ["updated_at"], + "selectedFields": [], + "fieldSelectionEnabled": False, + "primaryKey": [["id"]], + "selected": True, + } + assert other == original_other + + replace_catalog.assert_called_once() + saved_catalog = replace_catalog.call_args.kwargs["configured_catalog_dict"] + assert saved_catalog == catalog + assert saved_catalog["streams"][1] == original_other + assert "sync_mode" not in saved_catalog["streams"][0]["config"] + assert "supported_sync_modes" not in saved_catalog["streams"][0]["stream"] + assert replace_catalog.call_args.kwargs["config_api_root"] == _CONFIG_API_ROOT + + +@pytest.mark.parametrize( + ( + "catalog", + "stream_name", + "stream_namespace", + "sync_mode", + "cursor_field", + "error_message", + "expected_context", + ), + [ + pytest.param( + _catalog(_stream_entry("users", namespace="public")), + "missing", + None, + SyncMode.full_refresh, + None, + "not found", + {"available_streams": ["public.users"]}, + id="missing-stream", + ), + pytest.param( + _catalog( + _stream_entry("users", namespace="public"), + _stream_entry("users", namespace="sales"), + ), + "users", + None, + SyncMode.full_refresh, + None, + "ambiguous", + {"matching_namespaces": ["public", "sales"]}, + id="ambiguous-name", + ), + pytest.param( + _catalog( + _stream_entry("users", namespace="public", sync_modes=["full_refresh"]) + ), + "users", + "public", + SyncMode.incremental, + None, + "not supported", + {"requested_sync_mode": "incremental"}, + id="unsupported-sync-mode", + ), + pytest.param( + _catalog( + _stream_entry( + "users", + namespace="public", + config={ + "syncMode": "full_refresh", + "destinationSyncMode": "append", + }, + ) + ), + "users", + "public", + SyncMode.incremental, + None, + "needs a cursor", + {"stream_namespace": "public"}, + id="incremental-without-cursor", + ), + ], +) +def test_set_stream_sync_mode_rejects_invalid_updates( + monkeypatch: pytest.MonkeyPatch, + catalog: dict[str, Any], + stream_name: str, + stream_namespace: str | None, + sync_mode: SyncMode, + cursor_field: str | None, + error_message: str, + expected_context: dict[str, Any], +) -> None: + """Input errors leave the Config API catalog untouched.""" + _, replace_catalog = _mock_catalog_api(monkeypatch, deepcopy(catalog)) + + with pytest.raises(AirbyteLibInputError, match=error_message) as exc_info: + _connection().set_stream_sync_mode( + stream_name=stream_name, + stream_namespace=stream_namespace, + sync_mode=sync_mode, + cursor_field=cursor_field, + ) + + assert exc_info.value.context is not None + for key, value in expected_context.items(): + assert exc_info.value.context[key] == value + replace_catalog.assert_not_called() + + +def test_set_stream_sync_mode_namespace_disambiguates_name( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A namespace selects the intended same-named stream.""" + public = _stream_entry( + "users", + namespace="public", + config={ + "syncMode": "incremental", + "destinationSyncMode": "append", + "cursorField": ["updated_at"], + }, + ) + sales = _stream_entry("users", namespace="sales") + catalog = _catalog(public, sales) + _, replace_catalog = _mock_catalog_api(monkeypatch, catalog) + + _connection().set_stream_sync_mode( + stream_name="users", + stream_namespace="sales", + sync_mode=SyncMode.full_refresh, + ) + + assert public["config"]["syncMode"] == "incremental" + assert sales["config"]["syncMode"] == "full_refresh" + replace_catalog.assert_called_once() + + +@pytest.mark.parametrize( + "stream_options", + [ + pytest.param({"source_defined_cursor": True}, id="source-defined-cursor"), + pytest.param( + {"default_cursor_field": ["updated_at"]}, id="default-cursor-field" + ), + ], +) +def test_incremental_sync_accepts_stream_defined_cursor( + monkeypatch: pytest.MonkeyPatch, + stream_options: dict[str, Any], +) -> None: + """Incremental sync can use a source-defined or default cursor without an override.""" + entry = _stream_entry( + "users", + namespace="public", + config={"syncMode": "full_refresh", "destinationSyncMode": "append"}, + **stream_options, + ) + _, replace_catalog = _mock_catalog_api(monkeypatch, _catalog(entry)) + + _connection().set_stream_sync_mode( + stream_name="users", + stream_namespace="public", + sync_mode=SyncMode.incremental, + ) + + assert entry["config"]["syncMode"] == "incremental" + assert "cursorField" not in entry["config"] + replace_catalog.assert_called_once() + + +def test_refresh_catalog_gets_and_replaces_catalog_with_config_api_root( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Refresh the source catalog and persist the raw result.""" + refreshed_catalog = _catalog(_stream_entry("users", namespace="public")) + get_refreshed = MagicMock(return_value={"syncCatalog": refreshed_catalog}) + replace_catalog = MagicMock(return_value={}) + monkeypatch.setattr(api_util, "get_refreshed_connection_catalog", get_refreshed) + monkeypatch.setattr(api_util, "replace_connection_catalog", replace_catalog) + + result = _connection().refresh_catalog() + + assert result == refreshed_catalog + get_refreshed.assert_called_once_with( + connection_id=_CONNECTION_ID, + api_root=_API_ROOT, + client_id=None, + client_secret=None, + bearer_token="token", + config_api_root=_CONFIG_API_ROOT, + ) + replace_catalog.assert_called_once_with( + connection_id=_CONNECTION_ID, + configured_catalog_dict=refreshed_catalog, + api_root=_API_ROOT, + client_id=None, + client_secret=None, + bearer_token="token", + config_api_root=_CONFIG_API_ROOT, + ) + + +@pytest.mark.parametrize( + "response", + [ + pytest.param({}, id="missing-sync-catalog"), + pytest.param({"syncCatalog": None}, id="null-sync-catalog"), + pytest.param({"syncCatalog": {"streams": []}}, id="empty-sync-catalog"), + ], +) +def test_refresh_catalog_rejects_empty_or_missing_catalog( + monkeypatch: pytest.MonkeyPatch, + response: dict[str, Any], +) -> None: + """An absent or empty refreshed catalog is never written.""" + get_refreshed = MagicMock(return_value=response) + replace_catalog = MagicMock(return_value={}) + monkeypatch.setattr(api_util, "get_refreshed_connection_catalog", get_refreshed) + monkeypatch.setattr(api_util, "replace_connection_catalog", replace_catalog) + + with pytest.raises( + AirbyteLibInputError, match="Refreshed catalog is empty or missing" + ): + _connection().refresh_catalog() + + replace_catalog.assert_not_called() + + +def test_get_refreshed_connection_catalog_forwards_config_api_root( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Request discovery through the Config API catalog endpoint.""" + request = MagicMock(return_value={"syncCatalog": {"streams": []}}) + monkeypatch.setattr(api_util, "_make_config_api_request", request) + + result = api_util.get_refreshed_connection_catalog( + connection_id=_CONNECTION_ID, + api_root=_API_ROOT, + client_id=None, + client_secret=None, + bearer_token="token", + config_api_root=_CONFIG_API_ROOT, + ) + + assert result == {"syncCatalog": {"streams": []}} + request.assert_called_once_with( + path="/web_backend/connections/get", + json={"connectionId": _CONNECTION_ID, "withRefreshedCatalog": True}, + api_root=_API_ROOT, + config_api_root=_CONFIG_API_ROOT, + client_id=None, + client_secret=None, + bearer_token="token", + ) From ce6518bac53778a7de00f68ad69cfdbde1596050 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 16 Oct 2026 02:44:49 +0000 Subject: [PATCH 2/7] fix(mcp): update trace fixture and default cursor Co-Authored-By: AJ Steers --- airbyte/cloud/connections.py | 36 ++++++++++--------- tests/unit_tests/fixtures/arg_trace_plan.json | 13 +++++++ .../test_cloud_connection_catalog_edits.py | 5 ++- 3 files changed, 37 insertions(+), 17 deletions(-) diff --git a/airbyte/cloud/connections.py b/airbyte/cloud/connections.py index 5eee92133..95e1bd75b 100644 --- a/airbyte/cloud/connections.py +++ b/airbyte/cloud/connections.py @@ -1034,24 +1034,28 @@ def set_stream_sync_mode( config = target_entry.get("config") or {} if sync_mode is SyncMode.incremental: - existing_cursor = config.get("cursorField") or stream_def.get("defaultCursorField") + existing_cursor = config.get("cursorField") source_defined_cursor = bool(stream_def.get("sourceDefinedCursor")) + default_cursor_field = stream_def.get("defaultCursorField") if cursor_field is None and not existing_cursor and not source_defined_cursor: - raise AirbyteLibInputError( - message=( - f"Stream '{stream_name}' needs a cursor field before switching " - "to incremental sync mode." - ), - guidance=( - "Provide `cursor_field` or use a stream with a default or " - "source-defined cursor." - ), - context={ - "connection_id": self.connection_id, - "stream_name": stream_name, - "stream_namespace": stream_namespace, - }, - ) + if default_cursor_field: + config["cursorField"] = list(default_cursor_field) + else: + raise AirbyteLibInputError( + message=( + f"Stream '{stream_name}' needs a cursor field before switching " + "to incremental sync mode." + ), + guidance=( + "Provide `cursor_field` or use a stream with a default or " + "source-defined cursor." + ), + context={ + "connection_id": self.connection_id, + "stream_name": stream_name, + "stream_namespace": stream_namespace, + }, + ) config["syncMode"] = sync_mode.value if destination_sync_mode is not None: diff --git a/tests/unit_tests/fixtures/arg_trace_plan.json b/tests/unit_tests/fixtures/arg_trace_plan.json index 23b989e52..e886071b2 100644 --- a/tests/unit_tests/fixtures/arg_trace_plan.json +++ b/tests/unit_tests/fixtures/arg_trace_plan.json @@ -271,6 +271,10 @@ "source_connector_name": "fingerprint", "stream_name": "fingerprint" }, + "refresh_cloud_connection_catalog": { + "connection_id": "hash", + "workspace_id": "omit" + }, "rename_cloud_connection": { "connection_id": "hash", "name": "fingerprint", @@ -300,6 +304,15 @@ "stream_names": "fingerprint list", "workspace_id": "omit" }, + "set_cloud_connection_stream_sync_mode": { + "connection_id": "hash", + "cursor_field": "fingerprint", + "destination_sync_mode": "value allowed=[\"append\",\"append_dedup\",\"overwrite\",\"soft_delete\",\"update\"]", + "stream_name": "fingerprint", + "stream_namespace": "hash", + "sync_mode": "value allowed=[\"full_refresh\",\"incremental\"]", + "workspace_id": "omit" + }, "set_cloud_connection_table_prefix": { "connection_id": "hash", "prefix": "fingerprint", diff --git a/tests/unit_tests/test_cloud_connection_catalog_edits.py b/tests/unit_tests/test_cloud_connection_catalog_edits.py index 369359040..861890e23 100644 --- a/tests/unit_tests/test_cloud_connection_catalog_edits.py +++ b/tests/unit_tests/test_cloud_connection_catalog_edits.py @@ -292,7 +292,10 @@ def test_incremental_sync_accepts_stream_defined_cursor( ) assert entry["config"]["syncMode"] == "incremental" - assert "cursorField" not in entry["config"] + if "default_cursor_field" in stream_options: + assert entry["config"]["cursorField"] == ["updated_at"] + else: + assert "cursorField" not in entry["config"] replace_catalog.assert_called_once() From 5acdc20f4f01202629f2f20f236b56754098cf72 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 16 Oct 2026 19:38:24 +0000 Subject: [PATCH 3/7] refactor: move set_stream_sync_mode to a stacked PR Co-Authored-By: AJ Steers --- airbyte/cloud/connections.py | 132 ---------- airbyte/mcp/cloud.py | 87 ------- tests/unit_tests/fixtures/arg_trace_plan.json | 9 - .../test_cloud_connection_catalog_edits.py | 242 +----------------- 4 files changed, 2 insertions(+), 468 deletions(-) diff --git a/airbyte/cloud/connections.py b/airbyte/cloud/connections.py index 95e1bd75b..f6a9e9ec2 100644 --- a/airbyte/cloud/connections.py +++ b/airbyte/cloud/connections.py @@ -7,8 +7,6 @@ from http import HTTPStatus from typing import TYPE_CHECKING, Any, Literal, overload -from airbyte_protocol.models import DestinationSyncMode, SyncMode - from airbyte._util import api_util from airbyte.cloud._connection_catalog import ( _denormalize_catalog_to_api, @@ -936,136 +934,6 @@ def refresh_catalog(self) -> dict[str, Any]: self.import_raw_catalog(refreshed_catalog) return refreshed_catalog - def set_stream_sync_mode( - self, - stream_name: str, - *, - sync_mode: SyncMode, - stream_namespace: str | None = None, - destination_sync_mode: DestinationSyncMode | None = None, - cursor_field: str | None = None, - ) -> None: - """Set the sync mode for one stream on this connection. - - Only the requested stream's config is changed. A `cursor_field` can also - be set with `full_refresh` to pre-configure it for a later incremental sync. - - Args: - stream_name: The name of the stream to modify. - sync_mode: The source sync mode to set. - stream_namespace: The namespace to select when stream names are ambiguous. - destination_sync_mode: The destination sync mode to set, if any. - cursor_field: The cursor field to set. Required for incremental mode if the - stream has no existing cursor, default cursor, or source-defined cursor. - - Raises: - AirbyteLibInputError: If the catalog or stream is missing, the stream name is - ambiguous, the sync mode is unsupported, or incremental mode has no cursor. - """ - catalog = self.dump_raw_catalog(normalize=False) - if ( - not isinstance(catalog, dict) - or not isinstance(catalog.get("streams"), list) - or not catalog["streams"] - ): - raise AirbyteLibInputError( - message="Connection catalog is empty or missing.", - context={"connection_id": self.connection_id}, - ) - - matching_entries: list[dict[str, Any]] = [] - for entry in catalog["streams"]: - stream_def = entry.get("stream", {}) - if stream_def.get("name") != stream_name: - continue - if stream_namespace is not None and stream_def.get("namespace") != stream_namespace: - continue - matching_entries.append(entry) - - if not matching_entries: - available_streams = [] - for entry in catalog["streams"]: - stream_def = entry.get("stream", {}) - name = stream_def.get("name") - namespace = stream_def.get("namespace") - available_streams.append(f"{namespace}.{name}" if namespace else name) - raise AirbyteLibInputError( - message=f"Stream '{stream_name}' not found in connection catalog.", - guidance="Use a stream name from `available_streams`.", - context={ - "connection_id": self.connection_id, - "stream_name": stream_name, - "stream_namespace": stream_namespace, - "available_streams": available_streams, - }, - ) - - if len(matching_entries) > 1: - matching_namespaces = [ - entry.get("stream", {}).get("namespace") for entry in matching_entries - ] - raise AirbyteLibInputError( - message=f"Stream name '{stream_name}' is ambiguous.", - guidance="Specify `stream_namespace` to select one of the matching streams.", - context={ - "connection_id": self.connection_id, - "stream_name": stream_name, - "matching_namespaces": matching_namespaces, - }, - ) - - target_entry = matching_entries[0] - stream_def = target_entry.get("stream", {}) - supported_sync_modes: list[str] = stream_def.get("supportedSyncModes", []) - if sync_mode.value not in supported_sync_modes: - raise AirbyteLibInputError( - message=( - f"Sync mode '{sync_mode.value}' is not supported by stream '{stream_name}'." - ), - guidance="Choose a mode listed in `supported_sync_modes`.", - context={ - "connection_id": self.connection_id, - "stream_name": stream_name, - "stream_namespace": stream_namespace, - "requested_sync_mode": sync_mode.value, - "supported_sync_modes": supported_sync_modes, - }, - ) - - config = target_entry.get("config") or {} - if sync_mode is SyncMode.incremental: - existing_cursor = config.get("cursorField") - source_defined_cursor = bool(stream_def.get("sourceDefinedCursor")) - default_cursor_field = stream_def.get("defaultCursorField") - if cursor_field is None and not existing_cursor and not source_defined_cursor: - if default_cursor_field: - config["cursorField"] = list(default_cursor_field) - else: - raise AirbyteLibInputError( - message=( - f"Stream '{stream_name}' needs a cursor field before switching " - "to incremental sync mode." - ), - guidance=( - "Provide `cursor_field` or use a stream with a default or " - "source-defined cursor." - ), - context={ - "connection_id": self.connection_id, - "stream_name": stream_name, - "stream_namespace": stream_namespace, - }, - ) - - config["syncMode"] = sync_mode.value - if destination_sync_mode is not None: - config["destinationSyncMode"] = destination_sync_mode.value - if cursor_field is not None: - config["cursorField"] = [cursor_field] - target_entry["config"] = config - - self.import_raw_catalog(catalog) - def rename(self, name: str) -> CloudConnection: """Rename the connection. diff --git a/airbyte/mcp/cloud.py b/airbyte/mcp/cloud.py index 0925464e4..f74a3240e 100644 --- a/airbyte/mcp/cloud.py +++ b/airbyte/mcp/cloud.py @@ -19,8 +19,6 @@ from fastmcp_extensions import TraceArg, get_mcp_config, register_mcp_tools from pydantic import BaseModel, ConfigDict, Field -from airbyte_protocol.models import DestinationSyncMode, SyncMode - from airbyte import Destination, Source, get_destination, get_source from airbyte._direct_connectors import connector_docs from airbyte._direct_connectors.models import ( @@ -4006,91 +4004,6 @@ def refresh_cloud_connection_catalog( ) -@mcp_tool( - destructive=True, -) -def set_cloud_connection_stream_sync_mode( - ctx: Context, - connection_id: Annotated[ - str, - Field(description="The ID of the connection to update."), - ], - stream_name: Annotated[ - str, - Field(description="The name of the stream to update."), - TraceArg.FINGERPRINT, - ], - sync_mode: Annotated[ - SyncMode, - Field(description="The source sync mode to set."), - ], - *, - stream_namespace: Annotated[ - str | None, - Field( - description="The namespace of the stream, required when its name is ambiguous.", - default=None, - ), - ], - destination_sync_mode: Annotated[ - DestinationSyncMode | None, - Field( - description=( - "The destination sync mode to set. If omitted, the current mode is preserved." - ), - default=None, - ), - ], - cursor_field: Annotated[ - str | None, - Field( - description=( - "The cursor field to set. Required for incremental mode when no existing, default, " - "or source-defined cursor is available. It can also be set during full-refresh " - "syncs to pre-configure a later incremental sync." - ), - default=None, - ), - TraceArg.FINGERPRINT, - ], - workspace_id: Annotated[ - str | None, - Field( - description=WORKSPACE_ID_TIP_TEXT, - default=None, - ), - TraceArg.OMIT, - ], -) -> str: - """Set the sync and destination modes for one stream in a connection.""" - check_guid_created_in_session(connection_id) - workspace: CloudWorkspace = _get_cloud_workspace(ctx, workspace_id) - connection = workspace.get_connection(connection_id=connection_id) - connection.set_stream_sync_mode( - stream_name=stream_name, - sync_mode=sync_mode, - stream_namespace=stream_namespace, - destination_sync_mode=destination_sync_mode, - cursor_field=cursor_field, - ) - - stream_identifier = ( - f"{stream_namespace}.{stream_name}" if stream_namespace is not None else stream_name - ) - updates = [] - if destination_sync_mode is not None: - updates.append(f"destination sync mode '{destination_sync_mode.value}'") - if cursor_field is not None: - updates.append(f"cursor field '{cursor_field}'") - update_details = f" with {', '.join(updates)}" if updates else "" - - return ( - f"Successfully set sync mode for stream '{stream_identifier}' " - f"on connection '{connection_id}' to '{sync_mode.value}'{update_details}. " - f"URL: {connection.connection_url}" - ) - - @mcp_tool( destructive=True, ) diff --git a/tests/unit_tests/fixtures/arg_trace_plan.json b/tests/unit_tests/fixtures/arg_trace_plan.json index e886071b2..1b6df0a72 100644 --- a/tests/unit_tests/fixtures/arg_trace_plan.json +++ b/tests/unit_tests/fixtures/arg_trace_plan.json @@ -304,15 +304,6 @@ "stream_names": "fingerprint list", "workspace_id": "omit" }, - "set_cloud_connection_stream_sync_mode": { - "connection_id": "hash", - "cursor_field": "fingerprint", - "destination_sync_mode": "value allowed=[\"append\",\"append_dedup\",\"overwrite\",\"soft_delete\",\"update\"]", - "stream_name": "fingerprint", - "stream_namespace": "hash", - "sync_mode": "value allowed=[\"full_refresh\",\"incremental\"]", - "workspace_id": "omit" - }, "set_cloud_connection_table_prefix": { "connection_id": "hash", "prefix": "fingerprint", diff --git a/tests/unit_tests/test_cloud_connection_catalog_edits.py b/tests/unit_tests/test_cloud_connection_catalog_edits.py index 861890e23..0833c0710 100644 --- a/tests/unit_tests/test_cloud_connection_catalog_edits.py +++ b/tests/unit_tests/test_cloud_connection_catalog_edits.py @@ -3,12 +3,10 @@ from __future__ import annotations -from copy import deepcopy from typing import Any from unittest.mock import MagicMock import pytest -from airbyte_protocol.models import DestinationSyncMode, SyncMode from airbyte._util import api_util from airbyte.cloud.connections import CloudConnection @@ -38,26 +36,17 @@ def _stream_entry( name: str, *, namespace: str | None = None, - sync_modes: list[str] | None = None, - source_defined_cursor: bool = False, - default_cursor_field: list[str] | None = None, - config: dict[str, Any] | None = None, ) -> dict[str, Any]: stream: dict[str, Any] = { "name": name, "jsonSchema": {"type": "object"}, - "supportedSyncModes": sync_modes or ["full_refresh", "incremental"], + "supportedSyncModes": ["full_refresh", "incremental"], } if namespace is not None: stream["namespace"] = namespace - if source_defined_cursor: - stream["sourceDefinedCursor"] = True - if default_cursor_field is not None: - stream["defaultCursorField"] = default_cursor_field return { "stream": stream, - "config": config - or { + "config": { "syncMode": "full_refresh", "destinationSyncMode": "append", "selectedFields": [], @@ -72,233 +61,6 @@ def _catalog(*entries: dict[str, Any]) -> dict[str, Any]: return {"streams": list(entries)} -def _mock_catalog_api( - monkeypatch: pytest.MonkeyPatch, - catalog: dict[str, Any], -) -> tuple[MagicMock, MagicMock]: - get_catalog = MagicMock(return_value={"syncCatalog": catalog}) - replace_catalog = MagicMock(return_value={}) - monkeypatch.setattr(api_util, "get_connection_catalog", get_catalog) - monkeypatch.setattr(api_util, "replace_connection_catalog", replace_catalog) - return get_catalog, replace_catalog - - -def test_set_stream_sync_mode_updates_only_target_and_saves_raw_catalog( - monkeypatch: pytest.MonkeyPatch, -) -> None: - """Update the target stream without normalizing or changing other catalog data.""" - target = _stream_entry("users", namespace="public") - other = _stream_entry( - "orders", - namespace="sales", - config={ - "syncMode": "incremental", - "destinationSyncMode": "append_dedup", - "cursorField": ["changed_at"], - "selectedFields": [{"fieldPath": ["id"]}], - "fieldSelectionEnabled": True, - "primaryKey": [["order_id"]], - "selected": False, - }, - ) - original_other = deepcopy(other) - catalog = _catalog(target, other) - get_catalog, replace_catalog = _mock_catalog_api(monkeypatch, catalog) - - _connection().set_stream_sync_mode( - stream_name="users", - stream_namespace="public", - sync_mode=SyncMode.incremental, - destination_sync_mode=DestinationSyncMode.append_dedup, - cursor_field="updated_at", - ) - - get_catalog.assert_called_once_with( - connection_id=_CONNECTION_ID, - api_root=_API_ROOT, - client_id=None, - client_secret=None, - bearer_token="token", - config_api_root=_CONFIG_API_ROOT, - ) - assert target["config"] == { - "syncMode": "incremental", - "destinationSyncMode": "append_dedup", - "cursorField": ["updated_at"], - "selectedFields": [], - "fieldSelectionEnabled": False, - "primaryKey": [["id"]], - "selected": True, - } - assert other == original_other - - replace_catalog.assert_called_once() - saved_catalog = replace_catalog.call_args.kwargs["configured_catalog_dict"] - assert saved_catalog == catalog - assert saved_catalog["streams"][1] == original_other - assert "sync_mode" not in saved_catalog["streams"][0]["config"] - assert "supported_sync_modes" not in saved_catalog["streams"][0]["stream"] - assert replace_catalog.call_args.kwargs["config_api_root"] == _CONFIG_API_ROOT - - -@pytest.mark.parametrize( - ( - "catalog", - "stream_name", - "stream_namespace", - "sync_mode", - "cursor_field", - "error_message", - "expected_context", - ), - [ - pytest.param( - _catalog(_stream_entry("users", namespace="public")), - "missing", - None, - SyncMode.full_refresh, - None, - "not found", - {"available_streams": ["public.users"]}, - id="missing-stream", - ), - pytest.param( - _catalog( - _stream_entry("users", namespace="public"), - _stream_entry("users", namespace="sales"), - ), - "users", - None, - SyncMode.full_refresh, - None, - "ambiguous", - {"matching_namespaces": ["public", "sales"]}, - id="ambiguous-name", - ), - pytest.param( - _catalog( - _stream_entry("users", namespace="public", sync_modes=["full_refresh"]) - ), - "users", - "public", - SyncMode.incremental, - None, - "not supported", - {"requested_sync_mode": "incremental"}, - id="unsupported-sync-mode", - ), - pytest.param( - _catalog( - _stream_entry( - "users", - namespace="public", - config={ - "syncMode": "full_refresh", - "destinationSyncMode": "append", - }, - ) - ), - "users", - "public", - SyncMode.incremental, - None, - "needs a cursor", - {"stream_namespace": "public"}, - id="incremental-without-cursor", - ), - ], -) -def test_set_stream_sync_mode_rejects_invalid_updates( - monkeypatch: pytest.MonkeyPatch, - catalog: dict[str, Any], - stream_name: str, - stream_namespace: str | None, - sync_mode: SyncMode, - cursor_field: str | None, - error_message: str, - expected_context: dict[str, Any], -) -> None: - """Input errors leave the Config API catalog untouched.""" - _, replace_catalog = _mock_catalog_api(monkeypatch, deepcopy(catalog)) - - with pytest.raises(AirbyteLibInputError, match=error_message) as exc_info: - _connection().set_stream_sync_mode( - stream_name=stream_name, - stream_namespace=stream_namespace, - sync_mode=sync_mode, - cursor_field=cursor_field, - ) - - assert exc_info.value.context is not None - for key, value in expected_context.items(): - assert exc_info.value.context[key] == value - replace_catalog.assert_not_called() - - -def test_set_stream_sync_mode_namespace_disambiguates_name( - monkeypatch: pytest.MonkeyPatch, -) -> None: - """A namespace selects the intended same-named stream.""" - public = _stream_entry( - "users", - namespace="public", - config={ - "syncMode": "incremental", - "destinationSyncMode": "append", - "cursorField": ["updated_at"], - }, - ) - sales = _stream_entry("users", namespace="sales") - catalog = _catalog(public, sales) - _, replace_catalog = _mock_catalog_api(monkeypatch, catalog) - - _connection().set_stream_sync_mode( - stream_name="users", - stream_namespace="sales", - sync_mode=SyncMode.full_refresh, - ) - - assert public["config"]["syncMode"] == "incremental" - assert sales["config"]["syncMode"] == "full_refresh" - replace_catalog.assert_called_once() - - -@pytest.mark.parametrize( - "stream_options", - [ - pytest.param({"source_defined_cursor": True}, id="source-defined-cursor"), - pytest.param( - {"default_cursor_field": ["updated_at"]}, id="default-cursor-field" - ), - ], -) -def test_incremental_sync_accepts_stream_defined_cursor( - monkeypatch: pytest.MonkeyPatch, - stream_options: dict[str, Any], -) -> None: - """Incremental sync can use a source-defined or default cursor without an override.""" - entry = _stream_entry( - "users", - namespace="public", - config={"syncMode": "full_refresh", "destinationSyncMode": "append"}, - **stream_options, - ) - _, replace_catalog = _mock_catalog_api(monkeypatch, _catalog(entry)) - - _connection().set_stream_sync_mode( - stream_name="users", - stream_namespace="public", - sync_mode=SyncMode.incremental, - ) - - assert entry["config"]["syncMode"] == "incremental" - if "default_cursor_field" in stream_options: - assert entry["config"]["cursorField"] == ["updated_at"] - else: - assert "cursorField" not in entry["config"] - replace_catalog.assert_called_once() - - def test_refresh_catalog_gets_and_replaces_catalog_with_config_api_root( monkeypatch: pytest.MonkeyPatch, ) -> None: From a4fd77d13b25cda314850d62ce585e161bbb43be Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 16 Oct 2026 19:43:57 +0000 Subject: [PATCH 4/7] feat(mcp): fold catalog refresh into update_cloud_connection as force_catalog_refresh Co-Authored-By: AJ Steers --- airbyte/cloud/connections.py | 1 + airbyte/mcp/cloud.py | 53 +++++++------------ tests/unit_tests/fixtures/arg_trace_plan.json | 5 +- .../test_cloud_connection_catalog_edits.py | 23 ++++++++ .../test_mcp_connection_schedules.py | 19 +++++++ 5 files changed, 62 insertions(+), 39 deletions(-) diff --git a/airbyte/cloud/connections.py b/airbyte/cloud/connections.py index f6a9e9ec2..e5272ffb8 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 refresh_catalog(self) -> dict[str, Any]: """Refresh the connection's catalog by re-discovering the source schema. diff --git a/airbyte/mcp/cloud.py b/airbyte/mcp/cloud.py index f74a3240e..20bb6158e 100644 --- a/airbyte/mcp/cloud.py +++ b/airbyte/mcp/cloud.py @@ -3971,39 +3971,6 @@ def set_cloud_connection_selected_streams( ) -@mcp_tool( - destructive=True, -) -def refresh_cloud_connection_catalog( - ctx: Context, - connection_id: Annotated[ - str, - Field(description="The ID of the connection to refresh the catalog for."), - ], - *, - workspace_id: Annotated[ - str | None, - Field( - description=WORKSPACE_ID_TIP_TEXT, - default=None, - ), - TraceArg.OMIT, - ], -) -> str: - """Refresh a connection's catalog by re-discovering its source schema.""" - check_guid_created_in_session(connection_id) - workspace: CloudWorkspace = _get_cloud_workspace(ctx, workspace_id) - connection = workspace.get_connection(connection_id=connection_id) - refreshed_catalog = connection.refresh_catalog() - stream_count = len(refreshed_catalog["streams"]) - - return ( - f"Successfully refreshed catalog for connection '{connection_id}'. " - f"Catalog now contains {stream_count} stream(s). " - f"URL: {connection.connection_url}" - ) - - @mcp_tool( destructive=True, ) @@ -4073,6 +4040,15 @@ def update_cloud_connection( default=None, ), ], + force_catalog_refresh: Annotated[ + bool, + Field( + description=( + "If true, re-discover the source schema and replace the connection's catalog " + "with the refreshed result (same as 'Refresh source schema' in the UI)." + ), + ), + ] = False, workspace_id: Annotated[ str | None, Field( @@ -4089,9 +4065,11 @@ def update_cloud_connection( - Set a cron schedule for automatic syncs - Set a basic interval schedule in whole hours between syncs - Switch to manual scheduling (no automatic syncs) + - Refresh the source schema and replace the connection's catalog At least one setting must be provided. `interval_hours`, `cron_expression`, and - `manual_schedule=True` are mutually exclusive. + `manual_schedule=True` are mutually exclusive. `force_catalog_refresh` runs after + other requested settings. """ check_guid_created_in_session(connection_id) @@ -4101,11 +4079,12 @@ def update_cloud_connection( and cron_expression is None and interval_hours is None and manual_schedule is None + and not force_catalog_refresh ): raise AirbyteLibInputError( message=( "At least one setting must be provided: 'enabled', 'interval_hours', " - "'cron_expression', or 'manual_schedule'." + "'cron_expression', 'manual_schedule', or 'force_catalog_refresh'." ), ) @@ -4143,6 +4122,10 @@ def update_cloud_connection( connection.set_manual_schedule() changes_made.append("schedule set to 'manual'") + if force_catalog_refresh: + refreshed_catalog = connection.refresh_catalog() + changes_made.append(f"refreshed catalog ({len(refreshed_catalog['streams'])} streams)") + changes_summary = ", ".join(changes_made) return ( f"Successfully updated connection '{connection_id}': {changes_summary}. " diff --git a/tests/unit_tests/fixtures/arg_trace_plan.json b/tests/unit_tests/fixtures/arg_trace_plan.json index 1b6df0a72..72ebcbf00 100644 --- a/tests/unit_tests/fixtures/arg_trace_plan.json +++ b/tests/unit_tests/fixtures/arg_trace_plan.json @@ -271,10 +271,6 @@ "source_connector_name": "fingerprint", "stream_name": "fingerprint" }, - "refresh_cloud_connection_catalog": { - "connection_id": "hash", - "workspace_id": "omit" - }, "rename_cloud_connection": { "connection_id": "hash", "name": "fingerprint", @@ -349,6 +345,7 @@ "connection_id": "hash", "cron_expression": "fingerprint", "enabled": "value allowed=[false,true]", + "force_catalog_refresh": "value allowed=[false,true]", "interval_hours": "hash", "manual_schedule": "value allowed=[false,true]", "workspace_id": "omit" diff --git a/tests/unit_tests/test_cloud_connection_catalog_edits.py b/tests/unit_tests/test_cloud_connection_catalog_edits.py index 0833c0710..714d541df 100644 --- a/tests/unit_tests/test_cloud_connection_catalog_edits.py +++ b/tests/unit_tests/test_cloud_connection_catalog_edits.py @@ -93,6 +93,29 @@ def test_refresh_catalog_gets_and_replaces_catalog_with_config_api_root( ) +def test_refresh_catalog_clears_cached_connection_info( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A refreshed 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 + refreshed_catalog = _catalog(_stream_entry("new-stream")) + monkeypatch.setattr( + api_util, + "get_refreshed_connection_catalog", + MagicMock(return_value={"syncCatalog": refreshed_catalog}), + ) + replace_catalog = MagicMock(return_value={}) + monkeypatch.setattr(api_util, "replace_connection_catalog", replace_catalog) + + connection.refresh_catalog() + + replace_catalog.assert_called_once() + assert connection._connection_info is None + + @pytest.mark.parametrize( "response", [ diff --git a/tests/unit_tests/test_mcp_connection_schedules.py b/tests/unit_tests/test_mcp_connection_schedules.py index c8f0d2bfa..8db50fa54 100644 --- a/tests/unit_tests/test_mcp_connection_schedules.py +++ b/tests/unit_tests/test_mcp_connection_schedules.py @@ -319,3 +319,22 @@ def test_update_interval_api_failure_is_reported( with pytest.raises(ToolError, match="HTTP 400"): asyncio.run(_update(server, interval_hours=24)) + + +def test_update_force_catalog_refresh( + schedule_backend: tuple[FastMCP, responses.RequestsMock, MagicMock, MagicMock], + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A serialized update can refresh the source catalog.""" + server, _http, _lookup, _patch = schedule_backend + refreshed_catalog = {"streams": [{"stream": {"name": "users"}}]} + get_refreshed = MagicMock(return_value={"syncCatalog": refreshed_catalog}) + replace_catalog = MagicMock(return_value={}) + monkeypatch.setattr(api_util, "get_refreshed_connection_catalog", get_refreshed) + monkeypatch.setattr(api_util, "replace_connection_catalog", replace_catalog) + + result = asyncio.run(_update(server, force_catalog_refresh=True)) + + get_refreshed.assert_called_once() + replace_catalog.assert_called_once() + assert "refreshed catalog (1 streams)" in result From 68bb1436f84eb7aab402f56c3c49a4b3d42e6e4b Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 16 Oct 2026 20:59:23 +0000 Subject: [PATCH 5/7] feat(mcp): add async Cloud source catalog discovery Co-Authored-By: AJ Steers --- airbyte/_util/api_util.py | 136 ++++- airbyte/cloud/__init__.py | 6 + airbyte/cloud/connections.py | 31 - airbyte/cloud/connectors.py | 151 +++++ airbyte/cloud/models.py | 36 ++ airbyte/cloud/workspaces.py | 77 +++ airbyte/mcp/_otel.py | 2 + airbyte/mcp/cloud.py | 154 ++++- tests/unit_tests/fixtures/arg_trace_plan.json | 10 +- tests/unit_tests/test_cloud_api_util.py | 85 +++ .../test_cloud_connection_catalog_edits.py | 101 +--- .../test_cloud_source_catalog_discovery.py | 551 ++++++++++++++++++ .../test_mcp_connection_schedules.py | 19 - tests/unit_tests/test_mcp_otel.py | 29 + 14 files changed, 1195 insertions(+), 193 deletions(-) create mode 100644 tests/unit_tests/test_cloud_source_catalog_discovery.py diff --git a/airbyte/_util/api_util.py b/airbyte/_util/api_util.py index dcf8c5c71..193a3730c 100644 --- a/airbyte/_util/api_util.py +++ b/airbyte/_util/api_util.py @@ -3620,7 +3620,7 @@ def replace_connection_state( raise -def get_refreshed_connection_catalog( +def get_connection_catalog( connection_id: str, *, api_root: str, @@ -3629,25 +3629,27 @@ def get_refreshed_connection_catalog( bearer_token: SecretString | None, config_api_root: str | None = None, ) -> dict[str, Any]: - """Get a refreshed catalog for a connection. + """Get the configured catalog for a connection. Uses the Config API endpoint: POST /v1/web_backend/connections/get - with `withRefreshedCatalog` enabled. + + This returns the full connection info including the syncCatalog field, + which contains the configured catalog with full stream schemas. Args: - connection_id: The connection ID to get a refreshed catalog for. - api_root: The API root URL. - client_id: OAuth client ID. - client_secret: OAuth client secret. + connection_id: The connection ID to get catalog for + api_root: The API root URL + client_id: OAuth client ID + client_secret: OAuth client secret bearer_token: Bearer token for authentication (alternative to client credentials). config_api_root: Optional explicit Config API root URL. Returns: - Dictionary containing the connection info with refreshed syncCatalog. + Dictionary containing the connection info with syncCatalog. """ return _make_config_api_request( path="/web_backend/connections/get", - json={"connectionId": connection_id, "withRefreshedCatalog": True}, + json={"connectionId": connection_id, "withRefreshedCatalog": False}, api_root=api_root, config_api_root=config_api_root, client_id=client_id, @@ -3656,8 +3658,9 @@ def get_refreshed_connection_catalog( ) -def get_connection_catalog( - connection_id: str, +def run_source_discover_command( + source_id: str, + discover_request_id: str, *, api_root: str, client_id: SecretString | None, @@ -3665,27 +3668,106 @@ def get_connection_catalog( bearer_token: SecretString | None, config_api_root: str | None = None, ) -> dict[str, Any]: - """Get the configured catalog for a connection. + """Start a source discover command. - Uses the Config API endpoint: POST /v1/web_backend/connections/get + 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, + ) - This returns the full connection info including the syncCatalog field, - which contains the configured catalog with full stream schemas. - Args: - connection_id: The connection ID to get catalog for - api_root: The API root URL - client_id: OAuth client ID - client_secret: OAuth client secret - bearer_token: Bearer token for authentication (alternative to client credentials). - config_api_root: Optional explicit Config API root URL. +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, + ) - Returns: - Dictionary containing the connection info with syncCatalog. - """ + +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="/web_backend/connections/get", - json={"connectionId": connection_id, "withRefreshedCatalog": False}, + 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, diff --git a/airbyte/cloud/__init__.py b/airbyte/cloud/__init__.py index 2a2043f35..47fc940ba 100644 --- a/airbyte/cloud/__init__.py +++ b/airbyte/cloud/__init__.py @@ -94,6 +94,9 @@ from airbyte.cloud.connections import CloudConnection from airbyte.cloud.models import ( CloudDefaultContextInfo, + CloudSourceCatalogDiffStreamSummary, + CloudSourceCatalogDiscoverResult, + CloudSourceCatalogStreamSummary, CloudWorkspaceInfo, ConnectionStatus, JobStatusEnum, @@ -133,6 +136,9 @@ "CloudOrganization", "CloudWorkspace", "CloudConnection", + "CloudSourceCatalogDiffStreamSummary", + "CloudSourceCatalogDiscoverResult", + "CloudSourceCatalogStreamSummary", "CloudClientConfig", "CloudDefaultContextInfo", "CloudWorkspaceInfo", diff --git a/airbyte/cloud/connections.py b/airbyte/cloud/connections.py index e5272ffb8..67435df9c 100644 --- a/airbyte/cloud/connections.py +++ b/airbyte/cloud/connections.py @@ -904,37 +904,6 @@ def import_raw_catalog(self, catalog: dict[str, Any]) -> None: ) self._connection_info = None - def refresh_catalog(self) -> dict[str, Any]: - """Refresh the connection's catalog by re-discovering the source schema. - - Triggers a discover operation on the connection's source connector, then - replaces the connection's catalog with the refreshed result. This is - equivalent to clicking "Refresh source schema" in the Airbyte UI. - - Returns: - The refreshed raw syncCatalog dict. - - Raises: - AirbyteLibInputError: If the refreshed catalog is empty or missing. - """ - refreshed_response = api_util.get_refreshed_connection_catalog( - connection_id=self.connection_id, - api_root=self.workspace.api_root, - client_id=self.workspace.client_id, - client_secret=self.workspace.client_secret, - bearer_token=self.workspace.bearer_token, - config_api_root=self.workspace.config_api_root, - ) - refreshed_catalog = refreshed_response.get("syncCatalog") - if not isinstance(refreshed_catalog, dict) or not refreshed_catalog.get("streams"): - raise AirbyteLibInputError( - message="Refreshed catalog is empty or missing.", - context={"connection_id": self.connection_id}, - ) - - self.import_raw_catalog(refreshed_catalog) - return refreshed_catalog - def rename(self, name: str) -> CloudConnection: """Rename the connection. diff --git a/airbyte/cloud/connectors.py b/airbyte/cloud/connectors.py index 0e08038c2..fb80e13e6 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,91 @@ """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}, + ) + + api_args = { + "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, + } + command_status = api_util.get_command_status( + discover_request_id, + **api_args, + ) + 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_args, + ) + 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 +1174,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..4e53fb1c9 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,41 @@ def from_api_response(cls, source: _SourceResponseLike) -> CloudSourceInfo: ) +class CloudSourceCatalogStreamSummary(BaseModel): + """Summary of a stream in a source catalog.""" + + stream_name: str + stream_namespace: str | None = None + supported_sync_modes: list[str] = Field(default_factory=list) + + +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 + streams: list[CloudSourceCatalogStreamSummary] | 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 6eb1e7296..fa42f63ff 100644 --- a/airbyte/mcp/cloud.py +++ b/airbyte/mcp/cloud.py @@ -45,6 +45,7 @@ CloudDefaultContextInfo, CloudDefaultWorkspaceUpdateInfo, CloudOrganizationInfo, + CloudSourceCatalogDiscoverResult, ConnectionSchedule, ConnectionStatus, ConnectorFeature, @@ -4054,15 +4055,6 @@ def update_cloud_connection( default=None, ), ], - force_catalog_refresh: Annotated[ - bool, - Field( - description=( - "If true, re-discover the source schema and replace the connection's catalog " - "with the refreshed result (same as 'Refresh source schema' in the UI)." - ), - ), - ] = False, workspace_id: Annotated[ str | None, Field( @@ -4079,11 +4071,9 @@ def update_cloud_connection( - Set a cron schedule for automatic syncs - Set a basic interval schedule in whole hours between syncs - Switch to manual scheduling (no automatic syncs) - - Refresh the source schema and replace the connection's catalog At least one setting must be provided. `interval_hours`, `cron_expression`, and - `manual_schedule=True` are mutually exclusive. `force_catalog_refresh` runs after - other requested settings. + `manual_schedule=True` are mutually exclusive. """ check_guid_created_in_session(connection_id) @@ -4093,12 +4083,11 @@ def update_cloud_connection( and cron_expression is None and interval_hours is None and manual_schedule is None - and not force_catalog_refresh ): raise AirbyteLibInputError( message=( "At least one setting must be provided: 'enabled', 'interval_hours', " - "'cron_expression', 'manual_schedule', or 'force_catalog_refresh'." + "'cron_expression', or 'manual_schedule'." ), ) @@ -4136,10 +4125,6 @@ def update_cloud_connection( connection.set_manual_schedule() changes_made.append("schedule set to 'manual'") - if force_catalog_refresh: - refreshed_catalog = connection.refresh_catalog() - changes_made.append(f"refreshed catalog ({len(refreshed_catalog['streams'])} streams)") - changes_summary = ", ".join(changes_made) return ( f"Successfully updated connection '{connection_id}': {changes_summary}. " @@ -4147,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 72ebcbf00..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", @@ -345,7 +354,6 @@ "connection_id": "hash", "cron_expression": "fingerprint", "enabled": "value allowed=[false,true]", - "force_catalog_refresh": "value allowed=[false,true]", "interval_hours": "hash", "manual_schedule": "value allowed=[false,true]", "workspace_id": "omit" 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 index 714d541df..ae6204631 100644 --- a/tests/unit_tests/test_cloud_connection_catalog_edits.py +++ b/tests/unit_tests/test_cloud_connection_catalog_edits.py @@ -1,5 +1,5 @@ # Copyright (c) 2026 Airbyte, Inc., all rights reserved. -"""Unit tests for editing and refreshing Cloud connection catalogs.""" +"""Unit tests for editing Cloud connection catalogs.""" from __future__ import annotations @@ -11,7 +11,6 @@ from airbyte._util import api_util from airbyte.cloud.connections import CloudConnection from airbyte.cloud.workspaces import CloudWorkspace -from airbyte.exceptions import AirbyteLibInputError _CONNECTION_ID = "connection-id" @@ -61,110 +60,18 @@ def _catalog(*entries: dict[str, Any]) -> dict[str, Any]: return {"streams": list(entries)} -def test_refresh_catalog_gets_and_replaces_catalog_with_config_api_root( +def test_import_raw_catalog_clears_cached_connection_info( monkeypatch: pytest.MonkeyPatch, ) -> None: - """Refresh the source catalog and persist the raw result.""" - refreshed_catalog = _catalog(_stream_entry("users", namespace="public")) - get_refreshed = MagicMock(return_value={"syncCatalog": refreshed_catalog}) - replace_catalog = MagicMock(return_value={}) - monkeypatch.setattr(api_util, "get_refreshed_connection_catalog", get_refreshed) - monkeypatch.setattr(api_util, "replace_connection_catalog", replace_catalog) - - result = _connection().refresh_catalog() - - assert result == refreshed_catalog - get_refreshed.assert_called_once_with( - connection_id=_CONNECTION_ID, - api_root=_API_ROOT, - client_id=None, - client_secret=None, - bearer_token="token", - config_api_root=_CONFIG_API_ROOT, - ) - replace_catalog.assert_called_once_with( - connection_id=_CONNECTION_ID, - configured_catalog_dict=refreshed_catalog, - api_root=_API_ROOT, - client_id=None, - client_secret=None, - bearer_token="token", - config_api_root=_CONFIG_API_ROOT, - ) - - -def test_refresh_catalog_clears_cached_connection_info( - monkeypatch: pytest.MonkeyPatch, -) -> None: - """A refreshed catalog invalidates previously cached connection information.""" + """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 - refreshed_catalog = _catalog(_stream_entry("new-stream")) - monkeypatch.setattr( - api_util, - "get_refreshed_connection_catalog", - MagicMock(return_value={"syncCatalog": refreshed_catalog}), - ) replace_catalog = MagicMock(return_value={}) monkeypatch.setattr(api_util, "replace_connection_catalog", replace_catalog) - connection.refresh_catalog() + connection.import_raw_catalog(_catalog(_stream_entry("new-stream"))) replace_catalog.assert_called_once() assert connection._connection_info is None - - -@pytest.mark.parametrize( - "response", - [ - pytest.param({}, id="missing-sync-catalog"), - pytest.param({"syncCatalog": None}, id="null-sync-catalog"), - pytest.param({"syncCatalog": {"streams": []}}, id="empty-sync-catalog"), - ], -) -def test_refresh_catalog_rejects_empty_or_missing_catalog( - monkeypatch: pytest.MonkeyPatch, - response: dict[str, Any], -) -> None: - """An absent or empty refreshed catalog is never written.""" - get_refreshed = MagicMock(return_value=response) - replace_catalog = MagicMock(return_value={}) - monkeypatch.setattr(api_util, "get_refreshed_connection_catalog", get_refreshed) - monkeypatch.setattr(api_util, "replace_connection_catalog", replace_catalog) - - with pytest.raises( - AirbyteLibInputError, match="Refreshed catalog is empty or missing" - ): - _connection().refresh_catalog() - - replace_catalog.assert_not_called() - - -def test_get_refreshed_connection_catalog_forwards_config_api_root( - monkeypatch: pytest.MonkeyPatch, -) -> None: - """Request discovery through the Config API catalog endpoint.""" - request = MagicMock(return_value={"syncCatalog": {"streams": []}}) - monkeypatch.setattr(api_util, "_make_config_api_request", request) - - result = api_util.get_refreshed_connection_catalog( - connection_id=_CONNECTION_ID, - api_root=_API_ROOT, - client_id=None, - client_secret=None, - bearer_token="token", - config_api_root=_CONFIG_API_ROOT, - ) - - assert result == {"syncCatalog": {"streams": []}} - request.assert_called_once_with( - path="/web_backend/connections/get", - json={"connectionId": _CONNECTION_ID, "withRefreshedCatalog": True}, - api_root=_API_ROOT, - config_api_root=_CONFIG_API_ROOT, - client_id=None, - client_secret=None, - bearer_token="token", - ) 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_connection_schedules.py b/tests/unit_tests/test_mcp_connection_schedules.py index 8db50fa54..c8f0d2bfa 100644 --- a/tests/unit_tests/test_mcp_connection_schedules.py +++ b/tests/unit_tests/test_mcp_connection_schedules.py @@ -319,22 +319,3 @@ def test_update_interval_api_failure_is_reported( with pytest.raises(ToolError, match="HTTP 400"): asyncio.run(_update(server, interval_hours=24)) - - -def test_update_force_catalog_refresh( - schedule_backend: tuple[FastMCP, responses.RequestsMock, MagicMock, MagicMock], - monkeypatch: pytest.MonkeyPatch, -) -> None: - """A serialized update can refresh the source catalog.""" - server, _http, _lookup, _patch = schedule_backend - refreshed_catalog = {"streams": [{"stream": {"name": "users"}}]} - get_refreshed = MagicMock(return_value={"syncCatalog": refreshed_catalog}) - replace_catalog = MagicMock(return_value={}) - monkeypatch.setattr(api_util, "get_refreshed_connection_catalog", get_refreshed) - monkeypatch.setattr(api_util, "replace_connection_catalog", replace_catalog) - - result = asyncio.run(_update(server, force_catalog_refresh=True)) - - get_refreshed.assert_called_once() - replace_catalog.assert_called_once() - assert "refreshed catalog (1 streams)" in result 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 From ccdbf4b8a00e9a2ea293490ab7bd2c1a8571c6b0 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 16 Oct 2026 21:07:16 +0000 Subject: [PATCH 6/7] fix(cloud): preserve API argument types for catalog discovery Co-Authored-By: AJ Steers --- airbyte/cloud/connectors.py | 19 ++++++++++--------- 1 file changed, 10 insertions(+), 9 deletions(-) diff --git a/airbyte/cloud/connectors.py b/airbyte/cloud/connectors.py index fb80e13e6..2035efb02 100644 --- a/airbyte/cloud/connectors.py +++ b/airbyte/cloud/connectors.py @@ -144,16 +144,13 @@ def _get_catalog_discover_result( context={"discover_request_id": discover_request_id}, ) - api_args = { - "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, - } command_status = api_util.get_command_status( discover_request_id, - **api_args, + 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": @@ -164,7 +161,11 @@ def _get_catalog_discover_result( output = api_util.get_discover_command_output( discover_request_id, - **api_args, + 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 From 1202ea4445e8f137c1007b86a3c1ecf11b7b4763 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 16 Oct 2026 23:00:12 +0000 Subject: [PATCH 7/7] refactor(cloud): remove source catalog stream summary Co-Authored-By: AJ Steers --- airbyte/cloud/__init__.py | 2 -- airbyte/cloud/models.py | 9 --------- 2 files changed, 11 deletions(-) diff --git a/airbyte/cloud/__init__.py b/airbyte/cloud/__init__.py index 47fc940ba..3eb25f5c4 100644 --- a/airbyte/cloud/__init__.py +++ b/airbyte/cloud/__init__.py @@ -96,7 +96,6 @@ CloudDefaultContextInfo, CloudSourceCatalogDiffStreamSummary, CloudSourceCatalogDiscoverResult, - CloudSourceCatalogStreamSummary, CloudWorkspaceInfo, ConnectionStatus, JobStatusEnum, @@ -138,7 +137,6 @@ "CloudConnection", "CloudSourceCatalogDiffStreamSummary", "CloudSourceCatalogDiscoverResult", - "CloudSourceCatalogStreamSummary", "CloudClientConfig", "CloudDefaultContextInfo", "CloudWorkspaceInfo", diff --git a/airbyte/cloud/models.py b/airbyte/cloud/models.py index 4e53fb1c9..9e1355b11 100644 --- a/airbyte/cloud/models.py +++ b/airbyte/cloud/models.py @@ -510,14 +510,6 @@ def from_api_response(cls, source: _SourceResponseLike) -> CloudSourceInfo: ) -class CloudSourceCatalogStreamSummary(BaseModel): - """Summary of a stream in a source catalog.""" - - stream_name: str - stream_namespace: str | None = None - supported_sync_modes: list[str] = Field(default_factory=list) - - class CloudSourceCatalogDiffStreamSummary(BaseModel): """Summary of one stream transformation in a catalog diff.""" @@ -540,7 +532,6 @@ class CloudSourceCatalogDiscoverResult(BaseModel): stream_count: int | None = None updated_at: datetime | None = None failure_reason: str | None = None - streams: list[CloudSourceCatalogStreamSummary] | None = None catalog_diff_summary: list[CloudSourceCatalogDiffStreamSummary] | None = None schema_change: Literal["no_change", "non_breaking", "breaking"] | None = None