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..0298b2db6b668 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 @@ -107,11 +114,24 @@ def create_api_session(self) -> AirbyteAPI: self.airbyte_conn_id, ) - client = None + timeout = self.timeout if self.timeout is not None else self.conn["timeout"] + error_message = ( + 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}" + ) + if timeout is not None: + try: + timeout = float(timeout) + except (TypeError, ValueError) as e: + raise ValueError(error_message) from e + if timeout <= 0: + raise ValueError(error_message) + + mounts: dict[str, httpx.HTTPTransport] = {} if self.conn["proxies"]: self.log.debug("Creating client proxy...") proxies = self.conn["proxies"] - mounts = {} if isinstance(proxies, dict): for scheme, proxy_url in proxies.items(): # httpx mount keys require a "://" suffix @@ -122,7 +142,18 @@ def create_api_session(self) -> AirbyteAPI: "http://": httpx.HTTPTransport(proxy=proxies), "https://": httpx.HTTPTransport(proxy=proxies), } - client = httpx.Client(mounts=mounts) + + client = None + if mounts or timeout is not None: + # The timeout must be set on the client rather than passed as the SDK's + # timeout_ms: the SDK's client-credentials hook sends the OAuth token + # request directly through the client, bypassing per-operation timeouts. + # follow_redirects matches the default client the SDK creates otherwise. + client = httpx.Client( + mounts=mounts, + timeout=timeout if timeout is not None else 5.0, + follow_redirects=True, + ) return AirbyteAPI( server_url=self.conn["host"], diff --git a/providers/airbyte/tests/unit/airbyte/hooks/test_airbyte.py b/providers/airbyte/tests/unit/airbyte/hooks/test_airbyte.py index f837a8bf77e71..40214dd6010d7 100644 --- a/providers/airbyte/tests/unit/airbyte/hooks/test_airbyte.py +++ b/providers/airbyte/tests/unit/airbyte/hooks/test_airbyte.py @@ -262,6 +262,67 @@ 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"), + [ + pytest.param(None, None, 5.0, id="default-unchanged"), + pytest.param(300, None, 300.0, id="hook-parameter"), + pytest.param(None, 120, 120.0, id="connection-extra"), + pytest.param(None, "60", 60.0, id="connection-extra-string"), + pytest.param(30.5, 120, 30.5, id="hook-parameter-overrides-extra"), + ], + ) + def test_create_api_session_timeout( + self, create_connection_without_db, hook_timeout, extra_timeout, expected_timeout + ): + 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) + # The timeout is set on the httpx client (not the SDK's timeout_ms) so that it + # also covers the OAuth token request sent directly through the client. + client = hook.airbyte_api.sdk_configuration.client + assert client.timeout == httpx.Timeout(expected_timeout) + assert client.follow_redirects is True + + def test_create_api_session_timeout_with_proxy(self, create_connection_without_db): + create_connection_without_db( + Connection( + conn_id="airbyte_conn_id_test_timeout_proxy", + conn_type=self.conn_type, + host=self.host, + port=self.port, + extra={**self._mock_proxy, "timeout": 90}, + ) + ) + hook = AirbyteHook(airbyte_conn_id="airbyte_conn_id_test_timeout_proxy") + client = hook.airbyte_api.sdk_configuration.client + assert client.timeout == httpx.Timeout(90.0) + default_transport = client._transport + for scheme in self._mock_proxy["proxies"]: + url = httpx.URL(f"{scheme}://example.com") + assert client._transport_for_url(url) is not default_transport + + @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