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 01/12] 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 02/12] 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 03/12] 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 fb45501082ea6572cd40f3852526a528dd373bf4 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:39:03 +0000 Subject: [PATCH 04/12] feat(mcp): add tool to set stream sync mode on Cloud connections 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, 468 insertions(+), 2 deletions(-) diff --git a/airbyte/cloud/connections.py b/airbyte/cloud/connections.py index f6a9e9ec2..95e1bd75b 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, @@ -934,6 +936,136 @@ 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 f74a3240e..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 ( @@ -4004,6 +4006,91 @@ 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 1b6df0a72..e886071b2 100644 --- a/tests/unit_tests/fixtures/arg_trace_plan.json +++ b/tests/unit_tests/fixtures/arg_trace_plan.json @@ -304,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 0833c0710..861890e23 100644 --- a/tests/unit_tests/test_cloud_connection_catalog_edits.py +++ b/tests/unit_tests/test_cloud_connection_catalog_edits.py @@ -3,10 +3,12 @@ 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 @@ -36,17 +38,26 @@ 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": ["full_refresh", "incremental"], + "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": config + or { "syncMode": "full_refresh", "destinationSyncMode": "append", "selectedFields": [], @@ -61,6 +72,233 @@ 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 05/12] 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 6e8144bc512f934576e399c3f5a386c28564dc8b 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:11:48 +0000 Subject: [PATCH 06/12] feat(mcp): add safe batch stream config updates Co-Authored-By: AJ Steers --- airbyte/_util/api_util.py | 4 +- airbyte/cloud/__init__.py | 6 + airbyte/cloud/connections.py | 440 ++++++++++++---- airbyte/cloud/models.py | 55 +- airbyte/exceptions.py | 5 + airbyte/mcp/_error_handling.py | 2 + airbyte/mcp/cloud.py | 371 +++++++++----- tests/unit_tests/fixtures/arg_trace_plan.json | 21 +- .../test_cloud_connection_catalog_edits.py | 482 ++++++++++++------ .../test_mcp_connection_schedules.py | 251 ++++++++- tests/unit_tests/test_mcp_tool_errors.py | 8 + 11 files changed, 1250 insertions(+), 395 deletions(-) diff --git a/airbyte/_util/api_util.py b/airbyte/_util/api_util.py index dcf8c5c71..3156a1633 100644 --- a/airbyte/_util/api_util.py +++ b/airbyte/_util/api_util.py @@ -3703,6 +3703,7 @@ def replace_connection_catalog( client_id: SecretString | None, client_secret: SecretString | None, bearer_token: SecretString | None, + skip_reset: bool = True, ) -> dict[str, Any]: """Replace the configured catalog for a connection. @@ -3719,6 +3720,7 @@ def replace_connection_catalog( client_secret: OAuth client secret. bearer_token: Bearer token for authentication (alternative to client credentials). config_api_root: Optional explicit Config API root URL. + skip_reset: Whether to skip the platform reset of changed streams. Returns: Dictionary containing the updated WebBackendConnectionRead response. @@ -3730,7 +3732,7 @@ def replace_connection_catalog( "syncCatalog": configured_catalog_dict, # Resets are destructive and cause customer-side data outage. # If a reset is desired, caller will need to decide & manage. - "skipReset": True, + "skipReset": skip_reset, }, api_root=api_root, config_api_root=config_api_root, diff --git a/airbyte/cloud/__init__.py b/airbyte/cloud/__init__.py index 2a2043f35..b3df57aa4 100644 --- a/airbyte/cloud/__init__.py +++ b/airbyte/cloud/__init__.py @@ -98,6 +98,9 @@ ConnectionStatus, JobStatusEnum, JobTypeEnum, + StreamConfigChange, + StreamConfigOverride, + StreamConfigUpdateResult, WorkspacePrivilegeScope, ) from airbyte.cloud.organizations import CloudOrganization @@ -136,6 +139,9 @@ "CloudClientConfig", "CloudDefaultContextInfo", "CloudWorkspaceInfo", + "StreamConfigOverride", + "StreamConfigChange", + "StreamConfigUpdateResult", "SyncResult", # Enums "ExternalApiReadOnlyAction", diff --git a/airbyte/cloud/connections.py b/airbyte/cloud/connections.py index 44d620373..1b9ef91bd 100644 --- a/airbyte/cloud/connections.py +++ b/airbyte/cloud/connections.py @@ -4,11 +4,11 @@ from __future__ import annotations import logging +from copy import deepcopy +from dataclasses import dataclass 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, @@ -31,6 +31,9 @@ ConnectionSchedule, ConnectionStatus, JobTypeEnum, + StreamConfigChange, + StreamConfigOverride, + StreamConfigUpdateResult, _ConnectionResponseLike, ) from airbyte.cloud.sync_results import SyncResult @@ -39,6 +42,7 @@ AirbyteConnectionSyncError, AirbyteLibInputError, AirbyteMissingResourceError, + AirbyteUnsafeStreamConfigChangeError, AirbyteWorkspaceMismatchError, ) @@ -49,6 +53,69 @@ _QUARTZ_CRON_MAX_FIELDS = 8 # 7 fields plus an optional trailing timezone ID +@dataclass +class _StreamConfigUpdatePlan: + catalog: dict[str, Any] + updated_streams: list[StreamConfigChange] + unsafe_streams: list[StreamConfigChange] + catalog_refreshed: bool + stream_count: int + + +def _normalized_cursor_field(value: list[str] | str | None) -> list[str]: + if value is None: + return [] + if isinstance(value, str): + return [value] + return value + + +def _changed_stream_config_fields( + current: dict[str, Any], + updated: dict[str, Any], +) -> list[str]: + fields = ("syncMode", "destinationSyncMode", "cursorField", "primaryKey") + return [ + field + for field in fields + if ( + _normalized_cursor_field(current.get(field)) + if field == "cursorField" + else current.get(field) + ) + != ( + _normalized_cursor_field(updated.get(field)) + if field == "cursorField" + else updated.get(field) + ) + ] + + +def _stream_has_saved_state( + state_messages: list[dict[str, Any]], + stream_name: str, + stream_namespace: str | None, +) -> bool: + for message in state_messages: + if not isinstance(message, dict): + return True + message_type = message.get("type") + if message_type in {"GLOBAL", "LEGACY"}: + return True + if message_type != "STREAM": + return True + + stream = message.get("stream") + descriptor = stream.get("stream_descriptor") if isinstance(stream, dict) else None + if not isinstance(descriptor, dict) or not isinstance(descriptor.get("name"), str): + return True + if descriptor.get("name") != stream_name: + continue + if (descriptor.get("namespace") or None) == (stream_namespace or None): + return True + return False + + def _validate_quartz_cron_expression(cron_expression: str) -> None: """Raise `AirbyteLibInputError` if the expression is not a plausible Quartz cron. @@ -873,7 +940,7 @@ def dump_raw_catalog( return _normalize_catalog_to_protocol(raw) return raw - def import_raw_catalog(self, catalog: dict[str, Any]) -> None: + def import_raw_catalog(self, catalog: dict[str, Any], *, skip_reset: bool = True) -> None: """Replace the configured catalog for this connection. > ⚠️ **WARNING:** Modifying the catalog directly is not recommended and @@ -891,6 +958,7 @@ def import_raw_catalog(self, catalog: dict[str, Any]) -> None: Args: catalog: The configured catalog dict in either format. + skip_reset: Whether to skip the platform reset of changed streams. """ if _is_protocol_catalog_format(catalog): catalog = _denormalize_catalog_to_api(catalog) @@ -903,6 +971,7 @@ def import_raw_catalog(self, catalog: dict[str, Any]) -> None: client_secret=self.workspace.client_secret, bearer_token=self.workspace.bearer_token, config_api_root=self.workspace.config_api_root, + skip_reset=skip_reset, ) self._connection_info = None @@ -937,135 +1006,308 @@ def refresh_catalog(self) -> dict[str, Any]: self.import_raw_catalog(refreshed_catalog) return refreshed_catalog - def set_stream_sync_mode( + def plan_stream_config_update( # noqa: PLR0912, PLR0914, PLR0915 self, - stream_name: str, + overrides: list[StreamConfigOverride], *, - 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. + force_catalog_refresh: bool = False, + allow_unsafe: bool = False, + ) -> _StreamConfigUpdatePlan: + """Validate overrides and plan a catalog update without writing it.""" + if not overrides: + raise AirbyteLibInputError( + message="At least one stream configuration override must be provided.", + ) - 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. + if force_catalog_refresh: + 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, + ) + catalog = refreshed_response.get("syncCatalog") + else: + catalog = self.dump_raw_catalog(normalize=False) - 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.", + message=( + "Refreshed catalog is empty or missing." + if force_catalog_refresh + else "Connection catalog is empty or missing." + ), context={"connection_id": self.connection_id}, ) - matching_entries: list[dict[str, Any]] = [] + old_catalog = deepcopy(catalog) + catalog = deepcopy(catalog) + available_streams: list[str] = [] + catalog_keys: set[tuple[str, str | None]] = set() + duplicate_catalog_streams: list[str] = [] 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) + stream_def = entry.get("stream") if isinstance(entry, dict) else None + if not isinstance(stream_def, dict): + raise AirbyteLibInputError( + message="Connection catalog contains an invalid stream entry.", + context={"connection_id": self.connection_id}, + ) + name = stream_def.get("name") + namespace = stream_def.get("namespace") + if not isinstance(name, str) or ( + namespace is not None and not isinstance(namespace, str) + ): + raise AirbyteLibInputError( + message="Connection catalog contains an invalid stream descriptor.", + context={"connection_id": self.connection_id}, + ) + available_streams.append(f"{namespace}.{name}" if namespace else name) + key = (name, namespace) + if key in catalog_keys: + duplicate_catalog_streams.append(f"{namespace}.{name}" if namespace else name) + catalog_keys.add(key) + + if duplicate_catalog_streams: + duplicate_names = ", ".join(duplicate_catalog_streams) raise AirbyteLibInputError( - message=f"Stream '{stream_name}' not found in connection catalog.", - guidance="Use a stream name from `available_streams`.", + message=f"Connection catalog contains duplicate stream entries: {duplicate_names}.", context={ "connection_id": self.connection_id, - "stream_name": stream_name, - "stream_namespace": stream_namespace, - "available_streams": available_streams, + "duplicate_streams": duplicate_catalog_streams, }, ) - if len(matching_entries) > 1: - matching_namespaces = [ - entry.get("stream", {}).get("namespace") for entry in matching_entries + validation_errors: list[str] = [] + updated_streams: list[StreamConfigChange] = [] + override_keys: set[tuple[str, str | None]] = set() + for index, override in enumerate(overrides, start=1): + label = f"override {index} for " f"'{override.stream_name}'" + ( + f" in namespace '{override.stream_namespace}'" + if override.stream_namespace is not None + else "" + ) + matching_entries = [ + entry + for entry in catalog["streams"] + if entry["stream"].get("name") == override.stream_name + and ( + override.stream_namespace is None + or entry["stream"].get("namespace") == override.stream_namespace + ) ] + + if not matching_entries: + validation_errors.append( + f"{label}: stream not found. Available streams: {available_streams}." + ) + continue + if len(matching_entries) > 1: + matching_namespaces = [ + entry["stream"].get("namespace") for entry in matching_entries + ] + validation_errors.append( + f"{label}: stream name is ambiguous; specify a namespace. " + f"Matching namespaces: {matching_namespaces}." + ) + continue + + target_entry = matching_entries[0] + stream_def = target_entry["stream"] + stream_namespace = stream_def.get("namespace") + target_key = (override.stream_name, stream_namespace) + if target_key in override_keys: + validation_errors.append(f"{label}: this stream is overridden more than once.") + continue + override_keys.add(target_key) + + current_config = target_entry.get("config") or {} + if not isinstance(current_config, dict): + validation_errors.append(f"{label}: stream configuration is invalid.") + continue + + candidate_config = deepcopy(current_config) + supported_sync_modes = stream_def.get("supportedSyncModes") or [] + if ( + override.sync_mode is not None + and override.sync_mode.value not in supported_sync_modes + ): + validation_errors.append( + f"{label}: sync mode '{override.sync_mode.value}' is unsupported. " + f"Supported modes: {supported_sync_modes}." + ) + continue + + if override.sync_mode is not None: + candidate_config["syncMode"] = override.sync_mode.value + if override.destination_sync_mode is not None: + candidate_config["destinationSyncMode"] = override.destination_sync_mode.value + if override.cursor_field is not None: + candidate_config["cursorField"] = ( + [override.cursor_field] + if isinstance(override.cursor_field, str) + else deepcopy(override.cursor_field) + ) + + if candidate_config.get("syncMode") == "incremental" and not _normalized_cursor_field( + candidate_config.get("cursorField") + ): + if stream_def.get("sourceDefinedCursor"): + pass + elif ( + isinstance(stream_def.get("defaultCursorField"), list) + and stream_def["defaultCursorField"] + ): + candidate_config["cursorField"] = deepcopy(stream_def["defaultCursorField"]) + else: + validation_errors.append(f"{label}: incremental sync requires a cursor field.") + continue + + target_entry["config"] = candidate_config + updated_streams.append( + StreamConfigChange( + name=override.stream_name, + namespace=stream_namespace, + changed_fields=_changed_stream_config_fields( + current_config, + candidate_config, + ), + ) + ) + + if validation_errors: + joined_errors = " ".join(validation_errors) raise AirbyteLibInputError( - message=f"Stream name '{stream_name}' is ambiguous.", - guidance="Specify `stream_namespace` to select one of the matching streams.", + message=f"Invalid stream configuration overrides: {joined_errors}", + guidance="Correct every listed override and try again.", context={ "connection_id": self.connection_id, - "stream_name": stream_name, - "matching_namespaces": matching_namespaces, + "validation_errors": validation_errors, }, ) - 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`.", + changed_streams = [stream for stream in updated_streams if stream.changed_fields] + unsafe_streams: list[StreamConfigChange] = [] + if changed_streams: + state_messages = self.dump_raw_state(normalize=True) + unsafe_streams = self._find_unsafe_stream_changes( + old_catalog, + catalog, + state_messages, + ) + + self._raise_for_unsafe_stream_config_changes( + unsafe_streams, + allow_unsafe=allow_unsafe, + ) + return _StreamConfigUpdatePlan( + catalog=catalog, + updated_streams=updated_streams, + unsafe_streams=unsafe_streams, + catalog_refreshed=force_catalog_refresh, + stream_count=len(catalog["streams"]), + ) + + @staticmethod + def _find_unsafe_stream_changes( + old_catalog: dict[str, Any], + new_catalog: dict[str, Any], + state_messages: list[dict[str, Any]], + ) -> list[StreamConfigChange]: + old_streams = { + ( + entry["stream"]["name"], + entry["stream"].get("namespace"), + ): entry + for entry in old_catalog["streams"] + } + unsafe_streams: list[StreamConfigChange] = [] + for entry in new_catalog["streams"]: + stream_def = entry["stream"] + name = stream_def["name"] + namespace = stream_def.get("namespace") + current_entry = old_streams.get((name, namespace)) + if current_entry is None: + continue + changed_fields = _changed_stream_config_fields( + current_entry.get("config") or {}, + entry.get("config") or {}, + ) + if changed_fields and _stream_has_saved_state( + state_messages, + name, + namespace, + ): + unsafe_streams.append( + StreamConfigChange( + name=name, + namespace=namespace, + changed_fields=changed_fields, + ) + ) + return unsafe_streams + + def _raise_for_unsafe_stream_config_changes( + self, + unsafe_streams: list[StreamConfigChange], + *, + allow_unsafe: bool, + ) -> None: + if unsafe_streams and not allow_unsafe: + raise AirbyteUnsafeStreamConfigChangeError( + message="Stream configuration changes affect streams with saved sync state.", 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, + "unsafe_streams": [stream.model_dump() for stream in unsafe_streams], }, ) - 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 + def apply_stream_config_update( + self, + plan: _StreamConfigUpdatePlan, + *, + allow_unsafe: bool = False, + defer_reset: bool = False, + ) -> StreamConfigUpdateResult: + """Write a planned catalog update and return its result.""" + self._raise_for_unsafe_stream_config_changes( + plan.unsafe_streams, + allow_unsafe=allow_unsafe, + ) + reset_requested = bool(plan.unsafe_streams and allow_unsafe and not defer_reset) + self.import_raw_catalog(plan.catalog, skip_reset=not reset_requested) + return StreamConfigUpdateResult( + updated_streams=plan.updated_streams, + unsafe_streams=plan.unsafe_streams, + reset_requested=reset_requested, + catalog_refreshed=plan.catalog_refreshed, + stream_count=plan.stream_count, + ) - self.import_raw_catalog(catalog) + def update_stream_configs( + self, + overrides: list[StreamConfigOverride], + *, + force_catalog_refresh: bool = False, + allow_unsafe: bool = False, + defer_reset: bool = False, + ) -> StreamConfigUpdateResult: + """Update one or more stream configurations in a single catalog write.""" + plan = self.plan_stream_config_update( + overrides, + force_catalog_refresh=force_catalog_refresh, + allow_unsafe=allow_unsafe, + ) + return self.apply_stream_config_update( + plan, + allow_unsafe=allow_unsafe, + defer_reset=defer_reset, + ) def rename(self, name: str) -> CloudConnection: """Rename the connection. diff --git a/airbyte/cloud/models.py b/airbyte/cloud/models.py index 59caccd83..b7b3358c5 100644 --- a/airbyte/cloud/models.py +++ b/airbyte/cloud/models.py @@ -8,7 +8,9 @@ from enum import Enum from typing import Any, Literal, Protocol -from pydantic import BaseModel, ConfigDict, Field +from pydantic import BaseModel, ConfigDict, Field, model_validator + +from airbyte_protocol.models import DestinationSyncMode, SyncMode # noqa: TC001 from airbyte._direct_connectors.models import ( DirectAccessGuidance, @@ -390,6 +392,54 @@ def __str__(self) -> str: return self.friendly_description +class StreamConfigOverride(BaseModel): + """Overrides for one stream's sync configuration.""" + + stream_name: str + """The stream name to update.""" + + stream_namespace: str | None = None + """The namespace to match, or `None` to match the name in any namespace.""" + + sync_mode: SyncMode | None = None + """The source sync mode to set, or `None` to keep the current mode.""" + + destination_sync_mode: DestinationSyncMode | None = None + """The destination sync mode to set, or `None` to keep the current mode.""" + + cursor_field: str | list[str] | None = None + """The cursor field path to set, or `None` to keep the current cursor.""" + + @model_validator(mode="after") + def require_setting(self) -> StreamConfigOverride: + """Require at least one sync configuration setting.""" + if ( + self.sync_mode is None + and self.destination_sync_mode is None + and self.cursor_field is None + ): + raise ValueError("At least one stream configuration setting must be provided.") + return self + + +class StreamConfigChange(BaseModel): + """A stream whose requested sync configuration was applied.""" + + name: str + namespace: str | None = None + changed_fields: list[str] + + +class StreamConfigUpdateResult(BaseModel): + """Result of updating stream sync configurations on a connection.""" + + updated_streams: list[StreamConfigChange] + unsafe_streams: list[StreamConfigChange] + reset_requested: bool + catalog_refreshed: bool + stream_count: int + + class CloudConnectionInfo(BaseModel): """Information about an Airbyte Cloud connection.""" @@ -599,6 +649,9 @@ def _enum_value(value: object) -> str: __all__ = [ "CloudConnectionInfo", + "StreamConfigOverride", + "StreamConfigChange", + "StreamConfigUpdateResult", "ConnectionStatus", "ConnectionSchedule", "CloudCustomSourceDefinitionInfo", diff --git a/airbyte/exceptions.py b/airbyte/exceptions.py index 94f413cd6..6e91b501d 100644 --- a/airbyte/exceptions.py +++ b/airbyte/exceptions.py @@ -217,6 +217,11 @@ class AirbyteLibInputError(AirbyteLibError, ValueError): input_value: str | None = None +@dataclass +class AirbyteUnsafeStreamConfigChangeError(AirbyteLibInputError): + """Updating stream configuration could invalidate saved sync state.""" + + @dataclass class AirbyteDeletionGuardError(AirbyteLibInputError): """Deletion rejected: the resource name must contain 'delete-me' or 'deleteme'. diff --git a/airbyte/mcp/_error_handling.py b/airbyte/mcp/_error_handling.py index 0360fece0..d2a68635e 100644 --- a/airbyte/mcp/_error_handling.py +++ b/airbyte/mcp/_error_handling.py @@ -41,6 +41,7 @@ AirbyteMissingResourceError, AirbytePipelineChangesDisabledError, AirbyteReadOnlyModeError, + AirbyteUnsafeStreamConfigChangeError, ) @@ -83,6 +84,7 @@ MCP_TOOL_USER_FACING_ERRORS: tuple[type[AirbyteLibError], ...] = ( AirbyteLibInputError, + AirbyteUnsafeStreamConfigChangeError, AirbyteMCPError, AirbyteLocalDemoModeError, AirbytePipelineChangesDisabledError, diff --git a/airbyte/mcp/cloud.py b/airbyte/mcp/cloud.py index 605559233..2310497bf 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 ( @@ -53,6 +51,7 @@ ConnectorType, JobTypeEnum, OrganizationFeature, + StreamConfigOverride, WorkspacePrivilegeScope, ) from airbyte.cloud.workspaces import CloudWorkspace @@ -76,6 +75,7 @@ AirbyteLibInputError, AirbyteMissingResourceError, AirbyteMissingWorkspaceContextError, + AirbyteUnsafeStreamConfigChangeError, ) from airbyte.mcp._arg_resolvers import ( resolve_api_args, @@ -95,6 +95,7 @@ AIRBYTE_CLOUD_WORKSPACE_ID_IS_SET, check_guid_created_in_session, delete_guard_disabled, + destructive_tools_allowed, mcp_tool, register_guid_created_in_session, ) @@ -3987,53 +3988,221 @@ def set_cloud_connection_selected_streams( ) +def _update_cloud_connection_impl( # noqa: PLR0913 + ctx: Context, + connection_id: str, + *, + enabled: bool | None, + cron_expression: str | None, + interval_hours: int | None, + manual_schedule: bool | None, + force_catalog_refresh: bool, + override_stream_configs: list[StreamConfigOverride] | None, + workspace_id: str | None, + allow_unsafe: bool, + defer_reset: bool, +) -> str: + if not any( + ( + enabled is not None, + cron_expression is not None, + interval_hours is not None, + manual_schedule is not None, + force_catalog_refresh, + bool(override_stream_configs), + ) + ): + raise AirbyteLibInputError( + message=( + "At least one setting must be provided: 'enabled', 'interval_hours', " + "'cron_expression', 'manual_schedule', 'force_catalog_refresh', or " + "'override_stream_configs'." + ), + ) + + if sum((interval_hours is not None, cron_expression is not None, manual_schedule is True)) > 1: + raise AirbyteLibInputError( + message=( + "Cannot specify more than one of 'interval_hours', 'cron_expression', " + "and 'manual_schedule=True'." + ), + ) + + if interval_hours is not None: + api_util.validate_interval_hours(interval_hours) + + workspace: CloudWorkspace = _get_cloud_workspace(ctx, workspace_id) + connection = workspace.get_connection(connection_id=connection_id) + + plan = None + if override_stream_configs: + try: + plan = connection.plan_stream_config_update( + override_stream_configs, + force_catalog_refresh=force_catalog_refresh, + allow_unsafe=allow_unsafe, + ) + except AirbyteUnsafeStreamConfigChangeError as error: + guidance = ( + "Ask the user to make this change in the Airbyte UI. " + f"Connection URL: {connection.connection_url}" + if destructive_tools_allowed(ctx) is False + else "Retry with `update_cloud_connection_unsafe`." + ) + raise AirbyteUnsafeStreamConfigChangeError( + message=error.get_message(), + context=error.context, + guidance=guidance, + ) from error + + changes_made: list[str] = [] + + if enabled is not None: + connection.set_enabled(enabled=enabled) + status_str = "enabled" if enabled else "disabled" + changes_made.append(f"status set to '{status_str}'") + + if cron_expression is not None: + connection.set_schedule(cron_expression=cron_expression) + changes_made.append(f"schedule set to '{cron_expression}'") + elif interval_hours is not None: + connection.set_schedule(interval_hours=interval_hours) + changes_made.append(f"schedule set to 'every {interval_hours} hours'") + elif manual_schedule is True: + connection.set_manual_schedule() + changes_made.append("schedule set to 'manual'") + + if plan is not None: + result = connection.apply_stream_config_update( + plan, + allow_unsafe=allow_unsafe, + defer_reset=defer_reset, + ) + updated_names = [ + f"{stream.namespace}.{stream.name}" if stream.namespace else stream.name + for stream in result.updated_streams + ] + changes_made.append(f"updated stream configs for {updated_names}") + if result.catalog_refreshed: + changes_made.append(f"refreshed catalog ({result.stream_count} streams)") + + if result.unsafe_streams: + unsafe_names = [ + f"{stream.namespace}.{stream.name}" if stream.namespace else stream.name + for stream in result.unsafe_streams + ] + if result.reset_requested: + changes_made.append( + f"unsafe streams {unsafe_names}; the platform will refresh or reset these " + "streams (refresh keeps data if the destination supports it; otherwise a " + "reset clears that stream's destination data)" + ) + else: + changes_made.append( + f"unsafe streams {unsafe_names}; no reset was run; these streams have stale " + "state and need a refresh/reset before their next sync" + ) + elif 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}. " + f"URL: {connection.connection_url}" + ) + + @mcp_tool( - destructive=True, + destructive=False, ) -def set_cloud_connection_stream_sync_mode( +def update_cloud_connection( # noqa: PLR0913 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[ + enabled: Annotated[ + bool | None, + Field( + description=( + "Set the connection's enabled status. " + "True enables the connection (status='active'), " + "False disables it (status='inactive'). " + "Leave unset to keep the current status." + ), + default=None, + ), + ], + cron_expression: Annotated[ str | None, Field( - description="The namespace of the stream, required when its name is ambiguous.", + description=( + "A Quartz cron expression defining when syncs should run. " + "Must have 6 or 7 space-separated fields " + "(seconds, minutes, hours, day-of-month, month, day-of-week[, year]), " + "optionally followed by a timezone ID. Standard 5-field Unix cron " + "expressions are rejected by the API, and schedules may run at most " + "once per hour (seconds and minutes cannot be '*'). " + "Examples: '0 0 0 * * ?' (daily at midnight UTC), " + "'0 0 */6 * * ?' (every 6 hours), " + "'0 0 0 ? * SUN' (weekly on Sunday at midnight UTC), " + "'0 0 9 ? * MON-FRI US/Pacific' (weekdays at 9am Pacific). " + "Leave unset to keep the current schedule. " + "Cannot be used together with 'interval_hours' or 'manual_schedule=True'." + ), + default=None, + ), + TraceArg.FINGERPRINT, + ], + interval_hours: Annotated[ + int | None, + Field( + description=( + "A positive whole number of hours between automatic syncs. " + "Uses a basic interval schedule: an active connection runs a first sync " + "as soon as scheduling is enabled, then measures each interval from " + "the previous sync. Leave unset to keep the current schedule. " + "Cannot be used together with 'cron_expression' or 'manual_schedule=True'." + ), default=None, + strict=True, + gt=0, ), ], - destination_sync_mode: Annotated[ - DestinationSyncMode | None, + manual_schedule: Annotated[ + bool | None, Field( description=( - "The destination sync mode to set. If omitted, the current mode is preserved." + "Set to True to disable automatic syncs (manual scheduling only). " + "Syncs will only run when manually triggered. " + "Cannot be used together with 'cron_expression' or 'interval_hours'." ), default=None, ), ], - cursor_field: Annotated[ - str | 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, + override_stream_configs: Annotated[ + list[StreamConfigOverride] | 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." + "Update one or more streams' sync mode, destination sync mode, or cursor field. " + "Unsafe changes to streams with saved state must be run with " + "`update_cloud_connection_unsafe`." ), default=None, ), TraceArg.FINGERPRINT, - ], + ] = None, workspace_id: Annotated[ str | None, Field( @@ -4041,41 +4210,35 @@ def set_cloud_connection_stream_sync_mode( default=None, ), TraceArg.OMIT, - ], + ] = None, ) -> 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 "" + """Update a connection's settings and stream configurations on Airbyte Cloud. - 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}" + This tool can update the connection status, schedule, source catalog, and per-stream + sync configuration. `interval_hours`, `cron_expression`, and `manual_schedule=True` + are mutually exclusive. When stream overrides are provided, they are validated before + any settings are written and saved in one catalog update. + """ + check_guid_created_in_session(connection_id) + return _update_cloud_connection_impl( + ctx, + connection_id, + enabled=enabled, + cron_expression=cron_expression, + interval_hours=interval_hours, + manual_schedule=manual_schedule, + force_catalog_refresh=force_catalog_refresh, + override_stream_configs=override_stream_configs, + workspace_id=workspace_id, + allow_unsafe=False, + defer_reset=False, ) @mcp_tool( destructive=True, ) -def update_cloud_connection( +def update_cloud_connection_unsafe( # noqa: PLR0913 ctx: Context, connection_id: Annotated[ str, @@ -4150,6 +4313,23 @@ def update_cloud_connection( ), ), ] = False, + override_stream_configs: Annotated[ + list[StreamConfigOverride] | None, + Field( + description="Update one or more streams' sync configuration, including unsafe changes.", + default=None, + ), + TraceArg.FINGERPRINT, + ] = None, + defer_reset: Annotated[ + bool, + Field( + description=( + "If true, do not request a platform reset. Streams with saved state will need " + "a refresh or reset before their next sync." + ), + ), + ] = False, workspace_id: Annotated[ str | None, Field( @@ -4157,80 +4337,27 @@ def update_cloud_connection( default=None, ), TraceArg.OMIT, - ], + ] = None, ) -> str: - """Update a connection's settings on Airbyte Cloud. - - This tool allows updating multiple connection settings in a single call: - - Enable or disable the 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. + """Update connection settings and apply stream overrides even when saved state is affected. + + Unsafe changes can make the destination's data inconsistent. By default, the platform + refreshes or resets affected streams; set `defer_reset=True` to leave stale state for a + later refresh or reset. """ check_guid_created_in_session(connection_id) - - # Validate that at least one setting is provided - if ( - enabled is None - 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'." - ), - ) - - # Validate mutually exclusive schedule options - if sum((interval_hours is not None, cron_expression is not None, manual_schedule is True)) > 1: - raise AirbyteLibInputError( - message=( - "Cannot specify more than one of 'interval_hours', 'cron_expression', " - "and 'manual_schedule=True'." - ), - ) - - if interval_hours is not None: - api_util.validate_interval_hours(interval_hours) - - workspace: CloudWorkspace = _get_cloud_workspace(ctx, workspace_id) - connection = workspace.get_connection(connection_id=connection_id) - - changes_made: list[str] = [] - - # Apply enabled status change - if enabled is not None: - connection.set_enabled(enabled=enabled) - status_str = "enabled" if enabled else "disabled" - changes_made.append(f"status set to '{status_str}'") - - # Apply schedule change - if cron_expression is not None: - connection.set_schedule(cron_expression=cron_expression) - changes_made.append(f"schedule set to '{cron_expression}'") - elif interval_hours is not None: - connection.set_schedule(interval_hours=interval_hours) - changes_made.append(f"schedule set to 'every {interval_hours} hours'") - elif manual_schedule is True: - 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}. " - f"URL: {connection.connection_url}" + return _update_cloud_connection_impl( + ctx, + connection_id, + enabled=enabled, + cron_expression=cron_expression, + interval_hours=interval_hours, + manual_schedule=manual_schedule, + force_catalog_refresh=force_catalog_refresh, + override_stream_configs=override_stream_configs, + workspace_id=workspace_id, + allow_unsafe=True, + defer_reset=defer_reset, ) diff --git a/tests/unit_tests/fixtures/arg_trace_plan.json b/tests/unit_tests/fixtures/arg_trace_plan.json index 0977d9549..ca3e15159 100644 --- a/tests/unit_tests/fixtures/arg_trace_plan.json +++ b/tests/unit_tests/fixtures/arg_trace_plan.json @@ -300,15 +300,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", @@ -357,6 +348,18 @@ "force_catalog_refresh": "value allowed=[false,true]", "interval_hours": "hash", "manual_schedule": "value allowed=[false,true]", + "override_stream_configs": "fingerprint list", + "workspace_id": "omit" + }, + "update_cloud_connection_unsafe": { + "connection_id": "hash", + "cron_expression": "fingerprint", + "defer_reset": "value allowed=[false,true]", + "enabled": "value allowed=[false,true]", + "force_catalog_refresh": "value allowed=[false,true]", + "interval_hours": "hash", + "manual_schedule": "value allowed=[false,true]", + "override_stream_configs": "fingerprint list", "workspace_id": "omit" }, "update_cloud_connector_config": { diff --git a/tests/unit_tests/test_cloud_connection_catalog_edits.py b/tests/unit_tests/test_cloud_connection_catalog_edits.py index 81a76d093..150d33e01 100644 --- a/tests/unit_tests/test_cloud_connection_catalog_edits.py +++ b/tests/unit_tests/test_cloud_connection_catalog_edits.py @@ -9,11 +9,16 @@ import pytest from airbyte_protocol.models import DestinationSyncMode, SyncMode +from pydantic import ValidationError from airbyte._util import api_util from airbyte.cloud.connections import CloudConnection +from airbyte.cloud.models import StreamConfigOverride from airbyte.cloud.workspaces import CloudWorkspace -from airbyte.exceptions import AirbyteLibInputError +from airbyte.exceptions import ( + AirbyteLibInputError, + AirbyteUnsafeStreamConfigChangeError, +) _CONNECTION_ID = "connection-id" @@ -75,228 +80,362 @@ def _catalog(*entries: dict[str, Any]) -> dict[str, Any]: def _mock_catalog_api( monkeypatch: pytest.MonkeyPatch, catalog: dict[str, Any], -) -> tuple[MagicMock, MagicMock]: +) -> tuple[MagicMock, MagicMock, MagicMock]: get_catalog = MagicMock(return_value={"syncCatalog": catalog}) replace_catalog = MagicMock(return_value={}) + get_state = MagicMock( + return_value={"stateType": "not_set", "connectionId": _CONNECTION_ID} + ) monkeypatch.setattr(api_util, "get_connection_catalog", get_catalog) monkeypatch.setattr(api_util, "replace_connection_catalog", replace_catalog) - return get_catalog, replace_catalog + monkeypatch.setattr(api_util, "get_connection_state", get_state) + return get_catalog, replace_catalog, get_state -def test_set_stream_sync_mode_updates_only_target_and_saves_raw_catalog( +def test_batch_overrides_write_once_and_preserve_untouched_catalog_data( monkeypatch: pytest.MonkeyPatch, ) -> None: - """Update the target stream without normalizing or changing other catalog data.""" target = _stream_entry("users", namespace="public") - other = _stream_entry( + other_target = _stream_entry( "orders", namespace="sales", + config={ + "syncMode": "full_refresh", + "destinationSyncMode": "append", + "selectedFields": [], + "fieldSelectionEnabled": False, + "primaryKey": [["order_id"]], + "selected": True, + }, + ) + untouched = _stream_entry( + "products", config={ "syncMode": "incremental", "destinationSyncMode": "append_dedup", - "cursorField": ["changed_at"], + "cursorField": ["updated_at"], "selectedFields": [{"fieldPath": ["id"]}], "fieldSelectionEnabled": True, - "primaryKey": [["order_id"]], + "primaryKey": [["sku"]], "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 + original_untouched = deepcopy(untouched) + catalog = _catalog(target, other_target, untouched) + _, replace_catalog, get_state = _mock_catalog_api(monkeypatch, catalog) + + result = _connection().update_stream_configs([ + StreamConfigOverride( + stream_name="users", + stream_namespace="public", + sync_mode=SyncMode.incremental, + destination_sync_mode=DestinationSyncMode.append_dedup, + cursor_field="updated_at", + ), + StreamConfigOverride( + stream_name="orders", + stream_namespace="sales", + destination_sync_mode=DestinationSyncMode.overwrite, + ), + ]) 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 saved_catalog["streams"][0]["config"]["cursorField"] == ["updated_at"] + assert saved_catalog["streams"][0]["config"]["selectedFields"] == [] + assert saved_catalog["streams"][0]["config"]["fieldSelectionEnabled"] is False + assert saved_catalog["streams"][0]["config"]["primaryKey"] == [["id"]] + assert saved_catalog["streams"][0]["config"]["selected"] is True + assert saved_catalog["streams"][2] == original_untouched 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 + assert replace_catalog.call_args.kwargs["skip_reset"] is True + assert result.stream_count == 3 + assert len(result.updated_streams) == 2 + get_state.assert_called_once() -@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", +def test_namespace_selects_same_named_stream_and_cursor_paths_are_preserved( + monkeypatch: pytest.MonkeyPatch, +) -> None: + public = _stream_entry("users", namespace="public") + sales = _stream_entry("users", namespace="sales") + _, replace_catalog, _ = _mock_catalog_api(monkeypatch, _catalog(public, sales)) + + _connection().update_stream_configs([ + StreamConfigOverride( + stream_name="users", + stream_namespace="sales", + cursor_field=["nested", "updated_at"], + ) + ]) + + saved = replace_catalog.call_args.kwargs["configured_catalog_dict"]["streams"] + assert saved[0]["config"].get("cursorField") is None + assert saved[1]["config"]["cursorField"] == ["nested", "updated_at"] + + +def test_ambiguous_bare_name_is_rejected_without_write( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _, replace_catalog, _ = _mock_catalog_api( + monkeypatch, + _catalog( + _stream_entry("users", namespace="public"), + _stream_entry("users", namespace="sales"), ), - 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", + ) + + with pytest.raises(AirbyteLibInputError, match="ambiguous"): + _connection().update_stream_configs([ + StreamConfigOverride(stream_name="users", sync_mode=SyncMode.full_refresh) + ]) + + replace_catalog.assert_not_called() + + +def test_duplicate_catalog_stream_is_rejected( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _, replace_catalog, _ = _mock_catalog_api( + monkeypatch, + _catalog(_stream_entry("users"), _stream_entry("users")), + ) + + with pytest.raises(AirbyteLibInputError, match="duplicate stream"): + _connection().update_stream_configs([ + StreamConfigOverride(stream_name="users", sync_mode=SyncMode.incremental) + ]) + + replace_catalog.assert_not_called() + + +def test_stream_override_requires_at_least_one_setting() -> None: + with pytest.raises(ValidationError, match="At least one"): + StreamConfigOverride(stream_name="users") + + +def test_all_invalid_overrides_are_reported_before_writing( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _, replace_catalog, get_state = _mock_catalog_api( + monkeypatch, + _catalog( + _stream_entry("users", sync_modes=["full_refresh"]), + _stream_entry("orders"), ), - pytest.param( - _catalog( - _stream_entry( - "users", - namespace="public", - config={ - "syncMode": "full_refresh", - "destinationSyncMode": "append", - }, - ) + ) + + with pytest.raises(AirbyteLibInputError) as exc_info: + _connection().update_stream_configs([ + StreamConfigOverride( + stream_name="missing", sync_mode=SyncMode.full_refresh ), - "users", - "public", - SyncMode.incremental, - None, - "needs a cursor", - {"stream_namespace": "public"}, - id="incremental-without-cursor", - ), + StreamConfigOverride(stream_name="users", sync_mode=SyncMode.incremental), + ]) + + message = exc_info.value.get_message() + assert "missing" in message + assert "unsupported" in message + assert len(exc_info.value.context["validation_errors"]) == 2 + replace_catalog.assert_not_called() + get_state.assert_not_called() + + +@pytest.mark.parametrize( + "state", + [ + { + "stateType": "stream", + "connectionId": _CONNECTION_ID, + "streamState": [ + { + "streamDescriptor": {"name": "users", "namespace": "public"}, + "streamState": {"cursor": "saved"}, + } + ], + }, + { + "stateType": "global", + "connectionId": _CONNECTION_ID, + "globalState": {"sharedState": {"cursor": "saved"}, "streamStates": []}, + }, ], + ids=["per-stream-state", "global-state"], ) -def test_set_stream_sync_mode_rejects_invalid_updates( +def test_saved_state_rejects_unsafe_config_changes( 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], + state: 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, - ) + _, replace_catalog, _ = _mock_catalog_api( + monkeypatch, + _catalog(_stream_entry("users", namespace="public")), + ) + monkeypatch.setattr( + api_util, + "get_connection_state", + MagicMock(return_value=state), + ) - assert exc_info.value.context is not None - for key, value in expected_context.items(): - assert exc_info.value.context[key] == value + with pytest.raises(AirbyteUnsafeStreamConfigChangeError) as exc_info: + _connection().update_stream_configs([ + StreamConfigOverride( + stream_name="users", + stream_namespace="public", + sync_mode=SyncMode.incremental, + cursor_field="updated_at", + ) + ]) + + unsafe = exc_info.value.context["unsafe_streams"] + assert unsafe[0]["name"] == "users" + assert unsafe[0]["namespace"] == "public" + assert unsafe[0]["changed_fields"] == ["syncMode", "cursorField"] replace_catalog.assert_not_called() -def test_set_stream_sync_mode_namespace_disambiguates_name( +def test_find_unsafe_stream_changes_compares_catalogs_to_normalized_state() -> None: + old_catalog = _catalog( + _stream_entry( + "users", + namespace="public", + config={"syncMode": "full_refresh", "destinationSyncMode": "append"}, + ), + _stream_entry( + "orders", + config={"syncMode": "full_refresh", "destinationSyncMode": "append"}, + ), + ) + new_catalog = deepcopy(old_catalog) + new_catalog["streams"][0]["config"].update({ + "syncMode": "incremental", + "cursorField": ["updated_at"], + }) + new_catalog["streams"][1]["config"]["destinationSyncMode"] = "overwrite" + state_messages = [ + { + "type": "STREAM", + "stream": { + "stream_descriptor": {"name": "users", "namespace": "public"}, + "stream_state": {"cursor": "saved"}, + }, + } + ] + + unsafe = CloudConnection._find_unsafe_stream_changes( + old_catalog, + new_catalog, + state_messages, + ) + + assert [ + (stream.name, stream.namespace, stream.changed_fields) for stream in unsafe + ] == [("users", "public", ["syncMode", "cursorField"])] + + +def test_no_state_lookup_when_override_does_not_change_config( 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"], - }, + _, replace_catalog, get_state = _mock_catalog_api( + monkeypatch, + _catalog(_stream_entry("users")), ) - 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, - ) + _connection().update_stream_configs([ + StreamConfigOverride(stream_name="users", sync_mode=SyncMode.full_refresh) + ]) - assert public["config"]["syncMode"] == "incremental" - assert sales["config"]["syncMode"] == "full_refresh" + get_state.assert_not_called() replace_catalog.assert_called_once() @pytest.mark.parametrize( - "stream_options", + "stream_options,expected_cursor", [ - pytest.param({"source_defined_cursor": True}, id="source-defined-cursor"), - pytest.param( - {"default_cursor_field": ["updated_at"]}, id="default-cursor-field" - ), + ({"source_defined_cursor": True}, None), + ({"default_cursor_field": ["updated_at"]}, ["updated_at"]), ], + ids=["source-defined-cursor", "default-cursor-field"], ) -def test_incremental_sync_accepts_stream_defined_cursor( +def test_incremental_sync_uses_default_but_not_source_defined_cursor( monkeypatch: pytest.MonkeyPatch, stream_options: dict[str, Any], + expected_cursor: list[str] | None, ) -> 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)) + _, replace_catalog, _ = _mock_catalog_api(monkeypatch, _catalog(entry)) + + _connection().update_stream_configs([ + StreamConfigOverride( + stream_name="users", + stream_namespace="public", + sync_mode=SyncMode.incremental, + ) + ]) + + saved_config = replace_catalog.call_args.kwargs["configured_catalog_dict"][ + "streams" + ][0]["config"] + assert saved_config.get("cursorField") == expected_cursor - _connection().set_stream_sync_mode( - stream_name="users", - stream_namespace="public", - sync_mode=SyncMode.incremental, + +def test_refresh_and_overrides_use_one_refreshed_catalog_write( + monkeypatch: pytest.MonkeyPatch, +) -> None: + refreshed_catalog = _catalog(_stream_entry("users", namespace="public")) + get_refreshed = MagicMock(return_value={"syncCatalog": refreshed_catalog}) + replace_catalog = MagicMock(return_value={}) + get_state = MagicMock( + return_value={"stateType": "not_set", "connectionId": _CONNECTION_ID} + ) + monkeypatch.setattr(api_util, "get_refreshed_connection_catalog", get_refreshed) + monkeypatch.setattr(api_util, "replace_connection_catalog", replace_catalog) + monkeypatch.setattr(api_util, "get_connection_state", get_state) + + result = _connection().update_stream_configs( + [ + StreamConfigOverride( + stream_name="users", + stream_namespace="public", + sync_mode=SyncMode.incremental, + cursor_field="updated_at", + ) + ], + force_catalog_refresh=True, ) - 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"] + get_refreshed.assert_called_once() replace_catalog.assert_called_once() + assert result.catalog_refreshed is True + assert result.stream_count == 1 + + +def test_bad_override_after_refresh_does_not_write( + monkeypatch: pytest.MonkeyPatch, +) -> None: + 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) + + with pytest.raises(AirbyteLibInputError, match="not found"): + _connection().update_stream_configs( + [ + StreamConfigOverride( + stream_name="missing", sync_mode=SyncMode.full_refresh + ) + ], + force_catalog_refresh=True, + ) + + get_refreshed.assert_called_once() + replace_catalog.assert_not_called() def test_refresh_catalog_gets_and_replaces_catalog_with_config_api_root( @@ -328,6 +467,7 @@ def test_refresh_catalog_gets_and_replaces_catalog_with_config_api_root( client_secret=None, bearer_token="token", config_api_root=_CONFIG_API_ROOT, + skip_reset=True, ) @@ -406,3 +546,23 @@ def test_get_refreshed_connection_catalog_forwards_config_api_root( client_secret=None, bearer_token="token", ) + + +def test_replace_connection_catalog_forwards_skip_reset( + monkeypatch: pytest.MonkeyPatch, +) -> None: + request = MagicMock(return_value={}) + monkeypatch.setattr(api_util, "_make_config_api_request", request) + + api_util.replace_connection_catalog( + connection_id=_CONNECTION_ID, + configured_catalog_dict={"streams": []}, + api_root=_API_ROOT, + config_api_root=_CONFIG_API_ROOT, + client_id=None, + client_secret=None, + bearer_token="token", + skip_reset=False, + ) + + assert request.call_args.kwargs["json"]["skipReset"] is False diff --git a/tests/unit_tests/test_mcp_connection_schedules.py b/tests/unit_tests/test_mcp_connection_schedules.py index 8db50fa54..7de7dbff9 100644 --- a/tests/unit_tests/test_mcp_connection_schedules.py +++ b/tests/unit_tests/test_mcp_connection_schedules.py @@ -17,14 +17,20 @@ from fastmcp.exceptions import ToolError from airbyte._util import api_util +from airbyte.cloud.connections import CloudConnection +from airbyte.cloud.models import StreamConfigOverride from airbyte.cloud.workspaces import CloudWorkspace from airbyte.constants import ( MCP_LOCAL_DEMO_MODE_ENV_VAR, MCP_READONLY_MODE_ENV_VAR, ) -from airbyte.exceptions import AirbyteLibInputError +from airbyte.exceptions import ( + AirbyteLibInputError, + AirbyteUnsafeStreamConfigChangeError, +) from airbyte.mcp import _tool_utils from airbyte.mcp import cloud as cloud_mcp +from airbyte_protocol.models import SyncMode def _connection_response() -> models.ConnectionResponse: @@ -102,6 +108,15 @@ async def _update(server: FastMCP, **arguments: object) -> str: return str(result) +async def _unsafe_update(server: FastMCP, **arguments: object) -> str: + async with Client(server) as client: + result = await client.call_tool( + "update_cloud_connection_unsafe", + {"connection_id": "12345678-1234-1234-1234-1234567890ab", **arguments}, + ) + return str(result) + + def test_interval_schema() -> None: """Advertise an optional positive whole-hour interval for the update tool.""" server = FastMCP("schedule-schema-tests") @@ -119,8 +134,12 @@ async def inspect_tools() -> None: assert interval["default"] is None assert "interval_hours" not in update.input_schema.get("required", []) assert update.annotations is not None - assert update.annotations.destructive_hint is True + assert update.annotations.destructive_hint is False assert update.annotations.read_only_hint is False + unsafe = tools["update_cloud_connection_unsafe"] + assert unsafe.annotations is not None + assert unsafe.annotations.destructive_hint is True + assert "set_cloud_connection_stream_sync_mode" not in tools asyncio.run(inspect_tools()) @@ -338,3 +357,231 @@ def test_update_force_catalog_refresh( get_refreshed.assert_called_once() replace_catalog.assert_called_once() assert "refreshed catalog (1 streams)" in result + + +@pytest.mark.parametrize( + ("defer_reset", "expected_skip_reset", "message"), + [ + (False, False, "the platform will refresh or reset these streams"), + (True, True, "these streams have stale state"), + ], +) +def test_unsafe_update_controls_catalog_reset( + schedule_backend: tuple[FastMCP, responses.RequestsMock, MagicMock, MagicMock], + monkeypatch: pytest.MonkeyPatch, + defer_reset: bool, + expected_skip_reset: bool, + message: str, +) -> None: + server, _http, _lookup, _patch = schedule_backend + monkeypatch.setattr( + api_util, + "get_connection_catalog", + MagicMock( + return_value={ + "syncCatalog": { + "streams": [ + { + "stream": { + "name": "users", + "namespace": "public", + "supportedSyncModes": ["full_refresh", "incremental"], + }, + "config": { + "syncMode": "full_refresh", + "destinationSyncMode": "append", + }, + } + ] + } + } + ), + ) + monkeypatch.setattr( + api_util, + "get_connection_state", + MagicMock( + return_value={ + "stateType": "stream", + "connectionId": "12345678-1234-1234-1234-1234567890ab", + "streamState": [ + { + "streamDescriptor": {"name": "users", "namespace": "public"}, + "streamState": {"cursor": "saved"}, + } + ], + } + ), + ) + replace_catalog = MagicMock(return_value={}) + monkeypatch.setattr(api_util, "replace_connection_catalog", replace_catalog) + + result = asyncio.run( + _unsafe_update( + server, + defer_reset=defer_reset, + override_stream_configs=[ + { + "stream_name": "users", + "stream_namespace": "public", + "sync_mode": "incremental", + "cursor_field": "updated_at", + } + ], + ) + ) + + assert replace_catalog.call_args.kwargs["skip_reset"] is expected_skip_reset + assert "users" in result + assert message in result + if defer_reset: + assert "no reset was run" in result + else: + assert "otherwise a reset clears that stream's destination data" in result + + +def test_bad_refreshed_override_prevents_schedule_write( + schedule_backend: tuple[FastMCP, responses.RequestsMock, MagicMock, MagicMock], + monkeypatch: pytest.MonkeyPatch, +) -> None: + server, _http, _lookup, patch = schedule_backend + get_refreshed = MagicMock( + return_value={"syncCatalog": {"streams": [{"stream": {"name": "users"}}]}} + ) + replace_catalog = MagicMock(return_value={}) + monkeypatch.setattr(api_util, "get_refreshed_connection_catalog", get_refreshed) + monkeypatch.setattr(api_util, "replace_connection_catalog", replace_catalog) + monkeypatch.setattr( + api_util, + "get_connection_state", + MagicMock( + return_value={ + "stateType": "not_set", + "connectionId": "12345678-1234-1234-1234-1234567890ab", + } + ), + ) + + with pytest.raises(ToolError, match="not found"): + asyncio.run( + _update( + server, + interval_hours=24, + force_catalog_refresh=True, + override_stream_configs=[ + {"stream_name": "missing", "sync_mode": "full_refresh"} + ], + ) + ) + + get_refreshed.assert_called_once() + replace_catalog.assert_not_called() + patch.assert_not_called() + + +def test_update_refreshes_and_overrides_in_one_catalog_write( + schedule_backend: tuple[FastMCP, responses.RequestsMock, MagicMock, MagicMock], + monkeypatch: pytest.MonkeyPatch, +) -> None: + server, _http, _lookup, _patch = schedule_backend + get_refreshed = MagicMock( + return_value={ + "syncCatalog": { + "streams": [ + { + "stream": { + "name": "users", + "namespace": "public", + "supportedSyncModes": ["full_refresh", "incremental"], + }, + "config": { + "syncMode": "full_refresh", + "destinationSyncMode": "append", + }, + } + ] + } + } + ) + replace_catalog = MagicMock(return_value={}) + monkeypatch.setattr(api_util, "get_refreshed_connection_catalog", get_refreshed) + monkeypatch.setattr(api_util, "replace_connection_catalog", replace_catalog) + monkeypatch.setattr( + api_util, + "get_connection_state", + MagicMock( + return_value={ + "stateType": "not_set", + "connectionId": "12345678-1234-1234-1234-1234567890ab", + } + ), + ) + + result = asyncio.run( + _update( + server, + force_catalog_refresh=True, + override_stream_configs=[ + { + "stream_name": "users", + "stream_namespace": "public", + "sync_mode": "incremental", + "cursor_field": "updated_at", + } + ], + ) + ) + + get_refreshed.assert_called_once() + replace_catalog.assert_called_once() + assert "refreshed catalog (1 streams)" in result + + +@pytest.mark.parametrize( + ("allowed", "expected_guidance"), + [ + (False, "Airbyte UI"), + (True, "update_cloud_connection_unsafe"), + (None, "update_cloud_connection_unsafe"), + ], +) +def test_unsafe_change_guidance_uses_destructive_policy( + schedule_backend: tuple[FastMCP, responses.RequestsMock, MagicMock, MagicMock], + monkeypatch: pytest.MonkeyPatch, + allowed: bool | None, + expected_guidance: str, +) -> None: + monkeypatch.setattr( + CloudConnection, + "plan_stream_config_update", + MagicMock( + side_effect=AirbyteUnsafeStreamConfigChangeError( + message="unsafe", + context={"unsafe_streams": [{"name": "users"}]}, + ) + ), + ) + monkeypatch.setattr(cloud_mcp, "destructive_tools_allowed", lambda _ctx: allowed) + + with pytest.raises(AirbyteUnsafeStreamConfigChangeError) as exc_info: + cloud_mcp.update_cloud_connection( + ctx=None, + connection_id="12345678-1234-1234-1234-1234567890ab", + enabled=None, + cron_expression=None, + interval_hours=None, + manual_schedule=None, + force_catalog_refresh=False, + override_stream_configs=[ + StreamConfigOverride( + stream_name="users", + sync_mode=SyncMode.incremental, + cursor_field="updated_at", + ) + ], + workspace_id=None, + ) + + assert expected_guidance in (exc_info.value.guidance or "") + if allowed is False: + assert "Connection URL:" in (exc_info.value.guidance or "") diff --git a/tests/unit_tests/test_mcp_tool_errors.py b/tests/unit_tests/test_mcp_tool_errors.py index 5a1aa2799..e13ef70be 100644 --- a/tests/unit_tests/test_mcp_tool_errors.py +++ b/tests/unit_tests/test_mcp_tool_errors.py @@ -22,6 +22,7 @@ AirbyteMissingResourceError, AirbyteLibError, AirbyteLibInputError, + AirbyteUnsafeStreamConfigChangeError, ) from airbyte.mcp._error_handling import ( MCP_TOOL_USER_FACING_ERRORS, @@ -67,6 +68,13 @@ async def call_tool() -> None: AirbyteLibInputError(message="bad", guidance="fix it"), id="input-error", ), + pytest.param( + AirbyteUnsafeStreamConfigChangeError( + message="unsafe", + guidance="review state", + ), + id="unsafe-stream-config", + ), pytest.param( AirbyteConnectorInUseError( message="Connector is in use.", From b3deff44efdf6d37be3192c9e5eb393e1cb9964b 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:41:56 +0000 Subject: [PATCH 07/12] refactor(mcp): remove catalog refresh from stream config updates Co-Authored-By: AJ Steers --- airbyte/_util/api_util.py | 36 ---- airbyte/cloud/connections.py | 81 ++------ airbyte/cloud/models.py | 3 +- airbyte/mcp/cloud.py | 52 ++--- tests/unit_tests/fixtures/arg_trace_plan.json | 2 - .../test_cloud_connection_catalog_edits.py | 184 +++--------------- .../test_mcp_connection_schedules.py | 141 ++------------ 7 files changed, 83 insertions(+), 416 deletions(-) diff --git a/airbyte/_util/api_util.py b/airbyte/_util/api_util.py index 3156a1633..e89871e74 100644 --- a/airbyte/_util/api_util.py +++ b/airbyte/_util/api_util.py @@ -3620,42 +3620,6 @@ 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 1b9ef91bd..4ca05116c 100644 --- a/airbyte/cloud/connections.py +++ b/airbyte/cloud/connections.py @@ -58,8 +58,6 @@ class _StreamConfigUpdatePlan: catalog: dict[str, Any] updated_streams: list[StreamConfigChange] unsafe_streams: list[StreamConfigChange] - catalog_refreshed: bool - stream_count: int def _normalized_cursor_field(value: list[str] | str | None) -> list[str]: @@ -975,42 +973,10 @@ def import_raw_catalog(self, catalog: dict[str, Any], *, skip_reset: bool = True ) 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 plan_stream_config_update( # noqa: PLR0912, PLR0914, PLR0915 self, overrides: list[StreamConfigOverride], *, - force_catalog_refresh: bool = False, allow_unsafe: bool = False, ) -> _StreamConfigUpdatePlan: """Validate overrides and plan a catalog update without writing it.""" @@ -1019,18 +985,7 @@ def plan_stream_config_update( # noqa: PLR0912, PLR0914, PLR0915 message="At least one stream configuration override must be provided.", ) - if force_catalog_refresh: - 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, - ) - catalog = refreshed_response.get("syncCatalog") - else: - catalog = self.dump_raw_catalog(normalize=False) + catalog = self.dump_raw_catalog(normalize=False) if ( not isinstance(catalog, dict) @@ -1038,11 +993,7 @@ def plan_stream_config_update( # noqa: PLR0912, PLR0914, PLR0915 or not catalog["streams"] ): raise AirbyteLibInputError( - message=( - "Refreshed catalog is empty or missing." - if force_catalog_refresh - else "Connection catalog is empty or missing." - ), + message="Connection catalog is empty or missing.", context={"connection_id": self.connection_id}, ) @@ -1084,6 +1035,7 @@ def plan_stream_config_update( # noqa: PLR0912, PLR0914, PLR0915 ) validation_errors: list[str] = [] + missing_streams = False updated_streams: list[StreamConfigChange] = [] override_keys: set[tuple[str, str | None]] = set() for index, override in enumerate(overrides, start=1): @@ -1103,9 +1055,8 @@ def plan_stream_config_update( # noqa: PLR0912, PLR0914, PLR0915 ] if not matching_entries: - validation_errors.append( - f"{label}: stream not found. Available streams: {available_streams}." - ) + validation_errors.append(f"{label}: stream not found.") + missing_streams = True continue if len(matching_entries) > 1: matching_namespaces = [ @@ -1182,13 +1133,18 @@ def plan_stream_config_update( # noqa: PLR0912, PLR0914, PLR0915 if validation_errors: joined_errors = " ".join(validation_errors) + message = f"Invalid stream configuration overrides: {joined_errors}" + context: dict[str, Any] = { + "connection_id": self.connection_id, + "validation_errors": validation_errors, + } + if missing_streams: + message += f" Available streams: {available_streams}." + context["available_streams"] = available_streams raise AirbyteLibInputError( - message=f"Invalid stream configuration overrides: {joined_errors}", + message=message, guidance="Correct every listed override and try again.", - context={ - "connection_id": self.connection_id, - "validation_errors": validation_errors, - }, + context=context, ) changed_streams = [stream for stream in updated_streams if stream.changed_fields] @@ -1209,8 +1165,6 @@ def plan_stream_config_update( # noqa: PLR0912, PLR0914, PLR0915 catalog=catalog, updated_streams=updated_streams, unsafe_streams=unsafe_streams, - catalog_refreshed=force_catalog_refresh, - stream_count=len(catalog["streams"]), ) @staticmethod @@ -1284,23 +1238,20 @@ def apply_stream_config_update( return StreamConfigUpdateResult( updated_streams=plan.updated_streams, unsafe_streams=plan.unsafe_streams, + reset_streams=[stream for stream in plan.updated_streams if stream.changed_fields], reset_requested=reset_requested, - catalog_refreshed=plan.catalog_refreshed, - stream_count=plan.stream_count, ) def update_stream_configs( self, overrides: list[StreamConfigOverride], *, - force_catalog_refresh: bool = False, allow_unsafe: bool = False, defer_reset: bool = False, ) -> StreamConfigUpdateResult: """Update one or more stream configurations in a single catalog write.""" plan = self.plan_stream_config_update( overrides, - force_catalog_refresh=force_catalog_refresh, allow_unsafe=allow_unsafe, ) return self.apply_stream_config_update( diff --git a/airbyte/cloud/models.py b/airbyte/cloud/models.py index b7b3358c5..603afe0c8 100644 --- a/airbyte/cloud/models.py +++ b/airbyte/cloud/models.py @@ -435,9 +435,8 @@ class StreamConfigUpdateResult(BaseModel): updated_streams: list[StreamConfigChange] unsafe_streams: list[StreamConfigChange] + reset_streams: list[StreamConfigChange] reset_requested: bool - catalog_refreshed: bool - stream_count: int class CloudConnectionInfo(BaseModel): diff --git a/airbyte/mcp/cloud.py b/airbyte/mcp/cloud.py index 2310497bf..6e8f6e202 100644 --- a/airbyte/mcp/cloud.py +++ b/airbyte/mcp/cloud.py @@ -3996,7 +3996,6 @@ def _update_cloud_connection_impl( # noqa: PLR0913 cron_expression: str | None, interval_hours: int | None, manual_schedule: bool | None, - force_catalog_refresh: bool, override_stream_configs: list[StreamConfigOverride] | None, workspace_id: str | None, allow_unsafe: bool, @@ -4008,15 +4007,13 @@ def _update_cloud_connection_impl( # noqa: PLR0913 cron_expression is not None, interval_hours is not None, manual_schedule is not None, - force_catalog_refresh, bool(override_stream_configs), ) ): raise AirbyteLibInputError( message=( "At least one setting must be provided: 'enabled', 'interval_hours', " - "'cron_expression', 'manual_schedule', 'force_catalog_refresh', or " - "'override_stream_configs'." + "'cron_expression', 'manual_schedule', or 'override_stream_configs'." ), ) @@ -4039,7 +4036,6 @@ def _update_cloud_connection_impl( # noqa: PLR0913 try: plan = connection.plan_stream_config_update( override_stream_configs, - force_catalog_refresh=force_catalog_refresh, allow_unsafe=allow_unsafe, ) except AirbyteUnsafeStreamConfigChangeError as error: @@ -4083,8 +4079,6 @@ def _update_cloud_connection_impl( # noqa: PLR0913 for stream in result.updated_streams ] changes_made.append(f"updated stream configs for {updated_names}") - if result.catalog_refreshed: - changes_made.append(f"refreshed catalog ({result.stream_count} streams)") if result.unsafe_streams: unsafe_names = [ @@ -4092,20 +4086,20 @@ def _update_cloud_connection_impl( # noqa: PLR0913 for stream in result.unsafe_streams ] if result.reset_requested: + reset_names = [ + f"{stream.namespace}.{stream.name}" if stream.namespace else stream.name + for stream in result.reset_streams + ] changes_made.append( - f"unsafe streams {unsafe_names}; the platform will refresh or reset these " - "streams (refresh keeps data if the destination supports it; otherwise a " - "reset clears that stream's destination data)" + f"unsafe streams {unsafe_names}; the platform will refresh or reset changed " + f"streams {reset_names} (refresh keeps data if the destination supports it; " + "otherwise a reset clears that stream's destination data)" ) else: changes_made.append( f"unsafe streams {unsafe_names}; no reset was run; these streams have stale " "state and need a refresh/reset before their next sync" ) - elif 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}. " @@ -4116,7 +4110,7 @@ def _update_cloud_connection_impl( # noqa: PLR0913 @mcp_tool( destructive=False, ) -def update_cloud_connection( # noqa: PLR0913 +def update_cloud_connection( ctx: Context, connection_id: Annotated[ str, @@ -4182,15 +4176,6 @@ def update_cloud_connection( # noqa: PLR0913 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, override_stream_configs: Annotated[ list[StreamConfigOverride] | None, Field( @@ -4214,10 +4199,10 @@ def update_cloud_connection( # noqa: PLR0913 ) -> str: """Update a connection's settings and stream configurations on Airbyte Cloud. - This tool can update the connection status, schedule, source catalog, and per-stream - sync configuration. `interval_hours`, `cron_expression`, and `manual_schedule=True` - are mutually exclusive. When stream overrides are provided, they are validated before - any settings are written and saved in one catalog update. + This tool can update the connection status, schedule, and per-stream sync configuration. + `interval_hours`, `cron_expression`, and `manual_schedule=True` are mutually exclusive. + Stream overrides are validated before any settings are written and saved in one catalog + update. """ check_guid_created_in_session(connection_id) return _update_cloud_connection_impl( @@ -4227,7 +4212,6 @@ def update_cloud_connection( # noqa: PLR0913 cron_expression=cron_expression, interval_hours=interval_hours, manual_schedule=manual_schedule, - force_catalog_refresh=force_catalog_refresh, override_stream_configs=override_stream_configs, workspace_id=workspace_id, allow_unsafe=False, @@ -4304,15 +4288,6 @@ def update_cloud_connection_unsafe( # noqa: PLR0913 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, override_stream_configs: Annotated[ list[StreamConfigOverride] | None, Field( @@ -4353,7 +4328,6 @@ def update_cloud_connection_unsafe( # noqa: PLR0913 cron_expression=cron_expression, interval_hours=interval_hours, manual_schedule=manual_schedule, - force_catalog_refresh=force_catalog_refresh, override_stream_configs=override_stream_configs, workspace_id=workspace_id, allow_unsafe=True, diff --git a/tests/unit_tests/fixtures/arg_trace_plan.json b/tests/unit_tests/fixtures/arg_trace_plan.json index ca3e15159..fdb293e67 100644 --- a/tests/unit_tests/fixtures/arg_trace_plan.json +++ b/tests/unit_tests/fixtures/arg_trace_plan.json @@ -345,7 +345,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]", "override_stream_configs": "fingerprint list", @@ -356,7 +355,6 @@ "cron_expression": "fingerprint", "defer_reset": "value allowed=[false,true]", "enabled": "value allowed=[false,true]", - "force_catalog_refresh": "value allowed=[false,true]", "interval_hours": "hash", "manual_schedule": "value allowed=[false,true]", "override_stream_configs": "fingerprint list", diff --git a/tests/unit_tests/test_cloud_connection_catalog_edits.py b/tests/unit_tests/test_cloud_connection_catalog_edits.py index 150d33e01..0cc1e9e8a 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 @@ -150,7 +150,6 @@ def test_batch_overrides_write_once_and_preserve_untouched_catalog_data( 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["skip_reset"] is True - assert result.stream_count == 3 assert len(result.updated_streams) == 2 get_state.assert_called_once() @@ -242,6 +241,34 @@ def test_all_invalid_overrides_are_reported_before_writing( get_state.assert_not_called() +def test_multiple_missing_streams_list_available_streams_once( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _, replace_catalog, _ = _mock_catalog_api( + monkeypatch, + _catalog(_stream_entry("users", namespace="public")), + ) + + with pytest.raises(AirbyteLibInputError) as exc_info: + _connection().update_stream_configs([ + StreamConfigOverride( + stream_name="missing-one", + sync_mode=SyncMode.full_refresh, + ), + StreamConfigOverride( + stream_name="missing-two", + sync_mode=SyncMode.full_refresh, + ), + ]) + + message = exc_info.value.get_message() + assert "missing-one" in message + assert "missing-two" in message + assert message.count("Available streams:") == 1 + assert exc_info.value.context["available_streams"] == ["public.users"] + replace_catalog.assert_not_called() + + @pytest.mark.parametrize( "state", [ @@ -384,170 +411,23 @@ def test_incremental_sync_uses_default_but_not_source_defined_cursor( assert saved_config.get("cursorField") == expected_cursor -def test_refresh_and_overrides_use_one_refreshed_catalog_write( - monkeypatch: pytest.MonkeyPatch, -) -> None: - refreshed_catalog = _catalog(_stream_entry("users", namespace="public")) - get_refreshed = MagicMock(return_value={"syncCatalog": refreshed_catalog}) - replace_catalog = MagicMock(return_value={}) - get_state = MagicMock( - return_value={"stateType": "not_set", "connectionId": _CONNECTION_ID} - ) - monkeypatch.setattr(api_util, "get_refreshed_connection_catalog", get_refreshed) - monkeypatch.setattr(api_util, "replace_connection_catalog", replace_catalog) - monkeypatch.setattr(api_util, "get_connection_state", get_state) - - result = _connection().update_stream_configs( - [ - StreamConfigOverride( - stream_name="users", - stream_namespace="public", - sync_mode=SyncMode.incremental, - cursor_field="updated_at", - ) - ], - force_catalog_refresh=True, - ) - - get_refreshed.assert_called_once() - replace_catalog.assert_called_once() - assert result.catalog_refreshed is True - assert result.stream_count == 1 - - -def test_bad_override_after_refresh_does_not_write( - monkeypatch: pytest.MonkeyPatch, -) -> None: - 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) - - with pytest.raises(AirbyteLibInputError, match="not found"): - _connection().update_stream_configs( - [ - StreamConfigOverride( - stream_name="missing", sync_mode=SyncMode.full_refresh - ) - ], - force_catalog_refresh=True, - ) - - get_refreshed.assert_called_once() - replace_catalog.assert_not_called() - - -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, - skip_reset=True, - ) - - -def test_refresh_catalog_clears_cached_connection_info( +def test_import_raw_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", - ) - - def test_replace_connection_catalog_forwards_skip_reset( monkeypatch: pytest.MonkeyPatch, ) -> None: diff --git a/tests/unit_tests/test_mcp_connection_schedules.py b/tests/unit_tests/test_mcp_connection_schedules.py index 7de7dbff9..94965305a 100644 --- a/tests/unit_tests/test_mcp_connection_schedules.py +++ b/tests/unit_tests/test_mcp_connection_schedules.py @@ -340,29 +340,10 @@ def test_update_interval_api_failure_is_reported( 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 - - @pytest.mark.parametrize( ("defer_reset", "expected_skip_reset", "message"), [ - (False, False, "the platform will refresh or reset these streams"), + (False, False, "the platform will refresh or reset changed streams"), (True, True, "these streams have stale state"), ], ) @@ -391,7 +372,18 @@ def test_unsafe_update_controls_catalog_reset( "syncMode": "full_refresh", "destinationSyncMode": "append", }, - } + }, + { + "stream": { + "name": "orders", + "namespace": "sales", + "supportedSyncModes": ["full_refresh", "incremental"], + }, + "config": { + "syncMode": "full_refresh", + "destinationSyncMode": "append", + }, + }, ] } } @@ -426,7 +418,12 @@ def test_unsafe_update_controls_catalog_reset( "stream_namespace": "public", "sync_mode": "incremental", "cursor_field": "updated_at", - } + }, + { + "stream_name": "orders", + "stream_namespace": "sales", + "destination_sync_mode": "overwrite", + }, ], ) ) @@ -437,106 +434,11 @@ def test_unsafe_update_controls_catalog_reset( if defer_reset: assert "no reset was run" in result else: + assert "unsafe streams ['public.users']" in result + assert "changed streams ['public.users', 'sales.orders']" in result assert "otherwise a reset clears that stream's destination data" in result -def test_bad_refreshed_override_prevents_schedule_write( - schedule_backend: tuple[FastMCP, responses.RequestsMock, MagicMock, MagicMock], - monkeypatch: pytest.MonkeyPatch, -) -> None: - server, _http, _lookup, patch = schedule_backend - get_refreshed = MagicMock( - return_value={"syncCatalog": {"streams": [{"stream": {"name": "users"}}]}} - ) - replace_catalog = MagicMock(return_value={}) - monkeypatch.setattr(api_util, "get_refreshed_connection_catalog", get_refreshed) - monkeypatch.setattr(api_util, "replace_connection_catalog", replace_catalog) - monkeypatch.setattr( - api_util, - "get_connection_state", - MagicMock( - return_value={ - "stateType": "not_set", - "connectionId": "12345678-1234-1234-1234-1234567890ab", - } - ), - ) - - with pytest.raises(ToolError, match="not found"): - asyncio.run( - _update( - server, - interval_hours=24, - force_catalog_refresh=True, - override_stream_configs=[ - {"stream_name": "missing", "sync_mode": "full_refresh"} - ], - ) - ) - - get_refreshed.assert_called_once() - replace_catalog.assert_not_called() - patch.assert_not_called() - - -def test_update_refreshes_and_overrides_in_one_catalog_write( - schedule_backend: tuple[FastMCP, responses.RequestsMock, MagicMock, MagicMock], - monkeypatch: pytest.MonkeyPatch, -) -> None: - server, _http, _lookup, _patch = schedule_backend - get_refreshed = MagicMock( - return_value={ - "syncCatalog": { - "streams": [ - { - "stream": { - "name": "users", - "namespace": "public", - "supportedSyncModes": ["full_refresh", "incremental"], - }, - "config": { - "syncMode": "full_refresh", - "destinationSyncMode": "append", - }, - } - ] - } - } - ) - replace_catalog = MagicMock(return_value={}) - monkeypatch.setattr(api_util, "get_refreshed_connection_catalog", get_refreshed) - monkeypatch.setattr(api_util, "replace_connection_catalog", replace_catalog) - monkeypatch.setattr( - api_util, - "get_connection_state", - MagicMock( - return_value={ - "stateType": "not_set", - "connectionId": "12345678-1234-1234-1234-1234567890ab", - } - ), - ) - - result = asyncio.run( - _update( - server, - force_catalog_refresh=True, - override_stream_configs=[ - { - "stream_name": "users", - "stream_namespace": "public", - "sync_mode": "incremental", - "cursor_field": "updated_at", - } - ], - ) - ) - - get_refreshed.assert_called_once() - replace_catalog.assert_called_once() - assert "refreshed catalog (1 streams)" in result - - @pytest.mark.parametrize( ("allowed", "expected_guidance"), [ @@ -571,7 +473,6 @@ def test_unsafe_change_guidance_uses_destructive_policy( cron_expression=None, interval_hours=None, manual_schedule=None, - force_catalog_refresh=False, override_stream_configs=[ StreamConfigOverride( stream_name="users", From 7e72381dcb04b525c52bad08b8006a29d0679ded Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 16 Oct 2026 22:17:56 +0000 Subject: [PATCH 08/12] refactor(cloud): make stream config planning private and flatten connection update tool Co-Authored-By: AJ Steers --- airbyte/cloud/connections.py | 448 ++++++++++-------- airbyte/cloud/models.py | 7 +- airbyte/mcp/cloud.py | 201 +++++--- .../test_cloud_connection_catalog_edits.py | 129 ++++- .../test_mcp_connection_schedules.py | 15 +- 5 files changed, 530 insertions(+), 270 deletions(-) diff --git a/airbyte/cloud/connections.py b/airbyte/cloud/connections.py index 4ca05116c..a7ca09bc9 100644 --- a/airbyte/cloud/connections.py +++ b/airbyte/cloud/connections.py @@ -5,10 +5,11 @@ import logging from copy import deepcopy -from dataclasses import dataclass from http import HTTPStatus from typing import TYPE_CHECKING, Any, Literal, overload +from airbyte_protocol.models import DestinationSyncMode + from airbyte._util import api_util from airbyte.cloud._connection_catalog import ( _denormalize_catalog_to_api, @@ -51,13 +52,9 @@ _QUARTZ_CRON_MIN_FIELDS = 6 _QUARTZ_CRON_MAX_FIELDS = 8 # 7 fields plus an optional trailing timezone ID - - -@dataclass -class _StreamConfigUpdatePlan: - catalog: dict[str, Any] - updated_streams: list[StreamConfigChange] - unsafe_streams: list[StreamConfigChange] +_DEDUP_DESTINATION_SYNC_MODES = { + mode.value for mode in DestinationSyncMode if mode.value.endswith("_dedup") +} def _normalized_cursor_field(value: list[str] | str | None) -> list[str]: @@ -114,6 +111,213 @@ def _stream_has_saved_state( return False +def _index_catalog_streams( + catalog: dict[str, Any] | None, + *, + connection_id: str, +) -> list[str]: + 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": connection_id}, + ) + + available_streams: list[str] = [] + catalog_keys: set[tuple[str, str | None]] = set() + duplicate_catalog_streams: list[str] = [] + for entry in catalog["streams"]: + stream_def = entry.get("stream") if isinstance(entry, dict) else None + if not isinstance(stream_def, dict): + raise AirbyteLibInputError( + message="Connection catalog contains an invalid stream entry.", + context={"connection_id": connection_id}, + ) + + name = stream_def.get("name") + namespace = stream_def.get("namespace") + if not isinstance(name, str) or (namespace is not None and not isinstance(namespace, str)): + raise AirbyteLibInputError( + message="Connection catalog contains an invalid stream descriptor.", + context={"connection_id": connection_id}, + ) + + available_streams.append(f"{namespace}.{name}" if namespace else name) + key = (name, namespace or None) + if key in catalog_keys: + duplicate_catalog_streams.append(f"{namespace}.{name}" if namespace else name) + catalog_keys.add(key) + + if duplicate_catalog_streams: + raise AirbyteLibInputError( + message=( + "Connection catalog contains duplicate stream entries: " + f"{', '.join(duplicate_catalog_streams)}." + ), + context={ + "connection_id": connection_id, + "duplicate_streams": duplicate_catalog_streams, + }, + ) + + return available_streams + + +def _resolve_incremental_cursor( + stream_def: dict[str, Any], + config: dict[str, Any], +) -> str | None: + if config.get("syncMode") != "incremental" or _normalized_cursor_field( + config.get("cursorField") + ): + return None + if stream_def.get("sourceDefinedCursor"): + return None + default_cursor_field = stream_def.get("defaultCursorField") + if isinstance(default_cursor_field, list) and default_cursor_field: + config["cursorField"] = deepcopy(default_cursor_field) + return None + return "incremental sync requires a cursor field." + + +def _stream_override_label( + override_index: int, + override: StreamConfigOverride, +) -> str: + return f"override {override_index} for '{override.stream_name}'" + ( + f" in namespace '{override.stream_namespace}'" + if override.stream_namespace is not None + else "" + ) + + +def _resolve_stream_override_target( + catalog: dict[str, Any], + override: StreamConfigOverride, + *, + override_index: int, + override_keys: set[tuple[str, str | None]], +) -> dict[str, Any] | str: + label = _stream_override_label( + override_index, + override, + ) + matching_entries = [ + entry + for entry in catalog["streams"] + if entry["stream"].get("name") == override.stream_name + and ( + override.stream_namespace is None + or (entry["stream"].get("namespace") or None) == (override.stream_namespace or None) + ) + ] + if not matching_entries: + return f"{label}: stream not found." + if len(matching_entries) > 1: + matching_namespaces = [entry["stream"].get("namespace") for entry in matching_entries] + return ( + f"{label}: stream name is ambiguous; specify a namespace. " + f"Matching namespaces: {matching_namespaces}." + ) + + target_entry = matching_entries[0] + stream_def = target_entry["stream"] + stream_namespace = stream_def.get("namespace") + target_key = (override.stream_name, stream_namespace or None) + if target_key in override_keys: + return f"{label}: this stream is overridden more than once." + override_keys.add(target_key) + return target_entry + + +def _apply_stream_override( + catalog: dict[str, Any], + override: StreamConfigOverride, + *, + override_index: int, + override_keys: set[tuple[str, str | None]], +) -> StreamConfigChange | str: + target_entry = _resolve_stream_override_target( + catalog, + override, + override_index=override_index, + override_keys=override_keys, + ) + if isinstance(target_entry, str): + return target_entry + + label = _stream_override_label(override_index, override) + stream_def = target_entry["stream"] + stream_namespace = stream_def.get("namespace") + current_config = target_entry.get("config") or {} + if not isinstance(current_config, dict): + return f"{label}: stream configuration is invalid." + + candidate_config = deepcopy(current_config) + supported_sync_modes = stream_def.get("supportedSyncModes") or [] + if override.sync_mode is not None and override.sync_mode.value not in supported_sync_modes: + return ( + f"{label}: sync mode '{override.sync_mode.value}' is unsupported. " + f"Supported modes: {supported_sync_modes}." + ) + + if override.sync_mode is not None: + candidate_config["syncMode"] = override.sync_mode.value + if override.destination_sync_mode is not None: + candidate_config["destinationSyncMode"] = override.destination_sync_mode.value + if override.cursor_field is not None: + candidate_config["cursorField"] = [override.cursor_field] + + cursor_error = _resolve_incremental_cursor(stream_def, candidate_config) + if cursor_error is not None: + return f"{label}: {cursor_error}" + + destination_sync_mode = candidate_config.get("destinationSyncMode") + if ( + destination_sync_mode in _DEDUP_DESTINATION_SYNC_MODES + and not candidate_config.get("primaryKey") + and not stream_def.get("sourceDefinedPrimaryKey") + ): + return ( + f"{label}: destination sync mode '{destination_sync_mode}' requires " "a primary key." + ) + + target_entry["config"] = candidate_config + return StreamConfigChange( + name=override.stream_name, + namespace=stream_namespace, + changed_fields=_changed_stream_config_fields(current_config, candidate_config), + ) + + +def _raise_for_override_errors( + validation_errors: list[str], + available_streams: list[str], + *, + connection_id: str, +) -> None: + if not validation_errors: + return + + message = f"Invalid stream configuration overrides: {' '.join(validation_errors)}" + context: dict[str, Any] = { + "connection_id": connection_id, + "validation_errors": validation_errors, + } + if any(": stream not found." in error for error in validation_errors): + message += f" Available streams: {available_streams}." + context["available_streams"] = available_streams + + raise AirbyteLibInputError( + message=message, + guidance="Correct every listed override and try again.", + context=context, + ) + + def _validate_quartz_cron_expression(cron_expression: str) -> None: """Raise `AirbyteLibInputError` if the expression is not a plausible Quartz cron. @@ -973,179 +1177,52 @@ def import_raw_catalog(self, catalog: dict[str, Any], *, skip_reset: bool = True ) self._connection_info = None - def plan_stream_config_update( # noqa: PLR0912, PLR0914, PLR0915 + def _evaluate_stream_config_updates( self, overrides: list[StreamConfigOverride], *, allow_unsafe: bool = False, - ) -> _StreamConfigUpdatePlan: - """Validate overrides and plan a catalog update without writing it.""" + ) -> tuple[dict[str, Any], list[StreamConfigChange], list[StreamConfigChange]]: + """Validate overrides and evaluate their effect without writing the catalog.""" if not overrides: raise AirbyteLibInputError( message="At least one stream configuration override must be provided.", ) catalog = self.dump_raw_catalog(normalize=False) - - if ( - not isinstance(catalog, dict) - or not isinstance(catalog.get("streams"), list) - or not catalog["streams"] - ): + if catalog is None: raise AirbyteLibInputError( message="Connection catalog is empty or missing.", context={"connection_id": self.connection_id}, ) - old_catalog = deepcopy(catalog) + available_streams = _index_catalog_streams( + catalog, + connection_id=self.connection_id, + ) + old_catalog = catalog catalog = deepcopy(catalog) - available_streams: list[str] = [] - catalog_keys: set[tuple[str, str | None]] = set() - duplicate_catalog_streams: list[str] = [] - for entry in catalog["streams"]: - stream_def = entry.get("stream") if isinstance(entry, dict) else None - if not isinstance(stream_def, dict): - raise AirbyteLibInputError( - message="Connection catalog contains an invalid stream entry.", - context={"connection_id": self.connection_id}, - ) - name = stream_def.get("name") - namespace = stream_def.get("namespace") - if not isinstance(name, str) or ( - namespace is not None and not isinstance(namespace, str) - ): - raise AirbyteLibInputError( - message="Connection catalog contains an invalid stream descriptor.", - context={"connection_id": self.connection_id}, - ) - available_streams.append(f"{namespace}.{name}" if namespace else name) - key = (name, namespace) - if key in catalog_keys: - duplicate_catalog_streams.append(f"{namespace}.{name}" if namespace else name) - catalog_keys.add(key) - - if duplicate_catalog_streams: - duplicate_names = ", ".join(duplicate_catalog_streams) - raise AirbyteLibInputError( - message=f"Connection catalog contains duplicate stream entries: {duplicate_names}.", - context={ - "connection_id": self.connection_id, - "duplicate_streams": duplicate_catalog_streams, - }, - ) validation_errors: list[str] = [] - missing_streams = False updated_streams: list[StreamConfigChange] = [] override_keys: set[tuple[str, str | None]] = set() for index, override in enumerate(overrides, start=1): - label = f"override {index} for " f"'{override.stream_name}'" + ( - f" in namespace '{override.stream_namespace}'" - if override.stream_namespace is not None - else "" - ) - matching_entries = [ - entry - for entry in catalog["streams"] - if entry["stream"].get("name") == override.stream_name - and ( - override.stream_namespace is None - or entry["stream"].get("namespace") == override.stream_namespace - ) - ] - - if not matching_entries: - validation_errors.append(f"{label}: stream not found.") - missing_streams = True - continue - if len(matching_entries) > 1: - matching_namespaces = [ - entry["stream"].get("namespace") for entry in matching_entries - ] - validation_errors.append( - f"{label}: stream name is ambiguous; specify a namespace. " - f"Matching namespaces: {matching_namespaces}." - ) - continue - - target_entry = matching_entries[0] - stream_def = target_entry["stream"] - stream_namespace = stream_def.get("namespace") - target_key = (override.stream_name, stream_namespace) - if target_key in override_keys: - validation_errors.append(f"{label}: this stream is overridden more than once.") - continue - override_keys.add(target_key) - - current_config = target_entry.get("config") or {} - if not isinstance(current_config, dict): - validation_errors.append(f"{label}: stream configuration is invalid.") - continue - - candidate_config = deepcopy(current_config) - supported_sync_modes = stream_def.get("supportedSyncModes") or [] - if ( - override.sync_mode is not None - and override.sync_mode.value not in supported_sync_modes - ): - validation_errors.append( - f"{label}: sync mode '{override.sync_mode.value}' is unsupported. " - f"Supported modes: {supported_sync_modes}." - ) - continue - - if override.sync_mode is not None: - candidate_config["syncMode"] = override.sync_mode.value - if override.destination_sync_mode is not None: - candidate_config["destinationSyncMode"] = override.destination_sync_mode.value - if override.cursor_field is not None: - candidate_config["cursorField"] = ( - [override.cursor_field] - if isinstance(override.cursor_field, str) - else deepcopy(override.cursor_field) - ) - - if candidate_config.get("syncMode") == "incremental" and not _normalized_cursor_field( - candidate_config.get("cursorField") - ): - if stream_def.get("sourceDefinedCursor"): - pass - elif ( - isinstance(stream_def.get("defaultCursorField"), list) - and stream_def["defaultCursorField"] - ): - candidate_config["cursorField"] = deepcopy(stream_def["defaultCursorField"]) - else: - validation_errors.append(f"{label}: incremental sync requires a cursor field.") - continue - - target_entry["config"] = candidate_config - updated_streams.append( - StreamConfigChange( - name=override.stream_name, - namespace=stream_namespace, - changed_fields=_changed_stream_config_fields( - current_config, - candidate_config, - ), - ) + result = _apply_stream_override( + catalog, + override, + override_index=index, + override_keys=override_keys, ) + if isinstance(result, str): + validation_errors.append(result) + else: + updated_streams.append(result) - if validation_errors: - joined_errors = " ".join(validation_errors) - message = f"Invalid stream configuration overrides: {joined_errors}" - context: dict[str, Any] = { - "connection_id": self.connection_id, - "validation_errors": validation_errors, - } - if missing_streams: - message += f" Available streams: {available_streams}." - context["available_streams"] = available_streams - raise AirbyteLibInputError( - message=message, - guidance="Correct every listed override and try again.", - context=context, - ) + _raise_for_override_errors( + validation_errors, + available_streams, + connection_id=self.connection_id, + ) changed_streams = [stream for stream in updated_streams if stream.changed_fields] unsafe_streams: list[StreamConfigChange] = [] @@ -1161,11 +1238,7 @@ def plan_stream_config_update( # noqa: PLR0912, PLR0914, PLR0915 unsafe_streams, allow_unsafe=allow_unsafe, ) - return _StreamConfigUpdatePlan( - catalog=catalog, - updated_streams=updated_streams, - unsafe_streams=unsafe_streams, - ) + return catalog, updated_streams, unsafe_streams @staticmethod def _find_unsafe_stream_changes( @@ -1213,35 +1286,23 @@ def _raise_for_unsafe_stream_config_changes( allow_unsafe: bool, ) -> None: if unsafe_streams and not allow_unsafe: + stream_details = ", ".join( + f"{stream.namespace}.{stream.name} ({', '.join(stream.changed_fields)})" + if stream.namespace + else f"{stream.name} ({', '.join(stream.changed_fields)})" + for stream in unsafe_streams + ) raise AirbyteUnsafeStreamConfigChangeError( - message="Stream configuration changes affect streams with saved sync state.", + message=( + "Stream configuration changes affect streams with saved sync state: " + f"{stream_details}." + ), context={ "connection_id": self.connection_id, "unsafe_streams": [stream.model_dump() for stream in unsafe_streams], }, ) - def apply_stream_config_update( - self, - plan: _StreamConfigUpdatePlan, - *, - allow_unsafe: bool = False, - defer_reset: bool = False, - ) -> StreamConfigUpdateResult: - """Write a planned catalog update and return its result.""" - self._raise_for_unsafe_stream_config_changes( - plan.unsafe_streams, - allow_unsafe=allow_unsafe, - ) - reset_requested = bool(plan.unsafe_streams and allow_unsafe and not defer_reset) - self.import_raw_catalog(plan.catalog, skip_reset=not reset_requested) - return StreamConfigUpdateResult( - updated_streams=plan.updated_streams, - unsafe_streams=plan.unsafe_streams, - reset_streams=[stream for stream in plan.updated_streams if stream.changed_fields], - reset_requested=reset_requested, - ) - def update_stream_configs( self, overrides: list[StreamConfigOverride], @@ -1250,14 +1311,21 @@ def update_stream_configs( defer_reset: bool = False, ) -> StreamConfigUpdateResult: """Update one or more stream configurations in a single catalog write.""" - plan = self.plan_stream_config_update( + catalog, updated_streams, unsafe_streams = self._evaluate_stream_config_updates( overrides, allow_unsafe=allow_unsafe, ) - return self.apply_stream_config_update( - plan, - allow_unsafe=allow_unsafe, - defer_reset=defer_reset, + reset_requested = bool(unsafe_streams and allow_unsafe and not defer_reset) + self.import_raw_catalog(catalog, skip_reset=not reset_requested) + return StreamConfigUpdateResult( + updated_streams=updated_streams, + unsafe_streams=unsafe_streams, + reset_streams=( + [stream for stream in updated_streams if stream.changed_fields] + if reset_requested + else [] + ), + reset_requested=reset_requested, ) def rename(self, name: str) -> CloudConnection: diff --git a/airbyte/cloud/models.py b/airbyte/cloud/models.py index 603afe0c8..139184422 100644 --- a/airbyte/cloud/models.py +++ b/airbyte/cloud/models.py @@ -399,7 +399,8 @@ class StreamConfigOverride(BaseModel): """The stream name to update.""" stream_namespace: str | None = None - """The namespace to match, or `None` to match the name in any namespace.""" + """The namespace to match; `None` matches any namespace and `""` matches + unnamespaced streams.""" sync_mode: SyncMode | None = None """The source sync mode to set, or `None` to keep the current mode.""" @@ -407,8 +408,8 @@ class StreamConfigOverride(BaseModel): destination_sync_mode: DestinationSyncMode | None = None """The destination sync mode to set, or `None` to keep the current mode.""" - cursor_field: str | list[str] | None = None - """The cursor field path to set, or `None` to keep the current cursor.""" + cursor_field: str | None = None + """The top-level cursor field name to set, or `None` to keep the current cursor.""" @model_validator(mode="after") def require_setting(self) -> StreamConfigOverride: diff --git a/airbyte/mcp/cloud.py b/airbyte/mcp/cloud.py index 6e8f6e202..95b56274f 100644 --- a/airbyte/mcp/cloud.py +++ b/airbyte/mcp/cloud.py @@ -33,6 +33,7 @@ ) from airbyte._util import api_util from airbyte.cloud.client import _MAX_WORKSPACES_TO_VALIDATE, CloudClient +from airbyte.cloud.connections import CloudConnection from airbyte.cloud.connectors import ( CheckResult, CloudConnector, @@ -51,7 +52,9 @@ ConnectorType, JobTypeEnum, OrganizationFeature, + StreamConfigChange, StreamConfigOverride, + StreamConfigUpdateResult, WorkspacePrivilegeScope, ) from airbyte.cloud.workspaces import CloudWorkspace @@ -3988,19 +3991,14 @@ def set_cloud_connection_selected_streams( ) -def _update_cloud_connection_impl( # noqa: PLR0913 - ctx: Context, - connection_id: str, +def _validate_update_connection_args( *, enabled: bool | None, cron_expression: str | None, interval_hours: int | None, manual_schedule: bool | None, override_stream_configs: list[StreamConfigOverride] | None, - workspace_id: str | None, - allow_unsafe: bool, - defer_reset: bool, -) -> str: +) -> None: if not any( ( enabled is not None, @@ -4028,78 +4026,137 @@ def _update_cloud_connection_impl( # noqa: PLR0913 if interval_hours is not None: api_util.validate_interval_hours(interval_hours) - workspace: CloudWorkspace = _get_cloud_workspace(ctx, workspace_id) - connection = workspace.get_connection(connection_id=connection_id) - plan = None - if override_stream_configs: - try: - plan = connection.plan_stream_config_update( - override_stream_configs, - allow_unsafe=allow_unsafe, - ) - except AirbyteUnsafeStreamConfigChangeError as error: - guidance = ( - "Ask the user to make this change in the Airbyte UI. " - f"Connection URL: {connection.connection_url}" - if destructive_tools_allowed(ctx) is False - else "Retry with `update_cloud_connection_unsafe`." - ) - raise AirbyteUnsafeStreamConfigChangeError( - message=error.get_message(), - context=error.context, - guidance=guidance, - ) from error +def _apply_stream_config_overrides( + ctx: Context, + connection: CloudConnection, + overrides: list[StreamConfigOverride] | None, + *, + allow_unsafe: bool, + defer_reset: bool, +) -> StreamConfigUpdateResult | None: + if not overrides: + return None - changes_made: list[str] = [] + try: + return connection.update_stream_configs( + overrides, + allow_unsafe=allow_unsafe, + defer_reset=defer_reset, + ) + except AirbyteUnsafeStreamConfigChangeError as error: + guidance = ( + "Ask the user to make this change in the Airbyte UI. " + f"Connection URL: {connection.connection_url}" + if destructive_tools_allowed(ctx) is False + else "Retry with `update_cloud_connection_unsafe`." + ) + raise AirbyteUnsafeStreamConfigChangeError( + message=error.get_message(), + context=error.context, + guidance=guidance, + ) from error - if enabled is not None: - connection.set_enabled(enabled=enabled) - status_str = "enabled" if enabled else "disabled" - changes_made.append(f"status set to '{status_str}'") +def _set_connection_schedule( + connection: CloudConnection, + *, + cron_expression: str | None, + interval_hours: int | None, + manual_schedule: bool | None, +) -> str | None: if cron_expression is not None: connection.set_schedule(cron_expression=cron_expression) - changes_made.append(f"schedule set to '{cron_expression}'") - elif interval_hours is not None: + return f"schedule set to '{cron_expression}'" + if interval_hours is not None: connection.set_schedule(interval_hours=interval_hours) - changes_made.append(f"schedule set to 'every {interval_hours} hours'") - elif manual_schedule is True: + return f"schedule set to 'every {interval_hours} hours'" + if manual_schedule is True: connection.set_manual_schedule() - changes_made.append("schedule set to 'manual'") + return "schedule set to 'manual'" + return None - if plan is not None: - result = connection.apply_stream_config_update( - plan, - allow_unsafe=allow_unsafe, - defer_reset=defer_reset, + +def _stream_label(stream: StreamConfigChange) -> str: + return f"{stream.namespace}.{stream.name}" if stream.namespace else stream.name + + +def _describe_stream_config_result(result: StreamConfigUpdateResult) -> list[str]: + updated_names = [_stream_label(stream) for stream in result.updated_streams] + changes = [f"updated stream configs for {updated_names}"] + if not result.unsafe_streams: + return changes + + unsafe_names = [_stream_label(stream) for stream in result.unsafe_streams] + if result.reset_requested: + reset_names = [_stream_label(stream) for stream in result.reset_streams] + changes.append( + f"unsafe streams {unsafe_names}; the platform will refresh or reset changed " + f"streams {reset_names} (refresh keeps data if the destination supports it; " + "otherwise a reset clears that stream's destination data)" ) - updated_names = [ - f"{stream.namespace}.{stream.name}" if stream.namespace else stream.name - for stream in result.updated_streams - ] - changes_made.append(f"updated stream configs for {updated_names}") + else: + changes.append( + f"unsafe streams {unsafe_names}; no reset was run; these streams have stale " + "state and need a refresh/reset before their next sync" + ) + return changes + + +def _update_cloud_connection_impl( # noqa: PLR0913 + ctx: Context, + connection_id: str, + *, + enabled: bool | None, + cron_expression: str | None, + interval_hours: int | None, + manual_schedule: bool | None, + override_stream_configs: list[StreamConfigOverride] | None, + workspace_id: str | None, + allow_unsafe: bool, + defer_reset: bool, +) -> str: + _validate_update_connection_args( + enabled=enabled, + cron_expression=cron_expression, + interval_hours=interval_hours, + manual_schedule=manual_schedule, + override_stream_configs=override_stream_configs, + ) + + workspace: CloudWorkspace = _get_cloud_workspace(ctx, workspace_id) + connection = workspace.get_connection(connection_id=connection_id) + + stream_config_result = _apply_stream_config_overrides( + ctx, + connection, + override_stream_configs, + allow_unsafe=allow_unsafe, + defer_reset=defer_reset, + ) + + if enabled is not None: + connection.set_enabled(enabled=enabled) + status_str = "enabled" if enabled else "disabled" + status_summary = f"status set to '{status_str}'" + else: + status_summary = None + + schedule_summary = _set_connection_schedule( + connection, + cron_expression=cron_expression, + interval_hours=interval_hours, + manual_schedule=manual_schedule, + ) + + changes_made = ( + _describe_stream_config_result(stream_config_result) if stream_config_result else [] + ) + if status_summary is not None: + changes_made.append(status_summary) + if schedule_summary is not None: + changes_made.append(schedule_summary) - if result.unsafe_streams: - unsafe_names = [ - f"{stream.namespace}.{stream.name}" if stream.namespace else stream.name - for stream in result.unsafe_streams - ] - if result.reset_requested: - reset_names = [ - f"{stream.namespace}.{stream.name}" if stream.namespace else stream.name - for stream in result.reset_streams - ] - changes_made.append( - f"unsafe streams {unsafe_names}; the platform will refresh or reset changed " - f"streams {reset_names} (refresh keeps data if the destination supports it; " - "otherwise a reset clears that stream's destination data)" - ) - else: - changes_made.append( - f"unsafe streams {unsafe_names}; no reset was run; these streams have stale " - "state and need a refresh/reset before their next sync" - ) changes_summary = ", ".join(changes_made) return ( f"Successfully updated connection '{connection_id}': {changes_summary}. " @@ -4180,7 +4237,8 @@ def update_cloud_connection( list[StreamConfigOverride] | None, Field( description=( - "Update one or more streams' sync mode, destination sync mode, or cursor field. " + "Update one or more streams' sync mode, destination sync mode, or top-level " + "cursor field. " "Unsafe changes to streams with saved state must be run with " "`update_cloud_connection_unsafe`." ), @@ -4291,7 +4349,10 @@ def update_cloud_connection_unsafe( # noqa: PLR0913 override_stream_configs: Annotated[ list[StreamConfigOverride] | None, Field( - description="Update one or more streams' sync configuration, including unsafe changes.", + description=( + "Update one or more streams' sync configuration, including a top-level cursor " + "field and unsafe changes." + ), default=None, ), TraceArg.FINGERPRINT, diff --git a/tests/unit_tests/test_cloud_connection_catalog_edits.py b/tests/unit_tests/test_cloud_connection_catalog_edits.py index 0cc1e9e8a..f905f4a48 100644 --- a/tests/unit_tests/test_cloud_connection_catalog_edits.py +++ b/tests/unit_tests/test_cloud_connection_catalog_edits.py @@ -154,7 +154,7 @@ def test_batch_overrides_write_once_and_preserve_untouched_catalog_data( get_state.assert_called_once() -def test_namespace_selects_same_named_stream_and_cursor_paths_are_preserved( +def test_namespace_selects_same_named_stream_and_sets_top_level_cursor( monkeypatch: pytest.MonkeyPatch, ) -> None: public = _stream_entry("users", namespace="public") @@ -165,13 +165,44 @@ def test_namespace_selects_same_named_stream_and_cursor_paths_are_preserved( StreamConfigOverride( stream_name="users", stream_namespace="sales", - cursor_field=["nested", "updated_at"], + cursor_field="updated_at", ) ]) saved = replace_catalog.call_args.kwargs["configured_catalog_dict"]["streams"] assert saved[0]["config"].get("cursorField") is None - assert saved[1]["config"]["cursorField"] == ["nested", "updated_at"] + assert saved[1]["config"]["cursorField"] == ["updated_at"] + + +def test_empty_namespace_selects_unnamespaced_stream( + monkeypatch: pytest.MonkeyPatch, +) -> None: + unnamespaced = _stream_entry("users") + public = _stream_entry("users", namespace="public") + _, replace_catalog, _ = _mock_catalog_api( + monkeypatch, + _catalog(unnamespaced, public), + ) + + _connection().update_stream_configs([ + StreamConfigOverride( + stream_name="users", + stream_namespace="", + destination_sync_mode=DestinationSyncMode.overwrite, + ) + ]) + + saved = replace_catalog.call_args.kwargs["configured_catalog_dict"]["streams"] + assert saved[0]["config"]["destinationSyncMode"] == "overwrite" + assert saved[1]["config"]["destinationSyncMode"] == "append" + + +def test_cursor_override_rejects_nested_paths() -> None: + with pytest.raises(ValidationError): + StreamConfigOverride( + stream_name="users", + cursor_field=["nested", "updated_at"], + ) def test_ambiguous_bare_name_is_rejected_without_write( @@ -214,6 +245,30 @@ def test_stream_override_requires_at_least_one_setting() -> None: StreamConfigOverride(stream_name="users") +def test_dedup_destination_mode_requires_a_primary_key( + monkeypatch: pytest.MonkeyPatch, +) -> None: + entry = _stream_entry( + "users", + config={ + "syncMode": "full_refresh", + "destinationSyncMode": "append", + "primaryKey": [], + }, + ) + _, replace_catalog, _ = _mock_catalog_api(monkeypatch, _catalog(entry)) + + with pytest.raises(AirbyteLibInputError, match="requires a primary key"): + _connection().update_stream_configs([ + StreamConfigOverride( + stream_name="users", + destination_sync_mode=DestinationSyncMode.append_dedup, + ) + ]) + + replace_catalog.assert_not_called() + + def test_all_invalid_overrides_are_reported_before_writing( monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -318,6 +373,7 @@ def test_saved_state_rejects_unsafe_config_changes( assert unsafe[0]["name"] == "users" assert unsafe[0]["namespace"] == "public" assert unsafe[0]["changed_fields"] == ["syncMode", "cursorField"] + assert "public.users (syncMode, cursorField)" in exc_info.value.get_message() replace_catalog.assert_not_called() @@ -368,12 +424,77 @@ def test_no_state_lookup_when_override_does_not_change_config( _catalog(_stream_entry("users")), ) - _connection().update_stream_configs([ + result = _connection().update_stream_configs([ StreamConfigOverride(stream_name="users", sync_mode=SyncMode.full_refresh) ]) get_state.assert_not_called() replace_catalog.assert_called_once() + assert result.reset_streams == [] + + +def test_changed_stream_without_state_is_safe_and_does_not_request_reset( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _, replace_catalog, get_state = _mock_catalog_api( + monkeypatch, + _catalog(_stream_entry("users")), + ) + + result = _connection().update_stream_configs([ + StreamConfigOverride( + stream_name="users", + destination_sync_mode=DestinationSyncMode.overwrite, + ) + ]) + + get_state.assert_called_once() + assert replace_catalog.call_args.kwargs["skip_reset"] is True + assert result.unsafe_streams == [] + assert result.reset_requested is False + assert result.reset_streams == [] + + +def test_deferred_unsafe_update_does_not_report_reset_streams( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _, replace_catalog, _ = _mock_catalog_api( + monkeypatch, + _catalog(_stream_entry("users", namespace="public")), + ) + monkeypatch.setattr( + api_util, + "get_connection_state", + MagicMock( + return_value={ + "stateType": "stream", + "connectionId": _CONNECTION_ID, + "streamState": [ + { + "streamDescriptor": {"name": "users", "namespace": "public"}, + "streamState": {"cursor": "saved"}, + } + ], + } + ), + ) + + result = _connection().update_stream_configs( + [ + StreamConfigOverride( + stream_name="users", + stream_namespace="public", + sync_mode=SyncMode.incremental, + cursor_field="updated_at", + ) + ], + allow_unsafe=True, + defer_reset=True, + ) + + assert replace_catalog.call_args.kwargs["skip_reset"] is True + assert result.reset_requested is False + assert result.reset_streams == [] @pytest.mark.parametrize( diff --git a/tests/unit_tests/test_mcp_connection_schedules.py b/tests/unit_tests/test_mcp_connection_schedules.py index 94965305a..7b04419cd 100644 --- a/tests/unit_tests/test_mcp_connection_schedules.py +++ b/tests/unit_tests/test_mcp_connection_schedules.py @@ -453,9 +453,12 @@ def test_unsafe_change_guidance_uses_destructive_policy( allowed: bool | None, expected_guidance: str, ) -> None: + set_enabled = MagicMock() + set_schedule = MagicMock() + replace_catalog = MagicMock() monkeypatch.setattr( CloudConnection, - "plan_stream_config_update", + "update_stream_configs", MagicMock( side_effect=AirbyteUnsafeStreamConfigChangeError( message="unsafe", @@ -463,15 +466,18 @@ def test_unsafe_change_guidance_uses_destructive_policy( ) ), ) + monkeypatch.setattr(CloudConnection, "set_enabled", set_enabled) + monkeypatch.setattr(CloudConnection, "set_schedule", set_schedule) + monkeypatch.setattr(api_util, "replace_connection_catalog", replace_catalog) monkeypatch.setattr(cloud_mcp, "destructive_tools_allowed", lambda _ctx: allowed) with pytest.raises(AirbyteUnsafeStreamConfigChangeError) as exc_info: cloud_mcp.update_cloud_connection( ctx=None, connection_id="12345678-1234-1234-1234-1234567890ab", - enabled=None, + enabled=True, cron_expression=None, - interval_hours=None, + interval_hours=24, manual_schedule=None, override_stream_configs=[ StreamConfigOverride( @@ -486,3 +492,6 @@ def test_unsafe_change_guidance_uses_destructive_policy( assert expected_guidance in (exc_info.value.guidance or "") if allowed is False: assert "Connection URL:" in (exc_info.value.guidance or "") + set_enabled.assert_not_called() + set_schedule.assert_not_called() + replace_catalog.assert_not_called() From 50e4faff2b7420affcb7505c77a842431bcb9dc2 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 16 Oct 2026 22:21:40 +0000 Subject: [PATCH 09/12] ci: retrigger checks Co-Authored-By: AJ Steers From 17245d586d89bf956905a156bafe5a9b4c52576a Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 16 Oct 2026 22:58:53 +0000 Subject: [PATCH 10/12] refactor(cloud): validate stream overrides with protocol models Co-Authored-By: AJ Steers --- airbyte/cloud/_connection_settings.py | 71 ++++ airbyte/cloud/_stream_config_overrides.py | 301 +++++++++++++++ airbyte/cloud/connections.py | 362 ++---------------- airbyte/cloud/models.py | 21 + airbyte/mcp/cloud.py | 117 +----- tests/unit_tests/fixtures/arg_trace_plan.json | 1 - .../test_cloud_connection_catalog_edits.py | 122 +++--- .../test_mcp_connection_schedules.py | 28 +- 8 files changed, 513 insertions(+), 510 deletions(-) create mode 100644 airbyte/cloud/_connection_settings.py create mode 100644 airbyte/cloud/_stream_config_overrides.py diff --git a/airbyte/cloud/_connection_settings.py b/airbyte/cloud/_connection_settings.py new file mode 100644 index 000000000..d8d4ffaae --- /dev/null +++ b/airbyte/cloud/_connection_settings.py @@ -0,0 +1,71 @@ +# Copyright (c) 2026 Airbyte, Inc., all rights reserved. +"""Private helpers for validating and updating Cloud connection settings.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from airbyte._util import api_util +from airbyte.exceptions import AirbyteLibInputError + + +if TYPE_CHECKING: + from airbyte.cloud.connections import CloudConnection + from airbyte.cloud.models import StreamConfigOverride + + +def _validate_update_connection_args( + *, + enabled: bool | None, + cron_expression: str | None, + interval_hours: int | None, + manual_schedule: bool | None, + override_stream_configs: list[StreamConfigOverride] | None, +) -> None: + """Validate the connection settings before any API writes.""" + if not any( + ( + enabled is not None, + cron_expression is not None, + interval_hours is not None, + manual_schedule is not None, + bool(override_stream_configs), + ) + ): + raise AirbyteLibInputError( + message=( + "At least one setting must be provided: 'enabled', 'interval_hours', " + "'cron_expression', 'manual_schedule', or 'override_stream_configs'." + ), + ) + + if sum((interval_hours is not None, cron_expression is not None, manual_schedule is True)) > 1: + raise AirbyteLibInputError( + message=( + "Cannot specify more than one of 'interval_hours', 'cron_expression', " + "and 'manual_schedule=True'." + ), + ) + + if interval_hours is not None: + api_util.validate_interval_hours(interval_hours) + + +def _set_connection_schedule( + connection: CloudConnection, + *, + cron_expression: str | None, + interval_hours: int | None, + manual_schedule: bool | None, +) -> str | None: + """Set the requested connection schedule and return its summary.""" + if cron_expression is not None: + connection.set_schedule(cron_expression=cron_expression) + return f"schedule set to '{cron_expression}'" + if interval_hours is not None: + connection.set_schedule(interval_hours=interval_hours) + return f"schedule set to 'every {interval_hours} hours'" + if manual_schedule is True: + connection.set_manual_schedule() + return "schedule set to 'manual'" + return None diff --git a/airbyte/cloud/_stream_config_overrides.py b/airbyte/cloud/_stream_config_overrides.py new file mode 100644 index 000000000..997d37c8f --- /dev/null +++ b/airbyte/cloud/_stream_config_overrides.py @@ -0,0 +1,301 @@ +# Copyright (c) 2026 Airbyte, Inc., all rights reserved. +"""Private helpers for evaluating Cloud stream configuration overrides.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +from airbyte_protocol.models import DestinationSyncMode, SyncMode + +from airbyte.cloud.models import StreamConfigChange +from airbyte.exceptions import AirbyteLibInputError, AirbyteUnsafeStreamConfigChangeError + + +if TYPE_CHECKING: + from airbyte_protocol.models import ( + AirbyteStream, + ConfiguredAirbyteCatalog, + ConfiguredAirbyteStream, + ) + + from airbyte.cloud.models import StreamConfigOverride + + +def _index_catalog_streams( + catalog: ConfiguredAirbyteCatalog, + *, + connection_id: str, +) -> list[str]: + if not catalog.streams: + raise AirbyteLibInputError( + message="Connection catalog is empty or missing.", + context={"connection_id": connection_id}, + ) + + available_streams: list[str] = [] + catalog_keys: set[tuple[str, str | None]] = set() + duplicate_catalog_streams: list[str] = [] + for entry in catalog.streams: + stream = entry.stream + name = stream.name + namespace = stream.namespace + available_streams.append(f"{namespace}.{name}" if namespace else name) + key = (name, namespace or None) + if key in catalog_keys: + duplicate_catalog_streams.append(f"{namespace}.{name}" if namespace else name) + catalog_keys.add(key) + + if duplicate_catalog_streams: + raise AirbyteLibInputError( + message=( + "Connection catalog contains duplicate stream entries: " + f"{', '.join(duplicate_catalog_streams)}." + ), + context={ + "connection_id": connection_id, + "duplicate_streams": duplicate_catalog_streams, + }, + ) + + return available_streams + + +def _resolve_stream_override_target( + catalog: ConfiguredAirbyteCatalog, + override: StreamConfigOverride, + *, + override_index: int, + override_keys: set[tuple[str, str | None]], +) -> ConfiguredAirbyteStream | str: + label = _stream_override_label(override_index, override) + matching_entries = [ + entry + for entry in catalog.streams + if entry.stream.name == override.stream_name + and ( + override.stream_namespace is None + or (entry.stream.namespace or None) == (override.stream_namespace or None) + ) + ] + if not matching_entries: + return f"{label}: stream not found." + if len(matching_entries) > 1: + matching_namespaces = [entry.stream.namespace for entry in matching_entries] + return ( + f"{label}: stream name is ambiguous; specify a namespace. " + f"Matching namespaces: {matching_namespaces}." + ) + + target_entry = matching_entries[0] + target_key = (override.stream_name, target_entry.stream.namespace or None) + if target_key in override_keys: + return f"{label}: this stream is overridden more than once." + override_keys.add(target_key) + return target_entry + + +def _apply_stream_override( + catalog: ConfiguredAirbyteCatalog, + override: StreamConfigOverride, + *, + override_index: int, + override_keys: set[tuple[str, str | None]], +) -> StreamConfigChange | str: + target_entry = _resolve_stream_override_target( + catalog, + override, + override_index=override_index, + override_keys=override_keys, + ) + if isinstance(target_entry, str): + return target_entry + + label = _stream_override_label(override_index, override) + candidate = target_entry.model_copy(deep=True) + stream = candidate.stream + if override.sync_mode is not None: + supported_sync_modes = [mode.value for mode in stream.supported_sync_modes] + if override.sync_mode.value not in supported_sync_modes: + return ( + f"{label}: sync mode '{override.sync_mode.value}' is unsupported. " + f"Supported modes: {supported_sync_modes}." + ) + candidate.sync_mode = override.sync_mode + + if override.destination_sync_mode is not None: + candidate.destination_sync_mode = override.destination_sync_mode + if override.cursor_field is not None: + candidate.cursor_field = [override.cursor_field] + + cursor_error = _resolve_incremental_cursor(stream, candidate) + if cursor_error is not None: + return f"{label}: {cursor_error}" + + if ( + candidate.destination_sync_mode == DestinationSyncMode.append_dedup + and not candidate.primary_key + and not stream.source_defined_primary_key + ): + return ( + f"{label}: destination sync mode " + f"'{candidate.destination_sync_mode.value}' requires a primary key." + ) + + changed_fields = _changed_stream_config_fields(target_entry, candidate) + if override.sync_mode is not None: + target_entry.sync_mode = candidate.sync_mode + if override.destination_sync_mode is not None: + target_entry.destination_sync_mode = candidate.destination_sync_mode + if override.cursor_field is not None or candidate.cursor_field != target_entry.cursor_field: + target_entry.cursor_field = candidate.cursor_field + + return StreamConfigChange( + name=stream.name, + namespace=stream.namespace, + changed_fields=changed_fields, + ) + + +def _resolve_incremental_cursor( + stream_def: AirbyteStream, + config: ConfiguredAirbyteStream, +) -> str | None: + if config.sync_mode != SyncMode.incremental or config.cursor_field: + return None + if stream_def.source_defined_cursor: + return None + default_cursor_field = stream_def.default_cursor_field + if default_cursor_field: + config.cursor_field = list(default_cursor_field) + return None + return "incremental sync requires a cursor field." + + +def _raise_for_override_errors( + validation_errors: list[str], + available_streams: list[str], + *, + connection_id: str, +) -> None: + if not validation_errors: + return + + message = f"Invalid stream configuration overrides: {' '.join(validation_errors)}" + context: dict[str, Any] = { + "connection_id": connection_id, + "validation_errors": validation_errors, + } + if any(": stream not found." in error for error in validation_errors): + message += f" Available streams: {available_streams}." + context["available_streams"] = available_streams + + raise AirbyteLibInputError( + message=message, + guidance="Correct every listed override and try again.", + context=context, + ) + + +def _stream_override_label( + override_index: int, + override: StreamConfigOverride, +) -> str: + return f"override {override_index} for '{override.stream_name}'" + ( + f" in namespace '{override.stream_namespace}'" + if override.stream_namespace is not None + else "" + ) + + +def _changed_stream_config_fields( + current: ConfiguredAirbyteStream, + updated: ConfiguredAirbyteStream, +) -> list[str]: + fields = ( + ("syncMode", "sync_mode"), + ("destinationSyncMode", "destination_sync_mode"), + ("cursorField", "cursor_field"), + ("primaryKey", "primary_key"), + ) + return [ + wire_name + for wire_name, field_name in fields + if getattr(current, field_name) != getattr(updated, field_name) + ] + + +def _stream_has_saved_state( + state_messages: list[dict[str, Any]], + stream_name: str, + stream_namespace: str | None, +) -> bool: + for message in state_messages: + if not isinstance(message, dict): + return True + message_type = message.get("type") + if message_type in {"GLOBAL", "LEGACY"}: + return True + if message_type != "STREAM": + return True + + stream = message.get("stream") + descriptor = stream.get("stream_descriptor") if isinstance(stream, dict) else None + if not isinstance(descriptor, dict) or not isinstance(descriptor.get("name"), str): + return True + if descriptor.get("name") != stream_name: + continue + if (descriptor.get("namespace") or None) == (stream_namespace or None): + return True + return False + + +def _find_unsafe_stream_changes( + old_catalog: ConfiguredAirbyteCatalog, + new_catalog: ConfiguredAirbyteCatalog, + state_messages: list[dict[str, Any]], +) -> list[StreamConfigChange]: + old_streams = { + (entry.stream.name, entry.stream.namespace): entry for entry in old_catalog.streams + } + unsafe_streams: list[StreamConfigChange] = [] + for entry in new_catalog.streams: + stream = entry.stream + current_entry = old_streams.get((stream.name, stream.namespace)) + if current_entry is None: + continue + changed_fields = _changed_stream_config_fields(current_entry, entry) + if changed_fields and _stream_has_saved_state( + state_messages, + stream.name, + stream.namespace, + ): + unsafe_streams.append( + StreamConfigChange( + name=stream.name, + namespace=stream.namespace, + changed_fields=changed_fields, + ) + ) + return unsafe_streams + + +def _raise_for_unsafe_stream_config_changes( + unsafe_streams: list[StreamConfigChange], + *, + allow_unsafe: bool, + connection_id: str, +) -> None: + if unsafe_streams and not allow_unsafe: + stream_details = ", ".join( + f"{stream.label} ({', '.join(stream.changed_fields)})" for stream in unsafe_streams + ) + raise AirbyteUnsafeStreamConfigChangeError( + message=( + "Stream configuration changes affect streams with saved sync state: " + f"{stream_details}." + ), + context={ + "connection_id": connection_id, + "unsafe_streams": [stream.model_dump() for stream in unsafe_streams], + }, + ) diff --git a/airbyte/cloud/connections.py b/airbyte/cloud/connections.py index 05ce68db7..a52fdc533 100644 --- a/airbyte/cloud/connections.py +++ b/airbyte/cloud/connections.py @@ -4,11 +4,10 @@ from __future__ import annotations import logging -from copy import deepcopy from http import HTTPStatus from typing import TYPE_CHECKING, Any, Literal, overload -from airbyte_protocol.models import DestinationSyncMode +from airbyte_protocol.models import ConfiguredAirbyteCatalog from airbyte._direct_connectors import api_util as agents_api_util from airbyte._util import api_util @@ -25,6 +24,13 @@ _match_stream, _normalize_state_to_protocol, ) +from airbyte.cloud._stream_config_overrides import ( + _apply_stream_override, + _find_unsafe_stream_changes, + _index_catalog_streams, + _raise_for_override_errors, + _raise_for_unsafe_stream_config_changes, +) from airbyte.cloud.connectors import CloudDestination, CloudSource from airbyte.cloud.constants import FINAL_STATUSES from airbyte.cloud.models import ( @@ -44,7 +50,6 @@ AirbyteConnectionSyncError, AirbyteLibInputError, AirbyteMissingResourceError, - AirbyteUnsafeStreamConfigChangeError, AirbyteWorkspaceMismatchError, ) @@ -53,270 +58,6 @@ _QUARTZ_CRON_MIN_FIELDS = 6 _QUARTZ_CRON_MAX_FIELDS = 8 # 7 fields plus an optional trailing timezone ID -_DEDUP_DESTINATION_SYNC_MODES = { - mode.value for mode in DestinationSyncMode if mode.value.endswith("_dedup") -} - - -def _normalized_cursor_field(value: list[str] | str | None) -> list[str]: - if value is None: - return [] - if isinstance(value, str): - return [value] - return value - - -def _changed_stream_config_fields( - current: dict[str, Any], - updated: dict[str, Any], -) -> list[str]: - fields = ("syncMode", "destinationSyncMode", "cursorField", "primaryKey") - return [ - field - for field in fields - if ( - _normalized_cursor_field(current.get(field)) - if field == "cursorField" - else current.get(field) - ) - != ( - _normalized_cursor_field(updated.get(field)) - if field == "cursorField" - else updated.get(field) - ) - ] - - -def _stream_has_saved_state( - state_messages: list[dict[str, Any]], - stream_name: str, - stream_namespace: str | None, -) -> bool: - for message in state_messages: - if not isinstance(message, dict): - return True - message_type = message.get("type") - if message_type in {"GLOBAL", "LEGACY"}: - return True - if message_type != "STREAM": - return True - - stream = message.get("stream") - descriptor = stream.get("stream_descriptor") if isinstance(stream, dict) else None - if not isinstance(descriptor, dict) or not isinstance(descriptor.get("name"), str): - return True - if descriptor.get("name") != stream_name: - continue - if (descriptor.get("namespace") or None) == (stream_namespace or None): - return True - return False - - -def _index_catalog_streams( - catalog: dict[str, Any] | None, - *, - connection_id: str, -) -> list[str]: - 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": connection_id}, - ) - - available_streams: list[str] = [] - catalog_keys: set[tuple[str, str | None]] = set() - duplicate_catalog_streams: list[str] = [] - for entry in catalog["streams"]: - stream_def = entry.get("stream") if isinstance(entry, dict) else None - if not isinstance(stream_def, dict): - raise AirbyteLibInputError( - message="Connection catalog contains an invalid stream entry.", - context={"connection_id": connection_id}, - ) - - name = stream_def.get("name") - namespace = stream_def.get("namespace") - if not isinstance(name, str) or (namespace is not None and not isinstance(namespace, str)): - raise AirbyteLibInputError( - message="Connection catalog contains an invalid stream descriptor.", - context={"connection_id": connection_id}, - ) - - available_streams.append(f"{namespace}.{name}" if namespace else name) - key = (name, namespace or None) - if key in catalog_keys: - duplicate_catalog_streams.append(f"{namespace}.{name}" if namespace else name) - catalog_keys.add(key) - - if duplicate_catalog_streams: - raise AirbyteLibInputError( - message=( - "Connection catalog contains duplicate stream entries: " - f"{', '.join(duplicate_catalog_streams)}." - ), - context={ - "connection_id": connection_id, - "duplicate_streams": duplicate_catalog_streams, - }, - ) - - return available_streams - - -def _resolve_incremental_cursor( - stream_def: dict[str, Any], - config: dict[str, Any], -) -> str | None: - if config.get("syncMode") != "incremental" or _normalized_cursor_field( - config.get("cursorField") - ): - return None - if stream_def.get("sourceDefinedCursor"): - return None - default_cursor_field = stream_def.get("defaultCursorField") - if isinstance(default_cursor_field, list) and default_cursor_field: - config["cursorField"] = deepcopy(default_cursor_field) - return None - return "incremental sync requires a cursor field." - - -def _stream_override_label( - override_index: int, - override: StreamConfigOverride, -) -> str: - return f"override {override_index} for '{override.stream_name}'" + ( - f" in namespace '{override.stream_namespace}'" - if override.stream_namespace is not None - else "" - ) - - -def _resolve_stream_override_target( - catalog: dict[str, Any], - override: StreamConfigOverride, - *, - override_index: int, - override_keys: set[tuple[str, str | None]], -) -> dict[str, Any] | str: - label = _stream_override_label( - override_index, - override, - ) - matching_entries = [ - entry - for entry in catalog["streams"] - if entry["stream"].get("name") == override.stream_name - and ( - override.stream_namespace is None - or (entry["stream"].get("namespace") or None) == (override.stream_namespace or None) - ) - ] - if not matching_entries: - return f"{label}: stream not found." - if len(matching_entries) > 1: - matching_namespaces = [entry["stream"].get("namespace") for entry in matching_entries] - return ( - f"{label}: stream name is ambiguous; specify a namespace. " - f"Matching namespaces: {matching_namespaces}." - ) - - target_entry = matching_entries[0] - stream_def = target_entry["stream"] - stream_namespace = stream_def.get("namespace") - target_key = (override.stream_name, stream_namespace or None) - if target_key in override_keys: - return f"{label}: this stream is overridden more than once." - override_keys.add(target_key) - return target_entry - - -def _apply_stream_override( - catalog: dict[str, Any], - override: StreamConfigOverride, - *, - override_index: int, - override_keys: set[tuple[str, str | None]], -) -> StreamConfigChange | str: - target_entry = _resolve_stream_override_target( - catalog, - override, - override_index=override_index, - override_keys=override_keys, - ) - if isinstance(target_entry, str): - return target_entry - - label = _stream_override_label(override_index, override) - stream_def = target_entry["stream"] - stream_namespace = stream_def.get("namespace") - current_config = target_entry.get("config") or {} - if not isinstance(current_config, dict): - return f"{label}: stream configuration is invalid." - - candidate_config = deepcopy(current_config) - supported_sync_modes = stream_def.get("supportedSyncModes") or [] - if override.sync_mode is not None and override.sync_mode.value not in supported_sync_modes: - return ( - f"{label}: sync mode '{override.sync_mode.value}' is unsupported. " - f"Supported modes: {supported_sync_modes}." - ) - - if override.sync_mode is not None: - candidate_config["syncMode"] = override.sync_mode.value - if override.destination_sync_mode is not None: - candidate_config["destinationSyncMode"] = override.destination_sync_mode.value - if override.cursor_field is not None: - candidate_config["cursorField"] = [override.cursor_field] - - cursor_error = _resolve_incremental_cursor(stream_def, candidate_config) - if cursor_error is not None: - return f"{label}: {cursor_error}" - - destination_sync_mode = candidate_config.get("destinationSyncMode") - if ( - destination_sync_mode in _DEDUP_DESTINATION_SYNC_MODES - and not candidate_config.get("primaryKey") - and not stream_def.get("sourceDefinedPrimaryKey") - ): - return ( - f"{label}: destination sync mode '{destination_sync_mode}' requires " "a primary key." - ) - - target_entry["config"] = candidate_config - return StreamConfigChange( - name=override.stream_name, - namespace=stream_namespace, - changed_fields=_changed_stream_config_fields(current_config, candidate_config), - ) - - -def _raise_for_override_errors( - validation_errors: list[str], - available_streams: list[str], - *, - connection_id: str, -) -> None: - if not validation_errors: - return - - message = f"Invalid stream configuration overrides: {' '.join(validation_errors)}" - context: dict[str, Any] = { - "connection_id": connection_id, - "validation_errors": validation_errors, - } - if any(": stream not found." in error for error in validation_errors): - message += f" Available streams: {available_streams}." - context["available_streams"] = available_streams - - raise AirbyteLibInputError( - message=message, - guidance="Correct every listed override and try again.", - context=context, - ) def _validate_quartz_cron_expression(cron_expression: str) -> None: @@ -1220,26 +961,29 @@ def _evaluate_stream_config_updates( overrides: list[StreamConfigOverride], *, allow_unsafe: bool = False, - ) -> tuple[dict[str, Any], list[StreamConfigChange], list[StreamConfigChange]]: + ) -> tuple[ConfiguredAirbyteCatalog, list[StreamConfigChange], list[StreamConfigChange]]: """Validate overrides and evaluate their effect without writing the catalog.""" if not overrides: raise AirbyteLibInputError( message="At least one stream configuration override must be provided.", ) - catalog = self.dump_raw_catalog(normalize=False) - if catalog is None: + raw_catalog = self.dump_raw_catalog(normalize=False) + if raw_catalog is None: raise AirbyteLibInputError( message="Connection catalog is empty or missing.", context={"connection_id": self.connection_id}, ) + catalog = ConfiguredAirbyteCatalog.model_validate( + _normalize_catalog_to_protocol(raw_catalog) + ) available_streams = _index_catalog_streams( catalog, connection_id=self.connection_id, ) old_catalog = catalog - catalog = deepcopy(catalog) + catalog = catalog.model_copy(deep=True) validation_errors: list[str] = [] updated_streams: list[StreamConfigChange] = [] @@ -1266,95 +1010,35 @@ def _evaluate_stream_config_updates( unsafe_streams: list[StreamConfigChange] = [] if changed_streams: state_messages = self.dump_raw_state(normalize=True) - unsafe_streams = self._find_unsafe_stream_changes( + unsafe_streams = _find_unsafe_stream_changes( old_catalog, catalog, state_messages, ) - self._raise_for_unsafe_stream_config_changes( + _raise_for_unsafe_stream_config_changes( unsafe_streams, allow_unsafe=allow_unsafe, + connection_id=self.connection_id, ) return catalog, updated_streams, unsafe_streams - @staticmethod - def _find_unsafe_stream_changes( - old_catalog: dict[str, Any], - new_catalog: dict[str, Any], - state_messages: list[dict[str, Any]], - ) -> list[StreamConfigChange]: - old_streams = { - ( - entry["stream"]["name"], - entry["stream"].get("namespace"), - ): entry - for entry in old_catalog["streams"] - } - unsafe_streams: list[StreamConfigChange] = [] - for entry in new_catalog["streams"]: - stream_def = entry["stream"] - name = stream_def["name"] - namespace = stream_def.get("namespace") - current_entry = old_streams.get((name, namespace)) - if current_entry is None: - continue - changed_fields = _changed_stream_config_fields( - current_entry.get("config") or {}, - entry.get("config") or {}, - ) - if changed_fields and _stream_has_saved_state( - state_messages, - name, - namespace, - ): - unsafe_streams.append( - StreamConfigChange( - name=name, - namespace=namespace, - changed_fields=changed_fields, - ) - ) - return unsafe_streams - - def _raise_for_unsafe_stream_config_changes( - self, - unsafe_streams: list[StreamConfigChange], - *, - allow_unsafe: bool, - ) -> None: - if unsafe_streams and not allow_unsafe: - stream_details = ", ".join( - f"{stream.namespace}.{stream.name} ({', '.join(stream.changed_fields)})" - if stream.namespace - else f"{stream.name} ({', '.join(stream.changed_fields)})" - for stream in unsafe_streams - ) - raise AirbyteUnsafeStreamConfigChangeError( - message=( - "Stream configuration changes affect streams with saved sync state: " - f"{stream_details}." - ), - context={ - "connection_id": self.connection_id, - "unsafe_streams": [stream.model_dump() for stream in unsafe_streams], - }, - ) - def update_stream_configs( self, overrides: list[StreamConfigOverride], *, allow_unsafe: bool = False, - defer_reset: bool = False, ) -> StreamConfigUpdateResult: """Update one or more stream configurations in a single catalog write.""" catalog, updated_streams, unsafe_streams = self._evaluate_stream_config_updates( overrides, allow_unsafe=allow_unsafe, ) - reset_requested = bool(unsafe_streams and allow_unsafe and not defer_reset) - self.import_raw_catalog(catalog, skip_reset=not reset_requested) + reset_requested = bool(unsafe_streams and allow_unsafe) + self.import_raw_catalog( + _denormalize_catalog_to_api(catalog.model_dump(mode="json", exclude_unset=True)), + skip_reset=not reset_requested, + ) return StreamConfigUpdateResult( updated_streams=updated_streams, unsafe_streams=unsafe_streams, diff --git a/airbyte/cloud/models.py b/airbyte/cloud/models.py index 139184422..2a8e0f0f7 100644 --- a/airbyte/cloud/models.py +++ b/airbyte/cloud/models.py @@ -430,6 +430,11 @@ class StreamConfigChange(BaseModel): namespace: str | None = None changed_fields: list[str] + @property + def label(self) -> str: + """Return the namespace-qualified stream label.""" + return f"{self.namespace}.{self.name}" if self.namespace else self.name + class StreamConfigUpdateResult(BaseModel): """Result of updating stream sync configurations on a connection.""" @@ -439,6 +444,22 @@ class StreamConfigUpdateResult(BaseModel): reset_streams: list[StreamConfigChange] reset_requested: bool + def summary_lines(self) -> list[str]: + """Return concise summaries of updated streams and any requested resets.""" + updated_names = [stream.label for stream in self.updated_streams] + changes = [f"updated stream configs for {updated_names}"] + if not self.unsafe_streams: + return changes + + unsafe_names = [stream.label for stream in self.unsafe_streams] + reset_names = [stream.label for stream in self.reset_streams] + changes.append( + f"unsafe streams {unsafe_names}; the platform will refresh or reset changed " + f"streams {reset_names} (refresh keeps data if the destination supports it; " + "otherwise a reset clears that stream's destination data)" + ) + return changes + class CloudConnectionInfo(BaseModel): """Information about an Airbyte Cloud connection.""" diff --git a/airbyte/mcp/cloud.py b/airbyte/mcp/cloud.py index 050dfeca4..b2f6a74a7 100644 --- a/airbyte/mcp/cloud.py +++ b/airbyte/mcp/cloud.py @@ -34,6 +34,10 @@ ExternalSearchType, ) from airbyte._util import api_util +from airbyte.cloud._connection_settings import ( + _set_connection_schedule, + _validate_update_connection_args, +) from airbyte.cloud.client import _MAX_WORKSPACES_TO_VALIDATE, CloudClient from airbyte.cloud.connections import CloudConnection from airbyte.cloud.connectors import ( @@ -54,7 +58,6 @@ ConnectorType, JobTypeEnum, OrganizationFeature, - StreamConfigChange, StreamConfigOverride, StreamConfigUpdateResult, WorkspacePrivilegeScope, @@ -132,6 +135,7 @@ def _manifest_yaml_source_hint() -> str | None: + """Return a trusted-execution hint for custom connector manifests.""" return MANIFEST_YAML_SOURCE_HINT if is_trusted_execution_enabled() else None @@ -2237,6 +2241,7 @@ def get_cloud_search_status( return status def matches(stream: ExternalSearchStreamStatus) -> bool: + """Check whether a stream matches both optional filters.""" return (stream_name is None or stream.name == stream_name) and ( namespace is None or stream.namespace == namespace ) @@ -3220,6 +3225,7 @@ def get_cloud_organization_billing_status( def _get_custom_source_definition_description( custom_source: CustomCloudSourceDefinition, ) -> str: + """Format the useful metadata for a custom source definition.""" return "\n".join( [ f" - Custom Source Name: {custom_source.name}", @@ -4094,50 +4100,14 @@ def set_cloud_connection_selected_streams( ) -def _validate_update_connection_args( - *, - enabled: bool | None, - cron_expression: str | None, - interval_hours: int | None, - manual_schedule: bool | None, - override_stream_configs: list[StreamConfigOverride] | None, -) -> None: - if not any( - ( - enabled is not None, - cron_expression is not None, - interval_hours is not None, - manual_schedule is not None, - bool(override_stream_configs), - ) - ): - raise AirbyteLibInputError( - message=( - "At least one setting must be provided: 'enabled', 'interval_hours', " - "'cron_expression', 'manual_schedule', or 'override_stream_configs'." - ), - ) - - if sum((interval_hours is not None, cron_expression is not None, manual_schedule is True)) > 1: - raise AirbyteLibInputError( - message=( - "Cannot specify more than one of 'interval_hours', 'cron_expression', " - "and 'manual_schedule=True'." - ), - ) - - if interval_hours is not None: - api_util.validate_interval_hours(interval_hours) - - def _apply_stream_config_overrides( ctx: Context, connection: CloudConnection, overrides: list[StreamConfigOverride] | None, *, allow_unsafe: bool, - defer_reset: bool, ) -> StreamConfigUpdateResult | None: + """Apply overrides before other connection writes and add safety guidance.""" if not overrides: return None @@ -4145,7 +4115,6 @@ def _apply_stream_config_overrides( return connection.update_stream_configs( overrides, allow_unsafe=allow_unsafe, - defer_reset=defer_reset, ) except AirbyteUnsafeStreamConfigChangeError as error: guidance = ( @@ -4161,51 +4130,6 @@ def _apply_stream_config_overrides( ) from error -def _set_connection_schedule( - connection: CloudConnection, - *, - cron_expression: str | None, - interval_hours: int | None, - manual_schedule: bool | None, -) -> str | None: - if cron_expression is not None: - connection.set_schedule(cron_expression=cron_expression) - return f"schedule set to '{cron_expression}'" - if interval_hours is not None: - connection.set_schedule(interval_hours=interval_hours) - return f"schedule set to 'every {interval_hours} hours'" - if manual_schedule is True: - connection.set_manual_schedule() - return "schedule set to 'manual'" - return None - - -def _stream_label(stream: StreamConfigChange) -> str: - return f"{stream.namespace}.{stream.name}" if stream.namespace else stream.name - - -def _describe_stream_config_result(result: StreamConfigUpdateResult) -> list[str]: - updated_names = [_stream_label(stream) for stream in result.updated_streams] - changes = [f"updated stream configs for {updated_names}"] - if not result.unsafe_streams: - return changes - - unsafe_names = [_stream_label(stream) for stream in result.unsafe_streams] - if result.reset_requested: - reset_names = [_stream_label(stream) for stream in result.reset_streams] - changes.append( - f"unsafe streams {unsafe_names}; the platform will refresh or reset changed " - f"streams {reset_names} (refresh keeps data if the destination supports it; " - "otherwise a reset clears that stream's destination data)" - ) - else: - changes.append( - f"unsafe streams {unsafe_names}; no reset was run; these streams have stale " - "state and need a refresh/reset before their next sync" - ) - return changes - - def _update_cloud_connection_impl( # noqa: PLR0913 ctx: Context, connection_id: str, @@ -4217,8 +4141,8 @@ def _update_cloud_connection_impl( # noqa: PLR0913 override_stream_configs: list[StreamConfigOverride] | None, workspace_id: str | None, allow_unsafe: bool, - defer_reset: bool, ) -> str: + """Validate and apply settings in the order that protects catalog updates.""" _validate_update_connection_args( enabled=enabled, cron_expression=cron_expression, @@ -4235,7 +4159,6 @@ def _update_cloud_connection_impl( # noqa: PLR0913 connection, override_stream_configs, allow_unsafe=allow_unsafe, - defer_reset=defer_reset, ) if enabled is not None: @@ -4252,9 +4175,7 @@ def _update_cloud_connection_impl( # noqa: PLR0913 manual_schedule=manual_schedule, ) - changes_made = ( - _describe_stream_config_result(stream_config_result) if stream_config_result else [] - ) + changes_made = stream_config_result.summary_lines() if stream_config_result else [] if status_summary is not None: changes_made.append(status_summary) if schedule_summary is not None: @@ -4376,14 +4297,13 @@ def update_cloud_connection( override_stream_configs=override_stream_configs, workspace_id=workspace_id, allow_unsafe=False, - defer_reset=False, ) @mcp_tool( destructive=True, ) -def update_cloud_connection_unsafe( # noqa: PLR0913 +def update_cloud_connection_unsafe( ctx: Context, connection_id: Annotated[ str, @@ -4460,15 +4380,6 @@ def update_cloud_connection_unsafe( # noqa: PLR0913 ), TraceArg.FINGERPRINT, ] = None, - defer_reset: Annotated[ - bool, - Field( - description=( - "If true, do not request a platform reset. Streams with saved state will need " - "a refresh or reset before their next sync." - ), - ), - ] = False, workspace_id: Annotated[ str | None, Field( @@ -4478,11 +4389,10 @@ def update_cloud_connection_unsafe( # noqa: PLR0913 TraceArg.OMIT, ] = None, ) -> str: - """Update connection settings and apply stream overrides even when saved state is affected. + """Update connection settings and override stream configurations despite saved state. Unsafe changes can make the destination's data inconsistent. By default, the platform - refreshes or resets affected streams; set `defer_reset=True` to leave stale state for a - later refresh or reset. + refreshes or resets every changed stream. """ check_guid_created_in_session(connection_id) return _update_cloud_connection_impl( @@ -4495,7 +4405,6 @@ def update_cloud_connection_unsafe( # noqa: PLR0913 override_stream_configs=override_stream_configs, workspace_id=workspace_id, allow_unsafe=True, - defer_reset=defer_reset, ) diff --git a/tests/unit_tests/fixtures/arg_trace_plan.json b/tests/unit_tests/fixtures/arg_trace_plan.json index afdfd9f7f..70f8cc4e0 100644 --- a/tests/unit_tests/fixtures/arg_trace_plan.json +++ b/tests/unit_tests/fixtures/arg_trace_plan.json @@ -359,7 +359,6 @@ "update_cloud_connection_unsafe": { "connection_id": "hash", "cron_expression": "fingerprint", - "defer_reset": "value allowed=[false,true]", "enabled": "value allowed=[false,true]", "interval_hours": "hash", "manual_schedule": "value allowed=[false,true]", diff --git a/tests/unit_tests/test_cloud_connection_catalog_edits.py b/tests/unit_tests/test_cloud_connection_catalog_edits.py index f905f4a48..f4818314e 100644 --- a/tests/unit_tests/test_cloud_connection_catalog_edits.py +++ b/tests/unit_tests/test_cloud_connection_catalog_edits.py @@ -8,10 +8,16 @@ from unittest.mock import MagicMock import pytest -from airbyte_protocol.models import DestinationSyncMode, SyncMode +from airbyte_protocol.models import ( + ConfiguredAirbyteCatalog, + DestinationSyncMode, + SyncMode, +) from pydantic import ValidationError from airbyte._util import api_util +from airbyte.cloud._connection_catalog import _normalize_catalog_to_protocol +from airbyte.cloud._stream_config_overrides import _find_unsafe_stream_changes from airbyte.cloud.connections import CloudConnection from airbyte.cloud.models import StreamConfigOverride from airbyte.cloud.workspaces import CloudWorkspace @@ -154,6 +160,66 @@ def test_batch_overrides_write_once_and_preserve_untouched_catalog_data( get_state.assert_called_once() +def test_stream_config_override_round_trips_api_only_catalog_fields( + monkeypatch: pytest.MonkeyPatch, +) -> None: + users = _stream_entry( + "users", + namespace="public", + config={ + "syncMode": "full_refresh", + "destinationSyncMode": "append", + "cursorField": [], + "primaryKey": [["id"]], + "selected": True, + "selectedFields": [], + "fieldSelectionEnabled": False, + "mappers": [{"name": "users_mapper", "config": {"field": "id"}}], + }, + ) + users["stream"].update({ + "aliasName": "users_alias", + "hashedFields": [{"path": ["id"]}], + "suggested": True, + "selected": True, + }) + orders = _stream_entry( + "orders", + namespace="sales", + config={ + "syncMode": "full_refresh", + "destinationSyncMode": "append", + "primaryKey": [["order_id"]], + "selected": False, + "selectedFields": [{"fieldPath": ["order_id"]}], + "fieldSelectionEnabled": True, + "mappers": [{"name": "orders_mapper", "config": {"field": "order_id"}}], + }, + ) + orders["stream"].update({ + "aliasName": "orders_alias", + "hashedFields": [{"path": ["order_id"]}], + "suggested": False, + "selected": False, + }) + catalog = _catalog(users, orders) + expected = deepcopy(catalog) + expected["streams"][0]["config"]["syncMode"] = "incremental" + expected["streams"][0]["config"]["cursorField"] = ["updated_at"] + _, replace_catalog, _ = _mock_catalog_api(monkeypatch, catalog) + + _connection().update_stream_configs([ + StreamConfigOverride( + stream_name="users", + stream_namespace="public", + sync_mode=SyncMode.incremental, + cursor_field="updated_at", + ) + ]) + + assert replace_catalog.call_args.kwargs["configured_catalog_dict"] == expected + + def test_namespace_selects_same_named_stream_and_sets_top_level_cursor( monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -405,9 +471,15 @@ def test_find_unsafe_stream_changes_compares_catalogs_to_normalized_state() -> N } ] - unsafe = CloudConnection._find_unsafe_stream_changes( - old_catalog, - new_catalog, + old_typed_catalog = ConfiguredAirbyteCatalog.model_validate( + _normalize_catalog_to_protocol(old_catalog) + ) + new_typed_catalog = ConfiguredAirbyteCatalog.model_validate( + _normalize_catalog_to_protocol(new_catalog) + ) + unsafe = _find_unsafe_stream_changes( + old_typed_catalog, + new_typed_catalog, state_messages, ) @@ -455,48 +527,6 @@ def test_changed_stream_without_state_is_safe_and_does_not_request_reset( assert result.reset_streams == [] -def test_deferred_unsafe_update_does_not_report_reset_streams( - monkeypatch: pytest.MonkeyPatch, -) -> None: - _, replace_catalog, _ = _mock_catalog_api( - monkeypatch, - _catalog(_stream_entry("users", namespace="public")), - ) - monkeypatch.setattr( - api_util, - "get_connection_state", - MagicMock( - return_value={ - "stateType": "stream", - "connectionId": _CONNECTION_ID, - "streamState": [ - { - "streamDescriptor": {"name": "users", "namespace": "public"}, - "streamState": {"cursor": "saved"}, - } - ], - } - ), - ) - - result = _connection().update_stream_configs( - [ - StreamConfigOverride( - stream_name="users", - stream_namespace="public", - sync_mode=SyncMode.incremental, - cursor_field="updated_at", - ) - ], - allow_unsafe=True, - defer_reset=True, - ) - - assert replace_catalog.call_args.kwargs["skip_reset"] is True - assert result.reset_requested is False - assert result.reset_streams == [] - - @pytest.mark.parametrize( "stream_options,expected_cursor", [ diff --git a/tests/unit_tests/test_mcp_connection_schedules.py b/tests/unit_tests/test_mcp_connection_schedules.py index 7b04419cd..16cfb1f7b 100644 --- a/tests/unit_tests/test_mcp_connection_schedules.py +++ b/tests/unit_tests/test_mcp_connection_schedules.py @@ -340,19 +340,9 @@ def test_update_interval_api_failure_is_reported( asyncio.run(_update(server, interval_hours=24)) -@pytest.mark.parametrize( - ("defer_reset", "expected_skip_reset", "message"), - [ - (False, False, "the platform will refresh or reset changed streams"), - (True, True, "these streams have stale state"), - ], -) -def test_unsafe_update_controls_catalog_reset( +def test_unsafe_update_always_requests_catalog_reset( schedule_backend: tuple[FastMCP, responses.RequestsMock, MagicMock, MagicMock], monkeypatch: pytest.MonkeyPatch, - defer_reset: bool, - expected_skip_reset: bool, - message: str, ) -> None: server, _http, _lookup, _patch = schedule_backend monkeypatch.setattr( @@ -366,6 +356,7 @@ def test_unsafe_update_controls_catalog_reset( "stream": { "name": "users", "namespace": "public", + "jsonSchema": {"type": "object"}, "supportedSyncModes": ["full_refresh", "incremental"], }, "config": { @@ -377,6 +368,7 @@ def test_unsafe_update_controls_catalog_reset( "stream": { "name": "orders", "namespace": "sales", + "jsonSchema": {"type": "object"}, "supportedSyncModes": ["full_refresh", "incremental"], }, "config": { @@ -411,7 +403,6 @@ def test_unsafe_update_controls_catalog_reset( result = asyncio.run( _unsafe_update( server, - defer_reset=defer_reset, override_stream_configs=[ { "stream_name": "users", @@ -428,15 +419,12 @@ def test_unsafe_update_controls_catalog_reset( ) ) - assert replace_catalog.call_args.kwargs["skip_reset"] is expected_skip_reset + assert replace_catalog.call_args.kwargs["skip_reset"] is False assert "users" in result - assert message in result - if defer_reset: - assert "no reset was run" in result - else: - assert "unsafe streams ['public.users']" in result - assert "changed streams ['public.users', 'sales.orders']" in result - assert "otherwise a reset clears that stream's destination data" in result + assert "the platform will refresh or reset changed streams" in result + assert "unsafe streams ['public.users']" in result + assert "changed streams ['public.users', 'sales.orders']" in result + assert "otherwise a reset clears that stream's destination data" in result @pytest.mark.parametrize( From 6b18a2195c3d9bcb64bfb0223ceb5e5e14d9a03f Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 20 Oct 2026 01:33:09 +0000 Subject: [PATCH 11/12] style: put each argument on its own line in multi-arg signatures Co-Authored-By: AJ Steers --- airbyte/cloud/connections.py | 7 ++++++- tests/unit_tests/test_mcp_connection_schedules.py | 5 ++++- 2 files changed, 10 insertions(+), 2 deletions(-) diff --git a/airbyte/cloud/connections.py b/airbyte/cloud/connections.py index a52fdc533..e5befc965 100644 --- a/airbyte/cloud/connections.py +++ b/airbyte/cloud/connections.py @@ -921,7 +921,12 @@ def dump_raw_catalog( return _normalize_catalog_to_protocol(raw) return raw - def import_raw_catalog(self, catalog: dict[str, Any], *, skip_reset: bool = True) -> None: + def import_raw_catalog( + self, + catalog: dict[str, Any], + *, + skip_reset: bool = True, + ) -> None: """Replace the configured catalog for this connection. > ⚠️ **WARNING:** Modifying the catalog directly is not recommended and diff --git a/tests/unit_tests/test_mcp_connection_schedules.py b/tests/unit_tests/test_mcp_connection_schedules.py index 16cfb1f7b..d84c599ed 100644 --- a/tests/unit_tests/test_mcp_connection_schedules.py +++ b/tests/unit_tests/test_mcp_connection_schedules.py @@ -108,7 +108,10 @@ async def _update(server: FastMCP, **arguments: object) -> str: return str(result) -async def _unsafe_update(server: FastMCP, **arguments: object) -> str: +async def _unsafe_update( + server: FastMCP, + **arguments: object, +) -> str: async with Client(server) as client: result = await client.call_tool( "update_cloud_connection_unsafe", From 0552371cb5857e6428bfc162fc02578745700265 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 20 Oct 2026 01:39:52 +0000 Subject: [PATCH 12/12] style(mcp): define tool param defaults only in Field() Co-Authored-By: AJ Steers --- airbyte/mcp/cloud.py | 8 ++++---- tests/unit_tests/test_mcp_connection_schedules.py | 1 + 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/airbyte/mcp/cloud.py b/airbyte/mcp/cloud.py index b2f6a74a7..6743094d8 100644 --- a/airbyte/mcp/cloud.py +++ b/airbyte/mcp/cloud.py @@ -4269,7 +4269,7 @@ def update_cloud_connection( default=None, ), TraceArg.FINGERPRINT, - ] = None, + ], workspace_id: Annotated[ str | None, Field( @@ -4277,7 +4277,7 @@ def update_cloud_connection( default=None, ), TraceArg.OMIT, - ] = None, + ], ) -> str: """Update a connection's settings and stream configurations on Airbyte Cloud. @@ -4379,7 +4379,7 @@ def update_cloud_connection_unsafe( default=None, ), TraceArg.FINGERPRINT, - ] = None, + ], workspace_id: Annotated[ str | None, Field( @@ -4387,7 +4387,7 @@ def update_cloud_connection_unsafe( default=None, ), TraceArg.OMIT, - ] = None, + ], ) -> str: """Update connection settings and override stream configurations despite saved state. diff --git a/tests/unit_tests/test_mcp_connection_schedules.py b/tests/unit_tests/test_mcp_connection_schedules.py index d84c599ed..b5e825cb4 100644 --- a/tests/unit_tests/test_mcp_connection_schedules.py +++ b/tests/unit_tests/test_mcp_connection_schedules.py @@ -181,6 +181,7 @@ def test_direct_update_rejects_invalid_intervals_before_status_change( interval_hours=interval_hours, cron_expression=None, manual_schedule=None, + override_stream_configs=None, workspace_id=None, )