feat: add with_role to Client, Query, Insert{er}

src/insert.rs

@@ -95,24 +95,22 @@ 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
             InsertState::Active { handle, .. } => InsertState::Terminated { handle },
             _ => unreachable!(),
         });
     }
-
-    fn with_option(&mut self, name: impl Into<String>, value: impl Into<String>) {
-        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,14 +183,44 @@ impl<T> Insert<T> {
         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<Item = impl Into<String>>) -> 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.
     ///
     /// # Panics
     /// If called after the request is started, e.g., after [`Insert::write`].
     #[track_caller]
     pub fn with_option(mut self, name: impl Into<String>, value: impl Into<String>) -> Self {
-        self.state.with_option(name, value);
+        self.state.expect_client_mut().add_option(name, value);
         self
     }
 

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<Item = impl Into<String>>) -> 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<String>, value: impl Into<String>) -> Self {
         self.client.add_option(name, value);
         self

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<String>,
     authentication: Authentication,
     compression: Compression,
+    roles: HashSet<String>,
     options: HashMap<String, String>,
     headers: HashMap<String, String>,
     products_info: Vec<ProductInfo>,
@@ -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<Item = impl Into<String>>) -> 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<Item = impl Into<String>>) {
+        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

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<Item = impl Into<String>>) -> 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<String>, value: impl Into<String>) -> Self {
         self.client.add_option(name, value);

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::<Vec<_>>();
+
+    let insert_foos = async |mut insert: Insert<Foo>| {
+        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`");
+}