From d72c1cba9efa878b302352d6ec48283e3f6ff5df Mon Sep 17 00:00:00 2001 From: Filipe Oliveira <144428646+filipeaaoliveira@users.noreply.github.com> Date: Tue, 4 Aug 2026 14:08:59 +0200 Subject: [PATCH] Add configurable request timeout for the Airbyte API Since apache-airflow-providers-airbyte 6.0.0 the provider uses airbyte-api 1.x, which switched from requests to httpx and applies a 5-second default request timeout. On self-hosted Airbyte deployments job creation can take far longer than that under load, so every trigger fails - and because the create-job call is not idempotent, the timed-out request still creates a job server-side, causing 409 errors and duplicate syncs on retry. Allow users to raise the timeout via a hook parameter or a connection extra. --- providers/airbyte/docs/connections.rst | 10 +++-- .../providers/airbyte/hooks/airbyte.py | 22 +++++++++++ .../tests/unit/airbyte/hooks/test_airbyte.py | 39 +++++++++++++++++++ 3 files changed, 68 insertions(+), 3 deletions(-) diff --git a/providers/airbyte/docs/connections.rst b/providers/airbyte/docs/connections.rst index 63733b08bbcca..4a8ab2a9ff592 100644 --- a/providers/airbyte/docs/connections.rst +++ b/providers/airbyte/docs/connections.rst @@ -57,8 +57,12 @@ Client Secret (optional) Leave blank for Airbyte OSS deployments without auth enabled. Extra (optional) - Specify the ``proxies`` key in JSON format to route traffic through an HTTP proxy. + Specify extra parameters as JSON. The following keys are supported: - * ``proxies`` + * ``proxies`` - Route traffic through an HTTP proxy. + * ``timeout`` - Request timeout, in seconds, for each call to the Airbyte API. + When not set, the underlying HTTP client's default of 5 seconds applies. + Can also be set programmatically via the ``timeout`` parameter of + ``AirbyteHook``, which takes precedence over this extra. - Example: ``{"proxies": {"http": "http://proxy.example.com:8080", "https": "http://proxy.example.com:8080"}}`` + Example: ``{"proxies": {"http": "http://proxy.example.com:8080", "https": "http://proxy.example.com:8080"}, "timeout": 60}`` diff --git a/providers/airbyte/src/airflow/providers/airbyte/hooks/airbyte.py b/providers/airbyte/src/airflow/providers/airbyte/hooks/airbyte.py index dd7cc04562695..68061caf34134 100644 --- a/providers/airbyte/src/airflow/providers/airbyte/hooks/airbyte.py +++ b/providers/airbyte/src/airflow/providers/airbyte/hooks/airbyte.py @@ -37,6 +37,10 @@ class AirbyteHook(BaseHook): :param airbyte_conn_id: Optional. The name of the Airflow connection to get connection information for Airbyte. Defaults to "airbyte_default". :param api_version: Optional. Airbyte API version. Defaults to "v1". + :param timeout: Optional. Request timeout, in seconds, for each call to the + Airbyte API. Overrides the ``timeout`` key in the connection's Extra field. + When neither is set, the underlying HTTP client's default of 5 seconds + applies. """ conn_name_attr = "airbyte_conn_id" @@ -48,10 +52,12 @@ def __init__( self, airbyte_conn_id: str = "airbyte_default", api_version: str = "v1", + timeout: float | None = None, ) -> None: super().__init__() self.api_version: str = api_version self.airbyte_conn_id = airbyte_conn_id + self.timeout = timeout self.conn = self.get_conn_params(self.airbyte_conn_id) self.airbyte_api = self.create_api_session() @@ -71,6 +77,7 @@ def get_conn_params(self, conn_id: str) -> Any: conn_params["client_secret"] = conn.password conn_params["token_url"] = conn.schema or "v1/applications/token" conn_params["proxies"] = conn.extra_dejson.get("proxies", None) + conn_params["timeout"] = conn.extra_dejson.get("timeout", None) return conn_params @@ -124,10 +131,25 @@ def create_api_session(self) -> AirbyteAPI: } client = httpx.Client(mounts=mounts) + timeout = self.timeout if self.timeout is not None else self.conn["timeout"] + timeout_ms: int | None = None + if timeout is not None: + try: + timeout_ms = int(float(timeout) * 1000) + except (TypeError, ValueError): + timeout_ms = 0 + if timeout_ms <= 0: + raise ValueError( + f"Invalid Airbyte API request timeout {timeout!r}: expected a positive number of " + f"seconds, set via the AirbyteHook 'timeout' parameter or the 'timeout' extra of " + f"connection {self.airbyte_conn_id!r}" + ) + return AirbyteAPI( server_url=self.conn["host"], security=security, client=client, + timeout_ms=timeout_ms, ) @classmethod diff --git a/providers/airbyte/tests/unit/airbyte/hooks/test_airbyte.py b/providers/airbyte/tests/unit/airbyte/hooks/test_airbyte.py index f837a8bf77e71..7071c7722ef52 100644 --- a/providers/airbyte/tests/unit/airbyte/hooks/test_airbyte.py +++ b/providers/airbyte/tests/unit/airbyte/hooks/test_airbyte.py @@ -262,6 +262,45 @@ def test_create_api_session_with_proxy(self): transport = client._transport_for_url(url) assert transport is not default_transport, f"Expected proxy transport for {scheme}" + @pytest.mark.parametrize( + ("hook_timeout", "extra_timeout", "expected_timeout_ms"), + [ + pytest.param(None, None, None, id="default-unchanged"), + pytest.param(300, None, 300_000, id="hook-parameter"), + pytest.param(None, 120, 120_000, id="connection-extra"), + pytest.param(None, "60", 60_000, id="connection-extra-string"), + pytest.param(30.5, 120, 30_500, id="hook-parameter-overrides-extra"), + ], + ) + def test_create_api_session_timeout( + self, create_connection_without_db, hook_timeout, extra_timeout, expected_timeout_ms + ): + create_connection_without_db( + Connection( + conn_id="airbyte_conn_id_test_timeout", + conn_type=self.conn_type, + host=self.host, + port=self.port, + extra={"timeout": extra_timeout} if extra_timeout is not None else None, + ) + ) + hook = AirbyteHook(airbyte_conn_id="airbyte_conn_id_test_timeout", timeout=hook_timeout) + assert hook.airbyte_api.sdk_configuration.timeout_ms == expected_timeout_ms + + @pytest.mark.parametrize("bad_timeout", ["6o", 0, -5, {"connect": 5}]) + def test_create_api_session_invalid_timeout_extra(self, create_connection_without_db, bad_timeout): + create_connection_without_db( + Connection( + conn_id="airbyte_conn_id_test_bad_timeout", + conn_type=self.conn_type, + host=self.host, + port=self.port, + extra={"timeout": bad_timeout}, + ) + ) + with pytest.raises(ValueError, match="Invalid Airbyte API request timeout"): + AirbyteHook(airbyte_conn_id="airbyte_conn_id_test_bad_timeout") + def test_create_api_session_without_credentials(self): """Test that a session without OAuth credentials creates an unauthenticated client.""" # The default connection (self.airbyte_conn_id) has no login/password