phonograph/phono-backends/src/client.rs

use phono_pestgros::escape_identifier;
use sqlx::{PgConnection, Postgres, Row as _, pool::PoolConnection, query};

/// Newtype to differentiate between workspace and application database
/// connections.
#[derive(Debug)]
pub struct WorkspaceClient {
    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
    }

    /// If the given role does not exist, create it and grant it to the
    /// `session_user`. Roles are created with the `createrole` option.
    pub async fn ensure_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?;
        }
        Ok(())
    }

    /// Runs the Postgres `SET ROLE` command for the underlying connection.
    /// Returns an error if the given role does not exist. Use
    /// [`WorkspaceClient::ensure_role`] to ensure the role exists on the
    /// cluster, if needed.
    ///
    /// 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 set_role(&mut self, rolname: &str) -> Result<(), sqlx::Error> {
        query(&format!("set role {}", escape_identifier(rolname)))
            .execute(&mut *self.conn)
            .await?;
        Ok(())
    }
}