Adding custom User-Agent to spice-rs (#48)

spiceai · Nov 21, 2024 · 4b2f65a · 4b2f65a
1 parent 9db5f08
commit 4b2f65a
Show file tree

Hide file tree

Showing 3 changed files with 32 additions and 9 deletions.
diff --git a/src/client.rs b/src/client.rs
@@ -41,7 +41,7 @@ pub struct SpiceClient {
 }
 
 impl SpiceClient {
-    /// Creates a new `SpiceClient` with the given API key.
+    /// Creates a new `SpiceClient` with the given API key and default user agent.
     /// ```
     /// use spiceai::Client;
     ///
@@ -58,8 +58,12 @@ impl SpiceClient {
         let config = SpiceClientConfig::load_from_default().await?;
 
         Ok(Self {
-            flight: SqlFlightClient::new(config.flight_channel, Some(api_key.to_string())),
-            firecache: SqlFlightClient::new(config.firecache_channel, Some(api_key.to_string())),
+            flight: SqlFlightClient::new(config.flight_channel, Some(api_key.to_string()), None),
+            firecache: SqlFlightClient::new(
+                config.firecache_channel,
+                Some(api_key.to_string()),
+                None,
+            ),
         })
     }
 
@@ -141,6 +145,7 @@ impl SpiceClient {
 ///
 pub struct SpiceClientBuilder {
     api_key: Option<String>,
+    user_agent: Option<String>,
     firecache_url: Option<String>,
     flight_url: Option<String>,
 }
@@ -156,6 +161,7 @@ impl SpiceClientBuilder {
     pub fn new() -> Self {
         Self {
             api_key: None,
+            user_agent: None,
             firecache_url: None,
             flight_url: None,
         }
@@ -168,6 +174,13 @@ impl SpiceClientBuilder {
         self
     }
 
+    /// Configures the `SpiceClient` to use the given custom user agent.
+    #[must_use]
+    pub fn user_agent(mut self, user_agent: &str) -> Self {
+        self.user_agent = Some(user_agent.to_string());
+        self
+    }
+
     /// Configures the `SpiceClient` to use the given Spice Firecache endpoint.
     #[must_use]
     pub fn firecache_url(mut self, firecache_url: &str) -> Self {
@@ -208,8 +221,16 @@ impl SpiceClientBuilder {
         };
 
         Ok(SpiceClient {
-            flight: SqlFlightClient::new(flight_channel, self.api_key.clone()),
-            firecache: SqlFlightClient::new(firecache_channel, self.api_key.clone()),
+            flight: SqlFlightClient::new(
+                flight_channel,
+                self.api_key.clone(),
+                self.user_agent.clone(),
+            ),
+            firecache: SqlFlightClient::new(
+                firecache_channel,
+                self.api_key.clone(),
+                self.user_agent.clone(),
+            ),
         })
     }
 }
diff --git a/src/config.rs b/src/config.rs
@@ -56,7 +56,7 @@ pub(crate) fn get_user_agent() -> String {
         .to_string();
 
     format!(
-        "spice-rs {} ({os_type}/{os_release} {os_arch})",
+        "spice-rs/{} ({os_type}/{os_release} {os_arch})",
         env!("CARGO_PKG_VERSION")
     )
 }
@@ -68,7 +68,7 @@ mod test {
     #[test]
     fn test_get_user_agent() {
         let matching_regex = regex::Regex::new(
-            r"spice-rs \d+\.\d+\.\d+ \((Linux|Windows|Darwin)/[\d\w\.\-\_]+ (x86_64|aarch64|i386)\)",
+            r"spice-rs/\d+\.\d+\.\d+ \((Linux|Windows|Darwin)/[\d\w\.\-\_]+ (x86_64|aarch64|i386)\)",
         )
         .expect("regex should be constructed");
 

diff --git a/src/flight.rs b/src/flight.rs
@@ -31,9 +31,11 @@ fn status_to_arrow_error(status: tonic::Status) -> ArrowError {
 }
 
 impl SqlFlightClient {
-    pub fn new(chan: Channel, api_key: Option<String>) -> Self {
+    pub fn new(chan: Channel, api_key: Option<String>, user_agent: Option<String>) -> Self {
+        let user_agent = user_agent.unwrap_or_else(get_user_agent);
+
         let mut headers = HashMap::new();
-        headers.insert("x-spice-user-agent".to_string(), get_user_agent());
+        headers.insert("User-Agent".to_string(), user_agent);
 
         SqlFlightClient {
             api_key,