Skip to content

Commit

Permalink
feat(client): set the default user-agent and allow to add product inf…
Browse files Browse the repository at this point in the history
…o(s) (#135)
  • Loading branch information
slvrtrn committed Sep 17, 2024
1 parent 6a87d25 commit f78cf31
Show file tree
Hide file tree
Showing 6 changed files with 189 additions and 13 deletions.
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));
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 Down Expand Up @@ -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 {
Expand Down Expand Up @@ -132,6 +145,7 @@ impl Client {
compression: Compression::default(),
options: HashMap::new(),
headers: HashMap::new(),
products_info: Vec::default(),
}
}

Expand Down Expand Up @@ -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
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 @@ -63,6 +63,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);
}

0 comments on commit f78cf31

Please sign in to comment.