client/session: Added timestamp generator to SessionConfig

Also added an ability to set it through Session Builder

Author: smoczy123

scylla/src/client/session.rs

@@ -31,6 +31,7 @@ use crate::policies::host_filter::HostFilter;
 use crate::policies::load_balancing::{self, RoutingInfo};
 use crate::policies::retry::{RequestInfo, RetryDecision, RetrySession};
 use crate::policies::speculative_execution;
+use crate::policies::timestamp_generator::TimestampGenerator;
 use crate::prepared_statement::{PartitionKeyError, PreparedStatement};
 use crate::query::Query;
 #[allow(deprecated)]
@@ -181,6 +182,10 @@ pub struct SessionConfig {
     /// Generally, this options is best left as default (false).
     pub disallow_shard_aware_port: bool,
 
+    ///  Timestamp generator used for generating timestamps on the client-side
+    ///  If None, server-side timestamps are used.
+    pub timestamp_generator: Option<Arc<dyn TimestampGenerator>>,
+
     /// If empty, fetch all keyspaces
     pub keyspaces_to_fetch: Vec<String>,
 
@@ -293,6 +298,7 @@ impl SessionConfig {
             connect_timeout: Duration::from_secs(5),
             connection_pool_size: Default::default(),
             disallow_shard_aware_port: false,
+            timestamp_generator: None,
             keyspaces_to_fetch: Vec::new(),
             fetch_schema_metadata: true,
             keepalive_interval: Some(Duration::from_secs(30)),

scylla/src/client/session_builder.rs

@@ -14,6 +14,7 @@ use crate::cloud::{CloudConfig, CloudConfigError};
 use crate::errors::NewSessionError;
 use crate::policies::address_translator::AddressTranslator;
 use crate::policies::host_filter::HostFilter;
+use crate::policies::timestamp_generator::TimestampGenerator;
 use crate::statement::Consistency;
 #[cfg(feature = "ssl")]
 use openssl::ssl::SslContext;
@@ -684,6 +685,28 @@ impl<K: SessionBuilderKind> GenericSessionBuilder<K> {
         self
     }
 
+    /// Set the timestamp generator that will generate timestamps on the client-side.
+    ///
+    /// # Example
+    /// ```
+    /// # use scylla::client::session::Session;
+    /// # use scylla::client::session_builder::SessionBuilder;
+    /// # use scylla::policies::timestamp_generator::SimpleTimestampGenerator;
+    /// # use std::sync::Arc;
+    /// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
+    /// let session: Session = SessionBuilder::new()
+    ///     .known_node("127.0.0.1:9042")
+    ///     .timestamp_generator(Arc::new(SimpleTimestampGenerator::new()))
+    ///     .build()
+    ///     .await?;
+    /// # Ok(())
+    /// # }
+    /// ```
+    pub fn timestamp_generator(mut self, timestamp_generator: Arc<dyn TimestampGenerator>) -> Self {
+        self.config.timestamp_generator = Some(timestamp_generator);
+        self
+    }
+
     /// Set the keyspaces to be fetched, to retrieve their strategy, and schema metadata if enabled
     /// No keyspaces, the default value, means all the keyspaces will be fetched.
     ///