Repository: scylladb/scylla-rust-driver
Pull request #389 (2022-01-27 19:27:13)

Add CDC partitioner support
Fixes: #194 

The following PR refactors partitioner logic and implements CDC Partitioner.
## Pre-review checklist

<!--
    Make sure you took care of the issues on the list.
    Put 'x' into those boxes which apply.
    You can also create the PR now and click on all relevant checkboxes.
    See CONTRIBUTING.md for more details.
-->

- [x] I have split my patch into logically separate commits.
- [x] All commit messages clearly explain what they change and why.
- [x] I added relevant tests for new features and bug fixes.
- [x] All commits compile, pass static checks and pass test.
- [x] PR description sums up the changes and reasons why they should be introduced.
- [x] I added appropriate `Fixes:` annotations to PR description.


Ponewor: > It's quite sad that this implementation performs a lookup in the metadata every time calculate_token is called, which is also every time a prepared statement is executed. Per-table partitioners will not change during the lifetime of a table, so please precompute this information when the statement is being prepared.

I am not sure how this could be achieved. `PreparedStatement`s are generated in `Connection::prepare`. `Connection` has no idea about cluster's metadata so there is no direct way to access the partitioner name. Perhaps it could passed top-down but it would require to break a few public interfaces.
piodul: > > It's quite sad that this implementation performs a lookup in the metadata every time calculate_token is called, which is also every time a prepared statement is executed. Per-table partitioners will not change during the lifetime of a table, so please precompute this information when the statement is being prepared.
> 
> I am not sure how this could be achieved. `PreparedStatement`s are generated in `Connection::prepare`. `Connection` has no idea about cluster's metadata so there is no direct way to access the partitioner name. Perhaps it could passed top-down but it would require to break a few public interfaces.

I think it will be sufficient if we do it in `Session::prepare`. There, you should have all the necessary information and can set the partitioner field before returning `PreparedStatement`.

The `Connection::prepare` can keep returning `PreparedStatement` without the partitioner set. Users are supposed to mainly use `Session` to interact with the cluster, and the partitioner is just an optimization. I think that in the future we will hide the `Connection` from the public interface anyway, and if not, let's not worry about setting partitioner there now.
Ponewor: > > Please also add a test which verifies that a prepared statement has the right partitioner set for CDC/non-CDC table.
> 
> The implementation looks good, but the test is still missing.

Right, it somehow has slipped my attention.

While trying to complete this task I have found some unexpected difficulties - the `prepared_metadata` in `PreparedStatement` is always empty (`PreparedMetadata { col_count: 0, pk_indexes: [], col_specs: [] }`). This results in `get_keyspace_name` and `get_table_name` returning `None`.

Is this a bug or is this metadata fetched later? If the latter then when does this happen? How to trigger it?
piodul: > > > Please also add a test which verifies that a prepared statement has the right partitioner set for CDC/non-CDC table.
> > 
> > 
> > The implementation looks good, but the test is still missing.
> 
> Right, it somehow has slipped my attention.
> 
> While trying to complete this task I have found some unexpected difficulties - the `prepared_metadata` in `PreparedStatement` is always empty (`PreparedMetadata { col_count: 0, pk_indexes: [], col_specs: [] }`). This results in `get_keyspace_name` and `get_table_name` returning `None`.
> 
> Is this a bug or is this metadata fetched later? If the latter then when does this happen? How to trigger it?

Could you paste the query string? It should have at least one `?` in order for the metadata to be non-empty.
piodul: > > > > > Please also add a test which verifies that a prepared statement has the right partitioner set for CDC/non-CDC table.
> > > > 
> > > > 
> > > > The implementation looks good, but the test is still missing.
> > > 
> > > 
> > > Right, it somehow has slipped my attention.
> > > While trying to complete this task I have found some unexpected difficulties - the `prepared_metadata` in `PreparedStatement` is always empty (`PreparedMetadata { col_count: 0, pk_indexes: [], col_specs: [] }`). This results in `get_keyspace_name` and `get_table_name` returning `None`.
> > > Is this a bug or is this metadata fetched later? If the latter then when does this happen? How to trigger it?
> > 
> > 
> > Could you paste the query string? It should have at least one `?` in order for the metadata to be non-empty.
> 
> The table is given by `CREATE TABLE IF NOT EXISTS t (a int primary key) WITH cdc = {{'enabled':true}}` and the queries I've tried are `SELECT a FROM t` `INSERT INTO t (a) VALUES (0)`

Right, your queries need bind markers. `PreparedMetadata` contains information about bind markers, so if you don't have any then the metadata will be empty. Try changing your queries to `SELECT a FROM t WHERE a = ?` and `INSERT INTO t (a) VALUES (?)`.
piodul: LGTM. Please rebase onto newest main as merging https://github.com/scylladb/scylla-rust-driver/pull/426 seems to have introduced a conflict in the current PR.

Patch:
diff --git a/examples/compare-tokens.rs b/examples/compare-tokens.rs
index d2f7ad04..efc835fb 100644
--- a/examples/compare-tokens.rs
+++ b/examples/compare-tokens.rs
@@ -1,6 +1,6 @@
 use anyhow::Result;
 use scylla::frame::value::ValueList;
-use scylla::routing::murmur3_token;
+use scylla::transport::partitioner::{Murmur3Partitioner, Partitioner};
 use scylla::{Session, SessionBuilder};
 use std::env;
 
@@ -29,7 +29,7 @@ async fn main() -> Result<()> {
             .await?;
 
         let serialized_pk = (pk,).serialized()?.into_owned();
-        let t = murmur3_token(prepared.compute_partition_key(&serialized_pk)?).value;
+        let t = Murmur3Partitioner::hash(prepared.compute_partition_key(&serialized_pk)?).value;
 
         let qt = session
             .query(format!("SELECT token(pk) FROM ks.t where pk = {}", pk), &[])
diff --git a/scylla/src/routing.rs b/scylla/src/routing.rs
index 9ecdb08f..3eada0ae 100644
--- a/scylla/src/routing.rs
+++ b/scylla/src/routing.rs
@@ -1,9 +1,8 @@
-use bytes::{Buf, Bytes};
 use rand::Rng;
 use std::collections::HashMap;
 use std::convert::TryFrom;
 use std::net::SocketAddr;
-use std::num::{NonZeroU16, Wrapping};
+use std::num::NonZeroU16;
 use thiserror::Error;
 
 #[derive(PartialEq, Eq, PartialOrd, Ord, Hash, Clone, Copy)]
@@ -41,12 +40,6 @@ impl std::str::FromStr for Token {
     }
 }
 
-pub fn murmur3_token(pk: Bytes) -> Token {
-    Token {
-        value: hash3_x64_128(&pk) as i64,
-    }
-}
-
 impl ShardInfo {
     pub fn new(shard: u16, nr_shards: ShardCount, msb_ignore: u8) -> Self {
         ShardInfo {
@@ -148,99 +141,6 @@ impl<'a> TryFrom<&'a HashMap<String, Vec<String>>> for ShardInfo {
     }
 }
 
-// An implementation of MurmurHash3 ported from Scylla. Please note that this
-// is not a "correct" implementation of MurmurHash3 - it replicates the same
-// bugs made in the original Cassandra implementation in order to be compatible.
-pub fn hash3_x64_128(mut data: &[u8]) -> i128 {
-    let length = data.len();
-
-    let c1: Wrapping<i64> = Wrapping(0x87c3_7b91_1142_53d5_u64 as i64);
-    let c2: Wrapping<i64> = Wrapping(0x4cf5_ad43_2745_937f_u64 as i64);
-
-    let mut h1 = Wrapping(0_i64);
-    let mut h2 = Wrapping(0_i64);
-
-    while data.len() >= 16 {
-        let mut k1 = Wrapping(data.get_i64_le());
-        let mut k2 = Wrapping(data.get_i64_le());
-
-        k1 *= c1;
-        k1 = rotl64(k1, 31);
-        k1 *= c2;
-        h1 ^= k1;
-
-        h1 = rotl64(h1, 27);
-        h1 += h2;
-        h1 = h1 * Wrapping(5) + Wrapping(0x52dce729);
-
-        k2 *= c2;
-        k2 = rotl64(k2, 33);
-        k2 *= c1;
-        h2 ^= k2;
-
-        h2 = rotl64(h2, 31);
-        h2 += h1;
-        h2 = h2 * Wrapping(5) + Wrapping(0x38495ab5);
-    }
-
-    let mut k1 = Wrapping(0_i64);
-    let mut k2 = Wrapping(0_i64);
-
-    debug_assert!(data.len() < 16);
-
-    if data.len() > 8 {
-        for i in (8..data.len()).rev() {
-            k2 ^= Wrapping(data[i] as i8 as i64) << ((i - 8) * 8);
-        }
-
-        k2 *= c2;
-        k2 = rotl64(k2, 33);
-        k2 *= c1;
-        h2 ^= k2;
-    }
-
-    if !data.is_empty() {
-        for i in (0..std::cmp::min(8, data.len())).rev() {
-            k1 ^= Wrapping(data[i] as i8 as i64) << (i * 8);
-        }
-
-        k1 *= c1;
-        k1 = rotl64(k1, 31);
-        k1 *= c2;
-        h1 ^= k1;
-    }
-
-    h1 ^= Wrapping(length as i64);
-    h2 ^= Wrapping(length as i64);
-
-    h1 += h2;
-    h2 += h1;
-
-    h1 = fmix(h1);
-    h2 = fmix(h2);
-
-    h1 += h2;
-    h2 += h1;
-
-    ((h2.0 as i128) << 64) | h1.0 as i128
-}
-
-#[inline]
-fn rotl64(v: Wrapping<i64>, n: u32) -> Wrapping<i64> {
-    Wrapping((v.0 << n) | (v.0 as u64 >> (64 - n)) as i64)
-}
-
-#[inline]
-fn fmix(mut k: Wrapping<i64>) -> Wrapping<i64> {
-    k ^= Wrapping((k.0 as u64 >> 33) as i64);
-    k *= Wrapping(0xff51afd7ed558ccd_u64 as i64);
-    k ^= Wrapping((k.0 as u64 >> 33) as i64);
-    k *= Wrapping(0xc4ceb9fe1a85ec53_u64 as i64);
-    k ^= Wrapping((k.0 as u64 >> 33) as i64);
-
-    k
-}
-
 #[cfg(test)]
 mod tests {
     use super::Token;
diff --git a/scylla/src/statement/prepared_statement.rs b/scylla/src/statement/prepared_statement.rs
index 152d315c..5654c925 100644
--- a/scylla/src/statement/prepared_statement.rs
+++ b/scylla/src/statement/prepared_statement.rs
@@ -7,6 +7,7 @@ use super::StatementConfig;
 use crate::frame::response::result::PreparedMetadata;
 use crate::frame::types::{Consistency, SerialConsistency};
 use crate::frame::value::SerializedValues;
+use crate::transport::partitioner::PartitionerName;
 use crate::transport::retry_policy::RetryPolicy;
 
 /// Represents a statement prepared on the server.
@@ -19,6 +20,7 @@ pub struct PreparedStatement {
     metadata: PreparedMetadata,
     statement: String,
     page_size: Option<i32>,
+    partitioner_name: PartitionerName,
 }
 
 impl PreparedStatement {
@@ -36,6 +38,7 @@ impl PreparedStatement {
             prepare_tracing_ids: Vec::new(),
             page_size,
             config,
+            partitioner_name: Default::default(),
         }
     }
 
@@ -131,6 +134,14 @@ impl PreparedStatement {
             .map(|col_spec| col_spec.table_spec.ks_name.as_str())
     }
 
+    /// Returns the name of the table this statement is operating on.
+    pub fn get_table_name(&self) -> Option<&str> {
+        self.metadata
+            .col_specs
+            .first()
+            .map(|col_spec| col_spec.table_spec.table_name.as_str())
+    }
+
     /// Sets the consistency to be used when executing this statement.
     pub fn set_consistency(&mut self, c: Consistency) {
         self.config.consistency = Some(c);
@@ -203,6 +214,21 @@ impl PreparedStatement {
     pub fn get_timestamp(&self) -> Option<i64> {
         self.config.timestamp
     }
+
+    /// Sets the name of the partitioner used for this statement.
+    pub(crate) fn set_partitioner_name(&mut self, partitioner_name: Option<&str>) {
+        self.partitioner_name = match partitioner_name {
+            Some(partitioner_name) if partitioner_name.ends_with("CDCPartitioner") => {
+                PartitionerName::CDC
+            }
+            _ => PartitionerName::Murmur3,
+        }
+    }
+
+    /// Get the name of the partitioner used for this statement.
+    pub(crate) fn get_partitioner_name(&self) -> &PartitionerName {
+        &self.partitioner_name
+    }
 }
 
 #[derive(Debug, Error, PartialEq, Eq, PartialOrd, Ord)]
diff --git a/scylla/src/transport/mod.rs b/scylla/src/transport/mod.rs
index 909a4cd2..38847585 100644
--- a/scylla/src/transport/mod.rs
+++ b/scylla/src/transport/mod.rs
@@ -7,6 +7,7 @@ pub mod iterator;
 pub mod load_balancing;
 pub(crate) mod metrics;
 mod node;
+pub mod partitioner;
 pub mod query_result;
 pub mod retry_policy;
 pub mod session;
diff --git a/scylla/src/transport/partitioner.rs b/scylla/src/transport/partitioner.rs
new file mode 100644
index 00000000..d875cb5f
--- /dev/null
+++ b/scylla/src/transport/partitioner.rs
@@ -0,0 +1,176 @@
+use bytes::{Buf, Bytes};
+use std::num::Wrapping;
+
+use crate::routing::Token;
+
+#[allow(clippy::upper_case_acronyms)]
+#[derive(Clone, PartialEq, Debug)]
+pub(crate) enum PartitionerName {
+    Murmur3,
+    CDC,
+}
+
+impl Default for PartitionerName {
+    fn default() -> Self {
+        PartitionerName::Murmur3
+    }
+}
+
+pub trait Partitioner {
+    fn hash(pk: Bytes) -> Token;
+}
+
+pub struct Murmur3Partitioner;
+pub struct CDCPartitioner;
+
+impl Murmur3Partitioner {
+    // An implementation of MurmurHash3 ported from Scylla. Please note that this
+    // is not a "correct" implementation of MurmurHash3 - it replicates the same
+    // bugs made in the original Cassandra implementation in order to be compatible.
+    fn hash3_x64_128(mut data: &[u8]) -> i128 {
+        let length = data.len();
+
+        let c1: Wrapping<i64> = Wrapping(0x87c3_7b91_1142_53d5_u64 as i64);
+        let c2: Wrapping<i64> = Wrapping(0x4cf5_ad43_2745_937f_u64 as i64);
+
+        let mut h1 = Wrapping(0_i64);
+        let mut h2 = Wrapping(0_i64);
+
+        while data.len() >= 16 {
+            let mut k1 = Wrapping(data.get_i64_le());
+            let mut k2 = Wrapping(data.get_i64_le());
+
+            k1 *= c1;
+            k1 = Self::rotl64(k1, 31);
+            k1 *= c2;
+            h1 ^= k1;
+
+            h1 = Self::rotl64(h1, 27);
+            h1 += h2;
+            h1 = h1 * Wrapping(5) + Wrapping(0x52dce729);
+
+            k2 *= c2;
+            k2 = Self::rotl64(k2, 33);
+            k2 *= c1;
+            h2 ^= k2;
+
+            h2 = Self::rotl64(h2, 31);
+            h2 += h1;
+            h2 = h2 * Wrapping(5) + Wrapping(0x38495ab5);
+        }
+
+        let mut k1 = Wrapping(0_i64);
+        let mut k2 = Wrapping(0_i64);
+
+        debug_assert!(data.len() < 16);
+
+        if data.len() > 8 {
+            for i in (8..data.len()).rev() {
+                k2 ^= Wrapping(data[i] as i8 as i64) << ((i - 8) * 8);
+            }
+
+            k2 *= c2;
+            k2 = Self::rotl64(k2, 33);
+            k2 *= c1;
+            h2 ^= k2;
+        }
+
+        if !data.is_empty() {
+            for i in (0..std::cmp::min(8, data.len())).rev() {
+                k1 ^= Wrapping(data[i] as i8 as i64) << (i * 8);
+            }
+
+            k1 *= c1;
+            k1 = Self::rotl64(k1, 31);
+            k1 *= c2;
+            h1 ^= k1;
+        }
+
+        h1 ^= Wrapping(length as i64);
+        h2 ^= Wrapping(length as i64);
+
+        h1 += h2;
+        h2 += h1;
+
+        h1 = Self::fmix(h1);
+        h2 = Self::fmix(h2);
+
+        h1 += h2;
+        h2 += h1;
+
+        ((h2.0 as i128) << 64) | h1.0 as i128
+    }
+
+    #[inline]
+    fn rotl64(v: Wrapping<i64>, n: u32) -> Wrapping<i64> {
+        Wrapping((v.0 << n) | (v.0 as u64 >> (64 - n)) as i64)
+    }
+
+    #[inline]
+    fn fmix(mut k: Wrapping<i64>) -> Wrapping<i64> {
+        k ^= Wrapping((k.0 as u64 >> 33) as i64);
+        k *= Wrapping(0xff51afd7ed558ccd_u64 as i64);
+        k ^= Wrapping((k.0 as u64 >> 33) as i64);
+        k *= Wrapping(0xc4ceb9fe1a85ec53_u64 as i64);
+        k ^= Wrapping((k.0 as u64 >> 33) as i64);
+
+        k
+    }
+}
+
+impl Partitioner for Murmur3Partitioner {
+    fn hash(pk: Bytes) -> Token {
+        Token {
+            value: Self::hash3_x64_128(&pk) as i64,
+        }
+    }
+}
+
+impl Partitioner for CDCPartitioner {
+    fn hash(mut pk: Bytes) -> Token {
+        let value = if pk.len() < 8 { i64::MIN } else { pk.get_i64() };
+        Token { value }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::{CDCPartitioner, Murmur3Partitioner, Partitioner};
+    use bytes::Bytes;
+
+    fn assert_correct_murmur3_hash(pk: &'static str, expected_hash: i64) {
+        let pk = Bytes::from(pk);
+        let hash = Murmur3Partitioner::hash(pk).value;
+        assert_eq!(hash, expected_hash);
+    }
+
+    #[test]
+    fn test_murmur3_partitioner() {
+        for s in [
+            ("test", -6017608668500074083),
+            ("xd", 4507812186440344727),
+            ("primary_key", -1632642444691073360),
+            ("kremówki", 4354931215268080151),
+        ] {
+            assert_correct_murmur3_hash(s.0, s.1);
+        }
+    }
+
+    fn assert_correct_cdc_hash(pk: &'static str, expected_hash: i64) {
+        let pk = Bytes::from(pk);
+        let hash = CDCPartitioner::hash(pk).value;
+        assert_eq!(hash, expected_hash);
+    }
+
+    #[test]
+    fn test_cdc_partitioner() {
+        for s in [
+            ("test", -9223372036854775808),
+            ("xd", -9223372036854775808),
+            ("primary_key", 8102654598100187487),
+            ("kremówki", 7742362231512463211),
+        ] {
+            assert_correct_cdc_hash(s.0, s.1);
+        }
+    }
+}
diff --git a/scylla/src/transport/session.rs b/scylla/src/transport/session.rs
index c21b7976..1cc4ab03 100644
--- a/scylla/src/transport/session.rs
+++ b/scylla/src/transport/session.rs
@@ -22,7 +22,7 @@ use crate::frame::response::result;
 use crate::frame::value::{BatchValues, SerializedValues, ValueList};
 use crate::prepared_statement::{PartitionKeyError, PreparedStatement};
 use crate::query::Query;
-use crate::routing::{murmur3_token, Token};
+use crate::routing::Token;
 use crate::statement::{Consistency, SerialConsistency};
 use crate::tracing::{GetTracingConfig, TracingEvent, TracingInfo};
 use crate::transport::cluster::{Cluster, ClusterData};
@@ -36,6 +36,9 @@ use crate::transport::load_balancing::{
 };
 use crate::transport::metrics::Metrics;
 use crate::transport::node::Node;
+use crate::transport::partitioner::{
+    CDCPartitioner, Murmur3Partitioner, Partitioner, PartitionerName,
+};
 use crate::transport::query_result::QueryResult;
 use crate::transport::retry_policy::{
     DefaultRetryPolicy, QueryInfo, RetryDecision, RetryPolicy, RetrySession,
@@ -598,9 +601,27 @@ impl Session {
                 .extend(statement.prepare_tracing_ids);
         }
 
+        prepared.set_partitioner_name(
+            self.extract_partitioner_name(&prepared, &self.cluster.get_data()),
+        );
+
         Ok(prepared)
     }
 
+    fn extract_partitioner_name<'a>(
+        &self,
+        prepared: &PreparedStatement,
+        cluster_data: &'a ClusterData,
+    ) -> Option<&'a str> {
+        cluster_data
+            .keyspaces
+            .get(prepared.get_keyspace_name()?)?
+            .tables
+            .get(prepared.get_table_name()?)?
+            .partitioner
+            .as_deref()
+    }
+
     /// Execute a prepared query. Requires a [PreparedStatement](crate::prepared_statement::PreparedStatement)
     /// generated using [`Session::prepare`](Session::prepare)\
     /// Returns only a single page of results, to receive multiple pages use [execute_iter](Session::execute_iter)
@@ -662,7 +683,7 @@ impl Session {
         let values_ref = &serialized_values;
         let paging_state_ref = &paging_state;
 
-        let token = calculate_token(prepared, &serialized_values)?;
+        let token = self.calculate_token(prepared, &serialized_values)?;
 
         let statement_info = Statement {
             token: Some(token),
@@ -740,7 +761,7 @@ impl Session {
         let prepared = prepared.into();
         let serialized_values = values.serialized()?;
 
-        let token = calculate_token(&prepared, &serialized_values)?;
+        let token = self.calculate_token(&prepared, &serialized_values)?;
 
         let retry_session = match &prepared.config.retry_policy {
             Some(policy) => policy.new_session(),
@@ -1255,30 +1276,36 @@ impl Session {
         })
         .await
     }
+
+    fn calculate_token(
+        &self,
+        prepared: &PreparedStatement,
+        serialized_values: &SerializedValues,
+    ) -> Result<Token, QueryError> {
+        let partitioner_name = prepared.get_partitioner_name();
+
+        let partition_key = calculate_partition_key(prepared, serialized_values)?;
+
+        Ok(match partitioner_name {
+            PartitionerName::Murmur3 => Murmur3Partitioner::hash(partition_key),
+            PartitionerName::CDC => CDCPartitioner::hash(partition_key),
+        })
+    }
 }
 
-fn calculate_token(
+fn calculate_partition_key(
     stmt: &PreparedStatement,
     values: &SerializedValues,
-) -> Result<Token, QueryError> {
-    // TODO: take the partitioner of the table that is being queried and calculate the token using
-    // that partitioner. The below logic gives correct token only for murmur3partitioner
-    let partition_key = match stmt.compute_partition_key(values) {
-        Ok(key) => key,
-        Err(PartitionKeyError::NoPkIndexValue(_, _)) => {
-            return Err(QueryError::ProtocolError(
-                "No pk indexes - can't calculate token",
-            ))
-        }
-        Err(PartitionKeyError::ValueTooLong(values_len)) => {
-            return Err(QueryError::BadQuery(BadQuery::ValuesTooLongForKey(
-                values_len,
-                u16::MAX.into(),
-            )))
-        }
-    };
-
-    Ok(murmur3_token(partition_key))
+) -> Result<Bytes, QueryError> {
+    match stmt.compute_partition_key(values) {
+        Ok(key) => Ok(key),
+        Err(PartitionKeyError::NoPkIndexValue(_, _)) => Err(QueryError::ProtocolError(
+            "No pk indexes - can't calculate token",
+        )),
+        Err(PartitionKeyError::ValueTooLong(values_len)) => Err(QueryError::BadQuery(
+            BadQuery::ValuesTooLongForKey(values_len, u16::MAX.into()),
+        )),
+    }
 }
 
 // Resolve the given hostname using a DNS lookup if necessary.


Test patch:
diff --git a/scylla/src/transport/session_test.rs b/scylla/src/transport/session_test.rs
index 46d499bd..cc3fba3d 100644
--- a/scylla/src/transport/session_test.rs
+++ b/scylla/src/transport/session_test.rs
@@ -3,11 +3,12 @@ use crate::batch::Batch;
 use crate::frame::response::result::Row;
 use crate::frame::value::ValueList;
 use crate::query::Query;
-use crate::routing::hash3_x64_128;
+use crate::routing::Token;
 use crate::statement::Consistency;
 use crate::tracing::TracingInfo;
 use crate::transport::connection::BatchResult;
 use crate::transport::errors::{BadKeyspaceName, BadQuery, DbError, QueryError};
+use crate::transport::partitioner::{Murmur3Partitioner, Partitioner, PartitionerName};
 use crate::transport::topology::Strategy::SimpleStrategy;
 use crate::transport::topology::{CollectionType, ColumnKind, CqlType, NativeType};
 use crate::QueryResult;
@@ -200,16 +201,18 @@ async fn test_prepared_statement() {
             .unwrap()
             .rows
             .unwrap();
-        let token: i64 = rs.first().unwrap().columns[0]
-            .as_ref()
-            .unwrap()
-            .as_bigint()
-            .unwrap();
-        let expected_token = hash3_x64_128(
-            &prepared_statement
+        let token = Token {
+            value: rs.first().unwrap().columns[0]
+                .as_ref()
+                .unwrap()
+                .as_bigint()
+                .unwrap(),
+        };
+        let expected_token = Murmur3Partitioner::hash(
+            prepared_statement
                 .compute_partition_key(&serialized_values)
                 .unwrap(),
-        ) as i64;
+        );
 
         assert_eq!(token, expected_token)
     }
@@ -220,16 +223,18 @@ async fn test_prepared_statement() {
             .unwrap()
             .rows
             .unwrap();
-        let token: i64 = rs.first().unwrap().columns[0]
-            .as_ref()
-            .unwrap()
-            .as_bigint()
-            .unwrap();
-        let expected_token = hash3_x64_128(
-            &prepared_complex_pk_statement
+        let token = Token {
+            value: rs.first().unwrap().columns[0]
+                .as_ref()
+                .unwrap()
+                .as_bigint()
+                .unwrap(),
+        };
+        let expected_token = Murmur3Partitioner::hash(
+            prepared_complex_pk_statement
                 .compute_partition_key(&serialized_values)
                 .unwrap(),
-        ) as i64;
+        );
 
         assert_eq!(token, expected_token)
     }
@@ -473,16 +478,18 @@ async fn test_token_calculation() {
             .unwrap()
             .rows
             .unwrap();
-        let token: i64 = rs.first().unwrap().columns[0]
-            .as_ref()
-            .unwrap()
-            .as_bigint()
-            .unwrap();
-        let expected_token = hash3_x64_128(
-            &prepared_statement
+        let token = Token {
+            value: rs.first().unwrap().columns[0]
+                .as_ref()
+                .unwrap()
+                .as_bigint()
+                .unwrap(),
+        };
+        let expected_token = Murmur3Partitioner::hash(
+            prepared_statement
                 .compute_partition_key(&serialized_values)
                 .unwrap(),
-        ) as i64;
+        );
         assert_eq!(token, expected_token)
     }
 }
@@ -1694,3 +1701,56 @@ async fn test_named_bind_markers() {
         assert!(session.execute(&prepared, &wrongmap).await.is_err());
     }
 }
+
+#[tokio::test]
+async fn test_prepared_partitioner() {
+    let uri = std::env::var("SCYLLA_URI").unwrap_or_else(|_| "127.0.0.1:9042".to_string());
+    let session = SessionBuilder::new().known_node(uri).build().await.unwrap();
+    let ks = unique_name();
+
+    session.query(format!("CREATE KEYSPACE IF NOT EXISTS {} WITH REPLICATION = {{'class' : 'SimpleStrategy', 'replication_factor' : 1}}", ks), &[]).await.unwrap();
+    session.use_keyspace(ks, false).await.unwrap();
+
+    session
+        .query("CREATE TABLE IF NOT EXISTS t1 (a int primary key)", &[])
+        .await
+        .unwrap();
+
+    session.await_schema_agreement().await.unwrap();
+    session.refresh_metadata().await.unwrap();
+
+    let prepared_statement_for_main_table = session
+        .prepare("INSERT INTO t1 (a) VALUES (?)")
+        .await
+        .unwrap();
+
+    assert_eq!(
+        prepared_statement_for_main_table.get_partitioner_name(),
+        &PartitionerName::Murmur3
+    );
+
+    if option_env!("CDC") == Some("disabled") {
+        return;
+    }
+
+    session
+        .query(
+            "CREATE TABLE IF NOT EXISTS t2 (a int primary key) WITH cdc = {'enabled':true}",
+            &[],
+        )
+        .await
+        .unwrap();
+
+    session.await_schema_agreement().await.unwrap();
+    session.refresh_metadata().await.unwrap();
+
+    let prepared_statement_for_cdc_log = session
+        .prepare("SELECT a FROM t2_scylla_cdc_log WHERE \"cdc$stream_id\" = ?")
+        .await
+        .unwrap();
+
+    assert_eq!(
+        prepared_statement_for_cdc_log.get_partitioner_name(),
+        &PartitionerName::CDC
+    );
+}