diff --git a/airbyte/_util/api_util.py b/airbyte/_util/api_util.py index 910896771..d73c95939 100644 --- a/airbyte/_util/api_util.py +++ b/airbyte/_util/api_util.py @@ -3937,6 +3937,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. @@ -3953,6 +3954,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. @@ -3964,7 +3966,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 9b39ce135..e3a6099fe 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 @@ -138,6 +141,9 @@ "CloudClientConfig", "CloudDefaultContextInfo", "CloudWorkspaceInfo", + "StreamConfigOverride", + "StreamConfigChange", + "StreamConfigUpdateResult", "SyncResult", # Enums "ExternalApiReadOnlyAction", 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 15d29f401..e5befc965 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 ConfiguredAirbyteCatalog + from airbyte._direct_connectors import api_util as agents_api_util from airbyte._util import api_util from airbyte.cloud._connection_catalog import ( @@ -22,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 ( @@ -30,6 +39,9 @@ ConnectionSchedule, ConnectionStatus, JobTypeEnum, + StreamConfigChange, + StreamConfigOverride, + StreamConfigUpdateResult, _ConnectionResponseLike, ) from airbyte.cloud.sync_results import SyncResult @@ -909,7 +921,12 @@ 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 @@ -927,6 +944,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) @@ -939,6 +957,102 @@ 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 + + def _evaluate_stream_config_updates( + self, + overrides: list[StreamConfigOverride], + *, + allow_unsafe: bool = False, + ) -> 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.", + ) + + 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 = catalog.model_copy(deep=True) + + validation_errors: list[str] = [] + updated_streams: list[StreamConfigChange] = [] + override_keys: set[tuple[str, str | None]] = set() + for index, override in enumerate(overrides, start=1): + 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) + + _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] = [] + if changed_streams: + state_messages = self.dump_raw_state(normalize=True) + unsafe_streams = _find_unsafe_stream_changes( + old_catalog, + catalog, + state_messages, + ) + + _raise_for_unsafe_stream_config_changes( + unsafe_streams, + allow_unsafe=allow_unsafe, + connection_id=self.connection_id, + ) + return catalog, updated_streams, unsafe_streams + + def update_stream_configs( + self, + overrides: list[StreamConfigOverride], + *, + allow_unsafe: 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) + 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, + 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 59caccd83..2a8e0f0f7 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,75 @@ 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; `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.""" + + destination_sync_mode: DestinationSyncMode | None = None + """The destination sync mode to set, or `None` to keep the current mode.""" + + 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: + """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] + + @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.""" + + updated_streams: list[StreamConfigChange] + unsafe_streams: list[StreamConfigChange] + 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.""" @@ -599,6 +670,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 1723055f1..7b95ce058 100644 --- a/airbyte/mcp/cloud.py +++ b/airbyte/mcp/cloud.py @@ -34,7 +34,12 @@ 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.connector_builder import CloudConnectorBuilderProject from airbyte.cloud.connectors import ( CheckResult, @@ -54,6 +59,8 @@ ConnectorType, JobTypeEnum, OrganizationFeature, + StreamConfigOverride, + StreamConfigUpdateResult, WorkspacePrivilegeScope, ) from airbyte.cloud.workspaces import CloudWorkspace @@ -78,6 +85,7 @@ AirbyteLocalDemoModeError, AirbyteMissingResourceError, AirbyteMissingWorkspaceContextError, + AirbyteUnsafeStreamConfigChangeError, ) from airbyte.mcp._arg_resolvers import ( resolve_api_args, @@ -97,6 +105,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, ) @@ -133,6 +142,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 @@ -2248,6 +2258,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 ) @@ -3240,6 +3251,7 @@ def _check_builder_project_created_in_session(project: CloudConnectorBuilderProj 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}", @@ -4085,8 +4097,96 @@ def set_cloud_connection_selected_streams( ) +def _apply_stream_config_overrides( + ctx: Context, + connection: CloudConnection, + overrides: list[StreamConfigOverride] | None, + *, + allow_unsafe: bool, +) -> StreamConfigUpdateResult | None: + """Apply overrides before other connection writes and add safety guidance.""" + if not overrides: + return None + + try: + return connection.update_stream_configs( + overrides, + 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 _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, +) -> str: + """Validate and apply settings in the order that protects catalog updates.""" + _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, + ) + + 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 = 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: + changes_made.append(schedule_summary) + + 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 update_cloud_connection( ctx: Context, @@ -4154,6 +4254,19 @@ def update_cloud_connection( default=None, ), ], + override_stream_configs: Annotated[ + list[StreamConfigOverride] | None, + Field( + description=( + "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`." + ), + default=None, + ), + TraceArg.FINGERPRINT, + ], workspace_id: Annotated[ str | None, Field( @@ -4163,71 +4276,132 @@ def update_cloud_connection( TraceArg.OMIT, ], ) -> str: - """Update a connection's settings on Airbyte Cloud. + """Update a connection's settings and stream configurations 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) - - At least one setting must be provided. `interval_hours`, `cron_expression`, and - `manual_schedule=True` are mutually exclusive. + 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( + ctx, + connection_id, + enabled=enabled, + cron_expression=cron_expression, + interval_hours=interval_hours, + manual_schedule=manual_schedule, + override_stream_configs=override_stream_configs, + workspace_id=workspace_id, + allow_unsafe=False, + ) - # 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 - ): - raise AirbyteLibInputError( - message=( - "At least one setting must be provided: 'enabled', 'interval_hours', " - "'cron_expression', or 'manual_schedule'." - ), - ) - # 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'." +@mcp_tool( + destructive=True, +) +def update_cloud_connection_unsafe( + ctx: Context, + connection_id: Annotated[ + str, + Field(description="The ID of the connection to update."), + ], + *, + 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." ), - ) - - 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'") + default=None, + ), + ], + cron_expression: Annotated[ + str | None, + Field( + 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, + ), + ], + manual_schedule: Annotated[ + bool | None, + Field( + description=( + "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, + ), + ], + override_stream_configs: Annotated[ + list[StreamConfigOverride] | None, + Field( + description=( + "Update one or more streams' sync configuration, including a top-level cursor " + "field and unsafe changes." + ), + default=None, + ), + TraceArg.FINGERPRINT, + ], + workspace_id: Annotated[ + str | None, + Field( + description=WORKSPACE_ID_TIP_TEXT, + default=None, + ), + TraceArg.OMIT, + ], +) -> str: + """Update connection settings and override stream configurations despite saved state. - changes_summary = ", ".join(changes_made) - return ( - f"Successfully updated connection '{connection_id}': {changes_summary}. " - f"URL: {connection.connection_url}" + Unsafe changes can make the destination's data inconsistent. By default, the platform + refreshes or resets every changed stream. + """ + 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, + 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 1a5b51832..a4fa8c2c3 100644 --- a/tests/unit_tests/fixtures/arg_trace_plan.json +++ b/tests/unit_tests/fixtures/arg_trace_plan.json @@ -412,6 +412,16 @@ "enabled": "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", + "enabled": "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 new file mode 100644 index 000000000..f4818314e --- /dev/null +++ b/tests/unit_tests/test_cloud_connection_catalog_edits.py @@ -0,0 +1,599 @@ +# Copyright (c) 2026 Airbyte, Inc., all rights reserved. +"""Unit tests for editing 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 ( + 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 +from airbyte.exceptions import ( + AirbyteLibInputError, + AirbyteUnsafeStreamConfigChangeError, +) + + +_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, 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) + monkeypatch.setattr(api_util, "get_connection_state", get_state) + return get_catalog, replace_catalog, get_state + + +def test_batch_overrides_write_once_and_preserve_untouched_catalog_data( + monkeypatch: pytest.MonkeyPatch, +) -> None: + target = _stream_entry("users", namespace="public") + 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": ["updated_at"], + "selectedFields": [{"fieldPath": ["id"]}], + "fieldSelectionEnabled": True, + "primaryKey": [["sku"]], + "selected": False, + }, + ) + 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["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["skip_reset"] is True + assert len(result.updated_streams) == 2 + 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: + 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="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"] == ["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( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _, replace_catalog, _ = _mock_catalog_api( + monkeypatch, + _catalog( + _stream_entry("users", namespace="public"), + _stream_entry("users", namespace="sales"), + ), + ) + + 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_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: + _, replace_catalog, get_state = _mock_catalog_api( + monkeypatch, + _catalog( + _stream_entry("users", sync_modes=["full_refresh"]), + _stream_entry("orders"), + ), + ) + + with pytest.raises(AirbyteLibInputError) as exc_info: + _connection().update_stream_configs([ + StreamConfigOverride( + stream_name="missing", sync_mode=SyncMode.full_refresh + ), + 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() + + +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", + [ + { + "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_saved_state_rejects_unsafe_config_changes( + monkeypatch: pytest.MonkeyPatch, + state: dict[str, Any], +) -> None: + _, replace_catalog, _ = _mock_catalog_api( + monkeypatch, + _catalog(_stream_entry("users", namespace="public")), + ) + monkeypatch.setattr( + api_util, + "get_connection_state", + MagicMock(return_value=state), + ) + + 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"] + assert "public.users (syncMode, cursorField)" in exc_info.value.get_message() + replace_catalog.assert_not_called() + + +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"}, + }, + } + ] + + 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, + ) + + 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: + _, replace_catalog, get_state = _mock_catalog_api( + monkeypatch, + _catalog(_stream_entry("users")), + ) + + 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 == [] + + +@pytest.mark.parametrize( + "stream_options,expected_cursor", + [ + ({"source_defined_cursor": True}, None), + ({"default_cursor_field": ["updated_at"]}, ["updated_at"]), + ], + ids=["source-defined-cursor", "default-cursor-field"], +) +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: + entry = _stream_entry( + "users", + namespace="public", + config={"syncMode": "full_refresh", "destinationSyncMode": "append"}, + **stream_options, + ) + _, 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 + + +def test_import_raw_catalog_clears_cached_connection_info( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Importing a catalog invalidates previously cached connection information.""" + connection = _connection() + stale_connection_info = MagicMock() + stale_connection_info.configurations.streams = [MagicMock(name="old-stream")] + connection._connection_info = stale_connection_info + replace_catalog = MagicMock(return_value={}) + monkeypatch.setattr(api_util, "replace_connection_catalog", replace_catalog) + + connection.import_raw_catalog(_catalog(_stream_entry("new-stream"))) + + replace_catalog.assert_called_once() + assert connection._connection_info is None + + +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 c8f0d2bfa..b5e825cb4 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,18 @@ 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 +137,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()) @@ -159,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, ) @@ -319,3 +342,148 @@ def test_update_interval_api_failure_is_reported( with pytest.raises(ToolError, match="HTTP 400"): asyncio.run(_update(server, interval_hours=24)) + + +def test_unsafe_update_always_requests_catalog_reset( + schedule_backend: tuple[FastMCP, responses.RequestsMock, MagicMock, MagicMock], + monkeypatch: pytest.MonkeyPatch, +) -> None: + server, _http, _lookup, _patch = schedule_backend + monkeypatch.setattr( + api_util, + "get_connection_catalog", + MagicMock( + return_value={ + "syncCatalog": { + "streams": [ + { + "stream": { + "name": "users", + "namespace": "public", + "jsonSchema": {"type": "object"}, + "supportedSyncModes": ["full_refresh", "incremental"], + }, + "config": { + "syncMode": "full_refresh", + "destinationSyncMode": "append", + }, + }, + { + "stream": { + "name": "orders", + "namespace": "sales", + "jsonSchema": {"type": "object"}, + "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, + override_stream_configs=[ + { + "stream_name": "users", + "stream_namespace": "public", + "sync_mode": "incremental", + "cursor_field": "updated_at", + }, + { + "stream_name": "orders", + "stream_namespace": "sales", + "destination_sync_mode": "overwrite", + }, + ], + ) + ) + + assert replace_catalog.call_args.kwargs["skip_reset"] is False + assert "users" 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( + ("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: + set_enabled = MagicMock() + set_schedule = MagicMock() + replace_catalog = MagicMock() + monkeypatch.setattr( + CloudConnection, + "update_stream_configs", + MagicMock( + side_effect=AirbyteUnsafeStreamConfigChangeError( + message="unsafe", + context={"unsafe_streams": [{"name": "users"}]}, + ) + ), + ) + 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=True, + cron_expression=None, + interval_hours=24, + manual_schedule=None, + 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 "") + set_enabled.assert_not_called() + set_schedule.assert_not_called() + replace_catalog.assert_not_called() 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.",