diff --git a/scylla-rust-wrapper/src/cluster.rs b/scylla-rust-wrapper/src/cluster.rs index 1288cdc2..06b84b7a 100644 --- a/scylla-rust-wrapper/src/cluster.rs +++ b/scylla-rust-wrapper/src/cluster.rs @@ -7,6 +7,7 @@ use crate::retry_policy::CassRetryPolicy; use crate::retry_policy::RetryPolicy::*; use crate::ssl::CassSsl; use crate::types::*; +use crate::uuid::CassUuid; use openssl::ssl::SslContextBuilder; use openssl_sys::SSL_CTX_up_ref; use scylla::execution_profile::ExecutionProfileBuilder; @@ -16,6 +17,7 @@ use scylla::load_balancing::{DefaultPolicyBuilder, LoadBalancingPolicy}; use scylla::retry_policy::RetryPolicy; use scylla::speculative_execution::SimpleSpeculativeExecutionPolicy; use scylla::statement::{Consistency, SerialConsistency}; +use scylla::transport::SelfIdentity; use scylla::{SessionBuilder, SessionConfig}; use std::collections::HashMap; use std::convert::TryInto; @@ -32,6 +34,9 @@ include!(concat!(env!("OUT_DIR"), "/cppdriver_compression_types.rs")); const DEFAULT_CONSISTENCY: Consistency = Consistency::LocalOne; const DEFAULT_REQUEST_TIMEOUT_MILLIS: u64 = 12000; +const DRIVER_NAME: &str = "ScyllaDB Cpp-Rust Driver"; +const DRIVER_VERSION: &str = env!("CARGO_PKG_VERSION"); + #[derive(Clone, Debug)] pub(crate) struct LoadBalancingConfig { pub(crate) token_awareness_enabled: bool, @@ -85,6 +90,8 @@ pub struct CassCluster { use_beta_protocol_version: bool, auth_username: Option, auth_password: Option, + + client_id: Option, } impl CassCluster { @@ -106,6 +113,11 @@ impl CassCluster { pub(crate) fn get_contact_points(&self) -> &[String] { &self.contact_points } + + #[inline] + pub(crate) fn get_client_id(&self) -> Option { + self.client_id + } } pub struct CassCustomPayload; @@ -140,8 +152,13 @@ pub unsafe extern "C" fn cass_cluster_new() -> *mut CassCluster { .consistency(DEFAULT_CONSISTENCY) .request_timeout(Some(Duration::from_millis(DEFAULT_REQUEST_TIMEOUT_MILLIS))); + // Set DRIVER_NAME and DRIVER_VERSION of cpp-rust driver. + let custom_identity = SelfIdentity::new() + .with_custom_driver_name(DRIVER_NAME) + .with_custom_driver_version(DRIVER_VERSION); + Box::into_raw(Box::new(CassCluster { - session_builder: SessionBuilder::new(), + session_builder: SessionBuilder::new().custom_identity(custom_identity), port: 9042, contact_points: Vec::new(), // Per DataStax documentation: Without additional configuration the C/C++ driver @@ -152,6 +169,7 @@ pub unsafe extern "C" fn cass_cluster_new() -> *mut CassCluster { default_execution_profile_builder, execution_profile_map: Default::default(), load_balancing_config: Default::default(), + client_id: None, })) } @@ -219,6 +237,74 @@ pub unsafe extern "C" fn cass_cluster_set_use_randomized_contact_points( CassError::CASS_OK } +#[no_mangle] +pub unsafe extern "C" fn cass_cluster_set_application_name( + cluster_raw: *mut CassCluster, + app_name: *const c_char, +) { + cass_cluster_set_application_name_n(cluster_raw, app_name, strlen(app_name)) +} + +#[no_mangle] +pub unsafe extern "C" fn cass_cluster_set_application_name_n( + cluster_raw: *mut CassCluster, + app_name: *const c_char, + app_name_len: size_t, +) { + let cluster = ptr_to_ref_mut(cluster_raw); + let app_name = ptr_to_cstr_n(app_name, app_name_len).unwrap().to_string(); + + cluster + .session_builder + .config + .identity + .set_application_name(app_name) +} + +#[no_mangle] +pub unsafe extern "C" fn cass_cluster_set_application_version( + cluster_raw: *mut CassCluster, + app_version: *const c_char, +) { + cass_cluster_set_application_version_n(cluster_raw, app_version, strlen(app_version)) +} + +#[no_mangle] +pub unsafe extern "C" fn cass_cluster_set_application_version_n( + cluster_raw: *mut CassCluster, + app_version: *const c_char, + app_version_len: size_t, +) { + let cluster = ptr_to_ref_mut(cluster_raw); + let app_version = ptr_to_cstr_n(app_version, app_version_len) + .unwrap() + .to_string(); + + cluster + .session_builder + .config + .identity + .set_application_version(app_version); +} + +#[no_mangle] +pub unsafe extern "C" fn cass_cluster_set_client_id( + cluster_raw: *mut CassCluster, + client_id: CassUuid, +) { + let cluster = ptr_to_ref_mut(cluster_raw); + + let client_uuid: uuid::Uuid = client_id.into(); + let client_uuid_str = client_uuid.to_string(); + + cluster.client_id = Some(client_uuid); + cluster + .session_builder + .config + .identity + .set_client_id(client_uuid_str) +} + #[no_mangle] pub unsafe extern "C" fn cass_cluster_set_use_schema( cluster_raw: *mut CassCluster, diff --git a/scylla-rust-wrapper/src/session.rs b/scylla-rust-wrapper/src/session.rs index 2060308f..9e54737a 100644 --- a/scylla-rust-wrapper/src/session.rs +++ b/scylla-rust-wrapper/src/session.rs @@ -13,6 +13,7 @@ use crate::query_result::{CassResult, CassResultData, CassRow, CassValue, Collec use crate::statement::CassStatement; use crate::statement::Statement; use crate::types::{cass_uint64_t, size_t}; +use crate::uuid::CassUuid; use scylla::frame::response::result::{CqlValue, Row}; use scylla::frame::types::Consistency; use scylla::query::Query; @@ -31,6 +32,7 @@ use tokio::sync::RwLock; pub struct CassSessionInner { session: Session, exec_profile_map: HashMap, + client_id: uuid::Uuid, } impl CassSessionInner { @@ -82,6 +84,10 @@ impl CassSessionInner { session_opt, session_builder, exec_profile_map, + cluster + .get_client_id() + // If user did not set a client id, generate a random uuid v4. + .unwrap_or_else(uuid::Uuid::new_v4), keyspace, )) } @@ -90,6 +96,7 @@ impl CassSessionInner { session_opt: &RwLock>, session_builder_fut: impl Future, exec_profile_builder_map: HashMap, + client_id: uuid::Uuid, keyspace: Option, ) -> CassFutureResult { // This can sleep for a long time, but only if someone connects/closes session @@ -119,6 +126,7 @@ impl CassSessionInner { *session_guard = Some(CassSessionInner { session, exec_profile_map, + client_id, }); Ok(CassResultValue::Empty) } @@ -562,6 +570,14 @@ pub unsafe extern "C" fn cass_session_close(session: *mut CassSession) -> *const }) } +#[no_mangle] +pub unsafe extern "C" fn cass_session_get_client_id(session: *const CassSession) -> CassUuid { + let cass_session = ptr_to_ref(session); + + let client_id: uuid::Uuid = cass_session.blocking_read().as_ref().unwrap().client_id; + client_id.into() +} + #[no_mangle] pub unsafe extern "C" fn cass_session_get_schema_meta( session: *const CassSession,