transport/connection: Added logic for setting a timestamp through tim…

…estamp generator if user did not provide one
scylladb · Dec 4, 2024 · 72d579f · 72d579f
1 parent 0eaa012
commit 72d579f
Showing 1 changed file with 35 additions and 3 deletions.
diff --git 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,14 +1123,25 @@ 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 {
                 consistency,
                 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,14 +1273,24 @@ 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),
             values,
             batch_type: batch.get_type(),
             consistency,
             serial_consistency,
-            timestamp: batch.get_timestamp(),
+            timestamp,
         };
 
         loop {