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

Author: abonander

src/insert.rs

@@ -95,24 +95,21 @@ 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
             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 +182,40 @@ impl<T> Insert<T> {
         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<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
+    ///
+    /// # 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.
     ///
     /// # 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

@@ -161,8 +161,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<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()`].
+    ///
+    /// 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<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::borrow::Cow;
 use std::{collections::HashMap, fmt::Display, sync::Arc};
 use tokio::sync::RwLock;
 
@@ -57,7 +58,7 @@ pub struct Client {
     database: Option<String>,
     authentication: Authentication,
     compression: Compression,
-    options: HashMap<String, String>,
+    options: HashMap<Cow<'static, str>, String>,
     headers: HashMap<String, String>,
     products_info: Vec<ProductInfo>,
     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<String>) -> 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<String>, value: impl Into<String>) -> 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<String>, value: impl Into<String>) {
-        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<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");
+
+        match role {
+            Some(role) => {
+                self.options.insert(key, role);
+            }
+            None => {
+                self.options.remove(&key);
+            }
+        }
     }
 
     /// Use a mock server for testing purposes.

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<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