feat: implement HTTP caching with mitmproxy's native format

airbyte/__init__.py

@@ -131,6 +131,8 @@
 from airbyte.datasets import CachedDataset
 from airbyte.destinations.base import Destination
 from airbyte.destinations.util import get_destination
+from airbyte.http_caching import AirbyteConnectorCache, HttpCacheMode
+from airbyte.http_caching.serialization import SerializationFormat
 from airbyte.records import StreamRecord
 from airbyte.results import ReadResult, WriteResult
 from airbyte.secrets import SecretSourceEnum, get_secret
@@ -154,6 +156,7 @@
         documents,
         exceptions,  # noqa: ICN001  # No 'exc' alias for top-level module
         experimental,
+        http_caching,
         logs,
         records,
         results,
@@ -174,6 +177,7 @@
     "documents",
     "exceptions",
     "experimental",
+    "http_caching",
     "logs",
     "records",
     "registry",
@@ -189,12 +193,15 @@
     "get_source",
     "new_local_cache",
     # Classes
+    "AirbyteConnectorCache",
     "BigQueryCache",
     "CachedDataset",
     "Destination",
     "DuckDBCache",
+    "HttpCacheMode",
     "ReadResult",
     "SecretSourceEnum",
+    "SerializationFormat",
     "Source",
     "StreamRecord",
     "WriteResult",

airbyte/_connector_base.py

@@ -399,6 +399,7 @@ def _execute(
         stdin: IO[str] | AirbyteMessageIterator | None = None,
         *,
         progress_tracker: ProgressTracker | None = None,
+        env: dict[str, str] | None = None,
     ) -> Generator[AirbyteMessage, None, None]:
         """Execute the connector with the given arguments.
 
@@ -432,7 +433,7 @@ def _execute(
         )
 
         try:
-            for line in self.executor.execute(args, stdin=stdin):
+            for line in self.executor.execute(args, stdin=stdin, env=env):
                 try:
                     message: AirbyteMessage = AirbyteMessage.model_validate_json(json_data=line)
                     if progress_tracker and message.record:

airbyte/_executors/base.py

@@ -62,6 +62,7 @@ def _stream_from_subprocess(
     *,
     stdin: IO[str] | AirbyteMessageIterator | None = None,
     log_file: IO[str] | None = None,
+    env: dict[str, str] | None = None,
 ) -> Generator[Iterable[str], None, None]:
     """Stream lines from a subprocess."""
     input_thread: Thread | None = None
@@ -74,6 +75,7 @@ def _stream_from_subprocess(
             stderr=log_file,
             universal_newlines=True,
             encoding="utf-8",
+            env=env,
         )
         input_thread = Thread(
             target=_pump_input,
@@ -102,6 +104,7 @@ def _stream_from_subprocess(
             stderr=log_file,
             universal_newlines=True,
             encoding="utf-8",
+            env=env,
         )
 
     if process.stdout is None:
@@ -198,6 +201,7 @@ 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.
 
@@ -207,6 +211,7 @@ def execute(
         with _stream_from_subprocess(
             [*self._cli, *mapped_args],
             stdin=stdin,
+            env=env,
         ) as stream_lines:
             yield from stream_lines
 

airbyte/_executors/declarative.py

@@ -99,9 +99,10 @@ def execute(
         args: list[str],
         *,
         stdin: IO[str] | AirbyteMessageIterator | None = None,
+        env: dict[str, str] | None = None,
     ) -> Iterator[str]:
         """Execute the declarative source."""
-        _ = stdin  # Not used
+        _ = stdin, env  # Not used
         source_entrypoint = AirbyteEntrypoint(self.declarative_source)
 
         mapped_args: list[str] = self.map_cli_args(args)

airbyte/constants.py

@@ -124,3 +124,27 @@ def _str_to_bool(value: str) -> bool:
 If not set, the default value is `False` for non-CI environments.
 If running in a CI environment ("CI" env var is set), then the default value is `True`.
 """
+
+DEFAULT_HTTP_CACHE_DIR: Path = (
+    Path(".airbyte-http-cache")
+    if "AIRBYTE_HTTP_CACHE_DIR" not in os.environ
+    else Path(os.environ["AIRBYTE_HTTP_CACHE_DIR"])
+)
+"""Default HTTP cache directory is `.airbyte-http-cache` in the current working directory.
+
+The default location can be overridden by setting the `AIRBYTE_HTTP_CACHE_DIR` environment variable.
+
+This directory is used to store cached HTTP requests and responses when using the HTTP caching
+functionality with connectors.
+"""
+
+DEFAULT_HTTP_CACHE_READ_DIR: Path | None = (
+    None
+    if "AIRBYTE_HTTP_CACHE_READ_DIR" not in os.environ
+    else Path(os.environ["AIRBYTE_HTTP_CACHE_READ_DIR"])
+)
+"""Optional separate directory for reading HTTP cache files.
+
+If set, this directory will be used for reading cached HTTP responses, while the
+`DEFAULT_HTTP_CACHE_DIR` will be used for writing new cache files.
+"""

airbyte/http_caching/__init__.py

@@ -0,0 +1,10 @@
+# Copyright (c) 2024 Airbyte, Inc., all rights reserved.
+"""HTTP caching module for Airbyte connectors."""
+
+from __future__ import annotations
+
+from airbyte.http_caching.cache import AirbyteConnectorCache, HttpCacheMode
+from airbyte.http_caching.serialization import SerializationFormat
+
+
+__all__ = ["AirbyteConnectorCache", "HttpCacheMode", "SerializationFormat"]

airbyte/http_caching/cache.py

@@ -0,0 +1,112 @@
+# Copyright (c) 2024 Airbyte, Inc., all rights reserved.
+"""HTTP caching for Airbyte connectors."""
+
+from __future__ import annotations
+
+import threading
+from pathlib import Path
+from typing import cast
+
+from mitmproxy import options
+from mitmproxy.tools.dump import DumpMaster
+
+from airbyte.constants import DEFAULT_HTTP_CACHE_DIR, DEFAULT_HTTP_CACHE_READ_DIR
+from airbyte.http_caching.proxy import HttpCacheMode, HttpCachingAddon
+from airbyte.http_caching.serialization import SerializationFormat
+
+
+class AirbyteConnectorCache:
+    """Cache for Airbyte connector HTTP requests and responses.
+
+    This class manages an HTTP proxy that intercepts requests from connectors and either
+    serves them from the cache or forwards them to the server based on the cache mode.
+    """
+
+    def __init__(
+        self,
+        cache_dir: str | Path | None = None,
+        read_dir: str | Path | None = None,
+        mode: str | HttpCacheMode = HttpCacheMode.READ_WRITE,
+        serialization_format: str | SerializationFormat = SerializationFormat.BINARY,
+    ) -> None:
+        """Initialize the cache.
+
+        Args:
+            cache_dir: The directory where cache files are stored for writing. If not provided,
+                       the default directory will be used.
+            read_dir: Optional separate directory for reading cached responses. If not provided,
+                      cache_dir will be used for both reading and writing.
+            mode: The cache mode to use. Can be one of 'read_only', 'write_only',
+                  'read_write', or 'read_only_fail_on_miss'.
+            serialization_format: The format to use for serializing cached data. Can be
+                                 'json' or 'binary'.
+        """
+        self.cache_dir = Path(cache_dir) if cache_dir else DEFAULT_HTTP_CACHE_DIR
+
+        if read_dir:
+            self.read_dir = Path(read_dir)
+        elif DEFAULT_HTTP_CACHE_READ_DIR:
+            self.read_dir = DEFAULT_HTTP_CACHE_READ_DIR
+        else:
+            self.read_dir = self.cache_dir
+
+        self.mode = HttpCacheMode(mode) if isinstance(mode, str) else mode
+
+        self.serialization_format = (
+            SerializationFormat(serialization_format)
+            if isinstance(serialization_format, str)
+            else serialization_format
+        )
+
+        self._proxy_port: int | None = None
+        self._proxy_thread: threading.Thread | None = None
+        self._proxy: DumpMaster | None = None
+        self._addon: HttpCachingAddon | None = None
+
+    def start(self) -> int:
+        """Start the HTTP proxy.
+
+        Returns:
+            The port number the proxy is listening on.
+        """
+        if self._proxy_port is not None:
+            return self._proxy_port
+
+        port = 0
+
+        opts = options.Options(
+            listen_host="127.0.0.1",
+            listen_port=port,
+            ssl_insecure=True,  # Allow self-signed certificates
+            confdir=str(self.cache_dir),  # Store certificates in the cache directory
+        )
+
+        addon = HttpCachingAddon(
+            cache_dir=self.cache_dir,
+            read_dir=self.read_dir,
+            mode=self.mode,
+            serialization_format=self.serialization_format,
+        )
+        self._addon = addon
+
+        proxy = DumpMaster(opts)
+        self._proxy = proxy
+        proxy.addons.add(addon)
+
+        thread = threading.Thread(target=proxy.run, daemon=True)
+        self._proxy_thread = thread
+        thread.start()
+
+        port_number = cast("int", proxy.server.address[1])  # type: ignore[attr-defined]
+        self._proxy_port = port_number
+
+        return port_number
+
+    def stop(self) -> None:
+        """Stop the HTTP proxy."""
+        if self._proxy is not None:
+            self._proxy.shutdown()
+            self._proxy_thread = None
+            self._proxy = None
+            self._addon = None
+            self._proxy_port = None