Feat(cloud): Add check feature for Cloud connectors

airbyte/_util/api_util.py

@@ -14,9 +14,11 @@
 from __future__ import annotations
 
 import json
-from typing import TYPE_CHECKING, Any
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, Literal
 
 import airbyte_api
+import requests
 from airbyte_api import api, models
 
 from airbyte.exceptions import (
@@ -26,6 +28,11 @@
     AirbyteMultipleResourcesError,
     PyAirbyteInputError,
 )
+from airbyte.secrets.base import SecretString
+
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
 
 
 if TYPE_CHECKING:
@@ -35,19 +42,40 @@
         DestinationConfiguration,
     )
 
-    from airbyte.secrets.base import SecretString
-
 
 JOB_WAIT_INTERVAL_SECS = 2.0
 JOB_WAIT_TIMEOUT_SECS_DEFAULT = 60 * 60  # 1 hour
 CLOUD_API_ROOT = "https://api.airbyte.com/v1"
+"""The Airbyte Cloud API root URL.
+
+This is the root URL for the Airbyte Cloud API. It is used to interact with the Airbyte Cloud API
+and is the default API root for the `CloudWorkspace` class.
+- https://reference.airbyte.com/reference/getting-started
+"""
+CLOUD_CONFIG_API_ROOT = "https://cloud.airbyte.com/api/v1"
+"""Internal-Use API Root, aka Airbyte "Config API".
+
+Documentation:
+- https://docs.airbyte.com/api-documentation#configuration-api-deprecated
+- https://github.com/airbytehq/airbyte-platform-internal/blob/master/oss/airbyte-api/server-api/src/main/openapi/config.yaml
+"""
+CLOUD_CONFIG_API_TEST_ROOT = "https://dev-cloud.airbyte.com/api/v1"
+"""Test API root"."""
 
 
 def status_ok(status_code: int) -> bool:
     """Check if a status code is OK."""
     return status_code >= 200 and status_code < 300  # noqa: PLR2004  # allow inline magic numbers
 
 
+def get_config_api_root(api_root: str) -> str:
+    """Get the configuration API root from the main API root."""
+    if api_root == CLOUD_API_ROOT:
+        return CLOUD_CONFIG_API_ROOT
+
+    raise NotImplementedError("Configuration API root not implemented for this API root.")
+
+
 def get_airbyte_server_instance(
     *,
     api_root: str,
@@ -78,7 +106,7 @@ def get_workspace(
     client_id: SecretString,
     client_secret: SecretString,
 ) -> models.WorkspaceResponse:
-    """Get a connection."""
+    """Get a workspace object."""
     airbyte_instance = get_airbyte_server_instance(
         api_root=api_root,
         client_id=client_id,
@@ -113,7 +141,7 @@ def list_connections(
     name: str | None = None,
     name_filter: Callable[[str], bool] | None = None,
 ) -> list[models.ConnectionResponse]:
-    """Get a connection."""
+    """List connections."""
     if name and name_filter:
         raise PyAirbyteInputError(message="You can provide name or name_filter, but not both.")
 
@@ -155,7 +183,7 @@ def list_workspaces(
     name: str | None = None,
     name_filter: Callable[[str], bool] | None = None,
 ) -> list[models.WorkspaceResponse]:
-    """Get a connection."""
+    """List workspaces."""
     if name and name_filter:
         raise PyAirbyteInputError(message="You can provide name or name_filter, but not both.")
 
@@ -196,7 +224,7 @@ def list_sources(
     name: str | None = None,
     name_filter: Callable[[str], bool] | None = None,
 ) -> list[models.SourceResponse]:
-    """Get a connection."""
+    """List sources."""
     if name and name_filter:
         raise PyAirbyteInputError(message="You can provide name or name_filter, but not both.")
 
@@ -234,7 +262,7 @@ def list_destinations(
     name: str | None = None,
     name_filter: Callable[[str], bool] | None = None,
 ) -> list[models.DestinationResponse]:
-    """Get a connection."""
+    """List destinations."""
     if name and name_filter:
         raise PyAirbyteInputError(message="You can provide name or name_filter, but not both.")
 
@@ -720,16 +748,116 @@ def delete_connection(
         )
 
 
-# Not yet implemented
+# Functions for leveraging the Airbyte Config API (may not be supported or stable)
+
+
+def get_bearer_token(
+    *,
+    client_id: SecretString,
+    client_secret: SecretString,
+    api_root: str = CLOUD_API_ROOT,
+) -> SecretString:
+    """Get a bearer token.
+
+    https://reference.airbyte.com/reference/createaccesstoken
+
+    """
+    response = requests.post(
+        url=api_root + "/applications/token",
+        headers={
+            "content-type": "application/json",
+            "accept": "application/json",
+        },
+        json={
+            "client_id": client_id,
+            "client_secret": client_secret,
+        },
+    )
+    if not status_ok(response.status_code):
+        response.raise_for_status()
+
+    return SecretString(response.json()["access_token"])
+
+
+def _make_config_api_request(
+    *,
+    api_root: str,
+    path: str,
+    json: dict[str, Any],
+    client_id: SecretString,
+    client_secret: SecretString,
+) -> dict[str, Any]:
+    config_api_root = get_config_api_root(api_root)
+    bearer_token = get_bearer_token(
+        client_id=client_id,
+        client_secret=client_secret,
+        api_root=api_root,
+    )
+    headers: dict[str, Any] = {
+        "Content-Type": "application/json",
+        "Authorization": f"Bearer {bearer_token}",
+        "User-Agent": "PyAirbyte Client",
+    }
+    response = requests.request(
+        method="POST",
+        url=config_api_root + path,
+        headers=headers,
+        json=json,
+    )
+    if not status_ok(response.status_code):
+        try:
+            response.raise_for_status()
+        except requests.HTTPError as ex:
+            raise AirbyteError(
+                context={
+                    "url": response.request.url,
+                    "body": response.request.body,
+                    "response": response.__dict__,
+                },
+            ) from ex
+
+    return response.json()
+
+
+def check_connector(
+    *,
+    actor_id: str,
+    connector_type: Literal["source", "destination"],
+    client_id: SecretString,
+    client_secret: SecretString,
+    workspace_id: str | None = None,
+    api_root: str = CLOUD_API_ROOT,
+) -> tuple[bool, str | None]:
+    """Check a source.
+
+    Raises an exception if the check fails. Uses one of these endpoints:
+
+    - /v1/sources/check_connection: https://github.com/airbytehq/airbyte-platform-internal/blob/10bb92e1745a282e785eedfcbed1ba72654c4e4e/oss/airbyte-api/server-api/src/main/openapi/config.yaml#L1409
+    - /v1/destinations/check_connection: https://github.com/airbytehq/airbyte-platform-internal/blob/10bb92e1745a282e785eedfcbed1ba72654c4e4e/oss/airbyte-api/server-api/src/main/openapi/config.yaml#L1995
+    """
+    _ = workspace_id  # Not used (yet)
+
+    json_result = _make_config_api_request(
+        path=f"/{connector_type}s/check_connection",
+        json={
+            f"{connector_type}Id": actor_id,
+        },
+        api_root=api_root,
+        client_id=client_id,
+        client_secret=client_secret,
+    )
+    result, message = json_result.get("status"), json_result.get("message")
+
+    if result == "succeeded":
+        return True, None
 
+    if result == "failed":
+        return False, message
 
-# def check_source(
-#     source_id: str,
-#     *,
-#     api_root: str,
-#     api_key: str,
-#     workspace_id: str | None = None,
-# ) -> api.SourceCheckResponse:
-#     """Check a source."""
-#     _ = source_id, workspace_id, api_root, api_key
-#     raise NotImplementedError
+    raise AirbyteError(
+        context={
+            "actor_id": actor_id,
+            "connector_type": connector_type,
+            "response": json_result,
+        },
+    )

airbyte/cloud/connectors.py

@@ -6,11 +6,44 @@
 import abc
 from typing import TYPE_CHECKING, ClassVar, Literal
 
+from attr import dataclass
+
+from airbyte._util import api_util
+
 
 if TYPE_CHECKING:
     from airbyte.cloud.workspaces import CloudWorkspace
 
 
+@dataclass
+class CheckResult:
+    """A cloud check result object."""
+
+    success: bool
+    """Whether the check result is valid."""
+
+    error_message: str | None = None
+    """None if the check was successful. Otherwise the failure message from the check result."""
+
+    internal_error: str | None = None
+    """None if the check was able to be run. Otherwise, this will describe the internal failure."""
+
+    def __bool__(self) -> bool:
+        """Truthy when check was successful."""
+        return self.success
+
+    def __str__(self) -> str:
+        """Get a string representation of the check result."""
+        return "Success" if self.success else f"Failed: {self.error_message}"
+
+    def __repr__(self) -> str:
+        """Get a string representation of the check result."""
+        return (
+            f"CloudCheckResult(success={self.success}, "
+            f"error_message={self.error_message or self.internal_error})"
+        )
+
+
 class CloudConnector(abc.ABC):
     """A cloud connector is a deployed source or destination on Airbyte Cloud.
 
@@ -43,6 +76,35 @@ def permanently_delete(self) -> None:
         else:
             self.workspace.permanently_delete_destination(self.connector_id)
 
+    def check(
+        self,
+        *,
+        raise_on_error: bool = True,
+    ) -> CheckResult:
+        """Check the connector.
+
+        Returns:
+            A `CheckResult` object containing the result. The object is truthy if the check was
+            successful and falsy otherwise. The error message is available in the `error_message`
+            or by converting the object to a string.
+        """
+        result = api_util.check_connector(
+            workspace_id=self.workspace.workspace_id,
+            connector_type=self.connector_type,
+            actor_id=self.connector_id,
+            api_root=self.workspace.api_root,
+            client_id=self.workspace.client_id,
+            client_secret=self.workspace.client_secret,
+        )
+        check_result = CheckResult(
+            success=result[0],
+            error_message=result[1],
+        )
+        if raise_on_error and not check_result:
+            raise ValueError(f"Check failed: {check_result}")
+
+        return check_result
+
 
 class CloudSource(CloudConnector):
     """A cloud source is a source that is deployed on Airbyte Cloud."""

tests/integration_tests/cloud/test_cloud_api_util.py

@@ -6,12 +6,20 @@
 """
 
 from __future__ import annotations
+from typing import Literal
 
 from airbyte_api.models import DestinationResponse, SourceResponse, WorkspaceResponse
 from airbyte._util import api_util, text_util
+from airbyte._util.api_util import (
+    get_bearer_token,
+    check_connector,
+    AirbyteError,
+    CLOUD_API_ROOT,
+)
 from airbyte_api.models import DestinationDuckdb, SourceFaker
 
 from airbyte.secrets.base import SecretString
+import pytest
 
 
 def test_get_workspace(
@@ -226,3 +234,51 @@ def test_create_and_delete_connection(
         client_id=airbyte_cloud_client_id,
         client_secret=airbyte_cloud_client_secret,
     )
+
+
+@pytest.mark.parametrize(
+    "api_root",
+    [
+        CLOUD_API_ROOT,
+    ],
+)
+def test_get_bearer_token(
+    airbyte_cloud_client_id,
+    airbyte_cloud_client_secret,
+    api_root: str,
+) -> None:
+    try:
+        token: SecretString = get_bearer_token(
+            client_id=airbyte_cloud_client_id,
+            client_secret=airbyte_cloud_client_secret,
+            api_root=api_root,
+        )
+        assert token is not None
+    except AirbyteError as e:
+        pytest.fail(f"API call failed: {e}")
+
+
+@pytest.mark.parametrize(
+    "connector_id, connector_type, expect_success",
+    [
+        ("f45dd701-d1f0-4e8e-97c4-2b89c40ac928", "source", True),
+        # ("......-....-....-............", "destination", True),
+    ],
+)
+def test_check_connector(
+    airbyte_cloud_client_id: SecretString,
+    airbyte_cloud_client_secret: SecretString,
+    connector_id: str,
+    connector_type: Literal["source", "destination"],
+    expect_success: bool,
+) -> None:
+    try:
+        result, error_message = check_connector(
+            actor_id=connector_id,
+            connector_type=connector_type,
+            client_id=airbyte_cloud_client_id,
+            client_secret=airbyte_cloud_client_secret,
+        )
+        assert result == expect_success
+    except AirbyteError as e:
+        pytest.fail(f"API call failed: {e}")