feat(cloud): Add workspace validation to CloudConnection (#929)

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

Author: aaronsteers

airbyte/cloud/connections.py

@@ -8,6 +8,7 @@
 from airbyte._util import api_util
 from airbyte.cloud.connectors import CloudDestination, CloudSource
 from airbyte.cloud.sync_results import SyncResult
+from airbyte.exceptions import AirbyteWorkspaceMismatchError
 
 
 if TYPE_CHECKING:
@@ -16,7 +17,7 @@
     from airbyte.cloud.workspaces import CloudWorkspace
 
 
-class CloudConnection:
+class CloudConnection:  # noqa: PLR0904  # Too many public methods
     """A connection is an extract-load (EL) pairing of a source and destination in Airbyte Cloud.
 
     You can use a connection object to run sync jobs, retrieve logs, and manage the connection.
@@ -54,9 +55,41 @@ def __init__(
         self._cloud_destination_object: CloudDestination | None = None
         """The destination object. (Cached.)"""
 
-    def _fetch_connection_info(self) -> ConnectionResponse:
-        """Populate the connection with data from the API."""
-        return api_util.get_connection(
+    def _fetch_connection_info(
+        self,
+        *,
+        force_refresh: bool = False,
+        verify: bool = True,
+    ) -> ConnectionResponse:
+        """Fetch and cache connection info from the API.
+
+        By default, this method will only fetch from the API if connection info is not
+        already cached. It also verifies that the connection belongs to the expected
+        workspace unless verification is explicitly disabled.
+
+        Args:
+            force_refresh: If True, always fetch from the API even if cached.
+                If False (default), only fetch if not already cached.
+            verify: If True (default), verify that the connection is valid (e.g., that
+                the workspace_id matches this object's workspace). Raises an error if
+                validation fails.
+
+        Returns:
+            The ConnectionResponse from the API.
+
+        Raises:
+            AirbyteWorkspaceMismatchError: If verify is True and the connection's
+                workspace_id doesn't match the expected workspace.
+            AirbyteMissingResourceError: If the connection doesn't exist.
+        """
+        if not force_refresh and self._connection_info is not None:
+            # Use cached info, but still verify if requested
+            if verify:
+                self._verify_workspace_match(self._connection_info)
+            return self._connection_info
+
+        # Fetch from API
+        connection_info = api_util.get_connection(
             workspace_id=self.workspace.workspace_id,
             connection_id=self.connection_id,
             api_root=self.workspace.api_root,
@@ -65,6 +98,51 @@ def _fetch_connection_info(self) -> ConnectionResponse:
             bearer_token=self.workspace.bearer_token,
         )
 
+        # Cache the result first (before verification may raise)
+        self._connection_info = connection_info
+
+        # Verify if requested
+        if verify:
+            self._verify_workspace_match(connection_info)
+
+        return connection_info
+
+    def _verify_workspace_match(self, connection_info: ConnectionResponse) -> None:
+        """Verify that the connection belongs to the expected workspace.
+
+        Raises:
+            AirbyteWorkspaceMismatchError: If the workspace IDs don't match.
+        """
+        if connection_info.workspace_id != self.workspace.workspace_id:
+            raise AirbyteWorkspaceMismatchError(
+                resource_type="connection",
+                resource_id=self.connection_id,
+                workspace=self.workspace,
+                expected_workspace_id=self.workspace.workspace_id,
+                actual_workspace_id=connection_info.workspace_id,
+                message=(
+                    f"Connection '{self.connection_id}' belongs to workspace "
+                    f"'{connection_info.workspace_id}', not '{self.workspace.workspace_id}'."
+                ),
+            )
+
+    def check_is_valid(self) -> bool:
+        """Check if this connection exists and belongs to the expected workspace.
+
+        This method fetches connection info from the API (if not already cached) and
+        verifies that the connection's workspace_id matches the workspace associated
+        with this CloudConnection object.
+
+        Returns:
+            True if the connection exists and belongs to the expected workspace.
+
+        Raises:
+            AirbyteWorkspaceMismatchError: If the connection belongs to a different workspace.
+            AirbyteMissingResourceError: If the connection doesn't exist.
+        """
+        self._fetch_connection_info(force_refresh=False, verify=True)
+        return True
+
     @classmethod
     def _from_connection_response(
         cls,

airbyte/exceptions.py

@@ -490,6 +490,27 @@ class AirbyteConnectionSyncError(AirbyteConnectionError):
     """An error occurred while executing the remote Airbyte job."""
 
 
+@dataclass
+class AirbyteWorkspaceMismatchError(AirbyteError):
+    """Resource does not belong to the expected workspace.
+
+    This error is raised when a resource (connection, source, or destination) is fetched
+    from the API and the workspace ID in the response does not match the expected workspace.
+    """
+
+    resource_type: str | None = None
+    """The type of resource (e.g., 'connection', 'source', 'destination')."""
+
+    resource_id: str | None = None
+    """The ID of the resource that was fetched."""
+
+    expected_workspace_id: str | None = None
+    """The workspace ID that was expected."""
+
+    actual_workspace_id: str | None = None
+    """The workspace ID returned by the API."""
+
+
 @dataclass
 class AirbyteConnectionSyncTimeoutError(AirbyteConnectionSyncError):
     """An timeout occurred while waiting for the remote Airbyte job to complete."""