forked from 2sys/phonograph
60 lines
2.1 KiB
Rust
60 lines
2.1 KiB
Rust
use sqlx::{PgConnection, Postgres, Row as _, pool::PoolConnection, query};
|
|
|
|
use crate::escape_identifier;
|
|
|
|
/// Newtype to differentiate between workspace and application database
|
|
/// connections.
|
|
pub struct WorkspaceClient {
|
|
pub(crate) conn: PoolConnection<Postgres>,
|
|
}
|
|
|
|
impl WorkspaceClient {
|
|
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`. Roles are created with the `createrole` option.
|
|
///
|
|
/// 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 {0} createrole",
|
|
escape_identifier(rolname),
|
|
))
|
|
.execute(&mut *self.conn)
|
|
.await?;
|
|
query(&format!(
|
|
"grant {0} to {1}",
|
|
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(())
|
|
}
|
|
}
|