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"

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,
     ))
 }