Merge pull request #1039 from wprzytula/self-identity-customisation

connection: expose setters for STARTUP options
scylladb · Aug 3, 2024 · 36d5edd · 36d5edd
2 parents a707bfd + 9aadac8
commit 36d5edd
Show file tree

Hide file tree

Showing 11 changed files with 482 additions and 100 deletions.
diff --git a/scylla-cql/src/frame/mod.rs b/scylla-cql/src/frame/mod.rs
@@ -48,15 +48,21 @@ pub enum Compression {
     Snappy,
 }
 
-impl Display for Compression {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+impl Compression {
+    pub fn as_str(&self) -> &'static str {
         match self {
-            Compression::Lz4 => f.write_str("lz4"),
-            Compression::Snappy => f.write_str("snappy"),
+            Compression::Lz4 => "lz4",
+            Compression::Snappy => "snappy",
         }
     }
 }
 
+impl Display for Compression {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.write_str(self.as_str())
+    }
+}
+
 pub struct SerializedRequest {
     data: Vec<u8>,
 }

diff --git a/scylla-cql/src/frame/protocol_features.rs b/scylla-cql/src/frame/protocol_features.rs
@@ -1,3 +1,4 @@
+use std::borrow::Cow;
 use std::collections::HashMap;
 
 const RATE_LIMIT_ERROR_EXTENSION: &str = "SCYLLA_RATE_LIMIT_ERROR";
@@ -51,19 +52,19 @@ impl ProtocolFeatures {
             .find_map(|v| v.as_str().strip_prefix(key)?.strip_prefix('='))
     }
 
-    pub fn add_startup_options(&self, options: &mut HashMap<String, String>) {
+    pub fn add_startup_options(&self, options: &mut HashMap<Cow<'_, str>, Cow<'_, str>>) {
         if self.rate_limit_error.is_some() {
-            options.insert(RATE_LIMIT_ERROR_EXTENSION.to_string(), String::new());
+            options.insert(Cow::Borrowed(RATE_LIMIT_ERROR_EXTENSION), Cow::Borrowed(""));
         }
         if let Some(mask) = self.lwt_optimization_meta_bit_mask {
             options.insert(
-                SCYLLA_LWT_ADD_METADATA_MARK_EXTENSION.to_string(),
-                format!("{}={}", LWT_OPTIMIZATION_META_BIT_MASK_KEY, mask),
+                Cow::Borrowed(SCYLLA_LWT_ADD_METADATA_MARK_EXTENSION),
+                Cow::Owned(format!("{}={}", LWT_OPTIMIZATION_META_BIT_MASK_KEY, mask)),
             );
         }
 
         if self.tablets_v1_supported {
-            options.insert(TABLETS_ROUTING_V1_KEY.to_string(), String::new());
+            options.insert(Cow::Borrowed(TABLETS_ROUTING_V1_KEY), Cow::Borrowed(""));
         }
     }
 

diff --git a/scylla-cql/src/frame/request/options.rs b/scylla-cql/src/frame/request/options.rs
@@ -11,3 +11,20 @@ impl SerializableRequest for Options {
         Ok(())
     }
 }
+
+/* Key names for options in SUPPORTED/STARTUP */
+pub const SCYLLA_SHARD_AWARE_PORT: &str = "SCYLLA_SHARD_AWARE_PORT";
+pub const SCYLLA_SHARD_AWARE_PORT_SSL: &str = "SCYLLA_SHARD_AWARE_PORT_SSL";
+
+pub const COMPRESSION: &str = "COMPRESSION";
+pub const CQL_VERSION: &str = "CQL_VERSION";
+pub const DRIVER_NAME: &str = "DRIVER_NAME";
+pub const DRIVER_VERSION: &str = "DRIVER_VERSION";
+pub const APPLICATION_NAME: &str = "APPLICATION_NAME";
+pub const APPLICATION_VERSION: &str = "APPLICATION_VERSION";
+pub const CLIENT_ID: &str = "CLIENT_ID";
+
+/* Value names for options in SUPPORTED/STARTUP */
+pub const DEFAULT_CQL_PROTOCOL_VERSION: &str = "4.0.0";
+pub const DEFAULT_DRIVER_NAME: &str = "ScyllaDB Rust Driver";
+pub const DEFAULT_DRIVER_VERSION: &str = env!("CARGO_PKG_VERSION");
diff --git a/scylla-cql/src/frame/request/startup.rs b/scylla-cql/src/frame/request/startup.rs
@@ -1,17 +1,17 @@
 use crate::frame::frame_errors::ParseError;
 
-use std::collections::HashMap;
+use std::{borrow::Cow, collections::HashMap};
 
 use crate::{
     frame::request::{RequestOpcode, SerializableRequest},
     frame::types,
 };
 
-pub struct Startup {
-    pub options: HashMap<String, String>,
+pub struct Startup<'a> {
+    pub options: HashMap<Cow<'a, str>, Cow<'a, str>>,
 }
 
-impl SerializableRequest for Startup {
+impl SerializableRequest for Startup<'_> {
     const OPCODE: RequestOpcode = RequestOpcode::Startup;
 
     fn serialize(&self, buf: &mut Vec<u8>) -> Result<(), ParseError> {

diff --git a/scylla-cql/src/frame/types.rs b/scylla-cql/src/frame/types.rs
@@ -428,14 +428,14 @@ pub fn read_string_map(
 }
 
 pub fn write_string_map(
-    v: &HashMap<String, String>,
+    v: &HashMap<impl AsRef<str>, impl AsRef<str>>,
     buf: &mut impl BufMut,
 ) -> Result<(), std::num::TryFromIntError> {
     let len = v.len();
     write_short_length(len, buf)?;
     for (key, val) in v.iter() {
-        write_string(key, buf)?;
-        write_string(val, buf)?;
+        write_string(key.as_ref(), buf)?;
+        write_string(val.as_ref(), buf)?;
     }
     Ok(())
 }