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
 
@@ -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

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