diff --git a/airbyte/_util/api_util.py b/airbyte/_util/api_util.py index e64460517..cf6dee683 100644 --- a/airbyte/_util/api_util.py +++ b/airbyte/_util/api_util.py @@ -2929,6 +2929,52 @@ def replace_connection_catalog( ) +def clear_connection( + connection_id: str, + *, + streams: list[dict[str, str]] | None = None, + api_root: str, + config_api_root: str | None = None, + client_id: SecretString | None, + client_secret: SecretString | None, + bearer_token: SecretString | None, +) -> dict[str, Any]: + """Clear the synced data and state of a connection, or of some of its streams. + + Uses the Config API endpoint: POST /v1/connections/clear, or + POST /v1/connections/clear/stream when `streams` is provided. + + Args: + connection_id: The connection ID to clear. + streams: Optional stream descriptors to clear, each with a `streamName` and an + optional `streamNamespace`. If `None`, all streams of the connection are cleared. + 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 JobInfoRead response for the clear job. + """ + if streams is None: + path = "/connections/clear" + json: dict[str, Any] = {"connectionId": connection_id} + else: + path = "/connections/clear/stream" + json = {"connectionId": connection_id, "streams": streams} + + return _make_config_api_request( + path=path, + json=json, + api_root=api_root, + config_api_root=config_api_root, + client_id=client_id, + client_secret=client_secret, + bearer_token=bearer_token, + ) + + def get_organization_info( organization_id: str, *, diff --git a/airbyte/cloud/connections.py b/airbyte/cloud/connections.py index 02b70b2e2..9993a4f5e 100644 --- a/airbyte/cloud/connections.py +++ b/airbyte/cloud/connections.py @@ -792,6 +792,101 @@ def set_stream_state( self.import_raw_state(full_state) + def clear( + self, + stream_names: list[str] | None = None, + ) -> SyncResult: + """Clear the synced data of this connection, for all streams or only the named ones. + + Clearing removes the records this connection wrote to the destination and resets + the saved state of the cleared streams, so the next sync reads them from the start. + This cannot be undone. + + The clear runs as a job on Airbyte Cloud. This method starts the job and returns + without waiting. Call `wait_for_completion()` on the result to wait for it. + + Args: + stream_names: Optional names of the selected streams to clear. If `None` + (default), all streams of the connection are cleared. + + Returns: + A `SyncResult` for the clear job. + + Raises: + PyAirbyteInputError: If `stream_names` is empty, names a stream that is not + selected in this connection, or names a stream that exists in more than + one namespace. + """ + streams: list[dict[str, str]] | None = None + if stream_names is not None: + streams = self._resolve_stream_descriptors(stream_names) + + job_info = api_util.clear_connection( + connection_id=self.connection_id, + streams=streams, + 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, + ) + return SyncResult( + workspace=self.workspace, + connection=self, + job_id=job_info["job"]["id"], + ) + + def _resolve_stream_descriptors(self, stream_names: list[str]) -> list[dict[str, str]]: + """Return the Config API stream descriptors of the named selected streams.""" + if not stream_names: + raise PyAirbyteInputError( + message="No stream names were provided to clear.", + guidance="Pass at least one stream name, or omit `stream_names` to clear all.", + context={"connection_id": self.connection_id}, + ) + + raw_catalog = self.dump_raw_catalog(normalize=False) or {} + namespaces_by_name: dict[str, list[str | None]] = {} + for entry in raw_catalog.get("streams", []): + if not entry.get("config", {}).get("selected", False): + continue + stream = entry.get("stream", {}) + namespaces_by_name.setdefault(stream["name"], []).append(stream.get("namespace")) + + unknown_names = [name for name in stream_names if name not in namespaces_by_name] + if unknown_names: + raise PyAirbyteInputError( + message="Some stream names are not selected in this connection.", + guidance="Use the names of streams that this connection syncs.", + context={ + "connection_id": self.connection_id, + "unknown_stream_names": unknown_names, + "selected_stream_names": sorted(namespaces_by_name), + }, + ) + + ambiguous_names = [name for name in stream_names if len(namespaces_by_name[name]) > 1] + if ambiguous_names: + raise PyAirbyteInputError( + message="Some stream names exist in more than one namespace.", + guidance="Clear all streams of the connection instead.", + context={ + "connection_id": self.connection_id, + "ambiguous_stream_names": ambiguous_names, + }, + ) + + descriptors: list[dict[str, str]] = [] + for name in dict.fromkeys(stream_names): + namespace = namespaces_by_name[name][0] + descriptors.append( + { + "streamName": name, + **({"streamNamespace": namespace} if namespace else {}), + } + ) + return descriptors + @deprecated("Use 'dump_raw_catalog()' instead.") def get_catalog_artifact(self) -> dict[str, Any] | None: """Get the configured catalog for this connection. diff --git a/airbyte/mcp/cloud.py b/airbyte/mcp/cloud.py index 225c825a4..50e852025 100644 --- a/airbyte/mcp/cloud.py +++ b/airbyte/mcp/cloud.py @@ -2878,6 +2878,78 @@ def permanently_delete_cloud_connection( return f"Successfully deleted connection '{actual_name}' (ID: {connection_id})" +@mcp_tool( + destructive=True, + open_world=True, + extra_help_text=CLOUD_AUTH_TIP_TEXT, +) +def clear_cloud_connection( + ctx: Context, + connection_id: Annotated[ + str, + Field(description="The ID of the connection to clear."), + ], + name: Annotated[ + str, + Field(description="The expected name of the connection (for verification)."), + ], + *, + stream_names: Annotated[ + str | list[str] | None, + Field( + description=( + "Optional stream name or list of stream names to clear. " + "If not provided, all streams of the connection are cleared." + ), + default=None, + ), + ] = None, + workspace_id: Annotated[ + str | None, + Field( + description=WORKSPACE_ID_TIP_TEXT, + default=None, + ), + ], +) -> str: + """Clear the synced data of a connection on Airbyte Cloud, for all or selected streams. + + This removes the records the connection wrote to the destination and resets the saved + state of the cleared streams, so the next sync reads them from the start. It cannot be + undone. The clear runs as a job; use `get_cloud_sync_status` with the returned job ID to + follow it. + + The provided name must match the actual name of the connection for the operation to proceed. + This is a safety measure to ensure you are clearing the correct connection. + """ + check_guid_created_in_session(connection_id) + workspace: CloudWorkspace = _get_cloud_workspace(ctx, workspace_id) + connection = workspace.get_connection(connection_id=connection_id) + actual_name: str = cast(str, connection.name) + + if actual_name != name: + raise PyAirbyteInputError( + message=( + f"Name mismatch: expected '{name}' but found '{actual_name}'. " + "The provided name must exactly match the connection's actual name. " + "This is a safety measure to prevent clearing the wrong connection." + ), + context={ + "connection_id": connection_id, + "expected_name": name, + "actual_name": actual_name, + }, + ) + + resolved_stream_names = resolve_list_of_strings(stream_names) + clear_result = connection.clear(stream_names=resolved_stream_names) + scope = "all streams" if resolved_stream_names is None else f"streams {resolved_stream_names}" + return ( + f"Clear started for {scope} of connection '{actual_name}' (ID: {connection_id}). " + f"Job ID is '{clear_result.job_id}' and job URL is: {clear_result.job_url}" + ) + + @mcp_tool( open_world=True, extra_help_text=CLOUD_AUTH_TIP_TEXT, diff --git a/tests/unit_tests/test_cloud_api_util.py b/tests/unit_tests/test_cloud_api_util.py index dfa04aa0b..0edcfa722 100644 --- a/tests/unit_tests/test_cloud_api_util.py +++ b/tests/unit_tests/test_cloud_api_util.py @@ -249,6 +249,30 @@ def test_get_user_id_from_bearer_token_rejects_invalid_tokens( {"workspaceId": "workspace-id", "organizationId": "org-id"}, id="workspace-get", ), + pytest.param( + api_util.clear_connection, + {"connection_id": "connection-id"}, + "/connections/clear", + {"connectionId": "connection-id"}, + {"job": {"id": 77}, "attempts": []}, + {"job": {"id": 77}, "attempts": []}, + id="clear-connection", + ), + pytest.param( + api_util.clear_connection, + { + "connection_id": "connection-id", + "streams": [{"streamName": "users", "streamNamespace": "public"}], + }, + "/connections/clear/stream", + { + "connectionId": "connection-id", + "streams": [{"streamName": "users", "streamNamespace": "public"}], + }, + {"job": {"id": 78}, "attempts": []}, + {"job": {"id": 78}, "attempts": []}, + id="clear-connection-streams", + ), ], ) def test_config_api_helpers_forward_requests( diff --git a/tests/unit_tests/test_cloud_connections.py b/tests/unit_tests/test_cloud_connections.py index bb3360621..3181e8c2e 100644 --- a/tests/unit_tests/test_cloud_connections.py +++ b/tests/unit_tests/test_cloud_connections.py @@ -4,6 +4,7 @@ from __future__ import annotations from dataclasses import dataclass +from typing import cast from unittest.mock import MagicMock import pytest @@ -467,3 +468,117 @@ def test_run_connection_wraps_sdk_error( ) assert exc_info.value.context["status_code"] == 409 + + +def _raw_catalog() -> dict[str, object]: + """Create a raw Config API catalog with selected and unselected streams.""" + return { + "streams": [ + { + "stream": {"name": "users", "namespace": "public"}, + "config": {"selected": True}, + }, + { + "stream": {"name": "orders"}, + "config": {"selected": True}, + }, + { + "stream": {"name": "events", "namespace": "public"}, + "config": {"selected": False}, + }, + ] + } + + +def _patch_clear_connection( + monkeypatch: pytest.MonkeyPatch, + captured_calls: list[dict[str, object]], +) -> None: + """Patch the API utility clear call and capture its arguments.""" + + def clear_connection(**kwargs: object) -> dict[str, object]: + """Capture clear arguments and return a minimal JobInfoRead response.""" + captured_calls.append(kwargs) + return {"job": {"id": 77, "configType": "clear"}, "attempts": []} + + monkeypatch.setattr(api_util, "clear_connection", clear_connection) + + +def test_clear_all_streams_starts_a_clear_job( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Verify `clear()` with no stream names clears the whole connection.""" + captured_calls: list[dict[str, object]] = [] + _patch_clear_connection(monkeypatch, captured_calls) + connection = _connection() + + result = connection.clear() + + assert result.job_id == 77 + assert len(captured_calls) == 1 + assert captured_calls[0]["connection_id"] == "connection-id" + assert captured_calls[0]["streams"] is None + + +def test_clear_named_streams_sends_their_namespaces( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Verify `clear()` sends one descriptor per named stream, with its namespace.""" + captured_calls: list[dict[str, object]] = [] + _patch_clear_connection(monkeypatch, captured_calls) + connection = _connection() + monkeypatch.setattr(connection, "dump_raw_catalog", lambda **_: _raw_catalog()) + + result = connection.clear(stream_names=["users", "orders", "users"]) + + assert result.job_id == 77 + assert captured_calls[0]["streams"] == [ + {"streamName": "users", "streamNamespace": "public"}, + {"streamName": "orders"}, + ] + + +@pytest.mark.parametrize( + ("stream_names", "expected_message"), + [ + pytest.param([], "No stream names", id="empty-list"), + pytest.param(["missing"], "not selected", id="unknown-stream"), + pytest.param(["events"], "not selected", id="unselected-stream"), + ], +) +def test_clear_rejects_streams_it_cannot_clear( + monkeypatch: pytest.MonkeyPatch, + stream_names: list[str], + expected_message: str, +) -> None: + """Verify `clear()` raises before any API call for streams it cannot clear.""" + captured_calls: list[dict[str, object]] = [] + _patch_clear_connection(monkeypatch, captured_calls) + connection = _connection() + monkeypatch.setattr(connection, "dump_raw_catalog", lambda **_: _raw_catalog()) + + with pytest.raises(PyAirbyteInputError, match=expected_message): + connection.clear(stream_names=stream_names) + + assert captured_calls == [] + + +def test_clear_rejects_a_stream_name_in_two_namespaces( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Verify `clear()` refuses to guess which namespace a stream name means.""" + captured_calls: list[dict[str, object]] = [] + _patch_clear_connection(monkeypatch, captured_calls) + connection = _connection() + catalog = _raw_catalog() + streams = cast(list[dict[str, object]], catalog["streams"]) + streams.append({ + "stream": {"name": "users", "namespace": "archive"}, + "config": {"selected": True}, + }) + monkeypatch.setattr(connection, "dump_raw_catalog", lambda **_: catalog) + + with pytest.raises(PyAirbyteInputError, match="more than one namespace"): + connection.clear(stream_names=["users"]) + + assert captured_calls == [] diff --git a/tests/unit_tests/test_mcp_cloud.py b/tests/unit_tests/test_mcp_cloud.py index 90a749383..fc595d2be 100644 --- a/tests/unit_tests/test_mcp_cloud.py +++ b/tests/unit_tests/test_mcp_cloud.py @@ -4,7 +4,7 @@ from __future__ import annotations import inspect -from dataclasses import dataclass +from dataclasses import dataclass, field from datetime import datetime, timezone from types import SimpleNamespace from typing import Callable, cast @@ -17,6 +17,8 @@ CloudWorkspaceInfo, JobStatusEnum, ) +from airbyte.exceptions import PyAirbyteInputError +from airbyte.mcp import _tool_utils from airbyte.mcp import cloud as cloud_mcp from airbyte.mcp.cloud import ( CloudConnectionResult, @@ -25,7 +27,10 @@ ConnectorCheckResult, SyncJobResult, ) +from airbyte.mcp._tool_utils import SafeModeError from fastmcp import Context +from fastmcp_extensions.annotations import DESTRUCTIVE_HINT, READ_ONLY_HINT +from fastmcp_extensions.decorators import _REGISTERED_TOOLS # noqa: PLC2701 @dataclass @@ -875,3 +880,126 @@ def test_deploy_noop_destination_to_cloud_params_have_descriptions() -> None: assert field_infos, param_name assert field_infos[0].description, param_name assert "name" not in signature.parameters + + +@dataclass +class _ClearResultLike: + """Subset of `SyncResult` used by clear tool tests.""" + + job_id: int + job_url: str + + +@dataclass +class _ClearableConnectionLike: + """Subset of `CloudConnection` used by clear tool tests.""" + + name: str + received_stream_names: list[list[str] | None] = field(default_factory=list) + + def clear(self, stream_names: list[str] | None = None) -> _ClearResultLike: + """Capture the stream names and return a clear job result.""" + self.received_stream_names.append(stream_names) + return _ClearResultLike(job_id=77, job_url="https://cloud.airbyte.com/jobs/77") + + +def _patch_clear_workspace( + monkeypatch: pytest.MonkeyPatch, + connection: _ClearableConnectionLike, + seen_workspace_ids: list[str | None], +) -> None: + """Patch the Cloud workspace lookup to return the clear test double.""" + + def fake_get_cloud_workspace( + ctx: object, + workspace_id: str | None = None, + ) -> SimpleNamespace: + seen_workspace_ids.append(workspace_id) + return SimpleNamespace(get_connection=lambda **_: connection) + + monkeypatch.setattr(cloud_mcp, "_get_cloud_workspace", fake_get_cloud_workspace) + monkeypatch.setattr(cloud_mcp, "check_guid_created_in_session", lambda _: None) + + +@pytest.mark.parametrize( + ("stream_names", "expected_stream_names", "expected_scope"), + [ + pytest.param(None, None, "all streams", id="all-streams"), + pytest.param("users,orders", ["users", "orders"], "streams", id="csv-streams"), + pytest.param(["users"], ["users"], "streams", id="list-streams"), + ], +) +def test_clear_cloud_connection_forwards_stream_names( + monkeypatch: pytest.MonkeyPatch, + stream_names: str | list[str] | None, + expected_stream_names: list[str] | None, + expected_scope: str, +) -> None: + """Verify the clear tool resolves stream names and reports the clear job.""" + connection = _ClearableConnectionLike(name="my-connection") + seen_workspace_ids: list[str | None] = [] + _patch_clear_workspace(monkeypatch, connection, seen_workspace_ids) + + result = cloud_mcp.clear_cloud_connection( + cast(Context, object()), + connection_id="connection-id", + name="my-connection", + stream_names=stream_names, + workspace_id="explicit-workspace-id", + ) + + assert seen_workspace_ids == ["explicit-workspace-id"] + assert connection.received_stream_names == [expected_stream_names] + assert expected_scope in result + assert "'77'" in result + assert "https://cloud.airbyte.com/jobs/77" in result + + +def test_clear_cloud_connection_rejects_a_name_mismatch( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Verify the clear tool refuses to run when the confirmation name does not match.""" + connection = _ClearableConnectionLike(name="my-connection") + _patch_clear_workspace(monkeypatch, connection, []) + + with pytest.raises(PyAirbyteInputError, match="Name mismatch"): + cloud_mcp.clear_cloud_connection( + cast(Context, object()), + connection_id="connection-id", + name="other-connection", + workspace_id=None, + ) + + assert connection.received_stream_names == [] + + +def test_clear_cloud_connection_respects_safe_mode( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Verify the clear tool refuses a connection not created in this session in safe mode.""" + monkeypatch.setattr(_tool_utils, "AIRBYTE_CLOUD_MCP_SAFE_MODE", True) + monkeypatch.setattr( + cloud_mcp, + "_get_cloud_workspace", + lambda ctx, workspace_id=None: pytest.fail("workspace must not be looked up"), + ) + + with pytest.raises(SafeModeError): + cloud_mcp.clear_cloud_connection( + cast(Context, object()), + connection_id="connection-not-created-in-session", + name="my-connection", + workspace_id=None, + ) + + +def test_clear_cloud_connection_is_marked_destructive() -> None: + """Verify the clear tool is registered with the destructive hint.""" + annotations = next( + tool_annotations + for func, tool_annotations in _REGISTERED_TOOLS + if func is cloud_mcp.clear_cloud_connection + ) + + assert annotations[DESTRUCTIVE_HINT] is True + assert annotations[READ_ONLY_HINT] is False