[SPARK-52428] Add configurable gRPC connection and request timeouts

crates/connect/src/client/builder.rs

@@ -20,6 +20,7 @@
 use std::collections::HashMap;
 use std::env;
 use std::str::FromStr;
+use std::time::Duration;
 
 use crate::errors::SparkError;
 
@@ -43,6 +44,8 @@ pub struct ChannelBuilder {
     pub(super) user_agent: Option<String>,
     pub(super) use_ssl: bool,
     pub(super) headers: Option<HashMap<String, String>>,
+    pub connect_timeout: Option<Duration>,
+    pub request_timeout: Option<Duration>,
 }
 
 impl Default for ChannelBuilder {
@@ -172,6 +175,8 @@ impl ChannelBuilder {
             user_agent: ChannelBuilder::create_user_agent(None),
             use_ssl: false,
             headers: None,
+            connect_timeout: None,
+            request_timeout: None,
         };
 
         if let Some(mut headers) = headers {

crates/connect/src/client/config.rs

@@ -115,6 +115,8 @@ impl From<Config> for ChannelBuilder {
             } else {
                 Some(headers)
             },
+            connect_timeout: None,
+            request_timeout: None,
         }
     }
 }

crates/connect/src/session.rs

@@ -19,6 +19,7 @@
 
 use std::collections::HashMap;
 use std::sync::Arc;
+use std::time::Duration;
 
 use crate::client::{ChannelBuilder, Config, HeadersLayer, SparkClient, SparkConnectClient};
 
@@ -99,10 +100,30 @@ impl SparkSessionBuilder {
         self
     }
 
+    /// Sets the timeout for establishing the initial gRPC connection.
+    pub fn connect_timeout(mut self, timeout: Duration) -> Self {
+        self.channel_builder.connect_timeout = Some(timeout);
+        self
+    }
+
+    /// Sets the timeout for individual gRPC requests.
+    pub fn request_timeout(mut self, timeout: Duration) -> Self {
+        self.channel_builder.request_timeout = Some(timeout);
+        self
+    }
+
     async fn create_client(&self) -> Result<SparkSession, SparkError> {
-        let channel = Channel::from_shared(self.channel_builder.endpoint())?
-            .connect()
-            .await?;
+        let mut endpoint = Channel::from_shared(self.channel_builder.endpoint())?;
+
+        if let Some(timeout) = self.channel_builder.connect_timeout {
+            endpoint = endpoint.connect_timeout(timeout);
+        }
+
+        if let Some(timeout) = self.channel_builder.request_timeout {
+            endpoint = endpoint.timeout(timeout);
+        }
+
+        let channel = endpoint.connect().await?;
 
         let channel = ServiceBuilder::new()
             .layer(HeadersLayer::new(