From c4671a18f7390b3db283de719160a0d26cf7225f Mon Sep 17 00:00:00 2001
From: smoczy123 <michal.hudobski@poczta.onet.pl>
Date: Wed, 4 Dec 2024 00:36:58 +0100
Subject: [PATCH] transport/session: Added timestamp generator to SessionConfig
 and an ability to provide it in SessionBuilder

---
 scylla/src/transport/session.rs         |  6 ++++++
 scylla/src/transport/session_builder.rs | 23 +++++++++++++++++++++++
 2 files changed, 29 insertions(+)

diff --git a/scylla/src/transport/session.rs b/scylla/src/transport/session.rs
index b3efa7e07..6d098cd29 100644
--- a/scylla/src/transport/session.rs
+++ b/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)),
diff --git a/scylla/src/transport/session_builder.rs b/scylla/src/transport/session_builder.rs
index 404e27733..ca9cf47dc 100644
--- a/scylla/src/transport/session_builder.rs
+++ b/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,28 @@ impl<K: SessionBuilderKind> GenericSessionBuilder<K> {
         self
     }
 
+    /// Set the timestamp generator that will generate timestamps on the client-side.
+    ///
+    /// # Example
+    /// ```
+    /// # use scylla::{Session, SessionBuilder};
+    /// # use scylla::transport::timestamp_generator::MonotonicTimestampGenerator;
+    /// # 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(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.
     ///