From 506d401065d682ae1394d3f41eef8dcf01dab796 Mon Sep 17 00:00:00 2001 From: Piotr Dulikowski <piodul@scylladb.com> 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 db67b4fab8..4faa8df501 100644 --- a/scylla-cql/src/frame/value.rs +++ b/scylla-cql/src/frame/value.rs @@ -16,6 +16,7 @@ use chrono::{DateTime, NaiveDate, NaiveTime, TimeZone, Utc}; use super::response::result::CqlValue; use super::types::vint_encode; +use super::types::RawValue; #[cfg(feature = "secret")] use secrecy::{ExposeSecret, Secret, Zeroize}; @@ -366,7 +367,7 @@ impl SerializedValues { Ok(()) } - pub fn iter(&self) -> impl Iterator<Item = Option<&[u8]>> { + pub fn iter(&self) -> impl Iterator<Item = RawValue> { SerializedValuesIterator { serialized_values: &self.serialized_values, contains_names: self.contains_names, @@ -410,7 +411,7 @@ impl SerializedValues { }) } - pub fn iter_name_value_pairs(&self) -> impl Iterator<Item = (Option<&str>, &[u8])> { + pub fn iter_name_value_pairs(&self) -> impl Iterator<Item = (Option<&str>, 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, @@ -418,7 +419,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) }) } @@ -431,7 +432,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<Self::Item> { if self.serialized_values.is_empty() { @@ -443,7 +444,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 88de8dcbed..0992297a64 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, CqlDate, CqlTime, CqlTimestamp, MaybeUnset, SerializeValuesError, @@ -421,7 +421,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<_>>(), vec![Some([8].as_ref())]); + assert_eq!( + values.iter().collect::<Vec<_>>(), + vec![RawValue::Value([8].as_ref())] + ); } // Add second value @@ -436,7 +439,10 @@ fn serialized_values() { assert_eq!( values.iter().collect::<Vec<_>>(), - vec![Some([8].as_ref()), Some([0, 16].as_ref())] + vec![ + RawValue::Value([8].as_ref()), + RawValue::Value([0, 16].as_ref()) + ] ); } @@ -468,7 +474,10 @@ fn serialized_values() { assert_eq!( values.iter().collect::<Vec<_>>(), - vec![Some([8].as_ref()), Some([0, 16].as_ref())] + vec![ + RawValue::Value([8].as_ref()), + RawValue::Value([0, 16].as_ref()) + ] ); } } @@ -498,9 +507,9 @@ fn slice_value_list() { assert_eq!( serialized.iter().collect::<Vec<_>>(), 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()) ] ); } @@ -515,9 +524,9 @@ fn vec_value_list() { assert_eq!( serialized.iter().collect::<Vec<_>>(), 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()) ] ); } @@ -530,7 +539,7 @@ fn tuple_value_list() { let serialized_vals: Vec<u8> = serialized .iter() - .map(|o: Option<&[u8]>| o.unwrap()[0]) + .map(|o: RawValue| o.as_value().unwrap()[0]) .collect(); let expected: Vec<u8> = expected.collect(); @@ -604,9 +613,9 @@ fn ref_value_list() { assert_eq!( serialized.iter().collect::<Vec<_>>(), 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<P: Partitioner>( 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()