Merge pull request #1223 from wprzytula/partitioner-visibility

Partitioner visibility

Author: wprzytula

scylla/benches/benchmark.rs

@@ -1,7 +1,7 @@
 use criterion::{criterion_group, criterion_main, Criterion};
 
 use bytes::BytesMut;
-use scylla::routing::partitioner::{calculate_token_for_partition_key, Murmur3Partitioner};
+use scylla::routing::partitioner::{calculate_token_for_partition_key, PartitionerName};
 use scylla_cql::frame::response::result::{ColumnType, NativeType};
 use scylla_cql::frame::types;
 use scylla_cql::serialize::row::SerializedValues;
@@ -84,7 +84,9 @@ fn calculate_token_bench(c: &mut Criterion) {
         .unwrap();
 
     c.bench_function("calculate_token_from_partition_key simple pk", |b| {
-        b.iter(|| calculate_token_for_partition_key(&serialized_simple_pk, &Murmur3Partitioner))
+        b.iter(|| {
+            calculate_token_for_partition_key(&serialized_simple_pk, &PartitionerName::Murmur3)
+        })
     });
 
     c.bench_function(
@@ -93,14 +95,16 @@ fn calculate_token_bench(c: &mut Criterion) {
             b.iter(|| {
                 calculate_token_for_partition_key(
                     &serialized_simple_pk_long_column,
-                    &Murmur3Partitioner,
+                    &PartitionerName::Murmur3,
                 )
             })
         },
     );
 
     c.bench_function("calculate_token_from_partition_key complex pk", |b| {
-        b.iter(|| calculate_token_for_partition_key(&serialized_complex_pk, &Murmur3Partitioner))
+        b.iter(|| {
+            calculate_token_for_partition_key(&serialized_complex_pk, &PartitionerName::Murmur3)
+        })
     });
 
     c.bench_function(
@@ -109,7 +113,7 @@ fn calculate_token_bench(c: &mut Criterion) {
             b.iter(|| {
                 calculate_token_for_partition_key(
                     &serialized_values_long_column,
-                    &Murmur3Partitioner,
+                    &PartitionerName::Murmur3,
                 )
             })
         },

scylla/src/client/pager.rs

@@ -555,6 +555,7 @@ where
 /// needs to be cast into a typed stream. This is done by use of `rows_stream()` method.
 /// As the method is generic over the target type, the turbofish syntax
 /// can come in handy there, e.g. `query_pager.rows_stream::<(i32, String, Uuid)>()`.
+#[derive(Debug)]
 pub struct QueryPager {
     current_page: RawRowLendingIterator,
     page_receiver: mpsc::Receiver<Result<ReceivedPage, NextPageError>>,
@@ -640,7 +641,7 @@ impl QueryPager {
     /// Type-checks the iterator against given type.
     ///
     /// This is automatically called upon transforming [QueryPager] into [TypedRowStream].
-    /// Can be used with `next()` for manual deserialization. See `next()` for an example.
+    // Can be used with `next()` for manual deserialization.
     #[inline]
     pub fn type_check<'frame, 'metadata, RowT: DeserializeRow<'frame, 'metadata>>(
         &self,
@@ -972,6 +973,17 @@ pub struct TypedRowStream<RowT: 'static> {
     _phantom: std::marker::PhantomData<RowT>,
 }
 
+// Manual implementation not to depend on RowT implementing Debug.
+// Explanation: automatic derive of Debug would impose the RowT: Debug
+// constaint for the Debug impl.
+impl<T> std::fmt::Debug for TypedRowStream<T> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("TypedRowStream")
+            .field("raw_row_lending_stream", &self.raw_row_lending_stream)
+            .finish()
+    }
+}
+
 impl<RowT> Unpin for TypedRowStream<RowT> {}
 
 impl<RowT> TypedRowStream<RowT>

scylla/src/client/session_test.rs

@@ -3044,7 +3044,7 @@ async fn test_manual_primary_key_computation() {
 
         let token_by_hand = calculate_token_for_partition_key(
             serialized_pk_values_in_pk_order,
-            &Murmur3Partitioner,
+            &PartitionerName::Murmur3,
         )
         .unwrap();
         println!(

scylla/src/routing/partitioner.rs

@@ -16,8 +16,9 @@ use std::num::Wrapping;
 use crate::{prepared_statement::TokenCalculationError, routing::Token};
 
 #[allow(clippy::upper_case_acronyms)]
-#[derive(Clone, PartialEq, Debug, Default)]
-pub(crate) enum PartitionerName {
+#[derive(Clone, PartialEq, Eq, Debug, Default)]
+#[non_exhaustive]
+pub enum PartitionerName {
     #[default]
     Murmur3,
     CDC,
@@ -75,11 +76,12 @@ impl PartitionerHasher for PartitionerHasherAny {
 /// The Partitioners' design is based on std::hash design: `Partitioner`
 /// corresponds to `HasherBuilder`, and `PartitionerHasher` to `Hasher`.
 /// See their documentation for more details.
-pub trait Partitioner {
+pub(crate) trait Partitioner {
     type Hasher: PartitionerHasher;
 
     fn build_hasher(&self) -> Self::Hasher;
 
+    #[allow(unused)] // Currently, no public API uses this.
     fn hash_one(&self, data: &[u8]) -> Token {
         let mut hasher = self.build_hasher();
         hasher.write(data);
@@ -92,12 +94,12 @@ pub trait Partitioner {
 /// Instances of this trait are created by a `Partitioner` and are stateful.
 /// At any point, one can call `finish()` and a `Token` will be computed
 /// based on values that has been fed so far.
-pub trait PartitionerHasher {
+pub(crate) trait PartitionerHasher {
     fn write(&mut self, pk_part: &[u8]);
     fn finish(&self) -> Token;
 }
 
-pub struct Murmur3Partitioner;
+pub(crate) struct Murmur3Partitioner;
 
 impl Partitioner for Murmur3Partitioner {
     type Hasher = Murmur3PartitionerHasher;
@@ -112,7 +114,7 @@ impl Partitioner for Murmur3Partitioner {
     }
 }
 
-pub struct Murmur3PartitionerHasher {
+pub(crate) struct Murmur3PartitionerHasher {
     total_len: usize,
     buf: [u8; Self::BUF_CAPACITY],
     h1: Wrapping<i64>,
@@ -278,9 +280,9 @@ enum CDCPartitionerHasherState {
     Computed(Token),
 }
 
-pub struct CDCPartitioner;
+pub(crate) struct CDCPartitioner;
 
-pub struct CDCPartitionerHasher {
+pub(crate) struct CDCPartitionerHasher {
     state: CDCPartitionerHasherState,
 }
 
@@ -346,9 +348,9 @@ impl PartitionerHasher for CDCPartitionerHasher {
 ///
 /// NOTE: the provided values must completely constitute partition key
 /// and be in the order defined in CREATE TABLE statement.
-pub fn calculate_token_for_partition_key<P: Partitioner>(
+pub fn calculate_token_for_partition_key(
     serialized_partition_key_values: &SerializedValues,
-    partitioner: &P,
+    partitioner: &PartitionerName,
 ) -> Result<Token, TokenCalculationError> {
     let mut partitioner_hasher = partitioner.build_hasher();
 

scylla/src/statement/prepared_statement.rs

@@ -418,7 +418,7 @@ impl PreparedStatement {
     }
 
     /// Get the name of the partitioner used for this statement.
-    pub(crate) fn get_partitioner_name(&self) -> &PartitionerName {
+    pub fn get_partitioner_name(&self) -> &PartitionerName {
         &self.partitioner_name
     }