Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Set the default user-agent and allow to add product info(s) #135

Merged
merged 10 commits into from
Sep 17, 2024
35 changes: 35 additions & 0 deletions src/headers.rs
Original file line number Diff line number Diff line change
@@ -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));
loyd marked this conversation as resolved.
Show resolved Hide resolved
builder
}
6 changes: 2 additions & 4 deletions src/insert.rs
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@ use tokio::{
};
use url::Url;

use crate::headers::with_request_headers;
use crate::{
error::{Error, Result},
request_body::{ChunkSender, RequestBody},
Expand Down Expand Up @@ -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);
Expand Down
73 changes: 68 additions & 5 deletions src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -34,6 +33,7 @@ pub mod watch;
mod buflist;
mod compression;
mod cursor;
mod headers;
mod http_client;
mod request_body;
mod response;
Expand All @@ -60,6 +60,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 {
Expand Down Expand Up @@ -105,6 +118,7 @@ impl Client {
compression: Compression::default(),
options: HashMap::new(),
headers: HashMap::new(),
products_info: Vec::default(),
}
}

Expand Down Expand Up @@ -194,6 +208,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
Expand Down
6 changes: 2 additions & 4 deletions src/query.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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},
Expand Down Expand Up @@ -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");
Expand Down
1 change: 1 addition & 0 deletions tests/it/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -61,6 +61,7 @@ mod ip;
mod nested;
mod query;
mod time;
mod user_agent;
mod uuid;
mod watch;

Expand Down
81 changes: 81 additions & 0 deletions tests/it/user_agent.rs
Original file line number Diff line number Diff line change
@@ -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);
}