From 72d579f03cc32b08d10fbf653383333e675c6894 Mon Sep 17 00:00:00 2001
From: smoczy123 <michal.hudobski@poczta.onet.pl>
Date: Wed, 4 Dec 2024 23:21:30 +0100
Subject: [PATCH] transport/connection: Added logic for setting a timestamp
 through timestamp generator if user did not provide one

---
 scylla/src/transport/connection.rs | 38 +++++++++++++++++++++++++++---
 1 file changed, 35 insertions(+), 3 deletions(-)

diff --git a/scylla/src/transport/connection.rs b/scylla/src/transport/connection.rs
index 18ea4847d..882029673 100644
--- a/scylla/src/transport/connection.rs
+++ b/scylla/src/transport/connection.rs
@@ -1050,6 +1050,17 @@ impl Connection {
         page_size: Option<PageSize>,
         paging_state: PagingState,
     ) -> Result<QueryResponse, UserRequestError> {
+        let mut timestamp = None;
+        if query.get_timestamp().is_none() {
+            if let Some(x) = self.config.timestamp_generator.clone() {
+                timestamp = Some(x.next_timestamp().await);
+            }
+        }
+
+        if timestamp.is_none() {
+            timestamp = query.get_timestamp()
+        }
+
         let query_frame = query::Query {
             contents: Cow::Borrowed(&query.contents),
             parameters: query::QueryParameters {
@@ -1059,7 +1070,7 @@ impl Connection {
                 page_size: page_size.map(Into::into),
                 paging_state,
                 skip_metadata: false,
-                timestamp: query.get_timestamp(),
+                timestamp,
             },
         };
 
@@ -1112,6 +1123,17 @@ impl Connection {
         page_size: Option<PageSize>,
         paging_state: PagingState,
     ) -> Result<QueryResponse, UserRequestError> {
+        let mut timestamp = None;
+        if prepared_statement.get_timestamp().is_none() {
+            if let Some(x) = self.config.timestamp_generator.clone() {
+                timestamp = Some(x.next_timestamp().await);
+            }
+        }
+
+        if timestamp.is_none() {
+            timestamp = prepared_statement.get_timestamp()
+        }
+
         let execute_frame = execute::Execute {
             id: prepared_statement.get_id().to_owned(),
             parameters: query::QueryParameters {
@@ -1119,7 +1141,7 @@ impl Connection {
                 serial_consistency,
                 values: Cow::Borrowed(values),
                 page_size: page_size.map(Into::into),
-                timestamp: prepared_statement.get_timestamp(),
+                timestamp,
                 skip_metadata: prepared_statement.get_use_cached_result_metadata(),
                 paging_state,
             },
@@ -1251,6 +1273,16 @@ impl Connection {
         });
 
         let values = RawBatchValuesAdapter::new(values, contexts);
+        let mut timestamp = None;
+        if batch.get_timestamp().is_none() {
+            if let Some(x) = self.config.timestamp_generator.clone() {
+                timestamp = Some(x.next_timestamp().await);
+            }
+        }
+
+        if timestamp.is_none() {
+            timestamp = batch.get_timestamp()
+        }
 
         let batch_frame = batch::Batch {
             statements: Cow::Borrowed(&batch.statements),
@@ -1258,7 +1290,7 @@ impl Connection {
             batch_type: batch.get_type(),
             consistency,
             serial_consistency,
-            timestamp: batch.get_timestamp(),
+            timestamp,
         };
 
         loop {