From 7f07d1ca23e31074ab002812e35f59e32d7edee6 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Wojciech=20Przytu=C5=82a?= <wojciech.przytula@scylladb.com>
Date: Wed, 10 Jul 2024 16:52:47 +0200
Subject: [PATCH] session_builder: expose custom_identity setter

The exposed option enabled setting custom identity that is advertised
in STARTUP message by the driver, upon opening each connection.

Settable identity include:
- DRIVER_NAME
- DRIVER_VERSION
- APPLICATION_NAME
- APPLICATION_VERSION
- CLIENT_ID

All five are visible in Cassandra's `system_views.clients` in
`client_options` column.

DRIVER_NAME & DRIVER_VERSION are visible in ScyllaDB's `system.clients`.
---
 scylla/src/transport/session.rs         |  8 ++++--
 scylla/src/transport/session_builder.rs | 33 +++++++++++++++++++++++++
 2 files changed, 39 insertions(+), 2 deletions(-)

diff --git a/scylla/src/transport/session.rs b/scylla/src/transport/session.rs
index dd9084468c..80a23dace8 100644
--- a/scylla/src/transport/session.rs
+++ b/scylla/src/transport/session.rs
@@ -289,6 +289,10 @@ pub struct SessionConfig {
     /// for e.g: if they do not want unexpected traffic
     /// or they expect the topology to change frequently.
     pub cluster_metadata_refresh_interval: Duration,
+
+    /// Driver and application self-identifying information,
+    /// to be sent to server in STARTUP message.
+    pub identity: SelfIdentity<'static>,
 }
 
 impl SessionConfig {
@@ -335,6 +339,7 @@ impl SessionConfig {
             tracing_info_fetch_interval: Duration::from_millis(3),
             tracing_info_fetch_consistency: Consistency::One,
             cluster_metadata_refresh_interval: Duration::from_secs(60),
+            identity: SelfIdentity::default(),
         }
     }
 
@@ -515,8 +520,7 @@ impl Session {
             keepalive_interval: config.keepalive_interval,
             keepalive_timeout: config.keepalive_timeout,
             tablet_sender: Some(tablet_sender),
-            // A temporary stub, removed in the next commit.
-            identity: SelfIdentity::default(),
+            identity: config.identity,
         };
 
         let pool_config = PoolConfig {
diff --git a/scylla/src/transport/session_builder.rs b/scylla/src/transport/session_builder.rs
index e697bb6154..998803793f 100644
--- a/scylla/src/transport/session_builder.rs
+++ b/scylla/src/transport/session_builder.rs
@@ -1,5 +1,6 @@
 //! SessionBuilder provides an easy way to create new Sessions
 
+use super::connection::SelfIdentity;
 use super::errors::NewSessionError;
 use super::execution_profile::ExecutionProfileHandle;
 use super::session::{AddressTranslator, Session, SessionConfig};
@@ -942,6 +943,38 @@ impl<K: SessionBuilderKind> GenericSessionBuilder<K> {
         self.config.cluster_metadata_refresh_interval = interval;
         self
     }
+
+    /// Set the custom identity of the driver/application/instance,
+    /// to be sent as options in STARTUP message.
+    ///
+    /// By default driver name and version are sent;
+    /// application name and version and client id are not sent.
+    ///
+    /// # Example
+    /// ```
+    /// # use scylla::{Session, SessionBuilder};
+    /// # use scylla::transport::SelfIdentity;
+    /// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
+    ///     let (app_major, app_minor, app_patch) = (2, 1, 3);
+    ///     let app_version = format!("{app_major}.{app_minor}.{app_patch}");
+    ///
+    ///     let session: Session = SessionBuilder::new()
+    ///         .known_node("127.0.0.1:9042")
+    ///         .custom_identity(
+    ///             SelfIdentity::new()
+    ///                 .with_custom_driver_version("0.13.0-custom_build_17")
+    ///                 .with_application_name("my-app")
+    ///                 .with_application_version(app_version)
+    ///         )
+    ///         .build()
+    ///         .await?;
+    /// #   Ok(())
+    /// # }
+    /// ```
+    pub fn custom_identity(mut self, identity: SelfIdentity<'static>) -> Self {
+        self.config.identity = identity;
+        self
+    }
 }
 
 /// Creates a [`SessionBuilder`] with default configuration, same as [`SessionBuilder::new`]