From 4c7f6d5c9b47c76048d23affa3dea64ff6bf4a16 Mon Sep 17 00:00:00 2001 From: Austin Bonander Date: Wed, 15 Oct 2025 14:44:56 -0700 Subject: [PATCH 1/6] feat: add `with_role` to `Client`, `Query`, `Insert{er}` --- src/insert.rs | 47 +++++++++++++++++++++++++++++++++++------------ src/inserter.rs | 38 ++++++++++++++++++++++++++++++++++++++ src/lib.rs | 46 +++++++++++++++++++++++++++++++++++++++++++--- src/query.rs | 22 ++++++++++++++++++++++ 4 files changed, 138 insertions(+), 15 deletions(-) diff --git a/src/insert.rs b/src/insert.rs index 89ad7ac3..7ef2557a 100644 --- a/src/insert.rs +++ b/src/insert.rs @@ -95,6 +95,14 @@ impl InsertState { } } + fn expect_client_mut(&mut self) -> &mut Client { + let Self::NotStarted { client, .. } = self else { + panic!("cannot modify client options while an insert is in-progress") + }; + + client + } + fn terminated(&mut self) { replace_with_or_abort(self, |_self| match _self { InsertState::NotStarted { .. } => InsertState::Completed, // empty insert @@ -102,17 +110,6 @@ impl InsertState { _ => unreachable!(), }); } - - fn with_option(&mut self, name: impl Into, value: impl Into) { - assert!(matches!(self, InsertState::NotStarted { .. })); - replace_with_or_abort(self, |_self| match _self { - InsertState::NotStarted { mut client, sql } => { - client.add_option(name, value); - InsertState::NotStarted { client, sql } - } - _ => unreachable!(), - }); - } } // It should be a regular function, but it decreases performance. @@ -185,6 +182,32 @@ impl Insert { self } + /// Set the [role] to use when executing `INSERT` statements. + /// + /// Overrides any role previously set by [`Client::with_role()`] or [`Client::with_option()`]. + /// + /// [role]: https://clickhouse.com/docs/operations/access-rights#role-management + /// + /// # Panics + /// If called after the request is started, e.g., after [`Insert::write`]. + pub fn with_role(mut self, role: impl Into) -> Self { + self.state.expect_client_mut().set_role(Some(role.into())); + self + } + + /// Perform inserts without any explicit [role] set. + /// + /// Overrides any role previously set by [`Client::with_role()`] or [`Client::with_option()`]. + /// + /// [role]: https://clickhouse.com/docs/operations/access-rights#role-management + /// + /// # Panics + /// If called after the request is started, e.g., after [`Insert::write`]. + pub fn with_default_role(mut self) -> Self { + self.state.expect_client_mut().set_role(None); + self + } + /// Similar to [`Client::with_option`], but for this particular INSERT /// statement only. /// @@ -192,7 +215,7 @@ impl Insert { /// If called after the request is started, e.g., after [`Insert::write`]. #[track_caller] pub fn with_option(mut self, name: impl Into, value: impl Into) -> Self { - self.state.with_option(name, value); + self.state.expect_client_mut().add_option(name, value); self } diff --git a/src/inserter.rs b/src/inserter.rs index ddd6af93..d30208b8 100644 --- a/src/inserter.rs +++ b/src/inserter.rs @@ -165,8 +165,46 @@ where self } + /// Set the [role] to use when executing `INSERT` statements. + /// + /// Overrides any role previously set by [`Client::with_role()`] or [`Client::with_option()`]. + /// + /// This does not take effect until the next `INSERT` statement begins + /// if one is already in-progress. + /// + /// If you have already begun writing data, you may call [`.force_commit()`][Self::force_commit] + /// to end the current `INSERT` so this takes effect on the next call to [`.write()`][Self::write]. + /// + /// [role]: https://clickhouse.com/docs/operations/access-rights#role-management + pub fn with_role(mut self, role: impl Into) -> Self { + self.client.set_role(Some(role.into())); + self + } + + /// Perform inserts without any explicit [role] set. + /// + /// Overrides any role previously set by [`Client::with_role()`] or [`Client::with_option()`]. + /// + /// This does not take effect until the next `INSERT` statement begins + /// if one is already in-progress. + /// + /// If you have already begun writing data, you may call [`.force_commit()`][Self::force_commit] + /// to end the current `INSERT` so this takes effect on the next call to [`.write()`][Self::write]. + /// + /// [role]: https://clickhouse.com/docs/operations/access-rights#role-management + pub fn with_default_role(mut self) -> Self { + self.client.set_role(None); + self + } + /// Similar to [`Client::with_option`], but for the INSERT statements /// generated by this [`Inserter`] only. + /// + /// This does not take effect until the next `INSERT` statement begins + /// if one is already in-progress. + /// + /// If you have already begun writing data, you may call [`.force_commit()`][Self::force_commit] + /// to end the current `INSERT` so this takes effect on the next call to [`.write()`][Self::write]. pub fn with_option(mut self, name: impl Into, value: impl Into) -> Self { self.client.add_option(name, value); self diff --git a/src/lib.rs b/src/lib.rs index e5c56b53..3233ce66 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -16,6 +16,7 @@ pub use clickhouse_macros::Row; use clickhouse_types::{Column, DataTypeNode}; use crate::_priv::row_insert_metadata_query; +use std::borrow::Cow; use std::{collections::HashMap, fmt::Display, sync::Arc}; use tokio::sync::RwLock; @@ -57,7 +58,7 @@ pub struct Client { database: Option, authentication: Authentication, compression: Compression, - options: HashMap, + options: HashMap, String>, headers: HashMap, products_info: Vec, validation: bool, @@ -227,6 +228,30 @@ impl Client { self } + /// Set the [role] to use when executing statements with this `Client` instance. + /// + /// Overrides any role previously set by [`Client::with_role()`] or [`Client::with_option()`]. + /// + /// This setting is copied into cloned clients. + /// + /// [role]: https://clickhouse.com/docs/operations/access-rights#role-management + pub fn with_role(mut self, role: impl Into) -> Self { + self.set_role(Some(role.into())); + self + } + + /// Execute subsequent statements with this `Client` instance without any explicit [role] set. + /// + /// Overrides any role previously set by [`Client::with_role()`] or [`Client::with_option()`]. + /// + /// This setting is copied into cloned clients. + /// + /// [role]: https://clickhouse.com/docs/operations/access-rights#role-management + pub fn with_default_role(mut self) -> Self { + self.set_role(None); + self + } + /// A JWT access token to authenticate with ClickHouse. /// JWT token authentication is supported in ClickHouse Cloud only. /// Should not be called after [`Client::with_user`] or @@ -278,7 +303,7 @@ impl Client { /// Client::default().with_option("allow_nondeterministic_mutations", "1"); /// ``` pub fn with_option(mut self, name: impl Into, value: impl Into) -> Self { - self.options.insert(name.into(), value.into()); + self.options.insert(Cow::Owned(name.into()), value.into()); self } @@ -447,7 +472,22 @@ impl Client { /// Used internally to modify the options map of an _already cloned_ /// [`Client`] instance. pub(crate) fn add_option(&mut self, name: impl Into, value: impl Into) { - self.options.insert(name.into(), value.into()); + self.options.insert(Cow::Owned(name.into()), value.into()); + } + + pub(crate) fn set_role(&mut self, role: Option) { + // By setting the role via `options`, we can be sure we've overwritten any role + // manually set by the user with `with_option()`. + let key = Cow::Borrowed("role"); + + match role { + Some(role) => { + self.options.insert(key, role); + } + None => { + self.options.remove(&key); + } + } } /// Use a mock server for testing purposes. diff --git a/src/query.rs b/src/query.rs index ee9c4b80..afd8e361 100644 --- a/src/query.rs +++ b/src/query.rs @@ -39,6 +39,28 @@ impl Query { &self.sql } + /// Perform this query with the given [role] set. + /// + /// Overrides any role previously set + /// + /// [role]: https://clickhouse.com/docs/operations/access-rights#role-management + pub fn with_role(self, role: impl Into) -> Self { + Self { + client: self.client.with_role(role), + ..self + } + } + + /// Perform this query without any explicit [role] set. + /// + /// [role]: https://clickhouse.com/docs/operations/access-rights#role-management + pub fn with_default_role(self) -> Self { + Self { + client: self.client.with_default_role(), + ..self + } + } + /// Binds `value` to the next `?` in the query. /// /// The `value`, which must either implement [`Serialize`] or be an From f1dcf48fb00ef8bd9a884a9ac44d9fa20077810c Mon Sep 17 00:00:00 2001 From: Austin Bonander Date: Fri, 17 Oct 2025 13:56:18 -0700 Subject: [PATCH 2/6] feat: add `with_roles` to `Client`, `Query`, `Insert{er}` --- src/insert.rs | 25 +++++++-------------- src/inserter.rs | 29 ++++++++---------------- src/lib.rs | 60 ++++++++++++++++++++++++------------------------- src/query.rs | 40 +++++++++++++++------------------ 4 files changed, 65 insertions(+), 89 deletions(-) diff --git a/src/insert.rs b/src/insert.rs index 7ef2557a..91ecd91c 100644 --- a/src/insert.rs +++ b/src/insert.rs @@ -95,6 +95,7 @@ impl InsertState { } } + #[inline] fn expect_client_mut(&mut self) -> &mut Client { let Self::NotStarted { client, .. } = self else { panic!("cannot modify client options while an insert is in-progress") @@ -182,29 +183,19 @@ impl Insert { self } - /// Set the [role] to use when executing `INSERT` statements. + /// Configure the [roles] to use when executing `INSERT` statements. /// - /// Overrides any role previously set by [`Client::with_role()`] or [`Client::with_option()`]. + /// Overrides any roles previously set by this method, [`Insert::with_option`], + /// [`Client::with_roles`] or [`Client::with_option`]. /// - /// [role]: https://clickhouse.com/docs/operations/access-rights#role-management + /// An empty iterator may be passed to clear the set roles. /// - /// # Panics - /// If called after the request is started, e.g., after [`Insert::write`]. - pub fn with_role(mut self, role: impl Into) -> Self { - self.state.expect_client_mut().set_role(Some(role.into())); - self - } - - /// Perform inserts without any explicit [role] set. - /// - /// Overrides any role previously set by [`Client::with_role()`] or [`Client::with_option()`]. - /// - /// [role]: https://clickhouse.com/docs/operations/access-rights#role-management + /// [roles]: https://clickhouse.com/docs/operations/access-rights#role-management /// /// # Panics /// If called after the request is started, e.g., after [`Insert::write`]. - pub fn with_default_role(mut self) -> Self { - self.state.expect_client_mut().set_role(None); + pub fn with_roles(mut self, roles: impl IntoIterator>) -> Self { + self.state.expect_client_mut().set_roles(roles); self } diff --git a/src/inserter.rs b/src/inserter.rs index d30208b8..f5020a1d 100644 --- a/src/inserter.rs +++ b/src/inserter.rs @@ -165,41 +165,30 @@ where self } - /// Set the [role] to use when executing `INSERT` statements. + /// Set the [roles] to use when executing `INSERT` statements. /// - /// Overrides any role previously set by [`Client::with_role()`] or [`Client::with_option()`]. + /// Overrides any roles previously set by this method, [`Inserter::with_option`], + /// [`Client::with_roles`] or [`Client::with_option`]. /// - /// This does not take effect until the next `INSERT` statement begins - /// if one is already in-progress. - /// - /// If you have already begun writing data, you may call [`.force_commit()`][Self::force_commit] - /// to end the current `INSERT` so this takes effect on the next call to [`.write()`][Self::write]. - /// - /// [role]: https://clickhouse.com/docs/operations/access-rights#role-management - pub fn with_role(mut self, role: impl Into) -> Self { - self.client.set_role(Some(role.into())); - self - } - - /// Perform inserts without any explicit [role] set. - /// - /// Overrides any role previously set by [`Client::with_role()`] or [`Client::with_option()`]. + /// An empty iterator may be passed to clear the set roles. /// + /// # Note /// This does not take effect until the next `INSERT` statement begins /// if one is already in-progress. /// /// If you have already begun writing data, you may call [`.force_commit()`][Self::force_commit] /// to end the current `INSERT` so this takes effect on the next call to [`.write()`][Self::write]. /// - /// [role]: https://clickhouse.com/docs/operations/access-rights#role-management - pub fn with_default_role(mut self) -> Self { - self.client.set_role(None); + /// [roles]: https://clickhouse.com/docs/operations/access-rights#role-management + pub fn with_roles(mut self, roles: impl IntoIterator>) -> Self { + self.client.set_roles(roles); self } /// Similar to [`Client::with_option`], but for the INSERT statements /// generated by this [`Inserter`] only. /// + /// # Note /// This does not take effect until the next `INSERT` statement begins /// if one is already in-progress. /// diff --git a/src/lib.rs b/src/lib.rs index 3233ce66..c5f70904 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -16,7 +16,7 @@ pub use clickhouse_macros::Row; use clickhouse_types::{Column, DataTypeNode}; use crate::_priv::row_insert_metadata_query; -use std::borrow::Cow; +use std::collections::HashSet; use std::{collections::HashMap, fmt::Display, sync::Arc}; use tokio::sync::RwLock; @@ -58,7 +58,8 @@ pub struct Client { database: Option, authentication: Authentication, compression: Compression, - options: HashMap, String>, + roles: HashSet, + options: HashMap, headers: HashMap, products_info: Vec, validation: bool, @@ -122,6 +123,7 @@ impl Client { database: None, authentication: Authentication::default(), compression: Compression::default(), + roles: HashSet::new(), options: HashMap::new(), headers: HashMap::new(), products_info: Vec::default(), @@ -228,27 +230,32 @@ impl Client { self } - /// Set the [role] to use when executing statements with this `Client` instance. + /// Configure the [roles] to use when executing statements with this `Client` instance. /// - /// Overrides any role previously set by [`Client::with_role()`] or [`Client::with_option()`]. + /// Overrides any roles previously set by this method or [`Client::with_option`]. + /// + /// An empty iterator may be passed to clear the set roles. /// /// This setting is copied into cloned clients. /// - /// [role]: https://clickhouse.com/docs/operations/access-rights#role-management - pub fn with_role(mut self, role: impl Into) -> Self { - self.set_role(Some(role.into())); - self - } - - /// Execute subsequent statements with this `Client` instance without any explicit [role] set. + /// [roles]: https://clickhouse.com/docs/operations/access-rights#role-management /// - /// Overrides any role previously set by [`Client::with_role()`] or [`Client::with_option()`]. + /// # Examples /// - /// This setting is copied into cloned clients. + /// ``` + /// # use clickhouse::Client; + /// + /// # Single role + /// let client = Client::default().with_roles(["foo"]); /// - /// [role]: https://clickhouse.com/docs/operations/access-rights#role-management - pub fn with_default_role(mut self) -> Self { - self.set_role(None); + /// # Multiple roles + /// let client = Client::default().with_roles(["foo", "bar", "baz"]); + /// + /// # Clear all previously set roles + /// let client = Client::default().with_roles([]); + /// ``` + pub fn with_roles(mut self, roles: impl IntoIterator>) -> Self { + self.set_roles(roles); self } @@ -303,7 +310,7 @@ impl Client { /// Client::default().with_option("allow_nondeterministic_mutations", "1"); /// ``` pub fn with_option(mut self, name: impl Into, value: impl Into) -> Self { - self.options.insert(Cow::Owned(name.into()), value.into()); + self.options.insert(name.into(), value.into()); self } @@ -472,22 +479,15 @@ impl Client { /// Used internally to modify the options map of an _already cloned_ /// [`Client`] instance. pub(crate) fn add_option(&mut self, name: impl Into, value: impl Into) { - self.options.insert(Cow::Owned(name.into()), value.into()); + self.options.insert(name.into(), value.into()); } - pub(crate) fn set_role(&mut self, role: Option) { - // By setting the role via `options`, we can be sure we've overwritten any role - // manually set by the user with `with_option()`. - let key = Cow::Borrowed("role"); + pub(crate) fn set_roles(&mut self, roles: impl IntoIterator>) { + // Make sure we overwrite any role manually set by the user via `with_option()`. + self.options.remove("role"); - match role { - Some(role) => { - self.options.insert(key, role); - } - None => { - self.options.remove(&key); - } - } + self.roles.clear(); + self.roles.extend(roles.into_iter().map(Into::into)); } /// Use a mock server for testing purposes. diff --git a/src/query.rs b/src/query.rs index afd8e361..e5cf45e8 100644 --- a/src/query.rs +++ b/src/query.rs @@ -39,28 +39,6 @@ impl Query { &self.sql } - /// Perform this query with the given [role] set. - /// - /// Overrides any role previously set - /// - /// [role]: https://clickhouse.com/docs/operations/access-rights#role-management - pub fn with_role(self, role: impl Into) -> Self { - Self { - client: self.client.with_role(role), - ..self - } - } - - /// Perform this query without any explicit [role] set. - /// - /// [role]: https://clickhouse.com/docs/operations/access-rights#role-management - pub fn with_default_role(self) -> Self { - Self { - client: self.client.with_default_role(), - ..self - } - } - /// Binds `value` to the next `?` in the query. /// /// The `value`, which must either implement [`Serialize`] or be an @@ -211,6 +189,9 @@ impl Query { for (name, value) in &self.client.options { pairs.append_pair(name, value); } + + pairs.extend_pairs(self.client.roles.iter().map(|role| ("role", role))); + drop(pairs); let mut builder = Request::builder().method(method).uri(url.as_str()); @@ -231,6 +212,21 @@ impl Query { Ok(Response::new(future, self.client.compression)) } + /// Configure the [roles] to use when executing this query. + /// + /// Overrides any roles previously set by this method, [`Query::with_option`], + /// [`Client::with_roles`] or [`Client::with_option`]. + /// + /// An empty iterator may be passed to clear the set roles. + /// + /// [roles]: https://clickhouse.com/docs/operations/access-rights#role-management + pub fn with_roles(self, roles: impl IntoIterator>) -> Self { + Self { + client: self.client.with_roles(roles), + ..self + } + } + /// Similar to [`Client::with_option`], but for this particular query only. pub fn with_option(mut self, name: impl Into, value: impl Into) -> Self { self.client.add_option(name, value); From d411e7a5844e57eb8347398ff98be7d3b95d1169 Mon Sep 17 00:00:00 2001 From: Austin Bonander Date: Fri, 17 Oct 2025 14:29:22 -0700 Subject: [PATCH 3/6] feat: test `Query::with_roles()` --- tests/it/main.rs | 57 +++++++++++++++++++++++++++++++++++++ tests/it/query.rs | 71 +++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 128 insertions(+) diff --git a/tests/it/main.rs b/tests/it/main.rs index 0b2923ed..c3ef3a63 100644 --- a/tests/it/main.rs +++ b/tests/it/main.rs @@ -314,6 +314,63 @@ async fn create_readonly_user(client: &Client, database: &str) -> Client { .with_database(database) } +/// Create a test user and role for `test_db_name` with no default grants. +/// +/// Returns a `Client` with the user configured, and the role name. +async fn create_user_and_role(client: &Client, test_db_name: &str) -> (Client, String) { + let username = format!("{test_db_name}__user"); + let password = format!("CHRS_{:X}", rand::random::()); + + client + .query( + "CREATE USER OR REPLACE ? \ + IDENTIFIED WITH sha256_password BY ? \ + DEFAULT DATABASE ?", + ) + .bind(&username) + .bind(&password) + .bind(Identifier(test_db_name)) + .execute() + .await + .unwrap(); + + client + .query("REVOKE ALL ON *.* FROM ?") + .bind(&username) + .execute() + .await + .unwrap(); + + let role = format!("{test_db_name}__role"); + + client + .query("CREATE ROLE OR REPLACE ?") + .bind(Identifier(&role)) + .execute() + .await + .unwrap(); + + client + .query("GRANT ? TO ?") + .bind(Identifier(&role)) + .bind(Identifier(&username)) + .execute() + .await + .unwrap(); + + client + .query("SET DEFAULT ROLE NONE TO ?") + .bind(Identifier(&username)) + .execute() + .await + .unwrap(); + + ( + client.clone().with_user(username).with_password(password), + role, + ) +} + mod _priv { use super::*; use std::time::SystemTime; diff --git a/tests/it/query.rs b/tests/it/query.rs index 8e282d22..dbddaf3c 100644 --- a/tests/it/query.rs +++ b/tests/it/query.rs @@ -1,5 +1,6 @@ use serde::{Deserialize, Serialize}; +use clickhouse::sql::Identifier; use clickhouse::{Row, error::Error}; #[tokio::test] @@ -263,3 +264,73 @@ async fn prints_query() { "SELECT ?fields FROM test WHERE a = ? AND b < ?" ); } + +#[tokio::test] +async fn query_with_role() { + let db_name = test_database_name!(); + + let admin_client = crate::_priv::prepare_database(&db_name).await; + + let (user_client, role) = crate::create_user_and_role(&admin_client, &db_name).await; + + admin_client + .query( + "CREATE TABLE foo(\ + bar DateTime DEFAULT now(), \ + baz String\ + ) \ + ENGINE = MergeTree \ + PRIMARY KEY(bar)", + ) + .execute() + .await + .unwrap(); + + admin_client + .query("INSERT INTO foo(baz) VALUES ('lorem ipsum'), ('dolor sit amet')") + .execute() + .await + .unwrap(); + + user_client + .query("SELECT * FROM foo") + .execute() + .await + .expect_err("user should not be able to query `foo`"); + + admin_client + .query("GRANT SELECT ON ?.foo TO ?") + .bind(Identifier(&db_name)) + .bind(Identifier(&role)) + .execute() + .await + .unwrap(); + + user_client + .query("SELECT * FROM foo") + .execute() + .await + .expect_err("user should not be able to query `foo`"); + + user_client + .clone() + .with_roles([&role]) + .query("SELECT * FROM foo") + .execute() + .await + .expect("user should be able to query `foo` now"); + + // Roles should not have propagated back to parent instance + user_client + .query("SELECT * FROM foo") + .execute() + .await + .expect_err("user should not be able to query `foo`"); + + user_client + .query("SELECT * FROM foo") + .with_roles([&role]) + .execute() + .await + .expect("user should be able to query `foo` now"); +} From 5c9a780d4b3bca13f54896cfc58e777581393420 Mon Sep 17 00:00:00 2001 From: Austin Bonander Date: Fri, 17 Oct 2025 15:40:39 -0700 Subject: [PATCH 4/6] feat: add `with_default_roles` to `Client`, `Query`, `Insert{er}` --- src/insert.rs | 14 ++++++++++++++ src/inserter.rs | 26 ++++++++++++++++++++++---- src/lib.rs | 27 +++++++++++++++++++-------- src/query.rs | 13 +++++++++++++ tests/it/query.rs | 10 ++++++++++ 5 files changed, 78 insertions(+), 12 deletions(-) diff --git a/src/insert.rs b/src/insert.rs index 91ecd91c..82449c56 100644 --- a/src/insert.rs +++ b/src/insert.rs @@ -199,6 +199,20 @@ impl Insert { self } + /// Clear any explicit [roles] previously set on this `Insert` or inherited from [`Client`]. + /// + /// Overrides any roles previously set by [`Insert::with_roles`], [`Insert::with_option`], + /// [`Client::with_roles`] or [`Client::with_option`]. + /// + /// [roles]: https://clickhouse.com/docs/operations/access-rights#role-management + /// + /// # Panics + /// If called after the request is started, e.g., after [`Insert::write`]. + pub fn with_default_roles(mut self) -> Self { + self.state.expect_client_mut().clear_roles(); + self + } + /// Similar to [`Client::with_option`], but for this particular INSERT /// statement only. /// diff --git a/src/inserter.rs b/src/inserter.rs index f5020a1d..071866f6 100644 --- a/src/inserter.rs +++ b/src/inserter.rs @@ -176,8 +176,8 @@ where /// This does not take effect until the next `INSERT` statement begins /// if one is already in-progress. /// - /// If you have already begun writing data, you may call [`.force_commit()`][Self::force_commit] - /// to end the current `INSERT` so this takes effect on the next call to [`.write()`][Self::write]. + /// If you have already begun writing data, you may call [`Inserter::force_commit`] + /// to end the current `INSERT` so this takes effect on the next call to [`Inserter::write`]. /// /// [roles]: https://clickhouse.com/docs/operations/access-rights#role-management pub fn with_roles(mut self, roles: impl IntoIterator>) -> Self { @@ -185,6 +185,24 @@ where self } + /// Clear any explicit [roles] previously set on this `Inserter` or inherited from [`Client`]. + /// + /// Overrides any roles previously set by [`Inserter::with_roles`], [`Inserter::with_option`], + /// [`Client::with_roles`] or [`Client::with_option`]. + /// + /// # Note + /// This does not take effect until the next `INSERT` statement begins + /// if one is already in-progress. + /// + /// If you have already begun writing data, you may call [`Inserter::force_commit`] + /// to end the current `INSERT` so this takes effect on the next call to [`Inserter::write`]. + /// + /// [roles]: https://clickhouse.com/docs/operations/access-rights#role-management + pub fn with_default_roles(mut self) -> Self { + self.client.clear_roles(); + self + } + /// Similar to [`Client::with_option`], but for the INSERT statements /// generated by this [`Inserter`] only. /// @@ -192,8 +210,8 @@ where /// This does not take effect until the next `INSERT` statement begins /// if one is already in-progress. /// - /// If you have already begun writing data, you may call [`.force_commit()`][Self::force_commit] - /// to end the current `INSERT` so this takes effect on the next call to [`.write()`][Self::write]. + /// If you have already begun writing data, you may call [`Inserter::force_commit`] + /// to end the current `INSERT` so this takes effect on the next call to [`Inserter::write`]. pub fn with_option(mut self, name: impl Into, value: impl Into) -> Self { self.client.add_option(name, value); self diff --git a/src/lib.rs b/src/lib.rs index c5f70904..126cd131 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -234,7 +234,7 @@ impl Client { /// /// Overrides any roles previously set by this method or [`Client::with_option`]. /// - /// An empty iterator may be passed to clear the set roles. + /// Call [`Client::with_default_roles`] to clear any explicitly set roles. /// /// This setting is copied into cloned clients. /// @@ -245,20 +245,27 @@ impl Client { /// ``` /// # use clickhouse::Client; /// - /// # Single role + /// // Single role /// let client = Client::default().with_roles(["foo"]); /// - /// # Multiple roles + /// // Multiple roles /// let client = Client::default().with_roles(["foo", "bar", "baz"]); - /// - /// # Clear all previously set roles - /// let client = Client::default().with_roles([]); /// ``` pub fn with_roles(mut self, roles: impl IntoIterator>) -> Self { self.set_roles(roles); self } + /// Clear any explicitly set [roles] from this `Client` instance. + /// + /// Overrides any roles previously set by [`Client::with_roles`] or [`Client::with_option`]. + /// + /// [roles]: https://clickhouse.com/docs/operations/access-rights#role-management + pub fn with_default_roles(mut self) -> Self { + self.clear_roles(); + self + } + /// A JWT access token to authenticate with ClickHouse. /// JWT token authentication is supported in ClickHouse Cloud only. /// Should not be called after [`Client::with_user`] or @@ -483,11 +490,15 @@ impl Client { } pub(crate) fn set_roles(&mut self, roles: impl IntoIterator>) { + self.clear_roles(); + self.roles.extend(roles.into_iter().map(Into::into)); + } + + #[inline] + pub(crate) fn clear_roles(&mut self) { // Make sure we overwrite any role manually set by the user via `with_option()`. self.options.remove("role"); - self.roles.clear(); - self.roles.extend(roles.into_iter().map(Into::into)); } /// Use a mock server for testing purposes. diff --git a/src/query.rs b/src/query.rs index e5cf45e8..f27e69be 100644 --- a/src/query.rs +++ b/src/query.rs @@ -227,6 +227,19 @@ impl Query { } } + /// Clear any explicit [roles] previously set on this `Query` or inherited from [`Client`]. + /// + /// Overrides any roles previously set by [`Query::with_roles`], [`Query::with_option`], + /// [`Client::with_roles`] or [`Client::with_option`]. + /// + /// [roles]: https://clickhouse.com/docs/operations/access-rights#role-management + pub fn with_default_roles(self) -> Self { + Self { + client: self.client.with_default_roles(), + ..self + } + } + /// Similar to [`Client::with_option`], but for this particular query only. pub fn with_option(mut self, name: impl Into, value: impl Into) -> Self { self.client.add_option(name, value); diff --git a/tests/it/query.rs b/tests/it/query.rs index dbddaf3c..0a359e78 100644 --- a/tests/it/query.rs +++ b/tests/it/query.rs @@ -327,6 +327,16 @@ async fn query_with_role() { .await .expect_err("user should not be able to query `foo`"); + // Test `with_default_roles()` + user_client + .clone() + .with_roles([&role]) + .query("SELECT * FROM foo") + .with_default_roles() + .execute() + .await + .expect_err("user should not be able to query `foo`"); + user_client .query("SELECT * FROM foo") .with_roles([&role]) From ad4b18828c338c39a4150a4a8e92504704a053aa Mon Sep 17 00:00:00 2001 From: Austin Bonander Date: Fri, 17 Oct 2025 17:37:06 -0700 Subject: [PATCH 5/6] feat: test `Insert::with_roles` --- tests/it/insert.rs | 100 +++++++++++++++++++++++++++++++++++++++++++++ tests/it/main.rs | 9 ++++ 2 files changed, 109 insertions(+) diff --git a/tests/it/insert.rs b/tests/it/insert.rs index 4f43e7c4..32bad18b 100644 --- a/tests/it/insert.rs +++ b/tests/it/insert.rs @@ -1,4 +1,5 @@ use crate::{SimpleRow, create_simple_table, fetch_rows, flush_query_log}; +use clickhouse::insert::Insert; use clickhouse::{Row, sql::Identifier}; use serde::{Deserialize, Serialize}; use std::panic::AssertUnwindSafe; @@ -424,3 +425,102 @@ async fn clear_cached_metadata() { assert_eq!(*rows, [Foo2 { bar: 1 }, Foo2 { bar: 3 }]); } + +#[tokio::test] +async fn insert_with_role() { + #[derive(serde::Serialize, serde::Deserialize, clickhouse::Row)] + struct Foo { + bar: u64, + baz: String, + } + + let db_name = test_database_name!(); + + let admin_client = crate::_priv::prepare_database(&db_name).await; + + let (user_client, role) = crate::create_user_and_role(&admin_client, &db_name).await; + + admin_client + .query( + "CREATE TABLE foo(\ + bar UInt64, \ + baz String\ + ) \ + ENGINE = MergeTree \ + PRIMARY KEY(bar)", + ) + .execute() + .await + .unwrap(); + + let foos = [ + "lorem ipsum", + "dolor sit amet", + "consectetur adipiscing elit", + ] + .into_iter() + .enumerate() + .map(|(bar, baz)| Foo { + bar: bar as u64, + baz: baz.to_string(), + }) + .collect::>(); + + let insert_foos = async |mut insert: Insert| { + for foo in &foos { + insert.write(foo).await?; + } + + insert.end().await + }; + + insert_foos(user_client.insert("foo").await.unwrap()) + .await + .expect_err("user should not be able to insert into `foo`"); + + admin_client + .query("GRANT INSERT ON ?.foo TO ?") + .bind(Identifier(&db_name)) + .bind(Identifier(&role)) + .execute() + .await + .unwrap(); + + // We haven't set the role yet + insert_foos(user_client.insert("foo").await.unwrap()) + .await + .expect_err("user should not be able to insert into `foo`"); + + insert_foos( + user_client + .clone() + .with_roles([&role]) + .insert("foo") + .await + .unwrap(), + ) + .await + .expect_err("user should be able to insert into `foo` now"); + + // Roles should not propagate back to the parent instance + insert_foos(user_client.insert("foo").await.unwrap()) + .await + .expect_err("user should not be able to insert into `foo`"); + + insert_foos(user_client.insert("foo").await.unwrap().with_roles([&role])) + .await + .expect_err("user should be able to insert into `foo` now"); + + // `with_default_roles` should clear the role + insert_foos( + user_client + .clone() + .with_roles([&role]) + .insert("foo") + .await + .unwrap() + .with_default_roles(), + ) + .await + .expect_err("user should not be able to insert into `foo`"); +} diff --git a/tests/it/main.rs b/tests/it/main.rs index c3ef3a63..3d02e709 100644 --- a/tests/it/main.rs +++ b/tests/it/main.rs @@ -341,6 +341,15 @@ async fn create_user_and_role(client: &Client, test_db_name: &str) -> (Client, S .await .unwrap(); + // Needed for metadata queries + client + .query("GRANT SHOW ON ?.* TO ?") + .bind(Identifier(test_db_name)) + .bind(&username) + .execute() + .await + .unwrap(); + let role = format!("{test_db_name}__role"); client From 23dc690cd709d808e7e2709e70df27452f73161f Mon Sep 17 00:00:00 2001 From: Austin Bonander Date: Mon, 20 Oct 2025 08:18:24 -0700 Subject: [PATCH 6/6] feat: test `Inserter::with_role()` --- tests/it/inserter.rs | 95 +++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 93 insertions(+), 2 deletions(-) diff --git a/tests/it/inserter.rs b/tests/it/inserter.rs index 3f7c1ffe..0c9487a6 100644 --- a/tests/it/inserter.rs +++ b/tests/it/inserter.rs @@ -4,9 +4,10 @@ use std::string::ToString; use serde::Serialize; -use clickhouse::{Client, Row, inserter::Quantities}; - use crate::{SimpleRow, create_simple_table, fetch_rows, flush_query_log}; +use clickhouse::inserter::Inserter; +use clickhouse::sql::Identifier; +use clickhouse::{Client, Row, inserter::Quantities}; #[derive(Debug, Row, Serialize)] struct MyRow { @@ -291,3 +292,93 @@ async fn overrides_client_options() { let rows = fetch_rows::(&client, table_name).await; assert_eq!(rows, vec!(row)) } + +#[tokio::test] +async fn inserter_with_role() { + #[derive(serde::Serialize, serde::Deserialize, clickhouse::Row)] + struct Foo { + bar: u64, + baz: String, + } + + let db_name = test_database_name!(); + + let admin_client = crate::_priv::prepare_database(&db_name).await; + + let (user_client, role) = crate::create_user_and_role(&admin_client, &db_name).await; + + admin_client + .query( + "CREATE TABLE foo(\ + bar UInt64, \ + baz String\ + ) \ + ENGINE = MergeTree \ + PRIMARY KEY(bar)", + ) + .execute() + .await + .unwrap(); + + let foos = [ + "lorem ipsum", + "dolor sit amet", + "consectetur adipiscing elit", + ] + .into_iter() + .enumerate() + .map(|(bar, baz)| Foo { + bar: bar as u64, + baz: baz.to_string(), + }) + .collect::>(); + + let insert_foos = async |mut inserter: Inserter| { + for foo in &foos { + inserter.write(foo).await?; + } + + inserter.end().await + }; + + insert_foos(user_client.inserter("foo")) + .await + .expect_err("user should not be able to insert into `foo`"); + + admin_client + .query("GRANT INSERT ON ?.foo TO ?") + .bind(Identifier(&db_name)) + .bind(Identifier(&role)) + .execute() + .await + .unwrap(); + + // We haven't set the role yet + insert_foos(user_client.inserter("foo")) + .await + .expect_err("user should not be able to insert into `foo`"); + + insert_foos(user_client.clone().with_roles([&role]).inserter("foo")) + .await + .expect_err("user should be able to insert into `foo` now"); + + // Roles should not propagate back to the parent instance + insert_foos(user_client.inserter("foo")) + .await + .expect_err("user should not be able to insert into `foo`"); + + insert_foos(user_client.inserter("foo").with_roles([&role])) + .await + .expect_err("user should be able to insert into `foo` now"); + + // `with_default_roles` should clear the role + insert_foos( + user_client + .clone() + .with_roles([&role]) + .inserter("foo") + .with_default_roles(), + ) + .await + .expect_err("user should not be able to insert into `foo`"); +}