Feat: Add docker image override for connectors in Cloud

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,12 +42,23 @@
         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:
@@ -78,7 +96,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 +131,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 +173,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 +214,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 +252,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 +738,219 @@ 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
+    """
+    path = api_root + "/applications/token"
+    headers: dict[str, str] = {
+        "content-type": "application/json",
+        "accept": "application/json",
+    }
+    request_body: dict[str, str] = {
+        "grant-type": "client_credentials",
+        "client_id": client_id,
+        "client_secret": client_secret,
+    }
+    response = requests.request(
+        method="POST",
+        url=path,
+        headers=headers,
+        json=request_body,
+    )
+    if not status_ok(response.status_code):
+        response.raise_for_status()
+
+    return SecretString(response.json()["access_token"])
+
+
+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,
+    config_api_root: str = CLOUD_CONFIG_API_ROOT,
+) -> tuple[bool, str | None]:
+    """Check a source.
+
+    Raises an exception if the check fails.
+
+    https://github.com/airbytehq/airbyte-platform-internal/blob/10bb92e1745a282e785eedfcbed1ba72654c4e4e/oss/airbyte-api/server-api/src/main/openapi/config.yaml#L1409
+    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)
+    bearer_token: SecretString = get_bearer_token(
+        client_id=client_id,
+        client_secret=client_secret,
+        api_root=api_root,
+    )
+    path = config_api_root + f"/{connector_type}s/check_connection"
+    headers: dict[str, Any] = {
+        "Content-Type": "application/json",
+        "Accept": "application/json",
+        "Authorization": f"Bearer {bearer_token}",
+    }
+    request_body: dict[str, str] = {
+        f"{connector_type}Id": actor_id,
+    }
+    response = requests.request(
+        method="POST",
+        url=path,
+        headers=headers,
+        json=request_body,
+    )
+    if not status_ok(response.status_code):
+        response.raise_for_status()
+
+    response_json = response.json()
+    result, message = response_json.get("status"), response_json.get("message")
+
+    if result == "succeeded":
+        return True, None
+
+    if result == "failed":
+        return False, message
+
+    raise AirbyteError(
+        context={
+            "actor_id": actor_id,
+            "connector_type": connector_type,
+            "response": response,
+        },
+    )
+
+
+@dataclass
+class DockerImageOverride:
+    """Defines a connector image override."""
+
+    docker_image_override: str
+    override_level: Literal["workspace", "actor"] = "actor"
+
+
+def set_actor_override(
+    *,
+    workspace_id: str,
+    actor_id: str,
+    actor_type: Literal["source", "destination"],
+    override: DockerImageOverride,
+    config_api_root: str = CLOUD_CONFIG_API_ROOT,
+    client_id: SecretString,
+    client_secret: SecretString,
+) -> None:
+    """Override the docker image and tag for a specific connector.
+
+    https://github.com/airbytehq/airbyte-platform-internal/blob/10bb92e1745a282e785eedfcbed1ba72654c4e4e/oss/airbyte-api/server-api/src/main/openapi/config.yaml#L5111
+    """
+    path = config_api_root + "/v1/scoped_configuration/create"
+    bearer_token: SecretString = get_bearer_token(
+        client_id=client_id,
+        client_secret=client_secret,
+        api_root=config_api_root,
+    )
+    headers: dict[str, Any] = {
+        "Content-Type": "application/json",
+        "Authorization": SecretString(f"Bearer {bearer_token}"),
+    }
+    request_body: dict[str, str] = {
+        # https://github.com/airbytehq/airbyte-platform-internal/blob/10bb92e1745a282e785eedfcbed1ba72654c4e4e/oss/airbyte-api/server-api/src/main/openapi/config.yaml#L7376
+        "value": override.docker_image_override,
+        "config_key": "docker_image",  # TODO: Fix this.
+        "scope_id": actor_id,
+        "scope_type": actor_type,
+        "resource_id": "",  # TODO: Need to call something like get_actor_definition
+        "resource_type": "ACTOR_DEFINITION",
+        "origin": "",  # TODO: Need to get user ID somehow or use another origin type
+        "origin_type": "USER",
+    }
+    response = requests.request(
+        method="POST",
+        url=path,
+        headers=headers,
+        json=request_body,
+    )
+    if not status_ok(response.status_code):
+        raise AirbyteError(
+            context={
+                "workspace_id": workspace_id,
+                "actor_id": actor_id,
+                "actor_type": actor_type,
+                "response": response,
+            },
+        )
+
+
+def get_connector_image_override(
+    *,
+    workspace_id: str,
+    actor_id: str,
+    actor_type: Literal["source", "destination"],
+    config_api_root: str = CLOUD_CONFIG_API_ROOT,
+    client_id: SecretString,
+    client_secret: SecretString,
+) -> DockerImageOverride | None:
+    """Get the docker image and tag for a specific connector.
 
+    Result is a tuple of two values:
+    - A boolean indicating if an override is set.
+    - The docker image and tag, either from the override if set, or from the .
 
-# 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
+    https://github.com/airbytehq/airbyte-platform-internal/blob/10bb92e1745a282e785eedfcbed1ba72654c4e4e/oss/airbyte-api/server-api/src/main/openapi/config.yaml#L5066
+    """
+    path = config_api_root + "/v1/scoped_configuration/list"
+    bearer_token = get_bearer_token(
+        client_id=client_id,
+        client_secret=client_secret,
+        api_root=config_api_root,
+    )
+    headers: dict[str, Any] = {
+        "Content-Type": "application",
+        "Authorization": SecretString(f"Bearer {bearer_token}"),
+    }
+    request_body: dict[str, str] = {
+        "config_key": "docker_image",  # TODO: Fix this.
+    }
+    response = requests.request(
+        method="POST",
+        url=path,
+        headers=headers,
+        json=request_body,
+    )
+    if not status_ok(response.status_code):
+        raise AirbyteError(
+            context={
+                "workspace_id": workspace_id,
+                "actor_id": actor_id,
+                "actor_type": actor_type,
+                "response": response,
+            },
+        )
+    if not response.json():
+        return None
+
+    overrides = [
+        DockerImageOverride(
+            docker_image_override=entry["value"],
+            override_level=entry["scope_type"],
+        )
+        for entry in response.json()
+    ]
+    if not overrides:
+        return None
+    if len(overrides) > 1:
+        raise NotImplementedError(
+            "Multiple overrides found. This is not yet supported.",
+        )
+    return overrides[0]