-
Notifications
You must be signed in to change notification settings - Fork 87
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
feat(client): set the default user-agent and allow to add product inf…
…o(s) (#135)
- Loading branch information
Showing
6 changed files
with
189 additions
and
13 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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 | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -63,6 +63,7 @@ mod ip; | |
mod nested; | ||
mod query; | ||
mod time; | ||
mod user_agent; | ||
mod uuid; | ||
mod watch; | ||
|
||
|
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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); | ||
} |