diff --git a/airbyte_cdk/sources/declarative/declarative_component_schema.yaml b/airbyte_cdk/sources/declarative/declarative_component_schema.yaml index dfabccd6b0..251bce895d 100644 --- a/airbyte_cdk/sources/declarative/declarative_component_schema.yaml +++ b/airbyte_cdk/sources/declarative/declarative_component_schema.yaml @@ -2778,6 +2778,27 @@ definitions: description: Enables stream requests caching. When set to true, repeated requests to the same URL will return cached responses. Parent streams automatically have caching enabled. Only set this to false if you are certain that caching should be disabled, as it may negatively impact performance when the same data is needed multiple times (e.g., for scroll-based pagination APIs where caching causes duplicate records). type: boolean default: false + connect_timeout_in_seconds: + title: Connect Timeout (seconds) + description: Maximum time to establish the connection (TCP + TLS handshake). When unset, no connect timeout is applied. + type: number + exclusiveMinimum: 0 + examples: + - 30 + read_timeout_in_seconds: + title: Read Timeout (seconds) + description: The maximum time to wait for the server to send any data (the first byte of the response, or the next chunk while the body is read). It is NOT a limit on the total request duration. When unset, no read timeout is applied. Only set it when every endpoint served by this requester responds well within the value; synchronous report or export endpoints can stay silent for minutes. Each timed-out attempt counts against the error handler's retry window (600 s for manifest connectors), so a value near 300 allows only a few retries. + type: number + exclusiveMinimum: 0 + examples: + - 300 + use_tcp_keepalive: + title: Use TCP Keepalive + description: Enables TCP keepalive (60 s idle, 10 s probe interval, 6 probes) on this requester's connections, so dead or NAT-dropped connections are detected and fail instead of hanging forever. Does not fail slow but alive servers; use read_timeout_in_seconds to bound those. + type: boolean + default: false + examples: + - true $parameters: type: object additionalProperties: true diff --git a/airbyte_cdk/sources/declarative/models/declarative_component_schema.py b/airbyte_cdk/sources/declarative/models/declarative_component_schema.py index 5581be3757..2c3d864cd3 100644 --- a/airbyte_cdk/sources/declarative/models/declarative_component_schema.py +++ b/airbyte_cdk/sources/declarative/models/declarative_component_schema.py @@ -3204,6 +3204,26 @@ class HttpRequester(BaseModelWithDeprecations): description="Enables stream requests caching. When set to true, repeated requests to the same URL will return cached responses. Parent streams automatically have caching enabled. Only set this to false if you are certain that caching should be disabled, as it may negatively impact performance when the same data is needed multiple times (e.g., for scroll-based pagination APIs where caching causes duplicate records).", title="Use Cache", ) + connect_timeout_in_seconds: Optional[float] = Field( + None, + description="Maximum time to establish the connection (TCP + TLS handshake). When unset, no connect timeout is applied.", + examples=[30], + gt=0, + title="Connect Timeout (seconds)", + ) + read_timeout_in_seconds: Optional[float] = Field( + None, + description="The maximum time to wait for the server to send any data (the first byte of the response, or the next chunk while the body is read). It is NOT a limit on the total request duration. When unset, no read timeout is applied. Only set it when every endpoint served by this requester responds well within the value; synchronous report or export endpoints can stay silent for minutes. Each timed-out attempt counts against the error handler's retry window (600 s for manifest connectors), so a value near 300 allows only a few retries.", + examples=[300], + gt=0, + title="Read Timeout (seconds)", + ) + use_tcp_keepalive: Optional[bool] = Field( + False, + description="Enables TCP keepalive (60 s idle, 10 s probe interval, 6 probes) on this requester's connections, so dead or NAT-dropped connections are detected and fail instead of hanging forever. Does not fail slow but alive servers; use read_timeout_in_seconds to bound those.", + examples=[True], + title="Use TCP Keepalive", + ) parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters") diff --git a/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py b/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py index 97a1f4ce2d..219e0f70ad 100644 --- a/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py +++ b/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py @@ -2803,6 +2803,9 @@ def create_http_requester( parameters=model.parameters or {}, message_repository=self._message_repository, use_cache=should_use_cache, + connect_timeout_in_seconds=model.connect_timeout_in_seconds, + read_timeout_in_seconds=model.read_timeout_in_seconds, + use_tcp_keepalive=model.use_tcp_keepalive or False, decoder=decoder, stream_response=decoder.is_stream_response() if decoder else False, ) diff --git a/airbyte_cdk/sources/declarative/requesters/http_requester.py b/airbyte_cdk/sources/declarative/requesters/http_requester.py index 3ce4c85408..bd1468bc7b 100644 --- a/airbyte_cdk/sources/declarative/requesters/http_requester.py +++ b/airbyte_cdk/sources/declarative/requesters/http_requester.py @@ -66,6 +66,9 @@ class HttpRequester(Requester): disable_retries: bool = False message_repository: MessageRepository = NoopMessageRepository() use_cache: bool = False + connect_timeout_in_seconds: Optional[float] = None + read_timeout_in_seconds: Optional[float] = None + use_tcp_keepalive: bool = False _exit_on_rate_limit: bool = False stream_response: bool = False decoder: Decoder = field(default_factory=lambda: JsonDecoder(parameters={})) @@ -114,6 +117,13 @@ def __post_init__(self, parameters: Mapping[str, Any]) -> None: backoff_strategy=backoff_strategies, disable_retries=self.disable_retries, message_repository=self.message_repository, + request_timeout=( + (self.connect_timeout_in_seconds, self.read_timeout_in_seconds) + if self.connect_timeout_in_seconds is not None + or self.read_timeout_in_seconds is not None + else None + ), + use_tcp_keepalive=self.use_tcp_keepalive, ) @property diff --git a/airbyte_cdk/sources/streams/http/http_client.py b/airbyte_cdk/sources/streams/http/http_client.py index 50d03039cc..8da499aa32 100644 --- a/airbyte_cdk/sources/streams/http/http_client.py +++ b/airbyte_cdk/sources/streams/http/http_client.py @@ -65,6 +65,7 @@ ResponseAwareAuthenticator, TokenRotatingAuthenticator, ) +from airbyte_cdk.sources.streams.http.tcp_keepalive import TcpKeepaliveHTTPAdapter from airbyte_cdk.sources.utils.types import JsonType from airbyte_cdk.utils.airbyte_secrets_utils import filter_secrets from airbyte_cdk.utils.constants import ENV_REQUEST_CACHE_PATH @@ -160,6 +161,8 @@ def __init__( error_message_parser: Optional[ErrorMessageParser] = None, disable_retries: bool = False, message_repository: Optional[MessageRepository] = None, + request_timeout: Optional[Union[float, Tuple[Optional[float], Optional[float]]]] = None, + use_tcp_keepalive: bool = False, ): self._name = name self._api_budget: APIBudget = api_budget or APIBudget(policies=[]) @@ -168,12 +171,21 @@ def __init__( else: self._use_cache = use_cache self._session = self._request_session() - self._session.mount( - "https://", - requests.adapters.HTTPAdapter( - pool_connections=MAX_CONNECTION_POOL_SIZE, pool_maxsize=MAX_CONNECTION_POOL_SIZE - ), - ) + if use_tcp_keepalive: + tcp_keepalive_adapter = TcpKeepaliveHTTPAdapter( + pool_connections=MAX_CONNECTION_POOL_SIZE, + pool_maxsize=MAX_CONNECTION_POOL_SIZE, + ) + self._session.mount("https://", tcp_keepalive_adapter) + self._session.mount("http://", tcp_keepalive_adapter) + else: + self._session.mount( + "https://", + requests.adapters.HTTPAdapter( + pool_connections=MAX_CONNECTION_POOL_SIZE, + pool_maxsize=MAX_CONNECTION_POOL_SIZE, + ), + ) if isinstance(authenticator, AuthBase): self._session.auth = authenticator self._logger = logger @@ -191,6 +203,7 @@ def __init__( self._disable_retries = disable_retries self._message_repository = message_repository self._authenticator_update_failed = False + self._request_timeout = request_timeout @property def cache_filename(self) -> str: @@ -453,6 +466,9 @@ def _send( response: Optional[requests.Response] = None exc: Optional[requests.RequestException] = None + if self._request_timeout is not None and "timeout" not in request_kwargs: + request_kwargs = {**request_kwargs, "timeout": self._request_timeout} + try: response = self._session.send(request, **request_kwargs) except requests.RequestException as e: diff --git a/airbyte_cdk/sources/streams/http/tcp_keepalive.py b/airbyte_cdk/sources/streams/http/tcp_keepalive.py new file mode 100644 index 0000000000..35400fcd68 --- /dev/null +++ b/airbyte_cdk/sources/streams/http/tcp_keepalive.py @@ -0,0 +1,89 @@ +# +# Copyright (c) 2026 Airbyte, Inc., all rights reserved. +# + +import functools +import logging +import socket +from typing import List, Tuple, cast + +import requests +import urllib3.connection + +TCP_KEEPALIVE_IDLE_SECONDS = 60 +TCP_KEEPALIVE_INTERVAL_SECONDS = 10 +TCP_KEEPALIVE_PROBE_COUNT = 6 + +logger = logging.getLogger("airbyte") + + +@functools.lru_cache(maxsize=None) +def tcp_keepalive_socket_options() -> Tuple[Tuple[int, int, int], ...]: + """Socket options enabling TCP keepalive, probed for what this OS supports. + + Returns urllib3's default options plus SO_KEEPALIVE and the platform's idle / + interval / probe-count tuning options (TCP_KEEPIDLE on Linux/Windows, + TCP_KEEPALIVE on macOS). Each option is tried once on a throwaway socket and + dropped if the OS rejects it, because urllib3 applies these at connect time + and an unsupported option would fail every connection. Never raises. + """ + base_options = cast( + List[Tuple[int, int, int]], + list(urllib3.connection.HTTPConnection.default_socket_options), + ) + + idle_option = getattr(socket, "TCP_KEEPIDLE", None) or getattr(socket, "TCP_KEEPALIVE", None) + candidates: List[Tuple[int, int, int]] = [ + (socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1), + ] + if idle_option is not None: + candidates.append((socket.IPPROTO_TCP, idle_option, TCP_KEEPALIVE_IDLE_SECONDS)) + for name, value in ( + ("TCP_KEEPINTVL", TCP_KEEPALIVE_INTERVAL_SECONDS), + ("TCP_KEEPCNT", TCP_KEEPALIVE_PROBE_COUNT), + ): + option = getattr(socket, name, None) + if option is not None: + candidates.append((socket.IPPROTO_TCP, option, value)) + + try: + probe = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + except OSError as error: + _warn_keepalive_unavailable(error) + return tuple(base_options) + try: + supported = base_options + for level, optname, value in candidates: + try: + probe.setsockopt(level, optname, value) + except OSError as error: + if (level, optname) == (socket.SOL_SOCKET, socket.SO_KEEPALIVE): + # Without SO_KEEPALIVE the tuning options are useless. + _warn_keepalive_unavailable(error) + return tuple(base_options) + continue + supported.append((level, optname, value)) + return tuple(supported) + finally: + probe.close() + + +def _warn_keepalive_unavailable(error: OSError) -> None: + # Logged once per process because tcp_keepalive_socket_options() is cached. + logger.warning( + "TCP keepalive is unavailable on this platform, so connections use plain sockets: %s", error + ) + + +class TcpKeepaliveHTTPAdapter(requests.adapters.HTTPAdapter): + """HTTPAdapter that enables TCP keepalive on every connection it pools.""" + + def init_poolmanager(self, *args, **pool_kwargs): # type: ignore[no-untyped-def] + pool_kwargs.setdefault("socket_options", tcp_keepalive_socket_options()) + super().init_poolmanager(*args, **pool_kwargs) # type: ignore[no-untyped-call] + + def proxy_manager_for(self, proxy, **proxy_kwargs): # type: ignore[no-untyped-def] + # requests does not forward init_poolmanager kwargs to proxy managers, + # so without this override proxied connections get no keepalive. + proxy_kwargs.setdefault("socket_options", tcp_keepalive_socket_options()) + return super().proxy_manager_for(proxy, **proxy_kwargs) # type: ignore[no-untyped-call] diff --git a/poetry.lock b/poetry.lock index e649b2c475..671ac34351 100644 --- a/poetry.lock +++ b/poetry.lock @@ -8489,4 +8489,4 @@ vector-db-based = ["cohere", "langchain_community", "langchain_core", "langchain [metadata] lock-version = "2.1" python-versions = ">=3.10,<3.15" -content-hash = "120f1861e4b05dba481bfd92ce881558bc74b80f935e5d8b3e156ff0fd57a6ea" +content-hash = "fbdcb7c18d8e19a3db34d119631e5dc0e89347558c2877e40cc472f60a4839b8" diff --git a/pyproject.toml b/pyproject.toml index c0310807a9..7dd5111d5a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -58,6 +58,7 @@ requests_cache = "*" rich = "*" # Use whichever version meets other constraints rich-click = "^1.8.8" typing-extensions = "*" # Transitive dependency used directly in code +urllib3 = "*" # Transitive dependency used directly in code wcmatch = "10.0" # Extras depedencies avro = { version = ">=1.11.2,<1.13.0", optional = true } # TODO: Move into dev dependencies if only used in tests diff --git a/unit_tests/sources/declarative/parsers/test_model_to_component_factory.py b/unit_tests/sources/declarative/parsers/test_model_to_component_factory.py index 56e1ad76f0..e54f23e8e8 100644 --- a/unit_tests/sources/declarative/parsers/test_model_to_component_factory.py +++ b/unit_tests/sources/declarative/parsers/test_model_to_component_factory.py @@ -274,6 +274,7 @@ from airbyte_cdk.sources.streams.http.requests_native_auth.oauth import ( SingleUseRefreshTokenOauth2Authenticator, ) +from airbyte_cdk.sources.streams.http.tcp_keepalive import TcpKeepaliveHTTPAdapter from airbyte_cdk.sources.types import StreamSlice from airbyte_cdk.utils import AirbyteTracedException from airbyte_cdk.utils.datetime_helpers import AirbyteDateTime, ab_datetime_now, ab_datetime_parse @@ -10096,3 +10097,88 @@ def test_create_file_uploader_with_a_combined_download_target_extractor(): ) assert isinstance(file_uploader.download_target_extractor, CombinedExtractor) + + +@pytest.mark.parametrize( + "manifest_fields, expected_timeout", + [ + pytest.param( + {"connect_timeout_in_seconds": 30, "read_timeout_in_seconds": 300}, + (30.0, 300.0), + id="both_timeouts", + ), + pytest.param( + {"read_timeout_in_seconds": 300}, + (None, 300.0), + id="read_only", + ), + pytest.param( + {"connect_timeout_in_seconds": 30}, + (30.0, None), + id="connect_only", + ), + pytest.param({}, None, id="no_timeouts"), + ], +) +def test_create_http_requester_with_request_timeouts(manifest_fields, expected_timeout): + requester = factory.create_component( + model_type=HttpRequesterModel, + component_definition={ + "type": "HttpRequester", + "url_base": "https://api.test.com", + "path": "/data", + "http_method": "GET", + **manifest_fields, + }, + config=input_config, + name="name", + decoder=None, + ) + + assert requester._http_client._request_timeout == expected_timeout + if expected_timeout is None: + assert ( + type(requester._http_client._session.adapters["https://"]) + is requests.adapters.HTTPAdapter + ) + assert ( + type(requester._http_client._session.adapters["http://"]) + is requests.adapters.HTTPAdapter + ) + + +def test_create_http_requester_with_tcp_keepalive(): + requester = factory.create_component( + model_type=HttpRequesterModel, + component_definition={ + "type": "HttpRequester", + "url_base": "https://api.test.com", + "path": "/data", + "http_method": "GET", + "use_tcp_keepalive": True, + }, + config=input_config, + name="name", + decoder=None, + ) + + assert isinstance(requester._http_client._session.adapters["https://"], TcpKeepaliveHTTPAdapter) + assert isinstance(requester._http_client._session.adapters["http://"], TcpKeepaliveHTTPAdapter) + + +@pytest.mark.parametrize( + "manifest_fields", + [ + pytest.param({"read_timeout_in_seconds": 0}, id="read_zero"), + pytest.param({"connect_timeout_in_seconds": -1}, id="connect_negative"), + pytest.param({"read_timeout_in_seconds": -5}, id="read_negative"), + ], +) +def test_http_requester_rejects_non_positive_timeouts(manifest_fields): + with pytest.raises(ValidationError): + HttpRequesterModel( + type="HttpRequester", + url_base="https://api.test.com", + path="/data", + **manifest_fields, + ) diff --git a/unit_tests/sources/declarative/requesters/test_http_requester.py b/unit_tests/sources/declarative/requesters/test_http_requester.py index 211b348908..88eabad54f 100644 --- a/unit_tests/sources/declarative/requesters/test_http_requester.py +++ b/unit_tests/sources/declarative/requesters/test_http_requester.py @@ -53,6 +53,8 @@ def factory( disable_retries: bool = False, message_repository: Optional[MessageRepository] = None, use_cache: bool = False, + connect_timeout_in_seconds: Optional[float] = None, + read_timeout_in_seconds: Optional[float] = None, ) -> HttpRequester: return HttpRequester( name=name, @@ -68,6 +70,8 @@ def factory( disable_retries=disable_retries, message_repository=message_repository or MagicMock(), use_cache=use_cache, + connect_timeout_in_seconds=connect_timeout_in_seconds, + read_timeout_in_seconds=read_timeout_in_seconds, ) return factory @@ -960,3 +964,37 @@ def test_http_requester_with_mock_api_budget(http_requester_factory, monkeypatch assert response.status_code == 200 assert mock_budget.acquire_call.call_count == 1 + + +def test_default_requester_sends_no_timeout_kwarg(http_requester_factory): + requester = http_requester_factory(url_base="https://example.com", path="test") + + assert requester._http_client._request_timeout is None + + dummy_response = requests.Response() + dummy_response.status_code = 200 + send_mock = MagicMock(return_value=dummy_response) + requester._http_client._session.send = send_mock + + requester.send_request() + + assert send_mock.call_count == 1 + assert "timeout" not in send_mock.call_args.kwargs + + +def test_requester_with_timeouts_sends_timeout_tuple(http_requester_factory): + requester = http_requester_factory( + url_base="https://example.com", + path="test", + connect_timeout_in_seconds=30, + read_timeout_in_seconds=300, + ) + + dummy_response = requests.Response() + dummy_response.status_code = 200 + send_mock = MagicMock(return_value=dummy_response) + requester._http_client._session.send = send_mock + + requester.send_request() + + assert send_mock.call_args.kwargs["timeout"] == (30, 300) diff --git a/unit_tests/sources/streams/http/test_http_client.py b/unit_tests/sources/streams/http/test_http_client.py index ca960f95bb..31d4ff9a1c 100644 --- a/unit_tests/sources/streams/http/test_http_client.py +++ b/unit_tests/sources/streams/http/test_http_client.py @@ -7,6 +7,7 @@ import threading import time from datetime import timedelta +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from unittest.mock import MagicMock, patch import pytest @@ -16,6 +17,7 @@ from airbyte_cdk.models import FailureType, Level from airbyte_cdk.sources.declarative.auth.oauth import DeclarativeOauth2Authenticator +from airbyte_cdk.sources.http_config import MAX_CONNECTION_POOL_SIZE from airbyte_cdk.sources.http_logger import format_http_message from airbyte_cdk.sources.message import InMemoryMessageRepository from airbyte_cdk.sources.streams.call_rate import CachedLimiterSession, LimiterSession @@ -43,6 +45,7 @@ Oauth2Authenticator, TokenAuthenticator, ) +from airbyte_cdk.sources.streams.http.tcp_keepalive import TcpKeepaliveHTTPAdapter from airbyte_cdk.utils.constants import ENV_REQUEST_CACHE_PATH from airbyte_cdk.utils.datetime_helpers import ab_datetime_now from airbyte_cdk.utils.traced_exception import AirbyteTracedException @@ -2011,3 +2014,117 @@ def test_given_split_request_window_action_then_log_the_response_as_an_auxiliary "Stream 'test' request window rejected, retrying with a smaller window" ) assert "no records" in logged[0]["http"]["description"] + + +@pytest.mark.parametrize("use_cache", [True, False]) +def test_default_session_mounts_plain_http_adapter(use_cache): + http_client = HttpClient(name="test", logger=MagicMock(), use_cache=use_cache) + + assert type(http_client._session.adapters["https://"]) is requests.adapters.HTTPAdapter + assert not isinstance(http_client._session.adapters["http://"], TcpKeepaliveHTTPAdapter) + assert type(http_client._session.adapters["http://"]) is requests.adapters.HTTPAdapter + assert http_client._session.adapters["https://"]._pool_connections == MAX_CONNECTION_POOL_SIZE + assert http_client._session.adapters["https://"]._pool_maxsize == MAX_CONNECTION_POOL_SIZE + + +@pytest.mark.parametrize("use_cache", [True, False]) +def test_opted_in_session_mounts_tcp_keepalive_adapter(use_cache): + http_client = HttpClient( + name="test", logger=MagicMock(), use_cache=use_cache, use_tcp_keepalive=True + ) + + for prefix in ("https://", "http://"): + adapter = http_client._session.adapters[prefix] + assert isinstance(adapter, TcpKeepaliveHTTPAdapter) + assert adapter._pool_connections == MAX_CONNECTION_POOL_SIZE + assert adapter._pool_maxsize == MAX_CONNECTION_POOL_SIZE + + +def test_passed_in_session_is_not_modified_even_with_tcp_keepalive(): + mocked_session = MagicMock(spec=requests.Session) + HttpClient( + name="test", + logger=MagicMock(), + session=mocked_session, + use_tcp_keepalive=True, + ) + + mocked_session.mount.assert_not_called() + + +def _session_with_200(): + mocked_session = MagicMock(spec=requests.Session) + mocked_response = requests.Response() + mocked_response.status_code = 200 + mocked_session.send.return_value = mocked_response + return mocked_session + + +def test_no_timeout_kwarg_by_default(): + mocked_session = _session_with_200() + http_client = HttpClient(name="test", logger=MagicMock(), session=mocked_session) + + http_client.send_request(http_method="get", url="https://airbyte.io", request_kwargs={}) + + assert "timeout" not in mocked_session.send.call_args.kwargs + + +def test_request_timeout_is_passed_to_session_send(): + mocked_session = _session_with_200() + http_client = HttpClient( + name="test", logger=MagicMock(), session=mocked_session, request_timeout=(5, 10) + ) + + http_client.send_request(http_method="get", url="https://airbyte.io", request_kwargs={}) + + assert mocked_session.send.call_args.kwargs["timeout"] == (5, 10) + + +def test_explicit_timeout_in_request_kwargs_wins(): + mocked_session = _session_with_200() + http_client = HttpClient( + name="test", logger=MagicMock(), session=mocked_session, request_timeout=(5, 10) + ) + + http_client.send_request( + http_method="get", + url="https://airbyte.io", + request_kwargs={"timeout": 1}, + ) + + assert mocked_session.send.call_args.kwargs["timeout"] == 1 + + +def test_read_timeout_is_retried_then_raises(): + attempts = [] + + class SlowHandler(BaseHTTPRequestHandler): + def do_GET(self): + attempts.append(1) + time.sleep(2) + self.send_response(200) + self.send_header("Content-Length", "0") + self.end_headers() + + def log_message(self, *args): + pass + + server = ThreadingHTTPServer(("127.0.0.1", 0), SlowHandler) + threading.Thread(target=server.serve_forever, daemon=True).start() + try: + http_client = HttpClient( + name="test", + logger=MagicMock(), + error_handler=HttpStatusErrorHandler(logger=MagicMock(), max_retries=1), + request_timeout=(1, 0.5), + ) + with pytest.raises(AirbyteTracedException): + http_client.send_request( + http_method="get", + url=f"http://127.0.0.1:{server.server_port}/", + request_kwargs={}, + ) + assert len(attempts) == 2 + finally: + server.shutdown() + server.server_close() diff --git a/unit_tests/sources/streams/http/test_tcp_keepalive.py b/unit_tests/sources/streams/http/test_tcp_keepalive.py new file mode 100644 index 0000000000..3f79772638 --- /dev/null +++ b/unit_tests/sources/streams/http/test_tcp_keepalive.py @@ -0,0 +1,211 @@ +# +# Copyright (c) 2026 Airbyte, Inc., all rights reserved. +# + +import logging +import socket +import threading +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from unittest.mock import MagicMock + +import pytest +import urllib3.connection + +from airbyte_cdk.sources.streams.http.http_client import HttpClient +from airbyte_cdk.sources.streams.http.tcp_keepalive import ( + TCP_KEEPALIVE_IDLE_SECONDS, + TCP_KEEPALIVE_INTERVAL_SECONDS, + TCP_KEEPALIVE_PROBE_COUNT, + TcpKeepaliveHTTPAdapter, + tcp_keepalive_socket_options, +) + + +@pytest.fixture(autouse=True) +def clear_options_cache(): + tcp_keepalive_socket_options.cache_clear() + yield + tcp_keepalive_socket_options.cache_clear() + + +def _ok_server(): + class OkHandler(BaseHTTPRequestHandler): + def do_GET(self): + self.send_response(200) + self.send_header("Content-Length", "0") + self.end_headers() + + def log_message(self, *args): + pass + + server = ThreadingHTTPServer(("127.0.0.1", 0), OkHandler) + threading.Thread(target=server.serve_forever, daemon=True).start() + return server + + +def _capture_socket_options(monkeypatch): + captured_options = [] + real_connect = urllib3.connection.HTTPConnection.connect + + def capture_connect(connection): + real_connect(connection) + # The socket may be closed by the time the caller inspects it, so read + # the options while the connection is still open. + captured_options.append( + ( + connection.sock.getsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE), + connection.sock.getsockopt(socket.IPPROTO_TCP, socket.TCP_KEEPIDLE), + ) + ) + + monkeypatch.setattr(urllib3.connection.HTTPConnection, "connect", capture_connect) + return captured_options + + +def test_socket_options_returns_tuple(): + assert isinstance(tcp_keepalive_socket_options(), tuple) + + +@pytest.mark.skipif(not hasattr(socket, "TCP_KEEPIDLE"), reason="requires TCP_KEEPIDLE") +def test_socket_options_contains_keepalive_tuning(): + options = tcp_keepalive_socket_options() + + assert (socket.IPPROTO_TCP, socket.TCP_NODELAY, 1) in options + assert (socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1) in options + assert (socket.IPPROTO_TCP, socket.TCP_KEEPIDLE, TCP_KEEPALIVE_IDLE_SECONDS) in options + assert ( + socket.IPPROTO_TCP, + socket.TCP_KEEPINTVL, + TCP_KEEPALIVE_INTERVAL_SECONDS, + ) in options + assert (socket.IPPROTO_TCP, socket.TCP_KEEPCNT, TCP_KEEPALIVE_PROBE_COUNT) in options + + +def test_socket_options_degrade_when_idle_option_missing(monkeypatch): + monkeypatch.delattr(socket, "TCP_KEEPIDLE", raising=False) + monkeypatch.delattr(socket, "TCP_KEEPALIVE", raising=False) + + options = tcp_keepalive_socket_options() + + assert (socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1) in options + assert not any(value == TCP_KEEPALIVE_IDLE_SECONDS for _, _, value in options) + if hasattr(socket, "TCP_KEEPINTVL"): + assert ( + socket.IPPROTO_TCP, + socket.TCP_KEEPINTVL, + TCP_KEEPALIVE_INTERVAL_SECONDS, + ) in options + if hasattr(socket, "TCP_KEEPCNT"): + assert ( + socket.IPPROTO_TCP, + socket.TCP_KEEPCNT, + TCP_KEEPALIVE_PROBE_COUNT, + ) in options + + +def test_socket_options_drop_rejected_option(monkeypatch): + if not hasattr(socket, "TCP_KEEPCNT"): + pytest.skip("requires TCP_KEEPCNT") + + real_socket_cls = socket.socket + + class FakeSocket: + def setsockopt(self, level, optname, value): + if optname == socket.TCP_KEEPCNT: + raise OSError("unsupported") + + def close(self): + pass + + monkeypatch.setattr(socket, "socket", lambda *a, **kw: FakeSocket()) + + options = tcp_keepalive_socket_options() + + assert (socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1) in options + assert not any(optname == socket.TCP_KEEPCNT for _, optname, _ in options) + monkeypatch.undo() + assert isinstance(real_socket_cls(socket.AF_INET, socket.SOCK_STREAM), real_socket_cls) + + +def test_socket_options_warn_when_keepalive_rejected(monkeypatch, caplog): + class FakeSocket: + def setsockopt(self, level, optname, value): + if (level, optname) == (socket.SOL_SOCKET, socket.SO_KEEPALIVE): + raise OSError("unsupported") + + def close(self): + pass + + monkeypatch.setattr(socket, "socket", lambda *a, **kw: FakeSocket()) + + with caplog.at_level(logging.WARNING, logger="airbyte"): + options = tcp_keepalive_socket_options() + + assert options == tuple(urllib3.connection.HTTPConnection.default_socket_options) + assert "TCP keepalive is unavailable" in caplog.text + + +def test_socket_options_warn_when_probe_socket_fails(monkeypatch, caplog): + def failing_socket(*args, **kwargs): + raise OSError("no sockets") + + monkeypatch.setattr(socket, "socket", failing_socket) + + with caplog.at_level(logging.WARNING, logger="airbyte"): + options = tcp_keepalive_socket_options() + + assert options == tuple(urllib3.connection.HTTPConnection.default_socket_options) + assert "TCP keepalive is unavailable" in caplog.text + + +def test_socket_options_do_not_warn_when_supported(caplog): + with caplog.at_level(logging.WARNING, logger="airbyte"): + tcp_keepalive_socket_options() + + assert "TCP keepalive is unavailable" not in caplog.text + + +def test_adapter_passes_options_to_pool_and_proxy_managers(): + adapter = TcpKeepaliveHTTPAdapter() + + expected = tcp_keepalive_socket_options() + assert adapter.poolmanager.connection_pool_kw["socket_options"] == expected + proxy_manager = adapter.proxy_manager_for("http://proxy.invalid:3128") + assert proxy_manager.connection_pool_kw["socket_options"] == expected + + +@pytest.mark.skipif(not hasattr(socket, "TCP_KEEPIDLE"), reason="requires Linux") +def test_opted_in_http_client_connection_has_keepalive_enabled(monkeypatch): + captured_options = _capture_socket_options(monkeypatch) + + server = _ok_server() + try: + http_client = HttpClient(name="test", logger=MagicMock(), use_tcp_keepalive=True) + http_client.send_request( + http_method="GET", + url=f"http://127.0.0.1:{server.server_port}/", + request_kwargs={"stream": True}, + ) + assert captured_options == [(1, TCP_KEEPALIVE_IDLE_SECONDS)] + finally: + server.shutdown() + server.server_close() + + +@pytest.mark.skipif(not hasattr(socket, "TCP_KEEPIDLE"), reason="requires Linux") +def test_default_http_client_connection_has_no_keepalive(monkeypatch): + captured_options = _capture_socket_options(monkeypatch) + + server = _ok_server() + try: + http_client = HttpClient(name="test", logger=MagicMock()) + http_client.send_request( + http_method="GET", + url=f"http://127.0.0.1:{server.server_port}/", + request_kwargs={"stream": True}, + ) + assert captured_options + assert captured_options[0][0] == 0 + finally: + server.shutdown() + server.server_close()