Add configurable request timeout for the Airbyte API

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}``

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"],

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