Feat(cloud): Add check feature for Cloud connectors (#562)

Author: aaronsteers

airbyte/_util/api_util.py

@@ -14,9 +14,10 @@
 from __future__ import annotations
 
 import json
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, Literal
 
 import airbyte_api
+import requests
 from airbyte_api import api, models
 
 from airbyte.exceptions import (
@@ -26,6 +27,11 @@
     AirbyteMultipleResourcesError,
     PyAirbyteInputError,
 )
+from airbyte.secrets.base import SecretString
+
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
 
 
 if TYPE_CHECKING:
@@ -35,19 +41,38 @@
         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
+"""
 
 
 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 +103,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 +138,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 +180,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 +221,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 +259,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 +745,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/__init__.py

@@ -5,7 +5,7 @@
 
 ## Examples
 
-### Basic Usage Example:
+### Basic Sync Example:
 
 ```python
 import airbyte as ab
@@ -47,12 +47,7 @@
 for record in dataset:
     print(record)
 ```
-
-ℹ️ **Experimental Features**
-
-You can use the `airbyte.cloud.experimental` module to access experimental features.
-These additional features are subject to change and may not be available in all environments.
-"""  # noqa: RUF002  # Allow emoji
+"""
 
 from __future__ import annotations
 

airbyte/cloud/connectors.py

@@ -1,16 +1,84 @@
 # Copyright (c) 2024 Airbyte, Inc., all rights reserved.
-"""Cloud connectors module for working with Cloud sources and destinations."""
+"""Cloud connectors module for working with Cloud sources and destinations.
+
+This module provides classes for working with Cloud sources and destinations. Rather
+than creating `CloudConnector` objects directly, it is recommended to use the
+`airbyte.cloud.workspaces` module to create and manage cloud connector objects.
+
+Classes:
+  - `CloudConnector`: A cloud connector object.
+  - `CloudSource`: A cloud source object.
+  - `CloudDestination`: A cloud destination object.
+
+## Usage Examples
+
+Obtain a cloud source object and run a `check` on it:
+
+```python
+from airbyte.cloud import CloudWorkspace
+
+workspace = CloudWorkspace(
+    workspace_id="...",
+    client_id="...",
+    client_secret="...",
+)
+
+# Get the cloud source object
+cloud_source = workspace.get_source("...")
+
+# Check the source configuration and credentials
+check_result = cloud_source.check()
+if check_result:
+    # Truthy if the check was successful
+    print("Check successful")
+else:
+    # Stringify the check result to get the error message
+    print(f"Check failed: {check_result}")
+```
+"""
 
 from __future__ import annotations
 
 import abc
+from dataclasses import dataclass
 from typing import TYPE_CHECKING, ClassVar, Literal
 
+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"CheckResult(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 +111,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."""