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.

Worth noting that those methods can now return serialization error, so
the error type had to be change to PartitionKeyError because does not
have a variant with SerializationError.

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,15 @@
 use crate::errors::ConnectionPoolError;
 use crate::network::{Connection, PoolConfig, VerifiedKeyspaceName};
 use crate::policies::host_filter::HostFilter;
-use crate::prepared_statement::TokenCalculationError;
+use crate::prepared_statement::PartitionKeyError;
 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 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 +203,22 @@ impl ClusterState {
         &self,
         keyspace: &str,
         table: &str,
-        partition_key: &SerializedValues,
-    ) -> Result<Token, TokenCalculationError> {
-        let partitioner = self
+        partition_key: &dyn SerializeRow,
+    ) -> Result<Token, PartitionKeyError> {
+        let table = self
             .keyspaces
             .get(keyspace)
-            .and_then(|k| k.tables.get(table))
+            .and_then(|k| k.tables.get(table));
+        let values = SerializedValues::from_serializable(
+            &RowSerializationContext::from_specs(table.unwrap().pk_column_specs.as_slice()),
+            partition_key,
+        )?;
+        let partitioner = table
             .and_then(|t| t.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(PartitionKeyError::from)
     }
 
     /// Access to replicas owning a given token
@@ -250,8 +255,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)>, PartitionKeyError> {
         let token = self.compute_token(keyspace, table, partition_key)?;
         Ok(self.get_token_endpoints(keyspace, table, token))
     }