Introduce support for Tablets

.github/workflows/rust.yml

@@ -30,10 +30,14 @@ jobs:
       run: cargo fmt --verbose --all -- --check
     - name: Clippy check
       run: cargo clippy --verbose --all-targets
+    - name: Clippy check with all features
+      run: cargo clippy --verbose --all-targets --all-features
     - name: Cargo check without features
       run: cargo check --all-targets --manifest-path "scylla/Cargo.toml" --features ""
     - name: Cargo check with all serialization features
       run: cargo check --all-targets --manifest-path "scylla/Cargo.toml" --features "full-serialization"
+    - name: Cargo check with all features
+      run: cargo check --all-targets --manifest-path "scylla/Cargo.toml" --all-features
     - name: Cargo check with secret feature
       run: cargo check --all-targets --manifest-path "scylla/Cargo.toml" --features "secret"
     - name: Cargo check with chrono feature

.github/workflows/tablets.yml

@@ -0,0 +1,38 @@
+name: Tablets
+
+on:
+  push:
+    branches:
+    - main
+    - 'branch-*'
+  pull_request:
+    branches:
+    - main
+    - 'branch-*'
+
+env:
+  CARGO_TERM_COLOR: always
+  RUST_BACKTRACE: full
+
+jobs:
+  build:
+    runs-on: ubuntu-latest
+    timeout-minutes: 60
+    steps:
+      - uses: actions/checkout@v3
+      - name: Install scylla-ccm
+        run: pip3 install https://github.com/scylladb/scylla-ccm/archive/master.zip
+
+      - name: Create cluster with tablets enabled
+        run: |
+          ccm create tablets -i 127.0.1. -n 3 --scylla -v 'unstable/master:2024-05-01T18:26:10Z'
+          ccm updateconf 'experimental_features: [consistent-topology-changes, tablets]'
+          ccm start
+
+      - name: Check
+        run: cargo check --verbose
+      - name: Run tablets tests
+        run: SCYLLA_URI=127.0.1.1:9042 SCYLLA_URI2=127.0.1.2:9042 SCYLLA_URI3=127.0.1.3:9042 cargo test --verbose
+
+      - name: Remove tablets cluster
+        run: ccm remove tablets

Makefile

@@ -4,10 +4,10 @@ COMPOSE := docker compose -f test/cluster/docker-compose.yml
 all: test
 
 .PHONY: ci
-ci: fmt-check check check-without-features clippy test build
+ci: fmt-check check check-without-features check-all-features clippy clippy-all-features test build
 
 .PHONY: dockerized-ci
-dockerized-ci: fmt-check check check-without-features clippy dockerized-test build
+dockerized-ci: fmt-check check check-without-features check-all-features clippy clippy-all-features dockerized-test build
 
 .PHONY: fmt
 fmt:
@@ -25,10 +25,18 @@ check:
 check-without-features:
 	cargo check --manifest-path "scylla/Cargo.toml" --features "" --all-targets
 
+.PHONY: check-all-features
+check-all-features:
+	cargo check --all-targets --all-features
+
 .PHONY: clippy
 clippy:
 	RUSTFLAGS=-Dwarnings cargo clippy --all-targets
 
+.PHONY: clippy-all-features
+clippy-all-features:
+	RUSTFLAGS=-Dwarnings cargo clippy --all-targets --all-features
+
 .PHONY: test
 test: up
 	SCYLLA_URI=172.42.0.2:9042 \

examples/compare-tokens.rs

@@ -39,7 +39,7 @@ async fn main() -> Result<()> {
             "Token endpoints for query: {:?}",
             session
                 .get_cluster_data()
-                .get_token_endpoints("examples_ks", Token::new(t))
+                .get_token_endpoints("examples_ks", "compare_tokens", Token::new(t))
                 .iter()
                 .map(|(node, _shard)| node.address)
                 .collect::<Vec<NodeAddr>>()

scylla-cql/src/frame/mod.rs

@@ -15,7 +15,7 @@ use thiserror::Error;
 use tokio::io::{AsyncRead, AsyncReadExt};
 use uuid::Uuid;
 
-use std::convert::TryFrom;
+use std::{collections::HashMap, convert::TryFrom};
 
 use request::SerializableRequest;
 use response::ResponseOpcode;
@@ -169,6 +169,7 @@ pub struct ResponseBodyWithExtensions {
     pub trace_id: Option<Uuid>,
     pub warnings: Vec<String>,
     pub body: Bytes,
+    pub custom_payload: Option<HashMap<String, Vec<u8>>>,
 }
 
 pub fn parse_response_body_extensions(
@@ -204,20 +205,22 @@ pub fn parse_response_body_extensions(
         Vec::new()
     };
 
-    if flags & FLAG_CUSTOM_PAYLOAD != 0 {
-        // TODO: Do something useful with the custom payload map
-        // For now, just skip it
+    let custom_payload = if flags & FLAG_CUSTOM_PAYLOAD != 0 {
         let body_len = body.len();
         let buf = &mut &*body;
-        types::read_bytes_map(buf)?;
+        let payload_map = types::read_bytes_map(buf)?;
         let buf_len = buf.len();
         body.advance(body_len - buf_len);
-    }
+        Some(payload_map)
+    } else {
+        None
+    };
 
     Ok(ResponseBodyWithExtensions {
         trace_id,
         warnings,
         body,
+        custom_payload,
     })
 }
 

scylla-cql/src/frame/protocol_features.rs

@@ -3,11 +3,14 @@ use std::collections::HashMap;
 const RATE_LIMIT_ERROR_EXTENSION: &str = "SCYLLA_RATE_LIMIT_ERROR";
 pub const SCYLLA_LWT_ADD_METADATA_MARK_EXTENSION: &str = "SCYLLA_LWT_ADD_METADATA_MARK";
 pub const LWT_OPTIMIZATION_META_BIT_MASK_KEY: &str = "LWT_OPTIMIZATION_META_BIT_MASK";
+const TABLETS_ROUTING_V1_KEY: &str = "TABLETS_ROUTING_V1";
+
 #[derive(Default, Clone, Copy, Debug, PartialEq, Eq)]
 #[non_exhaustive]
 pub struct ProtocolFeatures {
     pub rate_limit_error: Option<i32>,
     pub lwt_optimization_meta_bit_mask: Option<u32>,
+    pub tablets_v1_supported: bool,
 }
 
 // TODO: Log information about options which failed to parse
@@ -19,6 +22,7 @@ impl ProtocolFeatures {
             lwt_optimization_meta_bit_mask: Self::maybe_parse_lwt_optimization_meta_bit_mask(
                 supported,
             ),
+            tablets_v1_supported: Self::check_tablets_routing_v1_support(supported),
         }
     }
 
@@ -37,6 +41,10 @@ impl ProtocolFeatures {
         mask_str.parse::<u32>().ok()
     }
 
+    fn check_tablets_routing_v1_support(supported: &HashMap<String, Vec<String>>) -> bool {
+        supported.contains_key(TABLETS_ROUTING_V1_KEY)
+    }
+
     // Looks up a field which starts with `key=` and returns the rest
     fn get_cql_extension_field<'a>(vals: &'a [String], key: &str) -> Option<&'a str> {
         vals.iter()
@@ -53,6 +61,10 @@ impl ProtocolFeatures {
                 format!("{}={}", LWT_OPTIMIZATION_META_BIT_MASK_KEY, mask),
             );
         }
+
+        if self.tablets_v1_supported {
+            options.insert(TABLETS_ROUTING_V1_KEY.to_string(), String::new());
+        }
     }
 
     pub fn prepared_flags_contain_lwt_mark(&self, flags: u32) -> bool {