Skip to content

Commit

Permalink
Merge pull request #1039 from wprzytula/self-identity-customisation
Browse files Browse the repository at this point in the history
connection: expose setters for STARTUP options
  • Loading branch information
wprzytula authored Aug 3, 2024
2 parents a707bfd + 9aadac8 commit 36d5edd
Show file tree
Hide file tree
Showing 11 changed files with 482 additions and 100 deletions.
14 changes: 10 additions & 4 deletions scylla-cql/src/frame/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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>,
}
Expand Down
11 changes: 6 additions & 5 deletions scylla-cql/src/frame/protocol_features.rs
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
use std::borrow::Cow;
use std::collections::HashMap;

const RATE_LIMIT_ERROR_EXTENSION: &str = "SCYLLA_RATE_LIMIT_ERROR";
Expand Down Expand Up @@ -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(""));
}
}

Expand Down
17 changes: 17 additions & 0 deletions scylla-cql/src/frame/request/options.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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");
8 changes: 4 additions & 4 deletions scylla-cql/src/frame/request/startup.rs
Original file line number Diff line number Diff line change
@@ -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> {
Expand Down
6 changes: 3 additions & 3 deletions scylla-cql/src/frame/types.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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(())
}
Expand Down
Loading

0 comments on commit 36d5edd

Please sign in to comment.