Feat: Accept docker image names, move DockerExecutor support to experimental module

airbyte/__init__.py

@@ -15,6 +15,7 @@
     datasets,
     documents,
     exceptions,  # noqa: ICN001  # No 'exc' alias for top-level module
+    experimental,
     results,
     secrets,
     sources,
@@ -39,6 +40,7 @@
     "datasets",
     "documents",
     "exceptions",
+    "experimental",
     "records",
     "registry",
     "results",

airbyte/experimental/__init__.py

@@ -0,0 +1,42 @@
+# Copyright (c) 2024 Airbyte, Inc., all rights reserved.
+"""Experimental features which may change.
+
+The experimental `get_source` implementation allows you to run sources
+using Docker containers. This feature is still in development and may
+change in the future.
+
+To use this feature, import `get_source` from this module and use it in place of the `get_source`
+function from the `airbyte` module.
+
+Instead of this:
+
+```python
+from airbyte import ab
+
+source = ab.get_source(...)
+```
+
+Use this:
+
+```python
+from airbyte.experimental import get_source
+
+source = get_source(...)
+```
+
+Experimental features may change without notice between minor versions of PyAirbyte. Although rare,
+they may also be entirely removed or refactored in future versions of PyAirbyte. Experimental
+features may also be less stable than other features, and may not be as well-tested.
+
+You can help improve this product by reporting issues and providing feedback for improvements in our
+[GitHub issue tracker](https://github.com/airbytehq/pyairbyte/issues).
+"""
+
+from __future__ import annotations
+
+from airbyte.sources.util import _get_source as get_source
+
+
+__all__ = [
+    "get_source",
+]

airbyte/sources/util.py

@@ -43,15 +43,17 @@ def get_connector(
     )
 
 
-def get_source(
+# This non-public function includes the `docker_image` parameter, which is not exposed in the
+# public API. See the `experimental` module for more info.
+def _get_source(  # noqa: PLR0912  # Too many branches
     name: str,
     config: dict[str, Any] | None = None,
     *,
     streams: str | list[str] | None = None,
     version: str | None = None,
     pip_url: str | None = None,
     local_executable: Path | str | None = None,
-    docker_executable: bool = False,
+    docker_image: str | bool = False,
     install_if_missing: bool = True,
 ) -> Source:
     """Get a connector by name and version.
@@ -71,14 +73,27 @@ def get_source(
         local_executable: If set, the connector will be assumed to already be installed and will be
             executed using this path or executable name. Otherwise, the connector will be installed
             automatically in a virtual environment.
+        docker_image: If set, the connector will be executed using Docker. You can specify `True`
+            to use the default image for the connector, or you can specify a custom image name.
+            If `version` is specified and your image name does not already contain a tag
+            (e.g. `my-image:latest`), the version will be appended as a tag (e.g. `my-image:0.1.0`).
         install_if_missing: Whether to install the connector if it is not available locally. This
             parameter is ignored when local_executable is set.
     """
+    if sum([bool(local_executable), bool(docker_image), bool(pip_url)]) > 1:
+        raise exc.PyAirbyteInputError(
+            message=(
+                "You can only specify one of the settings: 'local_executable', 'docker_image', "
+                "or 'pip_url'."
+            ),
+            context={
+                "local_executable": local_executable,
+                "docker_image": docker_image,
+                "pip_url": pip_url,
+            },
+        )
+
     if local_executable:
-        if pip_url:
-            raise exc.PyAirbyteInputError(
-                message="Param 'pip_url' is not supported when 'local_executable' is set."
-            )
         if version:
             raise exc.PyAirbyteInputError(
                 message="Param 'version' is not supported when 'local_executable' is set."
@@ -117,8 +132,23 @@ def get_source(
             ),
         )
 
-    if docker_executable:
-        version = version or "latest"
+    if docker_image:
+        if docker_image is True:
+            # Use the default image name for the connector
+            docker_image = f"airbyte/{name}"
+
+        if version is not None and ":" in docker_image:
+            raise exc.PyAirbyteInputError(
+                message="The 'version' parameter is not supported when a tag is already set in the "
+                "'docker_image' parameter.",
+                context={
+                    "docker_image": docker_image,
+                    "version": version,
+                },
+            )
+
+        if ":" not in docker_image:
+            docker_image = f"{docker_image}:{version or 'latest'}"
 
         temp_dir = tempfile.gettempdir()
         return Source(
@@ -134,7 +164,7 @@ def get_source(
                     "-i",
                     "--volume",
                     f"{temp_dir}:{temp_dir}",
-                    f"airbyte/{name}:{version}",
+                    docker_image,
                 ],
             ),
         )
@@ -171,6 +201,50 @@ def get_source(
         raise
 
 
+# This thin wrapper exposes only non-experimental functions.
+# Aka, exclude the `docker_image` parameter for now.
+# See the `experimental` module for more info.
+def get_source(
+    name: str,
+    config: dict[str, Any] | None = None,
+    *,
+    streams: str | list[str] | None = None,
+    version: str | None = None,
+    pip_url: str | None = None,
+    local_executable: Path | str | None = None,
+    install_if_missing: bool = True,
+) -> Source:
+    """Get a connector by name and version.
+
+    Args:
+        name: connector name
+        config: connector config - if not provided, you need to set it later via the set_config
+            method.
+        streams: list of stream names to select for reading. If set to "*", all streams will be
+            selected. If not provided, you can set it later via the `select_streams()` or
+            `select_all_streams()` method.
+        version: connector version - if not provided, the currently installed version will be used.
+            If no version is installed, the latest available version will be used. The version can
+            also be set to "latest" to force the use of the latest available version.
+        pip_url: connector pip URL - if not provided, the pip url will be inferred from the
+            connector name.
+        local_executable: If set, the connector will be assumed to already be installed and will be
+            executed using this path or executable name. Otherwise, the connector will be installed
+            automatically in a virtual environment.
+        install_if_missing: Whether to install the connector if it is not available locally. This
+            parameter is ignored when local_executable is set.
+    """
+    return _get_source(
+        name=name,
+        config=config,
+        streams=streams,
+        version=version,
+        pip_url=pip_url,
+        local_executable=local_executable,
+        install_if_missing=install_if_missing,
+    )
+
+
 __all__ = [
     "get_source",
 ]