From bc1503baa61ac8e69904dc2dac32397f78dde2d2 Mon Sep 17 00:00:00 2001 From: Piotr Dulikowski Date: Fri, 20 Oct 2023 17:10:16 +0200 Subject: [PATCH] frame: adjust serialized values iterator to return RawValue Currently, the SerializedValues' `iter()` method treats both null and unset values as None, and `iter_name_value_pairs()` just assumes that values are never null/unset and panics if they are. Make the interface more correct by adjusting both methods to return RawValue. The iterators will be used in the next commit to implement the fallback that allows to implement `SerializeRow`/`SerializeCql` via legacy `ValueList`/`Value` traits. --- scylla-cql/src/frame/value.rs | 11 ++++--- scylla-cql/src/frame/value_tests.rs | 37 ++++++++++++++-------- scylla/src/statement/prepared_statement.rs | 3 +- scylla/src/transport/partitioner.rs | 8 +++-- 4 files changed, 37 insertions(+), 22 deletions(-) diff --git a/scylla-cql/src/frame/value.rs b/scylla-cql/src/frame/value.rs index 17b75ea855..aeac47d4f0 100644 --- a/scylla-cql/src/frame/value.rs +++ b/scylla-cql/src/frame/value.rs @@ -15,6 +15,7 @@ use uuid::Uuid; use super::response::result::CqlValue; use super::types::vint_encode; +use super::types::RawValue; #[cfg(feature = "secret")] use secrecy::{ExposeSecret, Secret, Zeroize}; @@ -176,7 +177,7 @@ impl SerializedValues { Ok(()) } - pub fn iter(&self) -> impl Iterator> { + pub fn iter(&self) -> impl Iterator { SerializedValuesIterator { serialized_values: &self.serialized_values, contains_names: self.contains_names, @@ -220,7 +221,7 @@ impl SerializedValues { }) } - pub fn iter_name_value_pairs(&self) -> impl Iterator, &[u8])> { + pub fn iter_name_value_pairs(&self) -> impl Iterator, RawValue)> { let mut buf = &self.serialized_values[..]; (0..self.values_num).map(move |_| { // `unwrap()`s here are safe, as we assume type-safety: if `SerializedValues` exits, @@ -228,7 +229,7 @@ impl SerializedValues { let name = self .contains_names .then(|| types::read_string(&mut buf).unwrap()); - let serialized = types::read_bytes(&mut buf).unwrap(); + let serialized = types::read_value(&mut buf).unwrap(); (name, serialized) }) } @@ -241,7 +242,7 @@ pub struct SerializedValuesIterator<'a> { } impl<'a> Iterator for SerializedValuesIterator<'a> { - type Item = Option<&'a [u8]>; + type Item = RawValue<'a>; fn next(&mut self) -> Option { if self.serialized_values.is_empty() { @@ -253,7 +254,7 @@ impl<'a> Iterator for SerializedValuesIterator<'a> { types::read_short_bytes(&mut self.serialized_values).expect("badly encoded value name"); } - Some(types::read_bytes_opt(&mut self.serialized_values).expect("badly encoded value")) + Some(types::read_value(&mut self.serialized_values).expect("badly encoded value")) } } diff --git a/scylla-cql/src/frame/value_tests.rs b/scylla-cql/src/frame/value_tests.rs index bd411f45a8..5d60346492 100644 --- a/scylla-cql/src/frame/value_tests.rs +++ b/scylla-cql/src/frame/value_tests.rs @@ -1,4 +1,4 @@ -use crate::frame::value::BatchValuesIterator; +use crate::frame::{types::RawValue, value::BatchValuesIterator}; use super::value::{ BatchValues, Date, MaybeUnset, SerializeValuesError, SerializedValues, Time, Timestamp, Unset, @@ -225,7 +225,10 @@ fn serialized_values() { values.write_to_request(&mut request); assert_eq!(request, vec![0, 1, 0, 0, 0, 1, 8]); - assert_eq!(values.iter().collect::>(), vec![Some([8].as_ref())]); + assert_eq!( + values.iter().collect::>(), + vec![RawValue::Value([8].as_ref())] + ); } // Add second value @@ -240,7 +243,10 @@ fn serialized_values() { assert_eq!( values.iter().collect::>(), - vec![Some([8].as_ref()), Some([0, 16].as_ref())] + vec![ + RawValue::Value([8].as_ref()), + RawValue::Value([0, 16].as_ref()) + ] ); } @@ -272,7 +278,10 @@ fn serialized_values() { assert_eq!( values.iter().collect::>(), - vec![Some([8].as_ref()), Some([0, 16].as_ref())] + vec![ + RawValue::Value([8].as_ref()), + RawValue::Value([0, 16].as_ref()) + ] ); } } @@ -302,9 +311,9 @@ fn slice_value_list() { assert_eq!( serialized.iter().collect::>(), vec![ - Some([0, 0, 0, 1].as_ref()), - Some([0, 0, 0, 2].as_ref()), - Some([0, 0, 0, 3].as_ref()) + RawValue::Value([0, 0, 0, 1].as_ref()), + RawValue::Value([0, 0, 0, 2].as_ref()), + RawValue::Value([0, 0, 0, 3].as_ref()) ] ); } @@ -319,9 +328,9 @@ fn vec_value_list() { assert_eq!( serialized.iter().collect::>(), vec![ - Some([0, 0, 0, 1].as_ref()), - Some([0, 0, 0, 2].as_ref()), - Some([0, 0, 0, 3].as_ref()) + RawValue::Value([0, 0, 0, 1].as_ref()), + RawValue::Value([0, 0, 0, 2].as_ref()), + RawValue::Value([0, 0, 0, 3].as_ref()) ] ); } @@ -334,7 +343,7 @@ fn tuple_value_list() { let serialized_vals: Vec = serialized .iter() - .map(|o: Option<&[u8]>| o.unwrap()[0]) + .map(|o: RawValue| o.as_value().unwrap()[0]) .collect(); let expected: Vec = expected.collect(); @@ -408,9 +417,9 @@ fn ref_value_list() { assert_eq!( serialized.iter().collect::>(), vec![ - Some([0, 0, 0, 1].as_ref()), - Some([0, 0, 0, 2].as_ref()), - Some([0, 0, 0, 3].as_ref()) + RawValue::Value([0, 0, 0, 1].as_ref()), + RawValue::Value([0, 0, 0, 2].as_ref()), + RawValue::Value([0, 0, 0, 3].as_ref()) ] ); } diff --git a/scylla/src/statement/prepared_statement.rs b/scylla/src/statement/prepared_statement.rs index 22f34e60a2..58d8b9ea3d 100644 --- a/scylla/src/statement/prepared_statement.rs +++ b/scylla/src/statement/prepared_statement.rs @@ -1,5 +1,6 @@ use bytes::{Bytes, BytesMut}; use scylla_cql::errors::{BadQuery, QueryError}; +use scylla_cql::frame::types::RawValue; use smallvec::{smallvec, SmallVec}; use std::convert::TryInto; use std::sync::Arc; @@ -399,7 +400,7 @@ impl<'ps> PartitionKey<'ps> { PartitionKeyExtractionError::NoPkIndexValue(pk_index.index, bound_values.len()) })?; // Add it in sequence order to pk_values - if let Some(v) = next_val { + if let RawValue::Value(v) = next_val { let spec = &prepared_metadata.col_specs[pk_index.index as usize]; pk_values[pk_index.sequence as usize] = Some((v, spec)); } diff --git a/scylla/src/transport/partitioner.rs b/scylla/src/transport/partitioner.rs index 9c8a542325..4526715ab2 100644 --- a/scylla/src/transport/partitioner.rs +++ b/scylla/src/transport/partitioner.rs @@ -1,4 +1,5 @@ use bytes::Buf; +use scylla_cql::frame::types::RawValue; use std::num::Wrapping; use crate::{ @@ -343,11 +344,14 @@ pub fn calculate_token_for_partition_key( if serialized_partition_key_values.len() == 1 { let val = serialized_partition_key_values.iter().next().unwrap(); - if let Some(val) = val { + if let RawValue::Value(val) = val { partitioner_hasher.write(val); } } else { - for val in serialized_partition_key_values.iter().flatten() { + for val in serialized_partition_key_values + .iter() + .filter_map(|rv| rv.as_value()) + { let val_len_u16: u16 = val .len() .try_into()