ClusterState: Accepts &dyn SerializeRow instead of SerializedValues

There is no easy way for the users to create SerializedValues, which
makes the current APIs cumbersome to use. Instead they should accept
`&dyn SerializeRow` and perform serialization based on table metadata.

This change means that those methods can now also return
SerializationError, and also need to handle a table missing from metadata,
preferably also returning an error in this case.
No existing error type fits here, so either we need to extend some
existing one, or create new one. First idea was extending
PartitionKeyError, but it needs to be convertible to ExecutionError,
in which we don't need those new variants.
For that reason I introduced a new error type for those methods,
called ClusterStateTokenError.

Author: Lorak-mmk

scylla/src/client/session_test.rs

@@ -218,9 +218,6 @@ async fn test_prepared_statement() {
         .unwrap();
 
     let values = (17_i32, 16_i32, "I'm prepared!!!");
-    let serialized_values_complex_pk = prepared_complex_pk_statement
-        .serialize_values(&values)
-        .unwrap();
 
     session
         .execute_unpaged(&prepared_statement, &values)
@@ -245,12 +242,9 @@ async fn test_prepared_statement() {
         let prepared_token = Murmur3Partitioner
             .hash_one(&prepared_statement.compute_partition_key(&values).unwrap());
         assert_eq!(token, prepared_token);
-        let mut pk = SerializedValues::new();
-        pk.add_value(&17_i32, &ColumnType::Native(NativeType::Int))
-            .unwrap();
         let cluster_state_token = session
             .get_cluster_state()
-            .compute_token(&ks, "t2", &pk)
+            .compute_token(&ks, "t2", &(values.0,))
             .unwrap();
         assert_eq!(token, cluster_state_token);
     }
@@ -272,7 +266,7 @@ async fn test_prepared_statement() {
         assert_eq!(token, prepared_token);
         let cluster_state_token = session
             .get_cluster_state()
-            .compute_token(&ks, "complex_pk", &serialized_values_complex_pk)
+            .compute_token(&ks, "complex_pk", &values)
             .unwrap();
         assert_eq!(token, cluster_state_token);
     }
@@ -608,7 +602,6 @@ async fn test_token_calculation() {
             s.push('a');
         }
         let values = (&s,);
-        let serialized_values = prepared_statement.serialize_values(&values).unwrap();
         session
             .execute_unpaged(&prepared_statement, &values)
             .await
@@ -631,7 +624,7 @@ async fn test_token_calculation() {
         assert_eq!(token, prepared_token);
         let cluster_state_token = session
             .get_cluster_state()
-            .compute_token(&ks, "t3", &serialized_values)
+            .compute_token(&ks, "t3", &values)
             .unwrap();
         assert_eq!(token, cluster_state_token);
     }

scylla/src/cluster/state.rs

@@ -1,15 +1,14 @@
-use crate::errors::ConnectionPoolError;
+use crate::errors::{ClusterStateTokenError, ConnectionPoolError};
 use crate::network::{Connection, PoolConfig, VerifiedKeyspaceName};
 use crate::policies::host_filter::HostFilter;
 use crate::routing::locator::tablets::{RawTablet, Tablet, TabletsInfo};
 use crate::routing::locator::ReplicaLocator;
 use crate::routing::partitioner::{calculate_token_for_partition_key, PartitionerName};
 use crate::routing::{Shard, Token};
-use crate::statement::prepared::TokenCalculationError;
 
 use itertools::Itertools;
 use scylla_cql::frame::response::result::TableSpec;
-use scylla_cql::serialize::row::SerializedValues;
+use scylla_cql::serialize::row::{RowSerializationContext, SerializeRow, SerializedValues};
 use std::collections::{HashMap, HashSet};
 use std::sync::Arc;
 use tracing::{debug, warn};
@@ -203,17 +202,29 @@ impl ClusterState {
         &self,
         keyspace: &str,
         table: &str,
-        partition_key: &SerializedValues,
-    ) -> Result<Token, TokenCalculationError> {
-        let partitioner = self
+        partition_key: &dyn SerializeRow,
+    ) -> Result<Token, ClusterStateTokenError> {
+        let Some(table) = self
             .keyspaces
             .get(keyspace)
             .and_then(|k| k.tables.get(table))
-            .and_then(|t| t.partitioner.as_deref())
+        else {
+            return Err(ClusterStateTokenError::UnknownTable {
+                keyspace: keyspace.to_owned(),
+                table: table.to_owned(),
+            });
+        };
+        let values = SerializedValues::from_serializable(
+            &RowSerializationContext::from_specs(table.pk_column_specs.as_slice()),
+            partition_key,
+        )?;
+        let partitioner = table
+            .partitioner
+            .as_deref()
             .and_then(PartitionerName::from_str)
             .unwrap_or_default();
-
-        calculate_token_for_partition_key(partition_key, &partitioner)
+        calculate_token_for_partition_key(&values, &partitioner)
+            .map_err(ClusterStateTokenError::TokenCalculation)
     }
 
     /// Access to replicas owning a given token
@@ -250,8 +261,8 @@ impl ClusterState {
         &self,
         keyspace: &str,
         table: &str,
-        partition_key: &SerializedValues,
-    ) -> Result<Vec<(Arc<Node>, Shard)>, TokenCalculationError> {
+        partition_key: &dyn SerializeRow,
+    ) -> Result<Vec<(Arc<Node>, Shard)>, ClusterStateTokenError> {
         let token = self.compute_token(keyspace, table, partition_key)?;
         Ok(self.get_token_endpoints(keyspace, table, token))
     }

scylla/src/errors.rs

@@ -12,6 +12,7 @@ use crate::frame::response;
 // Re-export error types from pager module.
 pub use crate::client::pager::{NextPageError, NextRowError};
 
+use crate::statement::prepared::TokenCalculationError;
 // Re-export error types from query_result module.
 pub use crate::response::query_result::{
     FirstRowError, IntoRowsResultError, MaybeFirstRowError, ResultNotRowsError, RowsError,
@@ -934,6 +935,23 @@ pub(crate) enum ResponseParseError {
     CqlResponseParseError(#[from] CqlResponseParseError),
 }
 
+/// Error returned from [ClusterState](crate::cluster::ClusterState) APIs.
+#[derive(Clone, Debug, Error)]
+#[non_exhaustive]
+pub enum ClusterStateTokenError {
+    /// Failed to calculate token.
+    #[error(transparent)]
+    TokenCalculation(#[from] TokenCalculationError),
+
+    /// Failed to serialize values required to compute partition key.
+    #[error(transparent)]
+    Serialization(#[from] SerializationError),
+
+    /// ClusterState doesn't currently have metadata for the requested table.
+    #[error("Can't find metadata for requested table ({keyspace}.{table}).")]
+    UnknownTable { keyspace: String, table: String },
+}
+
 #[cfg(test)]
 mod tests {
     use scylla_cql::Consistency;