refactor http_cache as property of executor

Author: aaronsteers

airbyte/_connector_base.py

@@ -433,7 +433,7 @@ def _execute(
         )
 
         try:
-            for line in self.executor.execute(args, stdin=stdin, env=env):
+            for line in self.executor.execute(args, stdin=stdin):
                 try:
                     message: AirbyteMessage = AirbyteMessage.model_validate_json(json_data=line)
                     if progress_tracker and message.record:
@@ -464,7 +464,7 @@ def _execute(
                 connector_name=self.name,
                 log_text=self._last_log_messages,
                 original_exception=e,
-            ) from None
+            ) from e
 
 
 __all__ = [

airbyte/_executors/base.py

@@ -1,6 +1,7 @@
 # Copyright (c) 2023 Airbyte, Inc., all rights reserved.
 from __future__ import annotations
 
+import os
 import subprocess
 from abc import ABC, abstractmethod
 from contextlib import contextmanager
@@ -14,6 +15,7 @@
 if TYPE_CHECKING:
     from collections.abc import Generator, Iterable, Iterator
 
+    from airbyte.http_caching.cache import AirbyteConnectorCache
     from airbyte.sources.registry import ConnectorMetadata
 
 
@@ -157,6 +159,7 @@ def __init__(
         name: str | None = None,
         metadata: ConnectorMetadata | None = None,
         target_version: str | None = None,
+        http_cache: AirbyteConnectorCache | None = None,
     ) -> None:
         """Initialize a connector executor.
 
@@ -170,6 +173,7 @@ def __init__(
         )  # metadata is not None here
         self.metadata: ConnectorMetadata | None = metadata
         self.enforce_version: bool = target_version is not None
+        self.http_cache: AirbyteConnectorCache | None = http_cache
 
         self.reported_version: str | None = None
         self.target_version: str | None = None
@@ -201,7 +205,6 @@ def execute(
         args: list[str],
         *,
         stdin: IO[str] | AirbyteMessageIterator | None = None,
-        env: dict[str, str] | None = None,
     ) -> Iterator[str]:
         """Execute a command and return an iterator of STDOUT lines.
 
@@ -211,10 +214,23 @@ def execute(
         with _stream_from_subprocess(
             [*self._cli, *mapped_args],
             stdin=stdin,
-            env=env,
+            env=self.env_vars,
         ) as stream_lines:
             yield from stream_lines
 
+    @property
+    def env_vars(self) -> dict[str, str]:
+        """Get the environment variables for the connector.
+
+        By default, this is an empty dict. Subclasses may override this method
+        to provide custom environment variables.
+        """
+        result = cast(dict[str,str], os.environ.copy())
+        if self.http_cache:
+            result.update(self.http_cache.get_env_vars())
+
+        return result
+
     @abstractmethod
     def ensure_installation(self, *, auto_fix: bool = True) -> None:
         _ = auto_fix

airbyte/_executors/declarative.py

@@ -22,6 +22,7 @@
     from collections.abc import Iterator
 
     from airbyte._message_iterators import AirbyteMessageIterator
+    from airbyte.http_caching.cache import AirbyteConnectorCache
 
 
 def _suppress_cdk_pydantic_deprecation_warnings() -> None:
@@ -43,6 +44,8 @@ def __init__(
         self,
         name: str,
         manifest: dict | Path,
+        *,
+        http_cache: AirbyteConnectorCache | None = None,
     ) -> None:
         """Initialize a declarative executor.
 
@@ -54,6 +57,8 @@ def __init__(
 
         self.name = name
         self._manifest_dict: dict
+        self.http_cache = http_cache
+
         if isinstance(manifest, Path):
             self._manifest_dict = cast("dict", yaml.safe_load(manifest.read_text()))
 
@@ -99,7 +104,6 @@ def execute(
         args: list[str],
         *,
         stdin: IO[str] | AirbyteMessageIterator | None = None,
-        env: dict[str, str] | None = None,
     ) -> Iterator[str]:
         """Execute the declarative source."""
         _ = stdin, env  # Not used

airbyte/_executors/docker.py

@@ -1,13 +1,17 @@
 # Copyright (c) 2023 Airbyte, Inc., all rights reserved.
 from __future__ import annotations
 
+import http
 import logging
 import shutil
+import tempfile
 from pathlib import Path
 from typing import NoReturn
 
 from airbyte import exceptions as exc
 from airbyte._executors.base import Executor
+from airbyte.constants import TEMP_DIR_OVERRIDE
+from airbyte.http_caching.cache import AirbyteConnectorCache
 
 
 logger = logging.getLogger("airbyte")
@@ -20,16 +24,26 @@
 class DockerExecutor(Executor):
     def __init__(
         self,
-        name: str | None = None,
+        name: str,
+        docker_image: str,
         *,
-        executable: list[str],
-        target_version: str | None = None,
-        volumes: dict[Path, str] | None = None,
+        use_host_network: bool = False,
+        http_cache: AirbyteConnectorCache | None = None,
     ) -> None:
-        self.executable: list[str] = executable
-        self.volumes: dict[Path, str] = volumes or {}
-        name = name or executable[0]
-        super().__init__(name=name, target_version=target_version)
+        self.docker_image: str = docker_image
+        self.use_host_network: bool = use_host_network
+
+        local_mount_dir = Path().absolute() / name
+        local_mount_dir.mkdir(exist_ok=True)
+        self.volumes = {
+            local_mount_dir: "/local",
+            (TEMP_DIR_OVERRIDE or Path(tempfile.gettempdir())): DEFAULT_AIRBYTE_CONTAINER_TEMP_DIR,
+        }
+
+        super().__init__(
+            name=name or self.docker_image,
+            http_cache=http_cache,
+        )
 
     def ensure_installation(
         self,
@@ -65,8 +79,30 @@ def uninstall(self) -> NoReturn:
 
     @property
     def _cli(self) -> list[str]:
-        """Get the base args of the CLI executable."""
-        return self.executable
+        """Get the executable CLI command.
+
+        For Docker executors, this is the `docker run` command with the necessary arguments
+        to run the connector in a Docker container. This will be extended (suffixed) by the
+        connector's CLI commands and arguments, such as 'spec' or 'check --config=...'.
+        """
+        result: list[str] = [
+            "docker",
+            "run",
+            "--rm",
+            "-i",
+        ]
+        for local_dir, container_dir in self.volumes.items():
+            result.extend(["--volume", f"{local_dir}:{container_dir}"])
+
+        if self.http_cache:
+            for key, value in self.http_cache.get_env_vars().items():
+                result.extend(["-e", f"{key}={value}"])
+
+        if self.use_host_network:
+            result.extend(["--network", "host"])
+
+        result.append(self.docker_image)
+        return result
 
     def map_cli_args(self, args: list[str]) -> list[str]:
         """Map local file paths to the container's volume paths."""

airbyte/_executors/util.py

@@ -1,7 +1,6 @@
 # Copyright (c) 2023 Airbyte, Inc., all rights reserved.
 from __future__ import annotations
 
-import tempfile
 from pathlib import Path
 from typing import TYPE_CHECKING, Literal, cast
 
@@ -11,18 +10,19 @@
 
 from airbyte import exceptions as exc
 from airbyte._executors.declarative import DeclarativeExecutor
-from airbyte._executors.docker import DEFAULT_AIRBYTE_CONTAINER_TEMP_DIR, DockerExecutor
+from airbyte._executors.docker import DockerExecutor
 from airbyte._executors.local import PathExecutor
 from airbyte._executors.python import VenvExecutor
 from airbyte._util.meta import which
 from airbyte._util.telemetry import EventState, log_install_state  # Non-public API
-from airbyte.constants import AIRBYTE_OFFLINE_MODE, TEMP_DIR_OVERRIDE
+from airbyte.constants import AIRBYTE_OFFLINE_MODE
 from airbyte.sources.registry import ConnectorMetadata, InstallType, get_connector_metadata
 from airbyte.version import get_version
 
 
 if TYPE_CHECKING:
     from airbyte._executors.base import Executor
+    from airbyte.http_caching.cache import AirbyteConnectorCache
 
 
 VERSION_LATEST = "latest"
@@ -124,7 +124,7 @@ def _get_local_executor(
     )
 
 
-def get_connector_executor(  # noqa: PLR0912, PLR0913, PLR0915 # Too many branches/arguments/statements
+def get_connector_executor(  # noqa: PLR0912, PLR0913 # Too many branches/arguments
     name: str,
     *,
     version: str | None = None,
@@ -135,6 +135,7 @@ def get_connector_executor(  # noqa: PLR0912, PLR0913, PLR0915 # Too many branch
     source_manifest: bool | dict | Path | str | None = None,
     install_if_missing: bool = True,
     install_root: Path | None = None,
+    http_cache: AirbyteConnectorCache | None = None,
 ) -> Executor:
     """This factory function creates an executor for a connector.
 
@@ -226,41 +227,19 @@ def get_connector_executor(  # noqa: PLR0912, PLR0913, PLR0915 # Too many branch
         if ":" not in docker_image:
             docker_image = f"{docker_image}:{version or 'latest'}"
 
-        host_temp_dir = TEMP_DIR_OVERRIDE or Path(tempfile.gettempdir())
-        container_temp_dir = DEFAULT_AIRBYTE_CONTAINER_TEMP_DIR
-
-        local_mount_dir = Path().absolute() / name
-        local_mount_dir.mkdir(exist_ok=True)
-
-        volumes = {
-            local_mount_dir: "/local",
-            host_temp_dir: container_temp_dir,
-        }
-        docker_cmd = [
-            "docker",
-            "run",
-            "--rm",
-            "-i",
-        ]
-        for local_dir, container_dir in volumes.items():
-            docker_cmd.extend(["--volume", f"{local_dir}:{container_dir}"])
-
-        if use_host_network is True:
-            docker_cmd.extend(["--network", "host"])
-
-        docker_cmd.extend([docker_image])
-
         return DockerExecutor(
             name=name,
-            executable=docker_cmd,
-            volumes=volumes,
+            docker_image=docker_image,
+            use_host_network=use_host_network,
+            http_cache=http_cache,
         )
 
     if source_manifest:
         if isinstance(source_manifest, dict | Path):
             return DeclarativeExecutor(
                 name=name,
                 manifest=source_manifest,
+                http_cache=http_cache,
             )
 
         if isinstance(source_manifest, str | bool):
@@ -273,6 +252,7 @@ def get_connector_executor(  # noqa: PLR0912, PLR0913, PLR0915 # Too many branch
             return DeclarativeExecutor(
                 name=name,
                 manifest=source_manifest,
+                http_cache=http_cache,
             )
 
     # else: we are installing a connector in a Python virtual environment:
@@ -284,6 +264,7 @@ def get_connector_executor(  # noqa: PLR0912, PLR0913, PLR0915 # Too many branch
             target_version=version,
             pip_url=pip_url,
             install_root=install_root,
+            http_cache=http_cache,
         )
         if install_if_missing:
             executor.ensure_installation()

airbyte/http_caching/cache.py

@@ -68,6 +68,21 @@ def __init__(
         self._proxy: DumpMaster | None = None
         self._addon: HttpCachingAddon | None = None
 
+    def get_env_vars(self) -> dict[str, str]:
+        """Get the environment variables to apply to processes using the HTTP proxy."""
+        env_vars = {}
+        if self._proxy_port:
+            proxy_url = f"http://127.0.0.1:{self._proxy_port}"
+            env_vars = {
+                "HTTP_PROXY": proxy_url,
+                "HTTPS_PROXY": proxy_url,
+                "http_proxy": proxy_url,
+                "https_proxy": proxy_url,
+                "NO_PROXY": "localhost,127.0.0.1",
+                "no_proxy": "localhost,127.0.0.1",
+            }
+        return env_vars
+
     def start(self) -> int:
         """Start the HTTP proxy.