phonograph/interim-pgtypes/src/client.rs

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.
    ///
    /// 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(())
    }
}