feat(cloud): add connection state and catalog management to CloudConnection (#975)

Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

Author: aaronsteers

airbyte/_util/api_util.py

@@ -24,6 +24,7 @@
 
 from airbyte.constants import CLOUD_API_ROOT, CLOUD_CONFIG_API_ROOT, CLOUD_CONFIG_API_ROOT_ENV_VAR
 from airbyte.exceptions import (
+    AirbyteConnectionSyncActiveError,
     AirbyteConnectionSyncError,
     AirbyteError,
     AirbyteMissingResourceError,
@@ -2092,6 +2093,70 @@ def get_connection_state(
     )
 
 
+def replace_connection_state(
+    connection_id: str,
+    connection_state_dict: dict[str, Any],
+    *,
+    api_root: str,
+    client_id: SecretString | None,
+    client_secret: SecretString | None,
+    bearer_token: SecretString | None,
+) -> dict[str, Any]:
+    """Replace the state for a connection.
+
+    Uses the Config API endpoint: POST /v1/state/create_or_update_safe
+
+    Returns HTTP 423 if a sync is currently running, preventing state
+    corruption from concurrent modifications.
+
+    Important: This endpoint replaces the ENTIRE connection state. It does not
+    perform per-stream deduplication or merging on the backend. Callers that need
+    to update a single stream must first fetch the current state, merge the desired
+    stream change into the full state object, and then send the complete state back.
+    See ``CloudConnection.set_stream_state()`` for this fetch-modify-push pattern.
+
+    The provided ``connection_id`` is injected into both the outer request
+    wrapper and the inner ``connection_state`` payload to ensure consistency.
+
+    Args:
+        connection_id: The connection ID to update state for.
+        connection_state_dict: The full ConnectionState dict to set. Must include:
+            - stateType: "global", "stream", or "legacy"
+            - One of: state (legacy), streamState (stream), globalState (global)
+            All streams must be included; any stream omitted will have its state dropped.
+        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).
+
+    Returns:
+        Dictionary containing the updated ConnectionState object.
+    """
+    try:
+        return _make_config_api_request(
+            path="/state/create_or_update_safe",
+            json={
+                "connectionId": connection_id,
+                "connectionState": {
+                    **connection_state_dict,
+                    "connectionId": connection_id,
+                },
+            },
+            api_root=api_root,
+            client_id=client_id,
+            client_secret=client_secret,
+            bearer_token=bearer_token,
+        )
+    except AirbyteError as ex:
+        if ex.context and ex.context.get("status_code") == HTTPStatus.LOCKED:
+            raise AirbyteConnectionSyncActiveError(
+                message="Cannot update connection state while a sync is running.",
+                connection_id=connection_id,
+                guidance="Wait for the current sync to complete before updating state.",
+            ) from ex
+        raise
+
+
 def get_connection_catalog(
     connection_id: str,
     *,
@@ -2127,6 +2192,49 @@ def get_connection_catalog(
     )
 
 
+def replace_connection_catalog(
+    connection_id: str,
+    configured_catalog_dict: dict[str, Any],
+    *,
+    api_root: str,
+    client_id: SecretString | None,
+    client_secret: SecretString | None,
+    bearer_token: SecretString | None,
+) -> dict[str, Any]:
+    """Replace the configured catalog for a connection.
+
+    Uses the Config API endpoint: POST /v1/web_backend/connections/update
+
+    This is a patch-style update that replaces the connection's entire syncCatalog
+    with the provided catalog. All other connection settings remain unchanged.
+
+    Args:
+        connection_id: The connection ID to update catalog for.
+        configured_catalog_dict: The configured catalog dict (``{"streams": [...]}``) to set.
+        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).
+
+    Returns:
+        Dictionary containing the updated WebBackendConnectionRead response.
+    """
+    return _make_config_api_request(
+        path="/web_backend/connections/update",
+        json={
+            "connectionId": connection_id,
+            "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,
+        },
+        api_root=api_root,
+        client_id=client_id,
+        client_secret=client_secret,
+        bearer_token=bearer_token,
+    )
+
+
 def get_organization_info(
     organization_id: str,
     *,

airbyte/cloud/__init__.py

@@ -63,7 +63,13 @@
 # Submodules imported here for documentation reasons: https://github.com/mitmproxy/pdoc/issues/757
 if TYPE_CHECKING:
     # ruff: noqa: TC004
-    from airbyte.cloud import client_config, connections, constants, sync_results, workspaces
+    from airbyte.cloud import (
+        client_config,
+        connections,
+        constants,
+        sync_results,
+        workspaces,
+    )
 
 
 __all__ = [

airbyte/cloud/_connection_state.py

@@ -0,0 +1,133 @@
+# Copyright (c) 2024 Airbyte, Inc., all rights reserved.
+"""Pydantic models for connection state management.
+
+These models represent the state of an Airbyte connection, which tracks sync progress
+for incremental syncs. The state can be one of several types:
+- stream: Per-stream state
+- global: Global state with optional per-stream states
+- legacy: Legacy state blob
+- not_set: No state has been set yet
+"""
+
+from __future__ import annotations
+
+from typing import Any, Literal
+
+from pydantic import BaseModel, ConfigDict, Field
+
+
+class StreamDescriptor(BaseModel):
+    """Descriptor for a single stream within a connection state."""
+
+    name: str = Field(description="The stream name")
+    namespace: str | None = Field(
+        default=None,
+        description="The stream namespace (optional)",
+    )
+
+
+class StreamState(BaseModel):
+    """State for a single stream."""
+
+    model_config = ConfigDict(populate_by_name=True)
+
+    stream_descriptor: StreamDescriptor = Field(
+        alias="streamDescriptor",
+        description="The stream descriptor (name and namespace)",
+    )
+    stream_state: dict[str, Any] | None = Field(
+        default=None,
+        alias="streamState",
+        description="The state blob for this stream",
+    )
+
+
+class GlobalState(BaseModel):
+    """Global state containing shared state and per-stream states."""
+
+    model_config = ConfigDict(populate_by_name=True)
+
+    shared_state: dict[str, Any] | None = Field(
+        default=None,
+        alias="sharedState",
+        description="Shared state across all streams",
+    )
+    stream_states: list[StreamState] = Field(
+        default_factory=list,
+        alias="streamStates",
+        description="Per-stream states within the global state",
+    )
+
+
+class ConnectionStateResponse(BaseModel):
+    """Response model for connection state operations.
+
+    Represents the state of a connection, which can be one of:
+    - stream: Per-stream state (streamState field populated)
+    - global: Global state with optional per-stream states (globalState field populated)
+    - legacy: Legacy state blob (state field populated)
+    - not_set: No state has been set yet
+    """
+
+    model_config = ConfigDict(populate_by_name=True)
+
+    state_type: Literal["global", "stream", "legacy", "not_set"] = Field(
+        alias="stateType",
+        description="The type of state: global, stream, legacy, or not_set",
+    )
+    connection_id: str = Field(
+        alias="connectionId",
+        description="The connection ID (UUID)",
+    )
+    state: dict[str, Any] | None = Field(
+        default=None,
+        description="Legacy state blob (populated when stateType is 'legacy')",
+    )
+    stream_state: list[StreamState] | None = Field(
+        default=None,
+        alias="streamState",
+        description="Per-stream states (populated when stateType is 'stream')",
+    )
+    global_state: GlobalState | None = Field(
+        default=None,
+        alias="globalState",
+        description="Global state (populated when stateType is 'global')",
+    )
+
+    def __str__(self) -> str:
+        """Return a string representation of the connection state."""
+        stream_count = 0
+        if self.stream_state:
+            stream_count = len(self.stream_state)
+        elif self.global_state and self.global_state.stream_states:
+            stream_count = len(self.global_state.stream_states)
+        return (
+            f"Connection {self.connection_id}: "
+            f"stateType={self.state_type}, streams={stream_count}"
+        )
+
+
+def _match_stream(
+    stream: StreamState,
+    stream_name: str,
+    stream_namespace: str | None = None,
+) -> bool:
+    """Check if a StreamState matches the given name and namespace.
+
+    Namespace matching treats None and "" as equivalent (both mean "no namespace").
+    There is no wildcard behavior: the caller must match the actual namespace.
+    """
+    if stream.stream_descriptor.name != stream_name:
+        return False
+    return (stream.stream_descriptor.namespace or None) == (stream_namespace or None)
+
+
+def _get_stream_list(
+    state: ConnectionStateResponse,
+) -> list[StreamState]:
+    """Extract the stream list from a ConnectionStateResponse regardless of state type."""
+    if state.state_type == "stream" and state.stream_state:
+        return state.stream_state
+    if state.state_type == "global" and state.global_state:
+        return state.global_state.stream_states
+    return []