Move timestamp test to query.rs and add missing cases

Author: sylwiaszunejko

scylla/tests/integration/main.rs

@@ -11,6 +11,7 @@ mod hygiene;
 mod large_batch_statements;
 mod lwt_optimisation;
 mod new_session;
+mod query;
 mod query_result;
 mod retries;
 mod self_identity;

scylla/tests/integration/query.rs

@@ -0,0 +1,300 @@
+use crate::utils::{create_new_session_builder, setup_tracing, unique_keyspace_name, PerformDDL};
+use scylla::{
+    batch::{Batch, BatchType},
+    client::session::Session,
+    observability::tracing::TracingInfo,
+    query::Query,
+    response::{query_result::QueryResult, PagingState, PagingStateResponse},
+};
+use scylla_cql::Consistency;
+use std::collections::HashSet;
+use std::sync::{Arc, Mutex};
+use uuid::Uuid;
+
+#[tokio::test]
+async fn test_timestamp() {
+    setup_tracing();
+    let session = create_new_session_builder().build().await.unwrap();
+    let ks = unique_keyspace_name();
+
+    session.ddl(format!("CREATE KEYSPACE IF NOT EXISTS {} WITH REPLICATION = {{'class' : 'NetworkTopologyStrategy', 'replication_factor' : 1}}", ks)).await.unwrap();
+    session
+        .ddl(format!(
+            "CREATE TABLE IF NOT EXISTS {}.t_timestamp (a text, b text, primary key (a))",
+            ks
+        ))
+        .await
+        .unwrap();
+
+    session.await_schema_agreement().await.unwrap();
+
+    let query_str = format!("INSERT INTO {}.t_timestamp (a, b) VALUES (?, ?)", ks);
+
+    // test regular query timestamps
+
+    let mut regular_query = Query::new(query_str.to_string());
+
+    regular_query.set_timestamp(Some(420));
+    session
+        .query_unpaged(regular_query.clone(), ("regular query", "higher timestamp"))
+        .await
+        .unwrap();
+
+    regular_query.set_timestamp(Some(42));
+    session
+        .query_unpaged(regular_query.clone(), ("regular query", "lower timestamp"))
+        .await
+        .unwrap();
+
+    // test regular query iter timestamps
+
+    let mut regular_query = Query::new(query_str.to_string());
+
+    regular_query.set_timestamp(Some(420));
+    session
+        .query_iter(
+            regular_query.clone(),
+            ("regular query iter", "higher timestamp"),
+        )
+        .await
+        .unwrap();
+
+    regular_query.set_timestamp(Some(42));
+    session
+        .query_iter(
+            regular_query.clone(),
+            ("regular query iter", "lower timestamp"),
+        )
+        .await
+        .unwrap();
+
+    // test regular query single page timestamps
+
+    let mut regular_query = Query::new(query_str.to_string());
+
+    regular_query.set_timestamp(Some(420));
+    session
+        .query_single_page(
+            regular_query.clone(),
+            ("regular query single page", "higher timestamp"),
+            PagingState::start(),
+        )
+        .await
+        .unwrap();
+
+    regular_query.set_timestamp(Some(42));
+    session
+        .query_single_page(
+            regular_query.clone(),
+            ("regular query single page", "lower timestamp"),
+            PagingState::start(),
+        )
+        .await
+        .unwrap();
+
+    // test prepared statement timestamps
+
+    let mut prepared_statement = session.prepare(query_str.clone()).await.unwrap();
+
+    prepared_statement.set_timestamp(Some(420));
+    session
+        .execute_unpaged(&prepared_statement, ("prepared query", "higher timestamp"))
+        .await
+        .unwrap();
+
+    prepared_statement.set_timestamp(Some(42));
+    session
+        .execute_unpaged(&prepared_statement, ("prepared query", "lower timestamp"))
+        .await
+        .unwrap();
+
+    // test prepared statement iter timestamps
+
+    let mut prepared_statement = session.prepare(query_str.clone()).await.unwrap();
+
+    prepared_statement.set_timestamp(Some(420));
+    session
+        .execute_iter(
+            prepared_statement.clone(),
+            ("prepared query iter", "higher timestamp"),
+        )
+        .await
+        .unwrap();
+
+    prepared_statement.set_timestamp(Some(42));
+    session
+        .execute_iter(
+            prepared_statement,
+            ("prepared query iter", "lower timestamp"),
+        )
+        .await
+        .unwrap();
+
+    // test prepared statement single page timestamps
+
+    let mut prepared_statement = session.prepare(query_str).await.unwrap();
+
+    prepared_statement.set_timestamp(Some(420));
+    session
+        .execute_single_page(
+            &prepared_statement,
+            ("prepared query single page", "higher timestamp"),
+            PagingState::start(),
+        )
+        .await
+        .unwrap();
+
+    prepared_statement.set_timestamp(Some(42));
+    session
+        .execute_single_page(
+            &prepared_statement,
+            ("prepared query single page", "lower timestamp"),
+            PagingState::start(),
+        )
+        .await
+        .unwrap();
+
+    // test batch statement timestamps
+
+    let mut batch: Batch = Default::default();
+    batch.append_statement(regular_query);
+    batch.append_statement(prepared_statement);
+
+    batch.set_timestamp(Some(420));
+    session
+        .batch(
+            &batch,
+            (
+                ("first query in batch", "higher timestamp"),
+                ("second query in batch", "higher timestamp"),
+            ),
+        )
+        .await
+        .unwrap();
+
+    batch.set_timestamp(Some(42));
+    session
+        .batch(
+            &batch,
+            (
+                ("first query in batch", "lower timestamp"),
+                ("second query in batch", "lower timestamp"),
+            ),
+        )
+        .await
+        .unwrap();
+
+    let query_rows_result = session
+        .query_unpaged(
+            format!("SELECT a, b, WRITETIME(b) FROM {}.t_timestamp", ks),
+            &[],
+        )
+        .await
+        .unwrap()
+        .into_rows_result()
+        .unwrap();
+
+    let mut results = query_rows_result
+        .rows::<(&str, &str, i64)>()
+        .unwrap()
+        .map(Result::unwrap)
+        .collect::<Vec<_>>();
+    results.sort();
+
+    let expected_results = [
+        ("first query in batch", "higher timestamp", 420),
+        ("prepared query", "higher timestamp", 420),
+        ("prepared query iter", "higher timestamp", 420),
+        ("prepared query single page", "higher timestamp", 420),
+        ("regular query", "higher timestamp", 420),
+        ("regular query iter", "higher timestamp", 420),
+        ("regular query single page", "higher timestamp", 420),
+        ("second query in batch", "higher timestamp", 420),
+    ]
+    .into_iter()
+    .collect::<Vec<_>>();
+
+    assert_eq!(results, expected_results);
+}
+
+#[tokio::test]
+async fn test_timestamp_generator() {
+    use rand::random;
+    use scylla::policies::timestamp_generator::TimestampGenerator;
+
+    setup_tracing();
+    struct LocalTimestampGenerator {
+        generated_timestamps: Arc<Mutex<HashSet<i64>>>,
+    }
+
+    impl TimestampGenerator for LocalTimestampGenerator {
+        fn next_timestamp(&self) -> i64 {
+            let timestamp = random::<i64>().abs();
+            self.generated_timestamps.lock().unwrap().insert(timestamp);
+            timestamp
+        }
+    }
+
+    let timestamps = Arc::new(Mutex::new(HashSet::new()));
+    let generator = LocalTimestampGenerator {
+        generated_timestamps: timestamps.clone(),
+    };
+
+    let session = create_new_session_builder()
+        .timestamp_generator(Arc::new(generator))
+        .build()
+        .await
+        .unwrap();
+    let ks = unique_keyspace_name();
+    session.ddl(format!("CREATE KEYSPACE IF NOT EXISTS {} WITH REPLICATION = {{'class' : 'NetworkTopologyStrategy', 'replication_factor' : 1}}", ks)).await.unwrap();
+    session
+        .ddl(format!(
+            "CREATE TABLE IF NOT EXISTS {}.t_generator (a int primary key, b int)",
+            ks
+        ))
+        .await
+        .unwrap();
+
+    let prepared = session
+        .prepare(format!(
+            "INSERT INTO {}.t_generator (a, b) VALUES (1, 1)",
+            ks
+        ))
+        .await
+        .unwrap();
+    session.execute_unpaged(&prepared, []).await.unwrap();
+
+    let unprepared = Query::new(format!(
+        "INSERT INTO {}.t_generator (a, b) VALUES (2, 2)",
+        ks
+    ));
+    session.query_unpaged(unprepared, []).await.unwrap();
+
+    let mut batch = Batch::new(BatchType::Unlogged);
+    let stmt = session
+        .prepare(format!(
+            "INSERT INTO {}.t_generator (a, b) VALUES (3, 3)",
+            ks
+        ))
+        .await
+        .unwrap();
+    batch.append_statement(stmt);
+    session.batch(&batch, &((),)).await.unwrap();
+
+    let query_rows_result = session
+        .query_unpaged(
+            format!("SELECT a, b, WRITETIME(b) FROM {}.t_generator", ks),
+            &[],
+        )
+        .await
+        .unwrap()
+        .into_rows_result()
+        .unwrap();
+
+    let timestamps_locked = timestamps.lock().unwrap();
+    assert!(query_rows_result
+        .rows::<(i32, i32, i64)>()
+        .unwrap()
+        .map(|row_result| row_result.unwrap())
+        .all(|(_a, _b, writetime)| timestamps_locked.contains(&writetime)));
+}