phonograph/interim-server/src/routes/workspaces_multi/add_handlers.rs

use axum::{extract::State, response::IntoResponse};
use interim_models::{
    client::AppDbClient, cluster::Cluster, user::User, workspace::Workspace,
    workspace_user_perm::WorkspaceMembership,
};
use interim_pgtypes::{client::WorkspaceClient, escape_identifier, rolnames::ROLE_PREFIX_USER};
use sqlx::{Connection as _, PgConnection, query};

use crate::{
    app::AppDbConn,
    errors::AppError,
    navigator::{Navigator, NavigatorPage as _},
    user::CurrentUser,
    workspace_pooler::{RoleAssignment, WorkspacePooler},
    workspace_utils::PHONO_TABLE_NAMESPACE,
};

/// HTTP POST handler for creating a new workspace. This handler does not expect
/// any arguments, as the backing database name is generated pseudo-randomly and
/// the human-friendly name may be changed later rather than specified
/// permenantly upon creation.
pub(super) async fn post(
    State(mut pooler): State<WorkspacePooler>,
    AppDbConn(mut app_db): AppDbConn,
    navigator: Navigator,
    CurrentUser(user): CurrentUser,
) -> Result<impl IntoResponse, AppError> {
    // FIXME: csrf

    let cluster = Cluster::fetch_only(&mut app_db).await?;

    const NAME_LEN_WORDS: usize = 3;
    // WARNING: `db_name` is injected directly into the `create database` SQL
    // command. It **must not** contain spaces or any other unsafe characters.
    // Additionally, it **must** be URL safe without percent encoding.
    let db_name = interim_namegen::default_generator()
        .with_separator('_')
        .generate_name(NAME_LEN_WORDS);
    {
        // No need to pool these connections, since we don't expect to be using them
        // often. One less thing to keep track of in application state.
        let mut workspace_creator_conn = PgConnection::connect(
            cluster
                .conn_str_for_db("postgres", None)?
                .expose_secret()
                .as_str(),
        )
        .await?;
        query(&format!(
            // `db_name` is an underscore-separated sequence of alphabetical words,
            // which should be safe to inject directly into the SQL statement.
            "create database {db_name}"
        ))
        .execute(&mut workspace_creator_conn)
        .await?;
        query(&format!("revoke connect on database {db_name} from public"))
            .execute(&mut workspace_creator_conn)
            .await?;

        // Close and drop `workspace_creator_conn` at end of block, then
        // reconnect using a pooled connection.
        workspace_creator_conn.close().await?;
    }

    let workspace = Workspace::insert()
        .owner_id(user.id)
        .cluster_id(cluster.id)
        .db_name(&db_name)
        .build()?
        .insert(&mut app_db)
        .await?;

    let mut workspace_root_conn = pooler
        .acquire_for(workspace.id, RoleAssignment::Root)
        .await?;

    // Initialize database user. Connection is not used and may be dropped
    // immediately.
    pooler
        .acquire_for(workspace.id, RoleAssignment::User(user.id))
        .await?;
    query(&format!(
        "create schema {nsp}",
        nsp = escape_identifier(PHONO_TABLE_NAMESPACE)
    ))
    .execute(workspace_root_conn.get_conn())
    .await?;
    grant_workspace_membership(
        &db_name,
        &mut app_db,
        &mut workspace_root_conn,
        &user,
        &workspace,
    )
    .await?;

    Ok(navigator
        .workspace_page()
        .workspace_id(workspace.id)
        .build()?
        .redirect_to())
}

async fn grant_workspace_membership(
    db_name: &str,
    app_db_client: &mut AppDbClient,
    workspace_root_client: &mut WorkspaceClient,
    user: &User,
    workspace: &Workspace,
) -> Result<(), AppError> {
    let rolname = format!("{ROLE_PREFIX_USER}{uid}", uid = user.id.simple());
    query(&format!(
        "grant connect on database {db_name} to {db_user} with grant option",
        db_user = escape_identifier(&rolname),
    ))
    .execute(workspace_root_client.get_conn())
    .await?;

    // TODO: There may be cases in which we will want to grant granular
    // workspace access which excludes privileges to create tables.
    query(&format!(
        "grant usage, create on schema {nsp} to {rolname} with grant option",
        nsp = escape_identifier(PHONO_TABLE_NAMESPACE),
        rolname = escape_identifier(&rolname)
    ))
    .execute(workspace_root_client.get_conn())
    .await?;

    WorkspaceMembership::insert()
        .workspace_id(workspace.id)
        .user_id(user.id)
        .build()?
        .execute(app_db_client)
        .await?;

    Ok(())
}