tusk/src-tauri/src/commands/management.rs

use crate::error::{TuskError, TuskResult};
use crate::models::management::*;
use crate::state::{AppState, DbFlavor};
use crate::utils::escape_ident;
use sqlx::Row;
use std::sync::Arc;
use tauri::State;

#[tauri::command]
pub async fn get_database_info(
    state: State<'_, Arc<AppState>>,
    connection_id: String,
) -> TuskResult<Vec<DatabaseInfo>> {
    let pools = state.pools.read().await;
    let pool = pools
        .get(&connection_id)
        .ok_or(TuskError::NotConnected(connection_id))?;

    let rows = sqlx::query(
        "SELECT d.datname, \
                pg_catalog.pg_get_userbyid(d.datdba) AS owner, \
                pg_catalog.pg_encoding_to_char(d.encoding) AS encoding, \
                d.datcollate, \
                d.datctype, \
                COALESCE(t.spcname, 'pg_default') AS tablespace, \
                d.datconnlimit, \
                pg_catalog.pg_size_pretty(pg_catalog.pg_database_size(d.datname)) AS size, \
                pg_catalog.shobj_description(d.oid, 'pg_database') AS description \
         FROM pg_catalog.pg_database d \
         LEFT JOIN pg_catalog.pg_tablespace t ON d.dattablespace = t.oid \
         WHERE NOT d.datistemplate \
         ORDER BY d.datname",
    )
    .fetch_all(pool)
    .await
    .map_err(TuskError::Database)?;

    let databases = rows
        .iter()
        .map(|row| DatabaseInfo {
            name: row.get("datname"),
            owner: row.get("owner"),
            encoding: row.get("encoding"),
            collation: row.get("datcollate"),
            ctype: row.get("datctype"),
            tablespace: row.get("tablespace"),
            connection_limit: row.get("datconnlimit"),
            size: row.get("size"),
            description: row.get("description"),
        })
        .collect();

    Ok(databases)
}

#[tauri::command]
pub async fn create_database(
    state: State<'_, Arc<AppState>>,
    connection_id: String,
    params: CreateDatabaseParams,
) -> TuskResult<()> {
    if state.is_read_only(&connection_id).await {
        return Err(TuskError::ReadOnly);
    }

    let pools = state.pools.read().await;
    let pool = pools
        .get(&connection_id)
        .ok_or(TuskError::NotConnected(connection_id))?;

    let mut sql = format!("CREATE DATABASE {}", escape_ident(&params.name));

    if let Some(ref owner) = params.owner {
        sql.push_str(&format!(" OWNER {}", escape_ident(owner)));
    }
    if let Some(ref template) = params.template {
        sql.push_str(&format!(" TEMPLATE {}", escape_ident(template)));
    }
    if let Some(ref encoding) = params.encoding {
        sql.push_str(&format!(" ENCODING '{}'", encoding.replace('\'', "''")));
    }
    if let Some(ref tablespace) = params.tablespace {
        sql.push_str(&format!(" TABLESPACE {}", escape_ident(tablespace)));
    }
    if let Some(limit) = params.connection_limit {
        sql.push_str(&format!(" CONNECTION LIMIT {}", limit));
    }

    sqlx::query(&sql)
        .execute(pool)
        .await
        .map_err(TuskError::Database)?;

    Ok(())
}

#[tauri::command]
pub async fn drop_database(
    state: State<'_, Arc<AppState>>,
    connection_id: String,
    name: String,
) -> TuskResult<()> {
    if state.is_read_only(&connection_id).await {
        return Err(TuskError::ReadOnly);
    }

    let pools = state.pools.read().await;
    let pool = pools
        .get(&connection_id)
        .ok_or(TuskError::NotConnected(connection_id))?;

    // Terminate active connections to the target database
    let terminate_sql = format!(
        "SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE datname = '{}' AND pid <> pg_backend_pid()",
        name.replace('\'', "''")
    );
    sqlx::query(&terminate_sql)
        .execute(pool)
        .await
        .map_err(TuskError::Database)?;

    let drop_sql = format!("DROP DATABASE {}", escape_ident(&name));
    sqlx::query(&drop_sql)
        .execute(pool)
        .await
        .map_err(TuskError::Database)?;

    Ok(())
}

#[tauri::command]
pub async fn list_roles(
    state: State<'_, Arc<AppState>>,
    connection_id: String,
) -> TuskResult<Vec<RoleInfo>> {
    let pools = state.pools.read().await;
    let pool = pools
        .get(&connection_id)
        .ok_or(TuskError::NotConnected(connection_id))?;

    let rows = sqlx::query(
        "SELECT r.rolname, \
                r.rolsuper, \
                r.rolcanlogin, \
                r.rolcreatedb, \
                r.rolcreaterole, \
                r.rolinherit, \
                r.rolreplication, \
                r.rolconnlimit, \
                r.rolpassword IS NOT NULL AS password_set, \
                r.rolvaliduntil::text, \
                COALESCE(( \
                    SELECT array_agg(g.rolname ORDER BY g.rolname) \
                    FROM pg_catalog.pg_auth_members m \
                    JOIN pg_catalog.pg_roles g ON m.roleid = g.oid \
                    WHERE m.member = r.oid \
                ), ARRAY[]::text[]) AS member_of, \
                COALESCE(( \
                    SELECT array_agg(m2.rolname ORDER BY m2.rolname) \
                    FROM pg_catalog.pg_auth_members am \
                    JOIN pg_catalog.pg_roles m2 ON am.member = m2.oid \
                    WHERE am.roleid = r.oid \
                ), ARRAY[]::text[]) AS members, \
                pg_catalog.shobj_description(r.oid, 'pg_authid') AS description \
         FROM pg_catalog.pg_roles r \
         WHERE r.rolname !~ '^pg_' \
         ORDER BY r.rolname",
    )
    .fetch_all(pool)
    .await
    .map_err(TuskError::Database)?;

    let roles = rows
        .iter()
        .map(|row| RoleInfo {
            name: row.get("rolname"),
            is_superuser: row.get("rolsuper"),
            can_login: row.get("rolcanlogin"),
            can_create_db: row.get("rolcreatedb"),
            can_create_role: row.get("rolcreaterole"),
            inherit: row.get("rolinherit"),
            is_replication: row.get("rolreplication"),
            connection_limit: row.get("rolconnlimit"),
            password_set: row.get("password_set"),
            valid_until: row.get("rolvaliduntil"),
            member_of: row.get("member_of"),
            members: row.get("members"),
            description: row.get("description"),
        })
        .collect();

    Ok(roles)
}

#[tauri::command]
pub async fn create_role(
    state: State<'_, Arc<AppState>>,
    connection_id: String,
    params: CreateRoleParams,
) -> TuskResult<()> {
    if state.is_read_only(&connection_id).await {
        return Err(TuskError::ReadOnly);
    }

    let pools = state.pools.read().await;
    let pool = pools
        .get(&connection_id)
        .ok_or(TuskError::NotConnected(connection_id))?;

    let mut sql = format!("CREATE ROLE {}", escape_ident(&params.name));

    let mut options = Vec::new();
    options.push(if params.login { "LOGIN" } else { "NOLOGIN" });
    options.push(if params.superuser {
        "SUPERUSER"
    } else {
        "NOSUPERUSER"
    });
    options.push(if params.createdb {
        "CREATEDB"
    } else {
        "NOCREATEDB"
    });
    options.push(if params.createrole {
        "CREATEROLE"
    } else {
        "NOCREATEROLE"
    });
    options.push(if params.inherit {
        "INHERIT"
    } else {
        "NOINHERIT"
    });
    options.push(if params.replication {
        "REPLICATION"
    } else {
        "NOREPLICATION"
    });

    if let Some(ref password) = params.password {
        options.push("PASSWORD");
        // Will be appended separately
        sql.push_str(&format!(" {}", options.join(" ")));
        sql.push_str(&format!(" '{}'", password.replace('\'', "''")));
    } else {
        sql.push_str(&format!(" {}", options.join(" ")));
    }

    if let Some(limit) = params.connection_limit {
        sql.push_str(&format!(" CONNECTION LIMIT {}", limit));
    }

    if let Some(ref valid_until) = params.valid_until {
        sql.push_str(&format!(
            " VALID UNTIL '{}'",
            valid_until.replace('\'', "''")
        ));
    }

    if !params.in_roles.is_empty() {
        let roles: Vec<String> = params.in_roles.iter().map(|r| escape_ident(r)).collect();
        sql.push_str(&format!(" IN ROLE {}", roles.join(", ")));
    }

    sqlx::query(&sql)
        .execute(pool)
        .await
        .map_err(TuskError::Database)?;

    Ok(())
}

#[tauri::command]
pub async fn alter_role(
    state: State<'_, Arc<AppState>>,
    connection_id: String,
    params: AlterRoleParams,
) -> TuskResult<()> {
    if state.is_read_only(&connection_id).await {
        return Err(TuskError::ReadOnly);
    }

    let pools = state.pools.read().await;
    let pool = pools
        .get(&connection_id)
        .ok_or(TuskError::NotConnected(connection_id))?;

    let mut options = Vec::new();

    if let Some(login) = params.login {
        options.push(if login {
            "LOGIN".to_string()
        } else {
            "NOLOGIN".to_string()
        });
    }
    if let Some(superuser) = params.superuser {
        options.push(if superuser {
            "SUPERUSER".to_string()
        } else {
            "NOSUPERUSER".to_string()
        });
    }
    if let Some(createdb) = params.createdb {
        options.push(if createdb {
            "CREATEDB".to_string()
        } else {
            "NOCREATEDB".to_string()
        });
    }
    if let Some(createrole) = params.createrole {
        options.push(if createrole {
            "CREATEROLE".to_string()
        } else {
            "NOCREATEROLE".to_string()
        });
    }
    if let Some(inherit) = params.inherit {
        options.push(if inherit {
            "INHERIT".to_string()
        } else {
            "NOINHERIT".to_string()
        });
    }
    if let Some(replication) = params.replication {
        options.push(if replication {
            "REPLICATION".to_string()
        } else {
            "NOREPLICATION".to_string()
        });
    }
    if let Some(ref password) = params.password {
        options.push(format!("PASSWORD '{}'", password.replace('\'', "''")));
    }
    if let Some(limit) = params.connection_limit {
        options.push(format!("CONNECTION LIMIT {}", limit));
    }
    if let Some(ref valid_until) = params.valid_until {
        options.push(format!("VALID UNTIL '{}'", valid_until.replace('\'', "''")));
    }

    if !options.is_empty() {
        let sql = format!(
            "ALTER ROLE {} {}",
            escape_ident(&params.name),
            options.join(" ")
        );
        sqlx::query(&sql)
            .execute(pool)
            .await
            .map_err(TuskError::Database)?;
    }

    if let Some(ref new_name) = params.rename_to {
        let sql = format!(
            "ALTER ROLE {} RENAME TO {}",
            escape_ident(&params.name),
            escape_ident(new_name)
        );
        sqlx::query(&sql)
            .execute(pool)
            .await
            .map_err(TuskError::Database)?;
    }

    Ok(())
}

#[tauri::command]
pub async fn drop_role(
    state: State<'_, Arc<AppState>>,
    connection_id: String,
    name: String,
) -> TuskResult<()> {
    if state.is_read_only(&connection_id).await {
        return Err(TuskError::ReadOnly);
    }

    let pools = state.pools.read().await;
    let pool = pools
        .get(&connection_id)
        .ok_or(TuskError::NotConnected(connection_id))?;

    let sql = format!("DROP ROLE {}", escape_ident(&name));
    sqlx::query(&sql)
        .execute(pool)
        .await
        .map_err(TuskError::Database)?;

    Ok(())
}

#[tauri::command]
pub async fn get_table_privileges(
    state: State<'_, Arc<AppState>>,
    connection_id: String,
    schema: String,
    table: String,
) -> TuskResult<Vec<TablePrivilege>> {
    let pools = state.pools.read().await;
    let pool = pools
        .get(&connection_id)
        .ok_or(TuskError::NotConnected(connection_id))?;

    let rows = sqlx::query(
        "SELECT grantee, table_schema, table_name, privilege_type, \
                is_grantable = 'YES' AS is_grantable \
         FROM information_schema.role_table_grants \
         WHERE table_schema = $1 AND table_name = $2 \
         ORDER BY grantee, privilege_type",
    )
    .bind(&schema)
    .bind(&table)
    .fetch_all(pool)
    .await
    .map_err(TuskError::Database)?;

    let privileges = rows
        .iter()
        .map(|row| TablePrivilege {
            grantee: row.get("grantee"),
            table_schema: row.get("table_schema"),
            table_name: row.get("table_name"),
            privilege_type: row.get("privilege_type"),
            is_grantable: row.get("is_grantable"),
        })
        .collect();

    Ok(privileges)
}

#[tauri::command]
pub async fn grant_revoke(
    state: State<'_, Arc<AppState>>,
    connection_id: String,
    params: GrantRevokeParams,
) -> TuskResult<()> {
    if state.is_read_only(&connection_id).await {
        return Err(TuskError::ReadOnly);
    }

    let pools = state.pools.read().await;
    let pool = pools
        .get(&connection_id)
        .ok_or(TuskError::NotConnected(connection_id))?;

    let privs = params.privileges.join(", ");
    let object_type = params.object_type.to_uppercase();
    let object_ref = escape_ident(&params.object_name);
    let role_ref = escape_ident(&params.role_name);

    let sql = if params.action.to_uppercase() == "GRANT" {
        let grant_option = if params.with_grant_option {
            " WITH GRANT OPTION"
        } else {
            ""
        };
        format!(
            "GRANT {} ON {} {} TO {}{}",
            privs, object_type, object_ref, role_ref, grant_option
        )
    } else {
        format!(
            "REVOKE {} ON {} {} FROM {}",
            privs, object_type, object_ref, role_ref
        )
    };

    sqlx::query(&sql)
        .execute(pool)
        .await
        .map_err(TuskError::Database)?;

    Ok(())
}

#[tauri::command]
pub async fn manage_role_membership(
    state: State<'_, Arc<AppState>>,
    connection_id: String,
    params: RoleMembershipParams,
) -> TuskResult<()> {
    if state.is_read_only(&connection_id).await {
        return Err(TuskError::ReadOnly);
    }

    let pools = state.pools.read().await;
    let pool = pools
        .get(&connection_id)
        .ok_or(TuskError::NotConnected(connection_id))?;

    let role_ref = escape_ident(&params.role_name);
    let member_ref = escape_ident(&params.member_name);

    let sql = if params.action.to_uppercase() == "GRANT" {
        format!("GRANT {} TO {}", role_ref, member_ref)
    } else {
        format!("REVOKE {} FROM {}", role_ref, member_ref)
    };

    sqlx::query(&sql)
        .execute(pool)
        .await
        .map_err(TuskError::Database)?;

    Ok(())
}

#[tauri::command]
pub async fn list_sessions(
    state: State<'_, Arc<AppState>>,
    connection_id: String,
) -> TuskResult<Vec<SessionInfo>> {
    let flavor = state.get_flavor(&connection_id).await;
    let pools = state.pools.read().await;
    let pool = pools
        .get(&connection_id)
        .ok_or(TuskError::NotConnected(connection_id))?;

    let sql = if flavor == DbFlavor::Greenplum {
        "SELECT pid, usename, datname, state, query, \
                query_start::text, NULL::text as wait_event_type, NULL::text as wait_event, \
                client_addr::text \
         FROM pg_stat_activity \
         WHERE datname IS NOT NULL \
         ORDER BY query_start DESC NULLS LAST"
    } else {
        "SELECT pid, usename, datname, state, query, \
                query_start::text, wait_event_type, wait_event, \
                client_addr::text \
         FROM pg_stat_activity \
         WHERE datname IS NOT NULL \
         ORDER BY query_start DESC NULLS LAST"
    };

    let rows = sqlx::query(sql)
        .fetch_all(pool)
        .await
        .map_err(TuskError::Database)?;

    let sessions = rows
        .iter()
        .map(|row| SessionInfo {
            pid: row.get("pid"),
            usename: row.get("usename"),
            datname: row.get("datname"),
            state: row.get("state"),
            query: row.get("query"),
            query_start: row.get("query_start"),
            wait_event_type: row.get("wait_event_type"),
            wait_event: row.get("wait_event"),
            client_addr: row.get("client_addr"),
        })
        .collect();

    Ok(sessions)
}

#[tauri::command]
pub async fn cancel_query(
    state: State<'_, Arc<AppState>>,
    connection_id: String,
    pid: i32,
) -> TuskResult<bool> {
    let pools = state.pools.read().await;
    let pool = pools
        .get(&connection_id)
        .ok_or(TuskError::NotConnected(connection_id))?;

    let row = sqlx::query("SELECT pg_cancel_backend($1)")
        .bind(pid)
        .fetch_one(pool)
        .await
        .map_err(TuskError::Database)?;

    Ok(row.get::<bool, _>(0))
}

#[tauri::command]
pub async fn terminate_backend(
    state: State<'_, Arc<AppState>>,
    connection_id: String,
    pid: i32,
) -> TuskResult<bool> {
    let pools = state.pools.read().await;
    let pool = pools
        .get(&connection_id)
        .ok_or(TuskError::NotConnected(connection_id))?;

    let row = sqlx::query("SELECT pg_terminate_backend($1)")
        .bind(pid)
        .fetch_one(pool)
        .await
        .map_err(TuskError::Database)?;

    Ok(row.get::<bool, _>(0))
}