Add posibility to validate number of returned rows for read queries

Cargo.toml

@@ -1,7 +1,7 @@
 [package]
 name = "latte-cli"
 description = "A database benchmarking tool for Apache Cassandra and ScyllaDB"
-version = "0.28.4-scylladb"
+version = "0.28.5-scylladb"
 authors = ["Piotr Kołaczkowski <[email protected]>"]
 edition = "2021"
 readme = "README.md"

README.md

@@ -333,6 +333,58 @@ As a result we will be able to get multi-row partitions in a requested size prop
 
 Number of presets is unlimited. Any rune script may use multiple different presets for different tables.
 
+### Validating number of rows for select queries
+
+It is possible to validate number of rows.
+
+It can be done using 2 methods.
+
+First is `execute_with_validation` which is extended version of the common `execute` function.
+Second is `execute_prepared_with_validation` which is extended version of the `execute_prepared` respectively.
+
+Each of these methods expects, as a last parameter, a vector of `Rune Value`s.
+Both support following combinations:
+- [Integer] -> Exact number of expected rows
+- [Integer, Integer] -> Range of expected rows, both values are inclusive.
+- [Integer, String] -> Exact number of expected rows and custom error message.
+- [Integer, Integer, String] -> Range of expected rows and custom error message.
+
+Example using three validation elements:
+```
+  pub async fn some_select_rune_function(db, i) {
+    ...
+    let elapsed = db.elapsed_secs();
+    let rows_min = if elapsed > 100.0 { 0 } else { 1 };
+    let rows_max = if elapsed < 150.0 { 1 } else { 0 };
+    let custom_err = "rows must have been deleted by TTL after 100s-200s";
+    db.execute_prepared_with_validation(
+        PREPARED_STATEMENT_NAME,
+        [pk],
+        [rows_min, rows_max, custom_err],
+    ).await?
+  }
+```
+
+Example using just one validation element (expected strict number of rows):
+```
+  pub async fn prepare(db) {
+    db.init_partition_row_distribution_preset(
+      "main", ROW_COUNT, ROWS_PER_PARTITION, PARTITION_SIZES).await?;
+    ...
+  }
+
+  pub async fn some_select_rune_function(db, i) {
+    let idx = i % ROW_COUNT + OFFSET;
+    let partition = db.get_partition_info("main", idx).await;
+    partition.idx += OFFSET;
+    db.execute_prepared_with_validation(
+      PREPARED_STATEMENT_NAME,
+      [pk],
+      [partition.rows_num], // precise matching to calculated partition rows number
+    ).await?
+  }
+```
+
 ### Mixing workloads
 
 It is possible to run more than one workload function at the same time.

src/config.rs

@@ -220,6 +220,20 @@ pub struct ConnectionConf {
         value_name = "MIN[,MAX]"
     )]
     pub retry_interval: RetryInterval,
+
+    /// Defines the strategy for 'select' queries validation errors.
+    /// Gets applied when 'execute_prepared_with_validation'
+    /// or 'execute_with_validation' context methods are used in rune scripts.
+    #[clap(long("validation-strategy"), required = false, default_value = "retry")]
+    pub validation_strategy: ValidationStrategy,
+}
+
+#[derive(Clone, Copy, Default, Debug, Eq, PartialEq, Serialize, Deserialize, ValueEnum)]
+pub enum ValidationStrategy {
+    #[default]
+    Retry, // Retry 'select' queries if rows number is unexpected.
+    FailFast, // Stop stress execution right after any 'select' query validation fails.
+    Ignore,   // Ignore validation errors - face, print, go on.
 }
 
 #[derive(Clone, Copy, Default, Debug, Eq, PartialEq, Serialize, Deserialize)]

src/scripting/cass_error.rs

@@ -3,7 +3,7 @@ use rune::alloc::fmt::TryWrite;
 use rune::runtime::{TypeInfo, VmResult};
 use rune::{vm_write, Any};
 use scylla::errors::{ExecutionError, NewSessionError, PrepareError};
-use scylla::response::query_result::IntoRowsResultError;
+use scylla::response::query_result::{FirstRowError, IntoRowsResultError};
 use scylla::value::{CqlValue, ValueOverflow};
 use std::fmt::{Display, Formatter};
 
@@ -27,6 +27,40 @@ impl CassError {
         CassError(kind)
     }
 
+    pub fn query_validation_error(
+        cql: &str,
+        params: &[CqlValue],
+        expected_rows_num_min: u64,
+        expected_rows_num_max: u64,
+        actual_rows_num: u64,
+        custom_err_msg: String,
+    ) -> CassError {
+        let query = QueryInfo {
+            cql: cql.to_string(),
+            params: params.iter().map(cql_value_obj_to_string).collect(),
+        };
+        CassError(CassErrorKind::QueryResponseValidationError(
+            query,
+            expected_rows_num_min,
+            expected_rows_num_max,
+            actual_rows_num,
+            custom_err_msg,
+        ))
+    }
+
+    pub fn query_response_validation_not_applicable_error(
+        cql: &str,
+        params: &[CqlValue],
+    ) -> CassError {
+        let query = QueryInfo {
+            cql: cql.to_string(),
+            params: params.iter().map(cql_value_obj_to_string).collect(),
+        };
+        CassError(CassErrorKind::QueryResponseValidationNotApplicableError(
+            query,
+        ))
+    }
+
     pub fn query_retries_exceeded(retry_number: u64) -> CassError {
         CassError(CassErrorKind::QueryRetriesExceeded(format!(
             "Max retry attempts ({}) reached",
@@ -57,7 +91,11 @@ pub enum CassErrorKind {
     InvalidQueryParamsObject(TypeInfo),
     Prepare(String, PrepareError),
     Overloaded(QueryInfo, ExecutionError),
+
     QueryExecution(QueryInfo, ExecutionError),
+    QueryResponseValidationError(QueryInfo, u64, u64, u64, String),
+    QueryResponseValidationNotApplicableError(QueryInfo),
+
     Error(String),
 }
 
@@ -116,6 +154,28 @@ impl CassError {
             CassErrorKind::QueryExecution(q, e) => {
                 write!(buf, "Failed to execute query {q}: {e}")
             }
+            CassErrorKind::QueryResponseValidationError(q, emin, emax, a, err) => {
+                let custom_err = if !err.is_empty() {
+                    format!(" . Custom error msg: {err}", err = err)
+                } else {
+                    "".to_string()
+                };
+                let expected = if emin == emax {
+                    format!("'{emin}' rows")
+                } else {
+                    format!("'{emin}<=N<={emax}' rows")
+                };
+                write!(
+                    buf,
+                    "Expected {expected} in the response, but got '{a}'. Query: {q}{custom_err}"
+                )
+            }
+            CassErrorKind::QueryResponseValidationNotApplicableError(q) => {
+                write!(
+                    buf,
+                    "Response rows can be validated only for 'SELECT' queries, Query: {q}"
+                )
+            }
             CassErrorKind::Error(s) => {
                 write!(buf, "Error: {s}")
             }
@@ -143,6 +203,13 @@ impl From<ValueOverflow> for CassError {
     }
 }
 
+// FirstRowError
+impl From<FirstRowError> for CassError {
+    fn from(e: FirstRowError) -> CassError {
+        CassError(CassErrorKind::Error(e.to_string()))
+    }
+}
+
 impl std::error::Error for CassError {}
 
 /// Transforms a CqlValue object to a string dedicated to be part of CassError message

src/scripting/connect.rs

@@ -61,6 +61,7 @@ pub async fn connect(conf: &ConnectionConf) -> Result<Context, CassError> {
         datacenter,
         conf.retry_number,
         conf.retry_interval,
+        conf.validation_strategy,
     ))
 }