feat(client): set the default user-agent and allow to add product info(s) (#135)

Author: slvrtrn

src/headers.rs

@@ -0,0 +1,35 @@
+use crate::ProductInfo;
+use hyper::header::USER_AGENT;
+use hyper::http::request::Builder;
+use std::collections::HashMap;
+use std::env::consts::OS;
+
+fn get_user_agent(products_info: &[ProductInfo]) -> String {
+    // See https://doc.rust-lang.org/cargo/reference/environment-variables.html#environment-variables-cargo-sets-for-crates
+    let pkg_ver = option_env!("CARGO_PKG_VERSION").unwrap_or("unknown");
+    let rust_ver = option_env!("CARGO_PKG_RUST_VERSION").unwrap_or("unknown");
+    let default_agent = format!("clickhouse-rs/{pkg_ver} (lv:rust/{rust_ver}, os:{OS})");
+    if products_info.is_empty() {
+        default_agent
+    } else {
+        let products = products_info
+            .iter()
+            .rev()
+            .map(|product_info| product_info.to_string())
+            .collect::<Vec<String>>()
+            .join(" ");
+        format!("{products} {default_agent}")
+    }
+}
+
+pub(crate) fn with_request_headers(
+    mut builder: Builder,
+    headers: &HashMap<String, String>,
+    products_info: &[ProductInfo],
+) -> Builder {
+    for (name, value) in headers {
+        builder = builder.header(name, value);
+    }
+    builder = builder.header(USER_AGENT.to_string(), get_user_agent(products_info));
+    builder
+}

src/insert.rs

@@ -10,6 +10,7 @@ use tokio::{
 };
 use url::Url;
 
+use crate::headers::with_request_headers;
 use crate::{
     error::{Error, Result},
     request_body::{ChunkSender, RequestBody},
@@ -351,10 +352,7 @@ impl<T> Insert<T> {
         drop(pairs);
 
         let mut builder = Request::post(url.as_str());
-
-        for (name, value) in &client.headers {
-            builder = builder.header(name, value);
-        }
+        builder = with_request_headers(builder, &client.headers, &client.products_info);
 
         if let Some(user) = &client.user {
             builder = builder.header("X-ClickHouse-User", user);

src/lib.rs

@@ -5,19 +5,18 @@
 #[macro_use]
 extern crate static_assertions;
 
-use std::{collections::HashMap, sync::Arc, time::Duration};
-
+pub use clickhouse_derive::Row;
 #[cfg(feature = "tls")]
 use hyper_tls::HttpsConnector;
 use hyper_util::{
     client::legacy::{connect::HttpConnector, Client as HyperClient},
     rt::TokioExecutor,
 };
-
-use self::{error::Result, http_client::HttpClient};
+use std::fmt::Display;
+use std::{collections::HashMap, sync::Arc, time::Duration};
 
 pub use self::{compression::Compression, row::Row};
-pub use clickhouse_derive::Row;
+use self::{error::Result, http_client::HttpClient};
 
 pub mod error;
 pub mod insert;
@@ -34,6 +33,7 @@ pub mod watch;
 mod buflist;
 mod compression;
 mod cursor;
+mod headers;
 mod http_client;
 mod request_body;
 mod response;
@@ -76,6 +76,19 @@ pub struct Client {
     compression: Compression,
     options: HashMap<String, String>,
     headers: HashMap<String, String>,
+    products_info: Vec<ProductInfo>,
+}
+
+#[derive(Clone)]
+struct ProductInfo {
+    name: String,
+    version: String,
+}
+
+impl Display for ProductInfo {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{}/{}", self.name, self.version)
+    }
 }
 
 impl Default for Client {
@@ -132,6 +145,7 @@ impl Client {
             compression: Compression::default(),
             options: HashMap::new(),
             headers: HashMap::new(),
+            products_info: Vec::default(),
         }
     }
 
@@ -221,6 +235,55 @@ impl Client {
         self
     }
 
+    /// Specifies the product name and version that will be included
+    /// in the default User-Agent header. Multiple products are supported.
+    /// This could be useful for the applications built on top of this client.
+    ///
+    /// # Examples
+    ///
+    /// Sample default User-Agent header:
+    ///
+    /// ```plaintext
+    /// clickhouse-rs/0.12.2 (lv:rust/1.67.0, os:macos)
+    /// ```
+    ///
+    /// Sample User-Agent with a single product information:
+    ///
+    /// ```
+    /// # use clickhouse::Client;
+    /// let client = Client::default().with_product_info("MyDataSource", "v1.0.0");
+    /// ```
+    ///
+    /// ```plaintext
+    /// MyDataSource/v1.0.0 clickhouse-rs/0.12.2 (lv:rust/1.67.0, os:macos)
+    /// ```
+    ///
+    /// Sample User-Agent with multiple products information
+    /// (NB: the products are added in the reverse order of [`Client::with_product_info`] calls,
+    /// which could be useful to add higher abstraction layers first):
+    ///
+    /// ```
+    /// # use clickhouse::Client;
+    /// let client = Client::default()
+    ///     .with_product_info("MyDataSource", "v1.0.0")
+    ///     .with_product_info("MyApp", "0.0.1");
+    /// ```
+    ///
+    /// ```plaintext
+    /// MyApp/0.0.1 MyDataSource/v1.0.0 clickhouse-rs/0.12.2 (lv:rust/1.67.0, os:macos)
+    /// ```
+    pub fn with_product_info(
+        mut self,
+        product_name: impl Into<String>,
+        product_version: impl Into<String>,
+    ) -> Self {
+        self.products_info.push(ProductInfo {
+            name: product_name.into(),
+            version: product_version.into(),
+        });
+        self
+    }
+
     /// Starts a new INSERT statement.
     ///
     /// # Panics

src/query.rs

@@ -2,6 +2,7 @@ use hyper::{header::CONTENT_LENGTH, Method, Request};
 use serde::Deserialize;
 use url::Url;
 
+use crate::headers::with_request_headers;
 use crate::{
     cursor::RowBinaryCursor,
     error::{Error, Result},
@@ -159,10 +160,7 @@ impl Query {
         drop(pairs);
 
         let mut builder = Request::builder().method(method).uri(url.as_str());
-
-        for (name, value) in &self.client.headers {
-            builder = builder.header(name, value);
-        }
+        builder = with_request_headers(builder, &self.client.headers, &self.client.products_info);
 
         if content_length == 0 {
             builder = builder.header(CONTENT_LENGTH, "0");

tests/it/main.rs

@@ -63,6 +63,7 @@ mod ip;
 mod nested;
 mod query;
 mod time;
+mod user_agent;
 mod uuid;
 mod watch;
 

tests/it/user_agent.rs

@@ -0,0 +1,81 @@
+use crate::{create_simple_table, flush_query_log, SimpleRow};
+use clickhouse::sql::Identifier;
+use clickhouse::Client;
+
+const PKG_VER: &str = env!("CARGO_PKG_VERSION");
+const RUST_VER: &str = env!("CARGO_PKG_RUST_VERSION");
+const OS: &str = std::env::consts::OS;
+
+#[tokio::test]
+async fn default_user_agent() {
+    let table_name = "chrs_default_user_agent";
+    let client = prepare_database!();
+    let expected_user_agent = format!("clickhouse-rs/{PKG_VER} (lv:rust/{RUST_VER}, os:{OS})");
+    assert_queries_user_agents(&client, table_name, &expected_user_agent).await;
+}
+
+#[tokio::test]
+async fn user_agent_with_single_product_info() {
+    let table_name = "chrs_user_agent_with_single_product_info";
+    let client = prepare_database!().with_product_info("my-app", "0.1.0");
+    let expected_user_agent =
+        format!("my-app/0.1.0 clickhouse-rs/{PKG_VER} (lv:rust/{RUST_VER}, os:{OS})");
+    assert_queries_user_agents(&client, table_name, &expected_user_agent).await;
+}
+
+#[tokio::test]
+async fn user_agent_with_multiple_product_info() {
+    let table_name = "chrs_user_agent_with_multiple_product_info";
+    let client = prepare_database!()
+        .with_product_info("my-datasource", "2.5.0")
+        .with_product_info("my-app", "0.1.0");
+    let expected_user_agent = format!(
+        "my-app/0.1.0 my-datasource/2.5.0 clickhouse-rs/{PKG_VER} (lv:rust/{RUST_VER}, os:{OS})"
+    );
+    assert_queries_user_agents(&client, table_name, &expected_user_agent).await;
+}
+
+async fn assert_queries_user_agents(client: &Client, table_name: &str, expected_user_agent: &str) {
+    let row = SimpleRow::new(42, "foo");
+
+    create_simple_table(client, table_name).await;
+
+    let mut insert = client.insert(table_name).unwrap();
+    insert.write(&row).await.unwrap();
+    insert.end().await.unwrap();
+
+    let rows = client
+        .query("SELECT ?fields FROM ?")
+        .bind(Identifier(table_name))
+        .fetch_all::<SimpleRow>()
+        .await
+        .unwrap();
+
+    assert_eq!(rows.len(), 1);
+    assert_eq!(rows[0], row);
+
+    flush_query_log(client).await;
+
+    let recorded_user_agents = client
+        .query(&format!(
+            "
+            SELECT http_user_agent
+            FROM system.query_log
+            WHERE type = 'QueryFinish'
+            AND (
+              query LIKE 'SELECT%FROM%{table_name}%'
+              OR
+              query LIKE 'INSERT%INTO%{table_name}%'
+            )
+            ORDER BY event_time_microseconds DESC
+            LIMIT 2
+            "
+        ))
+        .fetch_all::<String>()
+        .await
+        .unwrap();
+
+    assert_eq!(recorded_user_agents.len(), 2);
+    assert_eq!(recorded_user_agents[0], expected_user_agent);
+    assert_eq!(recorded_user_agents[1], expected_user_agent);
+}