diff --git a/src/headers.rs b/src/headers.rs new file mode 100644 index 0000000..d758854 --- /dev/null +++ b/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::>() + .join(" "); + format!("{products} {default_agent}") + } +} + +pub(crate) fn with_request_headers( + mut builder: Builder, + headers: &HashMap, + 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 +} diff --git a/src/insert.rs b/src/insert.rs index 902e38f..42208a3 100644 --- a/src/insert.rs +++ b/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 Insert { 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); diff --git a/src/lib.rs b/src/lib.rs index b364a91..692211f 100644 --- a/src/lib.rs +++ b/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, headers: HashMap, + products_info: Vec, +} + +#[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, + product_version: impl Into, + ) -> Self { + self.products_info.push(ProductInfo { + name: product_name.into(), + version: product_version.into(), + }); + self + } + /// Starts a new INSERT statement. /// /// # Panics diff --git a/src/query.rs b/src/query.rs index 551e7a0..5a03d17 100644 --- a/src/query.rs +++ b/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"); diff --git a/tests/it/main.rs b/tests/it/main.rs index 001cbd6..7534917 100644 --- a/tests/it/main.rs +++ b/tests/it/main.rs @@ -63,6 +63,7 @@ mod ip; mod nested; mod query; mod time; +mod user_agent; mod uuid; mod watch; diff --git a/tests/it/user_agent.rs b/tests/it/user_agent.rs new file mode 100644 index 0000000..9db71ff --- /dev/null +++ b/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::() + .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::() + .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); +}