Added request_timeout field to statements

StatementConfig had this field added, and PreparedStatement and Query
got getters and setters for that.

Author: wprzytula

scylla/src/statement/mod.rs

@@ -1,4 +1,4 @@
-use std::sync::Arc;
+use std::{sync::Arc, time::Duration};
 
 use crate::transport::retry_policy::RetryPolicy;
 use crate::transport::speculative_execution::SpeculativeExecutionPolicy;
@@ -21,6 +21,7 @@ pub struct StatementConfig {
 
     pub tracing: bool,
     pub timestamp: Option<i64>,
+    pub request_timeout: Option<Duration>,
 }
 
 impl Default for StatementConfig {
@@ -33,23 +34,20 @@ impl Default for StatementConfig {
             speculative_execution_policy: None,
             tracing: false,
             timestamp: None,
+            request_timeout: None,
         }
     }
 }
 
 impl Clone for StatementConfig {
     fn clone(&self) -> Self {
         Self {
-            consistency: self.consistency,
-            serial_consistency: self.serial_consistency,
-            is_idempotent: self.is_idempotent,
             retry_policy: self
                 .retry_policy
                 .as_ref()
                 .map(|policy| policy.clone_boxed()),
             speculative_execution_policy: self.speculative_execution_policy.clone(),
-            tracing: self.tracing,
-            timestamp: self.timestamp,
+            ..*self
         }
     }
 }

scylla/src/statement/prepared_statement.rs

@@ -1,6 +1,7 @@
 use bytes::{BufMut, Bytes, BytesMut};
 use smallvec::{smallvec, SmallVec};
 use std::convert::TryInto;
+use std::time::Duration;
 use thiserror::Error;
 use uuid::Uuid;
 
@@ -247,6 +248,19 @@ impl PreparedStatement {
         self.config.timestamp
     }
 
+    /// Sets the client-side timeout for this statement.
+    /// If not None, the driver will stop waiting for the request
+    /// to finish after `timeout` passed.
+    /// Otherwise, default session client timeout will be applied.
+    pub fn set_request_timeout(&mut self, timeout: Option<Duration>) {
+        self.config.request_timeout = timeout
+    }
+
+    /// Gets client timeout associated with this query
+    pub fn get_request_timeout(&self) -> Option<Duration> {
+        self.config.request_timeout
+    }
+
     /// Sets the name of the partitioner used for this statement.
     pub(crate) fn set_partitioner_name(&mut self, partitioner_name: Option<&str>) {
         self.partitioner_name = match partitioner_name {

scylla/src/statement/query.rs

@@ -1,6 +1,7 @@
 use super::StatementConfig;
 use crate::frame::types::{Consistency, SerialConsistency};
 use crate::transport::retry_policy::RetryPolicy;
+use std::time::Duration;
 
 /// CQL query statement.
 ///
@@ -117,6 +118,19 @@ impl Query {
     pub fn get_timestamp(&self) -> Option<i64> {
         self.config.timestamp
     }
+
+    /// Sets the client-side timeout for this statement.
+    /// If not None, the driver will stop waiting for the request
+    /// to finish after `timeout` passed.
+    /// Otherwise, default session client timeout will be applied.
+    pub fn set_request_timeout(&mut self, timeout: Option<Duration>) {
+        self.config.request_timeout = timeout
+    }
+
+    /// Gets client timeout associated with this query
+    pub fn get_request_timeout(&self) -> Option<Duration> {
+        self.config.request_timeout
+    }
 }
 
 impl From<String> for Query {