errors: clean up errors related to pk and token extraction

I marked `PartitionKeyError` as non_exhaustive. Narrowed the return type
of remaining funtions that returned `QueryError` to `PartitionKeyError`.
Added an explicit conversion function PartitionKeyError -> QueryError.

Author: muzarski

scylla/src/client/pager.rs

@@ -803,7 +803,9 @@ impl QueryPager {
                 ) {
                 Ok(res) => res.unzip(),
                 Err(err) => {
-                    let (proof, _res) = ProvingSender::from(sender).send(Err(err)).await;
+                    let (proof, _res) = ProvingSender::from(sender)
+                        .send(Err(err.into_query_error()))
+                        .await;
                     return proof;
                 }
             };

scylla/src/client/session.rs

@@ -31,7 +31,7 @@ use crate::policies::host_filter::HostFilter;
 use crate::policies::load_balancing::{self, RoutingInfo};
 use crate::policies::retry::{RequestInfo, RetryDecision, RetrySession};
 use crate::policies::speculative_execution;
-use crate::prepared_statement::PreparedStatement;
+use crate::prepared_statement::{PartitionKeyError, PreparedStatement};
 use crate::query::Query;
 #[allow(deprecated)]
 use crate::response::legacy_query_result::LegacyQueryResult;
@@ -1394,7 +1394,8 @@ where
         let paging_state_ref = &paging_state;
 
         let (partition_key, token) = prepared
-            .extract_partition_key_and_calculate_token(prepared.get_partitioner_name(), values_ref)?
+            .extract_partition_key_and_calculate_token(prepared.get_partitioner_name(), values_ref)
+            .map_err(PartitionKeyError::into_query_error)?
             .unzip();
 
         let execution_profile = prepared

scylla/src/statement/batch.rs

@@ -217,6 +217,7 @@ pub(crate) mod batch_values {
     use scylla_cql::types::serialize::{RowWriter, SerializationError};
 
     use crate::errors::QueryError;
+    use crate::prepared_statement::PartitionKeyError;
     use crate::routing::Token;
 
     use super::BatchStatement;
@@ -247,7 +248,9 @@ pub(crate) mod batch_values {
                         .map(|o| o.is_some())
                 })?;
                 if did_write {
-                    let token = ps.calculate_token_untyped(&first_values)?;
+                    let token = ps
+                        .calculate_token_untyped(&first_values)
+                        .map_err(PartitionKeyError::into_query_error)?;
                     (token, Some(first_values))
                 } else {
                     (None, None)

scylla/src/statement/prepared_statement.rs

@@ -232,25 +232,13 @@ impl PreparedStatement {
         &'ps self,
         partitioner_name: &'ps PartitionerName,
         serialized_values: &'ps SerializedValues,
-    ) -> Result<Option<(PartitionKey<'ps>, Token)>, QueryError> {
+    ) -> Result<Option<(PartitionKey<'ps>, Token)>, PartitionKeyError> {
         if !self.is_token_aware() {
             return Ok(None);
         }
 
-        let partition_key =
-            self.extract_partition_key(serialized_values)
-                .map_err(|err| match err {
-                    PartitionKeyExtractionError::NoPkIndexValue(_, _) => {
-                        ProtocolError::PartitionKeyExtraction
-                    }
-                })?;
-        let token = partition_key
-            .calculate_token(partitioner_name)
-            .map_err(|err| match err {
-                TokenCalculationError::ValueTooLong(values_len) => {
-                    QueryError::BadQuery(BadQuery::ValuesTooLongForKey(values_len, u16::MAX.into()))
-                }
-            })?;
+        let partition_key = self.extract_partition_key(serialized_values)?;
+        let token = partition_key.calculate_token(partitioner_name)?;
 
         Ok(Some((partition_key, token)))
     }
@@ -262,7 +250,10 @@ impl PreparedStatement {
     // As this function creates a `PartitionKey`, it is intended rather for external usage (by users).
     // For internal purposes, `PartitionKey::calculate_token()` is preferred, as `PartitionKey`
     // is either way used internally, among others for display in traces.
-    pub fn calculate_token(&self, values: &impl SerializeRow) -> Result<Option<Token>, QueryError> {
+    pub fn calculate_token(
+        &self,
+        values: &impl SerializeRow,
+    ) -> Result<Option<Token>, PartitionKeyError> {
         self.calculate_token_untyped(&self.serialize_values(values)?)
     }
 
@@ -271,7 +262,7 @@ impl PreparedStatement {
     pub(crate) fn calculate_token_untyped(
         &self,
         values: &SerializedValues,
-    ) -> Result<Option<Token>, QueryError> {
+    ) -> Result<Option<Token>, PartitionKeyError> {
         self.extract_partition_key_and_calculate_token(&self.partitioner_name, values)
             .map(|opt| opt.map(|(_pk, token)| token))
     }
@@ -485,31 +476,37 @@ pub enum TokenCalculationError {
     ValueTooLong(usize),
 }
 
+/// An error returned by [`PreparedStatement::compute_partition_key()`].
 #[derive(Clone, Debug, Error)]
+#[non_exhaustive]
 pub enum PartitionKeyError {
+    /// Failed to extract partition key.
     #[error(transparent)]
-    PartitionKeyExtraction(PartitionKeyExtractionError),
-    #[error(transparent)]
-    TokenCalculation(TokenCalculationError),
-    #[error(transparent)]
-    Serialization(SerializationError),
-}
+    PartitionKeyExtraction(#[from] PartitionKeyExtractionError),
 
-impl From<PartitionKeyExtractionError> for PartitionKeyError {
-    fn from(err: PartitionKeyExtractionError) -> Self {
-        Self::PartitionKeyExtraction(err)
-    }
-}
+    /// Failed to calculate token.
+    #[error(transparent)]
+    TokenCalculation(#[from] TokenCalculationError),
 
-impl From<TokenCalculationError> for PartitionKeyError {
-    fn from(err: TokenCalculationError) -> Self {
-        Self::TokenCalculation(err)
-    }
+    /// Failed to serialize values required to compute partition key.
+    #[error(transparent)]
+    Serialization(#[from] SerializationError),
 }
 
-impl From<SerializationError> for PartitionKeyError {
-    fn from(err: SerializationError) -> Self {
-        Self::Serialization(err)
+impl PartitionKeyError {
+    /// Converts the error to [`QueryError`].
+    pub fn into_query_error(self) -> QueryError {
+        match self {
+            PartitionKeyError::PartitionKeyExtraction(_) => {
+                QueryError::ProtocolError(ProtocolError::PartitionKeyExtraction)
+            }
+            PartitionKeyError::TokenCalculation(TokenCalculationError::ValueTooLong(
+                values_len,
+            )) => QueryError::BadQuery(BadQuery::ValuesTooLongForKey(values_len, u16::MAX.into())),
+            PartitionKeyError::Serialization(err) => {
+                QueryError::BadQuery(BadQuery::SerializationError(err))
+            }
+        }
     }
 }