Feat: Add docker image override for connectors in Cloud

airbyte/_util/api_util.py

@@ -13,8 +13,11 @@
 
 from __future__ import annotations
 
+from dataclasses import dataclass
 import json
 from typing import TYPE_CHECKING, Any
+import requests
+from typing_extensions import Literal
 
 import airbyte_api
 from airbyte_api import api, models
@@ -26,6 +29,11 @@
     AirbyteMultipleResourcesError,
     PyAirbyteInputError,
 )
+from airbyte.secrets.base import SecretString
+
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
 
 
 if TYPE_CHECKING:
@@ -41,6 +49,19 @@
 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:
@@ -146,6 +167,9 @@ def list_connections(
     ]
 
 
+# Get and run connections
+
+
 def list_workspaces(
     workspace_id: str,
     *,
@@ -155,19 +179,17 @@ def list_workspaces(
     name: str | None = None,
     name_filter: Callable[[str], bool] | None = None,
 ) -> list[models.WorkspaceResponse]:
-    """Get a connection."""
     if name and name_filter:
         raise PyAirbyteInputError(message="You can provide name or name_filter, but not both.")
 
     name_filter = (lambda n: n == name) if name else name_filter or (lambda _: True)
 
     _ = workspace_id  # Not used (yet)
-    airbyte_instance: airbyte_api.AirbyteAPI = get_airbyte_server_instance(
+    airbyte_instance = get_airbyte_server_instance(
         client_id=client_id,
         client_secret=client_secret,
         api_root=api_root,
     )
-
     response: api.ListWorkspacesResponse = airbyte_instance.workspaces.list_workspaces(
         api.ListWorkspacesRequest(
             workspace_ids=[workspace_id],
@@ -234,7 +256,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,6 +742,119 @@ def delete_connection(
         )
 
 
+# Functions for leveraging the Airbyte Config API (may not be supported or stable)
+
+
+@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,
+    api_key: str | SecretString,
+) -> None:
+    """Override the docker image and tag for a specific connector.
+
+    https://github.com/airbytehq/airbyte-platform-internal/blob/master/oss/airbyte-api/server-api/src/main/openapi/config.yaml#L7234
+
+    """
+    path = config_api_root + "/v1/scoped_configuration/create"
+    headers: dict[str, Any] = {
+        "Content-Type": "application",
+        "Authorization": SecretString(f"Bearer {api_key}"),
+    }
+    request_body: dict[str, str] = {
+        "config_key": "docker_image",  # TODO: Fix this.
+        "value": override.docker_image_override,
+        "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,
+    api_key: str,
+) -> 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 .
+    """
+    path = config_api_root + "/v1/scoped_configuration/list"
+    headers: dict[str, Any] = {
+        "Content-Type": "application",
+        "Authorization": SecretString(f"Bearer {api_key}"),
+    }
+    request_body: dict[str, str] = {
+        "config_key": "docker_image",  # TODO: Fix this.
+    }
+    response = requests.request(
+        method="GET",
+        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]
+
+
 # Not yet implemented
 
 

airbyte/cloud/connections.py

@@ -89,6 +89,14 @@ def source(self) -> CloudSource:
         )
         return self._cloud_source_object
 
+    @property
+    def source(self) -> CloudSource:
+        """Get the source object."""
+        return CloudSource(
+            workspace=self.workspace,
+            connector_id=self.source_id,
+        )
+
     @property
     def destination_id(self) -> str:
         """The ID of the destination."""
@@ -112,6 +120,14 @@ def destination(self) -> CloudDestination:
         )
         return self._cloud_destination_object
 
+    @property
+    def destination(self) -> CloudDestination:
+        """Get the source object."""
+        return CloudDestination(
+            workspace=self.workspace,
+            connector_id=self.destination_id,
+        )
+
     @property
     def stream_names(self) -> list[str]:
         """The stream names."""
@@ -189,7 +205,7 @@ def get_previous_sync_logs(
             SyncResult(
                 workspace=self.workspace,
                 connection=self,
-                job_id=sync_log.job_id,
+                job_id=str(sync_log.job_id),
                 _latest_job_info=sync_log,
             )
             for sync_log in sync_logs

airbyte/cloud/workspaces.py

@@ -14,13 +14,15 @@
 from airbyte._util import api_util, text_util
 from airbyte.cloud.connections import CloudConnection
 from airbyte.cloud.connectors import CloudDestination, CloudSource
+from airbyte.cloud.sync_results import SyncResult
 from airbyte.destinations.base import Destination
 from airbyte.secrets.base import SecretString
 
 
 if TYPE_CHECKING:
     from collections.abc import Callable
 
+    from airbyte.destinations.base import Destination
     from airbyte.sources.base import Source
 
 
@@ -111,6 +113,10 @@ def deploy_source(
             workspace=self,
             connector_id=deployed_source.source_id,
         )
+        return CloudSource(
+            workspace=self,
+            connector_id=deployed_source.source_id,
+        )
 
     def deploy_destination(
         self,