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

Author: abonander

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<T> Insert<T> {
         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<String>) -> 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<Item = impl Into<String>>) -> Self {
+        self.state.expect_client_mut().set_roles(roles);
         self
     }
 

src/inserter.rs

@@ -161,41 +161,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<String>) -> 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<Item = impl Into<String>>) -> 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.
     ///

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<String>,
     authentication: Authentication,
     compression: Compression,
-    options: HashMap<Cow<'static, str>, String>,
+    roles: HashSet<String>,
+    options: HashMap<String, String>,
     headers: HashMap<String, String>,
     products_info: Vec<ProductInfo>,
     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<String>) -> 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<Item = impl Into<String>>) -> 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<String>, value: impl Into<String>) -> 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<String>, value: impl Into<String>) {
-        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<String>) {
-        // 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<Item = impl Into<String>>) {
+        // 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.

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