diff --git a/src/insert.rs b/src/insert.rs index 89ad7ac3..82449c56 100644 --- a/src/insert.rs +++ b/src/insert.rs @@ -95,6 +95,15 @@ 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") + }; + + client + } + fn terminated(&mut self) { replace_with_or_abort(self, |_self| match _self { InsertState::NotStarted { .. } => InsertState::Completed, // empty insert @@ -102,17 +111,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 +183,36 @@ impl Insert { self } + /// Configure the [roles] to use when executing `INSERT` statements. + /// + /// Overrides any roles previously set by this method, [`Insert::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 + /// + /// # Panics + /// If called after the request is started, e.g., after [`Insert::write`]. + pub fn with_roles(mut self, roles: impl IntoIterator>) -> Self { + self.state.expect_client_mut().set_roles(roles); + 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. /// @@ -192,7 +220,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..071866f6 100644 --- a/src/inserter.rs +++ b/src/inserter.rs @@ -165,8 +165,53 @@ where self } + /// Set the [roles] to use when executing `INSERT` statements. + /// + /// Overrides any roles previously set by this method, [`Inserter::with_option`], + /// [`Client::with_roles`] 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 [`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 { + self.client.set_roles(roles); + 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. + /// + /// # 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`]. 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..126cd131 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::collections::HashSet; use std::{collections::HashMap, fmt::Display, sync::Arc}; use tokio::sync::RwLock; @@ -57,6 +58,7 @@ pub struct Client { database: Option, authentication: Authentication, compression: Compression, + roles: HashSet, options: HashMap, headers: HashMap, products_info: Vec, @@ -121,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(), @@ -227,6 +230,42 @@ impl Client { self } + /// Configure the [roles] to use when executing statements with this `Client` instance. + /// + /// Overrides any roles previously set by this method or [`Client::with_option`]. + /// + /// Call [`Client::with_default_roles`] to clear any explicitly set roles. + /// + /// This setting is copied into cloned clients. + /// + /// [roles]: https://clickhouse.com/docs/operations/access-rights#role-management + /// + /// # Examples + /// + /// ``` + /// # use clickhouse::Client; + /// + /// // Single role + /// let client = Client::default().with_roles(["foo"]); + /// + /// // Multiple roles + /// let client = Client::default().with_roles(["foo", "bar", "baz"]); + /// ``` + 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 @@ -450,6 +489,18 @@ impl Client { self.options.insert(name.into(), value.into()); } + 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(); + } + /// Use a mock server for testing purposes. /// /// # Note diff --git a/src/query.rs b/src/query.rs index ee9c4b80..f27e69be 100644 --- a/src/query.rs +++ b/src/query.rs @@ -189,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()); @@ -209,6 +212,34 @@ 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 + } + } + + /// 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/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/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`"); +} diff --git a/tests/it/main.rs b/tests/it/main.rs index 0b2923ed..3d02e709 100644 --- a/tests/it/main.rs +++ b/tests/it/main.rs @@ -314,6 +314,72 @@ 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(); + + // 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 + .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..0a359e78 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,83 @@ 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`"); + + // 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]) + .execute() + .await + .expect("user should be able to query `foo` now"); +}