session: added timestamp generator to SessionConfig and an ability to provide it in SessionBuilder

Author: smoczy123

scylla/src/transport/session.rs

@@ -53,6 +53,7 @@ use super::node::{InternalKnownNode, KnownNode};
 use super::partitioner::PartitionerName;
 use super::query_result::MaybeFirstRowError;
 use super::query_result::RowsError;
+use super::timestamp_generator::TimestampGenerator;
 use super::topology::UntranslatedPeer;
 use super::{NodeRef, SelfIdentity};
 use crate::frame::response::result;
@@ -270,6 +271,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>,
 
@@ -382,6 +387,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/transport/session_builder.rs

@@ -7,6 +7,7 @@ use super::session::{
     AddressTranslator, CurrentDeserializationApi, GenericSession, LegacyDeserializationApi,
     SessionConfig,
 };
+use super::timestamp_generator::TimestampGenerator;
 use super::Compression;
 
 #[cfg(feature = "cloud")]
@@ -663,6 +664,26 @@ impl<K: SessionBuilderKind> GenericSessionBuilder<K> {
         self
     }
 
+    /// Set the timestamp generator that will generate timestamps on the client-side.
+    ///
+    /// # Example
+    /// ```
+    /// # use scylla::{MonotonicTimestampGenerator, Session, SessionBuilder};
+    /// # 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(MonotonicTimestampGenerator::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.
     ///