Repository: scylladb/scylla-rust-driver
Pull request #549 (2022-09-09 12:24:45)

Add support for rate limit exceeded exception
This PR adds support for a new Scylla-specific type of error returned from the database: RateLimitExceeded. It was introduced as a part of the _per-partition rate limit_ feature, upcoming in Scylla 5.1 (PR in Scylla: https://github.com/scylladb/scylladb/pull/9810). The RateLimitExceeded error is sent when Scylla detects that the per-partition rate limit is exceeded.

In order for Scylla to be able to send this error, the driver must tell the database during connection handshake that it is able to interpret it (otherwise Scylla falls back to the Config_error error). This PR adds support for Scylla-specific protocol extensions, negotiates the `SCYLLA_RATE_LIMIT_ERROR` during handshake and adds support for the RateLimitExceeded error deserialization.

## 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.
- [ ] ~I added appropriate `Fixes:` annotations to PR description.~


psarna: I assume it was tested manually - if so, please add a note in the PR description
psarna: Oh wait, there are automatic tests too, nvm
psarna: Nope, still valid - please let me know if you saw the new code successfully set up a connection to new Scylla and negotiated the new error type (e.g. via wireshark or logs)

Patch:
diff --git a/scylla-cql/src/errors.rs b/scylla-cql/src/errors.rs
index b12b434a..c046c4d5 100644
--- a/scylla-cql/src/errors.rs
+++ b/scylla-cql/src/errors.rs
@@ -205,11 +205,32 @@ pub enum DbError {
     #[error("Invalid protocol message received from the driver")]
     ProtocolError,
 
+    /// Rate limit was exceeded for a partition affected by the request.
+    /// (Scylla-specific)
+    /// TODO: Should this have a "Scylla" prefix?
+    #[error("Rate limit was exceeded for a partition affected by the request")]
+    RateLimitReached {
+        /// Type of the operation rejected by rate limiting.
+        op_type: OperationType,
+        /// Whether the operation was rate limited on the coordinator or not.
+        /// Writes rejected on the coordinator are guaranteed not to be applied
+        /// on any replica.
+        rejected_by_coordinator: bool,
+    },
+
     /// Other error code not specified in the specification
     #[error("Other error not specified in the specification. Error code: {0}")]
     Other(i32),
 }
 
+/// Type of the operation rejected by rate limiting
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub enum OperationType {
+    Read,
+    Write,
+    Other(u8),
+}
+
 /// Type of write operation requested
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub enum WriteType {
@@ -406,6 +427,16 @@ impl QueryError {
     }
 }
 
+impl From<u8> for OperationType {
+    fn from(operation_type: u8) -> OperationType {
+        match operation_type {
+            0 => OperationType::Read,
+            1 => OperationType::Write,
+            other => OperationType::Other(other),
+        }
+    }
+}
+
 impl From<&str> for WriteType {
     fn from(write_type_str: &str) -> WriteType {
         match write_type_str {
diff --git a/scylla-cql/src/frame/mod.rs b/scylla-cql/src/frame/mod.rs
index d409e154..f7c1ab7d 100644
--- a/scylla-cql/src/frame/mod.rs
+++ b/scylla-cql/src/frame/mod.rs
@@ -1,4 +1,5 @@
 pub mod frame_errors;
+pub mod protocol_features;
 pub mod request;
 pub mod response;
 pub mod server_event_type;
diff --git a/scylla-cql/src/frame/protocol_features.rs b/scylla-cql/src/frame/protocol_features.rs
new file mode 100644
index 00000000..b084462f
--- /dev/null
+++ b/scylla-cql/src/frame/protocol_features.rs
@@ -0,0 +1,37 @@
+use std::collections::HashMap;
+
+const RATE_LIMIT_ERROR_EXTENSION: &str = "SCYLLA_RATE_LIMIT_ERROR";
+
+#[derive(Default, Clone, Copy, Debug, PartialEq, Eq)]
+#[non_exhaustive]
+pub struct ProtocolFeatures {
+    pub rate_limit_error: Option<i32>,
+}
+
+// TODO: Log information about options which failed to parse
+
+impl ProtocolFeatures {
+    pub fn parse_from_supported(supported: &HashMap<String, Vec<String>>) -> Self {
+        Self {
+            rate_limit_error: Self::maybe_parse_rate_limit_error(supported),
+        }
+    }
+
+    fn maybe_parse_rate_limit_error(supported: &HashMap<String, Vec<String>>) -> Option<i32> {
+        let vals = supported.get(RATE_LIMIT_ERROR_EXTENSION)?;
+        let code_str = Self::get_cql_extension_field(vals.as_slice(), "ERROR_CODE")?;
+        code_str.parse::<i32>().ok()
+    }
+
+    // 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()
+            .find_map(|v| v.as_str().strip_prefix(key)?.strip_prefix('='))
+    }
+
+    pub fn add_startup_options(&self, options: &mut HashMap<String, String>) {
+        if self.rate_limit_error.is_some() {
+            options.insert(RATE_LIMIT_ERROR_EXTENSION.to_string(), String::new());
+        }
+    }
+}
diff --git a/scylla-cql/src/frame/response/error.rs b/scylla-cql/src/frame/response/error.rs
index 61338ea6..3984fe09 100644
--- a/scylla-cql/src/frame/response/error.rs
+++ b/scylla-cql/src/frame/response/error.rs
@@ -1,5 +1,6 @@
-use crate::errors::{DbError, QueryError, WriteType};
+use crate::errors::{DbError, OperationType, QueryError, WriteType};
 use crate::frame::frame_errors::ParseError;
+use crate::frame::protocol_features::ProtocolFeatures;
 use crate::frame::types;
 use byteorder::ReadBytesExt;
 use bytes::Bytes;
@@ -11,7 +12,7 @@ pub struct Error {
 }
 
 impl Error {
-    pub fn deserialize(buf: &mut &[u8]) -> Result<Self, ParseError> {
+    pub fn deserialize(features: &ProtocolFeatures, buf: &mut &[u8]) -> Result<Self, ParseError> {
         let code = types::read_int(buf)?;
         let reason = types::read_string(buf)?.to_owned();
 
@@ -69,6 +70,10 @@ impl Error {
             0x2500 => DbError::Unprepared {
                 statement_id: Bytes::from(types::read_short_bytes(buf)?.to_owned()),
             },
+            code if Some(code) == features.rate_limit_error => DbError::RateLimitReached {
+                op_type: OperationType::from(buf.read_u8()?),
+                rejected_by_coordinator: buf.read_u8()? != 0,
+            },
             _ => DbError::Other(code),
         };
 
@@ -85,7 +90,8 @@ impl From<Error> for QueryError {
 #[cfg(test)]
 mod tests {
     use super::Error;
-    use crate::errors::{DbError, WriteType};
+    use crate::errors::{DbError, OperationType, WriteType};
+    use crate::frame::protocol_features::ProtocolFeatures;
     use crate::frame::types::LegacyConsistency;
     use crate::Consistency;
     use bytes::Bytes;
@@ -121,9 +127,11 @@ mod tests {
             (0x1234, DbError::Other(0x1234)),
         ];
 
+        let features = ProtocolFeatures::default();
+
         for (error_code, expected_error) in &simple_error_mappings {
             let bytes: Vec<u8> = make_error_request_bytes(*error_code, "simple message");
-            let error: Error = Error::deserialize(&mut bytes.as_slice()).unwrap();
+            let error: Error = Error::deserialize(&features, &mut bytes.as_slice()).unwrap();
             assert_eq!(error.error, *expected_error);
             assert_eq!(error.reason, "simple message");
         }
@@ -131,12 +139,14 @@ mod tests {
 
     #[test]
     fn deserialize_unavailable() {
+        let features = ProtocolFeatures::default();
+
         let mut bytes = make_error_request_bytes(0x1000, "message 2");
         bytes.extend(&1_i16.to_be_bytes());
         bytes.extend(&2_i32.to_be_bytes());
         bytes.extend(&3_i32.to_be_bytes());
 
-        let error: Error = Error::deserialize(&mut bytes.as_slice()).unwrap();
+        let error: Error = Error::deserialize(&features, &mut bytes.as_slice()).unwrap();
 
         assert_eq!(
             error.error,
@@ -151,6 +161,8 @@ mod tests {
 
     #[test]
     fn deserialize_write_timeout() {
+        let features = ProtocolFeatures::default();
+
         let mut bytes = make_error_request_bytes(0x1100, "message 2");
         bytes.extend(&0x0004_i16.to_be_bytes());
         bytes.extend(&(-5_i32).to_be_bytes());
@@ -161,7 +173,7 @@ mod tests {
         bytes.extend(&write_type_str_len.to_be_bytes());
         bytes.extend(write_type_str.as_bytes());
 
-        let error: Error = Error::deserialize(&mut bytes.as_slice()).unwrap();
+        let error: Error = Error::deserialize(&features, &mut bytes.as_slice()).unwrap();
 
         assert_eq!(
             error.error,
@@ -177,13 +189,15 @@ mod tests {
 
     #[test]
     fn deserialize_read_timeout() {
+        let features = ProtocolFeatures::default();
+
         let mut bytes = make_error_request_bytes(0x1200, "message 2");
         bytes.extend(&0x0002_i16.to_be_bytes());
         bytes.extend(&8_i32.to_be_bytes());
         bytes.extend(&32_i32.to_be_bytes());
         bytes.push(0_u8);
 
-        let error: Error = Error::deserialize(&mut bytes.as_slice()).unwrap();
+        let error: Error = Error::deserialize(&features, &mut bytes.as_slice()).unwrap();
 
         assert_eq!(
             error.error,
@@ -199,6 +213,8 @@ mod tests {
 
     #[test]
     fn deserialize_read_failure() {
+        let features = ProtocolFeatures::default();
+
         let mut bytes = make_error_request_bytes(0x1300, "message 2");
         bytes.extend(&0x0003_i16.to_be_bytes());
         bytes.extend(&4_i32.to_be_bytes());
@@ -206,7 +222,7 @@ mod tests {
         bytes.extend(&6_i32.to_be_bytes());
         bytes.push(123_u8); // Any non-zero value means data_present is true
 
-        let error: Error = Error::deserialize(&mut bytes.as_slice()).unwrap();
+        let error: Error = Error::deserialize(&features, &mut bytes.as_slice()).unwrap();
 
         assert_eq!(
             error.error,
@@ -223,6 +239,8 @@ mod tests {
 
     #[test]
     fn deserialize_function_failure() {
+        let features = ProtocolFeatures::default();
+
         let mut bytes = make_error_request_bytes(0x1400, "message 2");
 
         let keyspace_name: &str = "keyspace_name";
@@ -247,7 +265,7 @@ mod tests {
         bytes.extend(&type2_len.to_be_bytes());
         bytes.extend(type2.as_bytes());
 
-        let error: Error = Error::deserialize(&mut bytes.as_slice()).unwrap();
+        let error: Error = Error::deserialize(&features, &mut bytes.as_slice()).unwrap();
 
         assert_eq!(
             error.error,
@@ -262,6 +280,8 @@ mod tests {
 
     #[test]
     fn deserialize_write_failure() {
+        let features = ProtocolFeatures::default();
+
         let mut bytes = make_error_request_bytes(0x1500, "message 2");
 
         bytes.extend(&0x0000_i16.to_be_bytes());
@@ -274,7 +294,7 @@ mod tests {
         bytes.extend(&write_type_str_len.to_be_bytes());
         bytes.extend(write_type_str.as_bytes());
 
-        let error: Error = Error::deserialize(&mut bytes.as_slice()).unwrap();
+        let error: Error = Error::deserialize(&features, &mut bytes.as_slice()).unwrap();
 
         assert_eq!(
             error.error,
@@ -291,6 +311,8 @@ mod tests {
 
     #[test]
     fn deserialize_already_exists() {
+        let features = ProtocolFeatures::default();
+
         let mut bytes = make_error_request_bytes(0x2400, "message 2");
 
         let keyspace_name: &str = "keyspace_name";
@@ -304,7 +326,7 @@ mod tests {
         bytes.extend(&table_name_len.to_be_bytes());
         bytes.extend(table_name.as_bytes());
 
-        let error: Error = Error::deserialize(&mut bytes.as_slice()).unwrap();
+        let error: Error = Error::deserialize(&features, &mut bytes.as_slice()).unwrap();
 
         assert_eq!(
             error.error,
@@ -318,12 +340,14 @@ mod tests {
 
     #[test]
     fn deserialize_unprepared() {
+        let features = ProtocolFeatures::default();
+
         let mut bytes = make_error_request_bytes(0x2500, "message 3");
         let statement_id = b"deadbeef";
         bytes.extend((statement_id.len() as i16).to_be_bytes());
         bytes.extend(statement_id);
 
-        let error: Error = Error::deserialize(&mut bytes.as_slice()).unwrap();
+        let error: Error = Error::deserialize(&features, &mut bytes.as_slice()).unwrap();
 
         assert_eq!(
             error.error,
@@ -333,4 +357,43 @@ mod tests {
         );
         assert_eq!(error.reason, "message 3");
     }
+
+    #[test]
+    fn deserialize_rate_limit_error() {
+        let features = ProtocolFeatures {
+            rate_limit_error: Some(0x4321),
+            ..Default::default()
+        };
+        let mut bytes = make_error_request_bytes(0x4321, "message 1");
+        bytes.extend([0u8]); // Read type
+        bytes.extend([1u8]); // Rejected by coordinator
+        let error = Error::deserialize(&features, &mut bytes.as_slice()).unwrap();
+
+        assert_eq!(
+            error.error,
+            DbError::RateLimitReached {
+                op_type: OperationType::Read,
+                rejected_by_coordinator: true,
+            }
+        );
+        assert_eq!(error.reason, "message 1");
+
+        let features = ProtocolFeatures {
+            rate_limit_error: Some(0x8765),
+            ..Default::default()
+        };
+        let mut bytes = make_error_request_bytes(0x8765, "message 2");
+        bytes.extend([1u8]); // Write type
+        bytes.extend([0u8]); // Not rejected by coordinator
+        let error = Error::deserialize(&features, &mut bytes.as_slice()).unwrap();
+
+        assert_eq!(
+            error.error,
+            DbError::RateLimitReached {
+                op_type: OperationType::Write,
+                rejected_by_coordinator: false,
+            }
+        );
+        assert_eq!(error.reason, "message 2");
+    }
 }
diff --git a/scylla-cql/src/frame/response/mod.rs b/scylla-cql/src/frame/response/mod.rs
index e53ae874..c8c4ec10 100644
--- a/scylla-cql/src/frame/response/mod.rs
+++ b/scylla-cql/src/frame/response/mod.rs
@@ -8,6 +8,7 @@ pub mod supported;
 use crate::{errors::QueryError, frame::frame_errors::ParseError};
 use num_enum::TryFromPrimitive;
 
+use crate::frame::protocol_features::ProtocolFeatures;
 pub use error::Error;
 pub use supported::Supported;
 
@@ -37,9 +38,13 @@ pub enum Response {
 }
 
 impl Response {
-    pub fn deserialize(opcode: ResponseOpcode, buf: &mut &[u8]) -> Result<Response, ParseError> {
+    pub fn deserialize(
+        features: &ProtocolFeatures,
+        opcode: ResponseOpcode,
+        buf: &mut &[u8],
+    ) -> Result<Response, ParseError> {
         let response = match opcode {
-            ResponseOpcode::Error => Response::Error(Error::deserialize(buf)?),
+            ResponseOpcode::Error => Response::Error(Error::deserialize(features, buf)?),
             ResponseOpcode::Ready => Response::Ready,
             ResponseOpcode::Authenticate => {
                 Response::Authenticate(authenticate::Authenticate::deserialize(buf)?)
diff --git a/scylla/src/transport/connection.rs b/scylla/src/transport/connection.rs
index 422eadf8..79cba517 100644
--- a/scylla/src/transport/connection.rs
+++ b/scylla/src/transport/connection.rs
@@ -29,6 +29,7 @@ use std::{
 use super::errors::{BadKeyspaceName, BadQuery, DbError, QueryError};
 
 use crate::batch::{Batch, BatchStatement};
+use crate::frame::protocol_features::ProtocolFeatures;
 use crate::frame::{
     self,
     request::{self, batch, execute, query, register, Request},
@@ -70,9 +71,8 @@ pub struct Connection {
     _worker_handle: RemoteHandle<()>,
 
     connect_address: SocketAddr,
-    shard_info: Option<ShardInfo>,
-    shard_aware_port: Option<u16>,
     config: ConnectionConfig,
+    features: ConnectionFeatures,
 
     // Each request send by `Connection::send_request` needs a unique request id.
     // This field is a monotonic generator of such ids.
@@ -85,6 +85,13 @@ pub struct Connection {
     orphan_notification_sender: mpsc::UnboundedSender<RequestId>,
 }
 
+#[derive(Default)]
+pub(crate) struct ConnectionFeatures {
+    shard_info: Option<ShardInfo>,
+    shard_aware_port: Option<u16>,
+    protocol_features: ProtocolFeatures,
+}
+
 type RequestId = u64;
 
 struct ResponseHandler {
@@ -297,9 +304,8 @@ impl Connection {
             submit_channel: sender,
             _worker_handle,
             config,
+            features: Default::default(),
             connect_address: addr,
-            shard_info: None,
-            shard_aware_port: None,
             request_id_generator: AtomicU64::new(0),
             orphan_notification_sender,
         };
@@ -825,12 +831,17 @@ impl Connection {
         // notification about orphaning.
         notifier.disable();
 
-        Self::parse_response(task_response?, self.config.compression)
+        Self::parse_response(
+            task_response?,
+            self.config.compression,
+            &self.features.protocol_features,
+        )
     }
 
     fn parse_response(
         task_response: TaskResponse,
         compression: Option<Compression>,
+        features: &ProtocolFeatures,
     ) -> Result<QueryResponse, QueryError> {
         let body_with_ext = frame::parse_response_body_extensions(
             task_response.params.flags,
@@ -842,7 +853,8 @@ impl Connection {
             warn!(warning = warn_description.as_str());
         }
 
-        let response = Response::deserialize(task_response.opcode, &mut &*body_with_ext.body)?;
+        let response =
+            Response::deserialize(features, task_response.opcode, &mut &*body_with_ext.body)?;
 
         Ok(QueryResponse {
             response,
@@ -1107,7 +1119,18 @@ impl Connection {
         compression: Option<Compression>,
         event_sender: &mpsc::Sender<Event>,
     ) -> Result<(), QueryError> {
-        let response = Self::parse_response(task_response, compression)?.response;
+        // Protocol features are negotiated during connection handshake.
+        // However, the router is already created and sent to a different tokio
+        // task before the handshake begins, therefore it's hard to cleanly
+        // update the protocol features in the router at this point.
+        // Making it possible would require restructuring the handshake process,
+        // or passing the negotiated features via a channel/mutex/etc.
+        // Fortunately, events do not need information about protocol features
+        // to be serialized (yet), therefore I'm leaving this problem for
+        // future implementors.
+        let features = ProtocolFeatures::default(); // TODO: Use the right features
+
+        let response = Self::parse_response(task_response, compression, &features)?.response;
         let event = match response {
             Response::Event(e) => e,
             _ => {
@@ -1125,19 +1148,15 @@ impl Connection {
     }
 
     pub fn get_shard_info(&self) -> &Option<ShardInfo> {
-        &self.shard_info
+        &self.features.shard_info
     }
 
     pub fn get_shard_aware_port(&self) -> Option<u16> {
-        self.shard_aware_port
-    }
-
-    fn set_shard_info(&mut self, shard_info: Option<ShardInfo>) {
-        self.shard_info = shard_info
+        self.features.shard_aware_port
     }
 
-    fn set_shard_aware_port(&mut self, shard_aware_port: Option<u16>) {
-        self.shard_aware_port = shard_aware_port;
+    fn set_features(&mut self, features: ConnectionFeatures) {
+        self.features = features;
     }
 
     pub fn get_connect_address(&self) -> SocketAddr {
@@ -1176,25 +1195,37 @@ pub async fn open_named_connection(
         false => "SCYLLA_SHARD_AWARE_PORT",
     };
 
-    let (shard_info, supported_compression, shard_aware_port) = match options_result {
-        Response::Supported(mut supported) => {
-            let shard_info = ShardInfo::try_from(&supported.options).ok();
-            let supported_compression = supported.options.remove("COMPRESSION").unwrap_or_default();
-            let shard_aware_port = supported
-                .options
-                .remove(shard_aware_port_key)
-                .unwrap_or_default()
-                .into_iter()
-                .next()
-                .and_then(|p| p.parse::<u16>().ok());
-            (shard_info, supported_compression, shard_aware_port)
+    let mut supported = match options_result {
+        Response::Supported(supported) => supported,
+        _ => {
+            return Err(QueryError::ProtocolError(
+                "Wrong response to OPTIONS message was received",
+            ));
         }
-        _ => (None, Vec::new(), None),
     };
-    connection.set_shard_info(shard_info);
-    connection.set_shard_aware_port(shard_aware_port);
+
+    let shard_info = ShardInfo::try_from(&supported.options).ok();
+    let supported_compression = supported.options.remove("COMPRESSION").unwrap_or_default();
+    let shard_aware_port = supported
+        .options
+        .remove(shard_aware_port_key)
+        .unwrap_or_default()
+        .into_iter()
+        .next()
+        .and_then(|p| p.parse::<u16>().ok());
+
+    let protocol_features = ProtocolFeatures::parse_from_supported(&supported.options);
 
     let mut options = HashMap::new();
+    protocol_features.add_startup_options(&mut options);
+
+    let features = ConnectionFeatures {
+        shard_info,
+        shard_aware_port,
+        protocol_features,
+    };
+    connection.set_features(features);
+
     options.insert("CQL_VERSION".to_string(), "4.0.0".to_string()); // FIXME: hardcoded values
     if let Some(name) = driver_name {
         options.insert("DRIVER_NAME".to_string(), name);


Test patch:
diff --git a/scylla/src/transport/session_test.rs b/scylla/src/transport/session_test.rs
index a1e66666..cb300cca 100644
--- a/scylla/src/transport/session_test.rs
+++ b/scylla/src/transport/session_test.rs
@@ -12,7 +12,7 @@ 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::utils::test_utils::unique_keyspace_name;
+use crate::utils::test_utils::{supports_feature, unique_keyspace_name};
 use crate::CachingSession;
 use crate::QueryResult;
 use crate::{IntoTypedRows, Session, SessionBuilder};
@@ -2241,3 +2241,51 @@ async fn test_refresh_metadata_after_schema_agreement() {
         ])
     );
 }
+
+#[tokio::test]
+async fn test_rate_limit_exceeded_exception() {
+    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();
+
+    // Typed errors in RPC were introduced along with per-partition rate limiting.
+    // There is no dedicated feature for per-partition rate limiting, so we are
+    // looking at the other one.
+    if !supports_feature(&session, "TYPED_ERRORS_IN_READ_RPC").await {
+        println!("Skipping because the cluster doesn't support per partition rate limiting");
+        return;
+    }
+
+    let ks = unique_keyspace_name();
+    session.query(format!("CREATE KEYSPACE IF NOT EXISTS {} WITH REPLICATION = {{'class' : 'SimpleStrategy', 'replication_factor' : 1}}", ks), &[]).await.unwrap();
+    session.use_keyspace(ks.clone(), false).await.unwrap();
+    session.query("CREATE TABLE tbl (pk int PRIMARY KEY, v int) WITH per_partition_rate_limit = {'max_writes_per_second': 1}", ()).await.unwrap();
+
+    let stmt = session
+        .prepare("INSERT INTO tbl (pk, v) VALUES (?, ?)")
+        .await
+        .unwrap();
+
+    // The rate limit is 1 write/s, so repeat the same query
+    // until an error occurs, it should happen quickly
+
+    let mut maybe_err = None;
+
+    for _ in 0..1000 {
+        match session.execute(&stmt, (123, 456)).await {
+            Ok(_) => {} // Try again
+            Err(err) => {
+                maybe_err = Some(err);
+                break;
+            }
+        }
+    }
+
+    use scylla_cql::errors::OperationType;
+
+    match maybe_err.expect("Rate limit error didn't occur") {
+        QueryError::DbError(DbError::RateLimitReached { op_type, .. }, _) => {
+            assert_eq!(op_type, OperationType::Write);
+        }
+        err => panic!("Unexpected error type received: {:?}", err),
+    }
+}
diff --git a/scylla/src/utils/test_utils.rs b/scylla/src/utils/test_utils.rs
index 082f158f..c15bba0c 100644
--- a/scylla/src/utils/test_utils.rs
+++ b/scylla/src/utils/test_utils.rs
@@ -3,6 +3,8 @@ use std::{
     time::{SystemTime, UNIX_EPOCH},
 };
 
+use crate::Session;
+
 static UNIQUE_COUNTER: AtomicUsize = AtomicUsize::new(0);
 
 pub(crate) fn unique_keyspace_name() -> String {
@@ -18,3 +20,33 @@ pub(crate) fn unique_keyspace_name() -> String {
     println!("Unique name: {}", name);
     name
 }
+
+pub(crate) async fn supports_feature(session: &Session, feature: &str) -> bool {
+    // Cassandra doesn't have a concept of features, so first detect
+    // if there is the `supported_features` column in system.local
+
+    let meta = session.get_cluster_data();
+    let system_local = meta
+        .keyspaces
+        .get("system")
+        .unwrap()
+        .tables
+        .get("local")
+        .unwrap();
+
+    if !system_local.columns.contains_key("supported_features") {
+        return false;
+    }
+
+    let (features,): (Option<String>,) = session
+        .query("SELECT supported_features FROM system.local", ())
+        .await
+        .unwrap()
+        .single_row_typed()
+        .unwrap();
+
+    features
+        .unwrap_or_default()
+        .split(',')
+        .any(|f| f == feature)
+}