56 lines
1.9 KiB
Rust
56 lines
1.9 KiB
Rust
|
use sqlx::{PgConnection, Postgres, Row as _, pool::PoolConnection, query};
|
||
|
|
|
||
|
|
use crate::escape_identifier;
|
||
|
|
|
||
|
|
pub struct BaseClient {
|
||
|
|
pub(crate) conn: PoolConnection<Postgres>,
|
||
|
|
}
|
||
|
|
|
||
|
|
impl BaseClient {
|
||
|
|
pub fn from_pool_conn(conn: PoolConnection<Postgres>) -> Self {
|
||
|
|
Self { conn }
|
||
|
|
}
|
||
|
|
|
||
|
|
pub fn get_conn(&mut self) -> &mut PgConnection {
|
||
|
|
&mut self.conn
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Runs the Postgres `set role` command for the underlying connection. If
|
||
|
|
/// the given role does not exist, it is created and granted to the
|
||
|
|
/// session_user.
|
||
|
|
///
|
||
|
|
/// Note that while using `set role` simulates impersonation for most data
|
||
|
|
/// access and RLS purposes, it is both incomplete and easily reversible:
|
||
|
|
/// some commands and system tables will still behave according to the
|
||
|
|
/// privileges of the session user, and clients relying on this abstraction
|
||
|
|
/// should **NEVER** execute untrusted SQL.
|
||
|
|
pub async fn init_role(&mut self, rolname: &str) -> Result<(), sqlx::Error> {
|
||
|
|
let session_user = query!("select session_user;")
|
||
|
|
.fetch_one(&mut *self.conn)
|
||
|
|
.await?
|
||
|
|
.session_user
|
||
|
|
.unwrap();
|
||
|
|
if !query("select exists(select 1 from pg_roles where rolname = $1)")
|
||
|
|
.bind(rolname)
|
||
|
|
.fetch_one(&mut *self.conn)
|
||
|
|
.await?
|
||
|
|
.try_get(0)?
|
||
|
|
{
|
||
|
|
query(&format!("create role {}", escape_identifier(rolname)))
|
||
|
|
.execute(&mut *self.conn)
|
||
|
|
.await?;
|
||
|
|
query(&format!(
|
||
|
|
"grant {} to {}",
|
||
|
|
escape_identifier(rolname),
|
||
|
|
escape_identifier(&session_user),
|
||
|
|
))
|
||
|
|
.execute(&mut *self.conn)
|
||
|
|
.await?;
|
||
|
|
}
|
||
|
|
query(&format!("set role {}", escape_identifier(rolname)))
|
||
|
|
.execute(&mut *self.conn)
|
||
|
|
.await?;
|
||
|
|
Ok(())
|
||
|
|
}
|
||
|
|
}
|