From b436bf7b2e9a181b07f334b904a71b58c3a476c0 Mon Sep 17 00:00:00 2001
From: Jonathan Giddy <jgiddy@cloudflare.com>
Date: Mon, 5 Aug 2024 19:03:55 +0100
Subject: [PATCH] Provide `with_header` method on client

This allows arbitrary headers to be included in requests to ClickHouse.
---
 src/insert.rs |  4 ++++
 src/lib.rs    | 18 ++++++++++++++++++
 src/query.rs  |  4 ++++
 3 files changed, 26 insertions(+)

diff --git a/src/insert.rs b/src/insert.rs
index e48ba9e..0ea351b 100644
--- a/src/insert.rs
+++ b/src/insert.rs
@@ -88,6 +88,10 @@ impl<T> Insert<T> {
 
         let mut builder = Request::post(url.as_str());
 
+        for (name, value) in &client.headers {
+            builder = builder.header(name, value);
+        }
+
         if let Some(user) = &client.user {
             builder = builder.header("X-ClickHouse-User", user);
         }
diff --git a/src/lib.rs b/src/lib.rs
index 02e4a5b..874ab85 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -59,6 +59,7 @@ pub struct Client {
     password: Option<String>,
     compression: Compression,
     options: HashMap<String, String>,
+    headers: HashMap<String, String>,
 }
 
 impl Default for Client {
@@ -103,6 +104,7 @@ impl Client {
             password: None,
             compression: Compression::default(),
             options: HashMap::new(),
+            headers: HashMap::new(),
         }
     }
 
@@ -180,6 +182,22 @@ impl Client {
         self
     }
 
+    /// Used to specify a header that will be passed to all queries.
+    ///
+    /// # Example
+    /// ```
+    /// # use clickhouse::Client;
+    /// Client::default().with_header(http::header::COOKIE, http::HeaderValue::from_static("A=1"));
+    /// ```
+    pub fn with_header(
+        mut self,
+        name: impl Into<String>,
+        value: impl Into<String>,
+    ) -> Self {
+        self.headers.insert(name.into(), value.into());
+        self
+    }
+
     /// Starts a new INSERT statement.
     ///
     /// # Panics
diff --git a/src/query.rs b/src/query.rs
index 20ed96b..28253f9 100644
--- a/src/query.rs
+++ b/src/query.rs
@@ -160,6 +160,10 @@ impl Query {
 
         let mut builder = Request::builder().method(method).uri(url.as_str());
 
+        for (name, value) in &self.client.headers {
+            builder = builder.header(name, value);
+        }
+
         if content_length == 0 {
             builder = builder.header(CONTENT_LENGTH, "0");
         } else {