lmrc_postgres/

user_db_management.rs

//! User and Database Management
//!
//! This module provides comprehensive user and database management operations including:
//! - List users and databases
//! - Create/drop users and databases
//! - Update user passwords
//! - Role management
//! - Granular permission management

use crate::error::{Error, Result};
use lmrc_ssh::SshClient;
use tracing::{debug, info};

/// PostgreSQL user information
#[derive(Debug, Clone)]
pub struct UserInfo {
    /// Username
    pub name: String,
    /// Is superuser
    pub is_superuser: bool,
    /// Can create databases
    pub can_create_db: bool,
    /// Can create roles
    pub can_create_role: bool,
    /// Connection limit (-1 for unlimited)
    pub connection_limit: i32,
}

/// PostgreSQL database information
#[derive(Debug, Clone)]
pub struct DatabaseInfo {
    /// Database name
    pub name: String,
    /// Owner username
    pub owner: String,
    /// Encoding (e.g., UTF8)
    pub encoding: String,
    /// Database size (human-readable)
    pub size: Option<String>,
}

/// PostgreSQL privilege types
#[derive(Debug, Clone, PartialEq)]
pub enum Privilege {
    /// SELECT privilege
    Select,
    /// INSERT privilege
    Insert,
    /// UPDATE privilege
    Update,
    /// DELETE privilege
    Delete,
    /// TRUNCATE privilege
    Truncate,
    /// REFERENCES privilege
    References,
    /// TRIGGER privilege
    Trigger,
    /// CREATE privilege
    Create,
    /// CONNECT privilege
    Connect,
    /// TEMPORARY privilege
    Temporary,
    /// EXECUTE privilege
    Execute,
    /// USAGE privilege
    Usage,
    /// ALL privileges
    All,
}

impl Privilege {
    /// Convert privilege to PostgreSQL keyword
    pub fn as_str(&self) -> &str {
        match self {
            Privilege::Select => "SELECT",
            Privilege::Insert => "INSERT",
            Privilege::Update => "UPDATE",
            Privilege::Delete => "DELETE",
            Privilege::Truncate => "TRUNCATE",
            Privilege::References => "REFERENCES",
            Privilege::Trigger => "TRIGGER",
            Privilege::Create => "CREATE",
            Privilege::Connect => "CONNECT",
            Privilege::Temporary => "TEMPORARY",
            Privilege::Execute => "EXECUTE",
            Privilege::Usage => "USAGE",
            Privilege::All => "ALL PRIVILEGES",
        }
    }
}

/// List all users in the PostgreSQL instance
pub async fn list_users(ssh: &mut SshClient) -> Result<Vec<UserInfo>> {
    debug!("Listing PostgreSQL users");

    let query = r#"
        SELECT
            usename,
            usesuper,
            usecreatedb,
            usecreaterole,
            useconnlimit
        FROM pg_user
        ORDER BY usename;
    "#;

    let cmd = format!(
        r#"sudo -u postgres psql -t -A -F '|' -c "{}""#,
        query.replace('\n', " ")
    );

    let output = ssh
        .execute(&cmd)
        .map_err(|e| Error::Configuration(format!("Failed to list users: {}", e)))?;

    let mut users = Vec::new();

    for line in output.stdout.lines() {
        let line = line.trim();
        if line.is_empty() {
            continue;
        }

        let parts: Vec<&str> = line.split('|').collect();
        if parts.len() >= 5 {
            users.push(UserInfo {
                name: parts[0].to_string(),
                is_superuser: parts[1] == "t",
                can_create_db: parts[2] == "t",
                can_create_role: parts[3] == "t",
                connection_limit: parts[4].parse().unwrap_or(-1),
            });
        }
    }

    info!("Found {} user(s)", users.len());
    Ok(users)
}

/// List all databases in the PostgreSQL instance
pub async fn list_databases(ssh: &mut SshClient) -> Result<Vec<DatabaseInfo>> {
    debug!("Listing PostgreSQL databases");

    let query = r#"
        SELECT
            d.datname,
            u.usename,
            pg_encoding_to_char(d.encoding),
            pg_size_pretty(pg_database_size(d.datname))
        FROM pg_database d
        JOIN pg_user u ON d.datdba = u.usesysid
        WHERE d.datistemplate = false
        ORDER BY d.datname;
    "#;

    let cmd = format!(
        r#"sudo -u postgres psql -t -A -F '|' -c "{}""#,
        query.replace('\n', " ")
    );

    let output = ssh
        .execute(&cmd)
        .map_err(|e| Error::Configuration(format!("Failed to list databases: {}", e)))?;

    let mut databases = Vec::new();

    for line in output.stdout.lines() {
        let line = line.trim();
        if line.is_empty() {
            continue;
        }

        let parts: Vec<&str> = line.split('|').collect();
        if parts.len() >= 4 {
            databases.push(DatabaseInfo {
                name: parts[0].to_string(),
                owner: parts[1].to_string(),
                encoding: parts[2].to_string(),
                size: Some(parts[3].to_string()),
            });
        }
    }

    info!("Found {} database(s)", databases.len());
    Ok(databases)
}

/// Drop a database
pub async fn drop_database(ssh: &mut SshClient, database_name: &str) -> Result<()> {
    info!("Dropping database: {}", database_name);

    // Terminate connections to the database first
    let terminate_cmd = format!(
        r#"sudo -u postgres psql -c "SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE datname = '{}' AND pid <> pg_backend_pid();" || true"#,
        database_name
    );

    ssh.execute(&terminate_cmd)
        .map_err(|e| Error::Configuration(format!("Failed to terminate connections: {}", e)))?;

    // Drop the database
    let drop_cmd = format!(
        r#"sudo -u postgres psql -c "DROP DATABASE IF EXISTS {};" || true"#,
        database_name
    );

    ssh.execute(&drop_cmd)
        .map_err(|e| Error::Configuration(format!("Failed to drop database: {}", e)))?;

    info!("✓ Database {} dropped successfully", database_name);
    Ok(())
}

/// Drop a user
pub async fn drop_user(ssh: &mut SshClient, username: &str) -> Result<()> {
    info!("Dropping user: {}", username);

    // Reassign owned objects first (to avoid dependency errors)
    let reassign_cmd = format!(
        r#"sudo -u postgres psql -c "REASSIGN OWNED BY {} TO postgres;" || true"#,
        username
    );

    ssh.execute(&reassign_cmd).ok(); // Don't fail if user owns nothing

    // Drop owned objects
    let drop_owned_cmd = format!(
        r#"sudo -u postgres psql -c "DROP OWNED BY {} CASCADE;" || true"#,
        username
    );

    ssh.execute(&drop_owned_cmd).ok(); // Don't fail if user owns nothing

    // Drop the user
    let drop_cmd = format!(
        r#"sudo -u postgres psql -c "DROP USER IF EXISTS {};" || true"#,
        username
    );

    ssh.execute(&drop_cmd)
        .map_err(|e| Error::Configuration(format!("Failed to drop user: {}", e)))?;

    info!("✓ User {} dropped successfully", username);
    Ok(())
}

/// Update user password
pub async fn update_user_password(
    ssh: &mut SshClient,
    username: &str,
    new_password: &str,
) -> Result<()> {
    info!("Updating password for user: {}", username);

    let update_cmd = format!(
        r#"sudo -u postgres psql -c "ALTER USER {} WITH PASSWORD '{}';" || true"#,
        username, new_password
    );

    ssh.execute(&update_cmd)
        .map_err(|e| Error::Configuration(format!("Failed to update password: {}", e)))?;

    info!("✓ Password updated for user {}", username);
    Ok(())
}

/// Grant specific privileges on a database to a user
pub async fn grant_privileges(
    ssh: &mut SshClient,
    database: &str,
    username: &str,
    privileges: &[Privilege],
) -> Result<()> {
    info!(
        "Granting privileges on database {} to user {}",
        database, username
    );

    // Grant database-level privileges
    let privs: Vec<&str> = privileges.iter().map(|p| p.as_str()).collect();
    let priv_string = privs.join(", ");

    let grant_cmd = format!(
        r#"sudo -u postgres psql -c "GRANT {} ON DATABASE {} TO {};" || true"#,
        priv_string, database, username
    );

    ssh.execute(&grant_cmd)
        .map_err(|e| Error::Configuration(format!("Failed to grant database privileges: {}", e)))?;

    // For schema-level privileges (tables, sequences, functions)
    if privileges.contains(&Privilege::All)
        || privileges.contains(&Privilege::Select)
        || privileges.contains(&Privilege::Insert)
        || privileges.contains(&Privilege::Update)
        || privileges.contains(&Privilege::Delete)
    {
        // Grant on all tables in public schema
        let grant_tables_cmd = format!(
            r#"sudo -u postgres psql -d {} -c "GRANT {} ON ALL TABLES IN SCHEMA public TO {};" || true"#,
            database, priv_string, username
        );

        ssh.execute(&grant_tables_cmd).ok();

        // Grant on all sequences
        let grant_sequences_cmd = format!(
            r#"sudo -u postgres psql -d {} -c "GRANT ALL ON ALL SEQUENCES IN SCHEMA public TO {};" || true"#,
            database, username
        );

        ssh.execute(&grant_sequences_cmd).ok();

        // Grant default privileges for future tables
        let grant_default_cmd = format!(
            r#"sudo -u postgres psql -d {} -c "ALTER DEFAULT PRIVILEGES IN SCHEMA public GRANT {} ON TABLES TO {};" || true"#,
            database, priv_string, username
        );

        ssh.execute(&grant_default_cmd).ok();
    }

    info!("✓ Privileges granted successfully");
    Ok(())
}

/// Revoke privileges on a database from a user
pub async fn revoke_privileges(
    ssh: &mut SshClient,
    database: &str,
    username: &str,
    privileges: &[Privilege],
) -> Result<()> {
    info!(
        "Revoking privileges on database {} from user {}",
        database, username
    );

    let privs: Vec<&str> = privileges.iter().map(|p| p.as_str()).collect();
    let priv_string = privs.join(", ");

    let revoke_cmd = format!(
        r#"sudo -u postgres psql -c "REVOKE {} ON DATABASE {} FROM {};" || true"#,
        priv_string, database, username
    );

    ssh.execute(&revoke_cmd)
        .map_err(|e| Error::Configuration(format!("Failed to revoke privileges: {}", e)))?;

    info!("✓ Privileges revoked successfully");
    Ok(())
}

/// Create a role (user without login)
pub async fn create_role(
    ssh: &mut SshClient,
    role_name: &str,
    can_login: bool,
    is_superuser: bool,
) -> Result<()> {
    info!("Creating role: {}", role_name);

    let mut options = Vec::new();

    if can_login {
        options.push("LOGIN");
    } else {
        options.push("NOLOGIN");
    }

    if is_superuser {
        options.push("SUPERUSER");
    }

    let options_str = options.join(" ");

    let create_cmd = format!(
        r#"sudo -u postgres psql -c "CREATE ROLE {} {};" || true"#,
        role_name, options_str
    );

    ssh.execute(&create_cmd)
        .map_err(|e| Error::Configuration(format!("Failed to create role: {}", e)))?;

    info!("✓ Role {} created successfully", role_name);
    Ok(())
}

/// Grant a role to a user
pub async fn grant_role(ssh: &mut SshClient, role_name: &str, username: &str) -> Result<()> {
    info!("Granting role {} to user {}", role_name, username);

    let grant_cmd = format!(
        r#"sudo -u postgres psql -c "GRANT {} TO {};" || true"#,
        role_name, username
    );

    ssh.execute(&grant_cmd)
        .map_err(|e| Error::Configuration(format!("Failed to grant role: {}", e)))?;

    info!("✓ Role granted successfully");
    Ok(())
}

/// Revoke a role from a user
pub async fn revoke_role(ssh: &mut SshClient, role_name: &str, username: &str) -> Result<()> {
    info!("Revoking role {} from user {}", role_name, username);

    let revoke_cmd = format!(
        r#"sudo -u postgres psql -c "REVOKE {} FROM {};" || true"#,
        role_name, username
    );

    ssh.execute(&revoke_cmd)
        .map_err(|e| Error::Configuration(format!("Failed to revoke role: {}", e)))?;

    info!("✓ Role revoked successfully");
    Ok(())
}

/// Check if a user exists
pub async fn user_exists(ssh: &mut SshClient, username: &str) -> Result<bool> {
    debug!("Checking if user {} exists", username);

    let check_cmd = format!(
        r#"sudo -u postgres psql -t -c "SELECT 1 FROM pg_user WHERE usename = '{}';" | grep -q 1"#,
        username
    );

    Ok(ssh.execute(&check_cmd).is_ok())
}

/// Check if a database exists
pub async fn database_exists(ssh: &mut SshClient, database: &str) -> Result<bool> {
    debug!("Checking if database {} exists", database);

    let check_cmd = format!(
        r#"sudo -u postgres psql -t -c "SELECT 1 FROM pg_database WHERE datname = '{}';" | grep -q 1"#,
        database
    );

    Ok(ssh.execute(&check_cmd).is_ok())
}

/// Create a database with options
pub async fn create_database_with_options(
    ssh: &mut SshClient,
    database_name: &str,
    owner: Option<&str>,
    encoding: Option<&str>,
    template: Option<&str>,
) -> Result<()> {
    info!("Creating database: {}", database_name);

    let mut options = Vec::new();

    if let Some(owner) = owner {
        options.push(format!("OWNER {}", owner));
    }

    if let Some(encoding) = encoding {
        options.push(format!("ENCODING '{}'", encoding));
    }

    if let Some(template) = template {
        options.push(format!("TEMPLATE {}", template));
    }

    let options_str = if options.is_empty() {
        String::new()
    } else {
        format!(" WITH {}", options.join(" "))
    };

    let create_cmd = format!(
        r#"sudo -u postgres psql -c "CREATE DATABASE {}{};" || true"#,
        database_name, options_str
    );

    ssh.execute(&create_cmd)
        .map_err(|e| Error::Configuration(format!("Failed to create database: {}", e)))?;

    info!("✓ Database {} created successfully", database_name);
    Ok(())
}

/// Create a user with full options
pub async fn create_user_with_options(
    ssh: &mut SshClient,
    username: &str,
    password: &str,
    is_superuser: bool,
    can_create_db: bool,
    can_create_role: bool,
    connection_limit: Option<i32>,
) -> Result<()> {
    info!("Creating user: {}", username);

    let mut options = Vec::new();

    options.push(format!("PASSWORD '{}'", password));

    if is_superuser {
        options.push("SUPERUSER".to_string());
    } else {
        options.push("NOSUPERUSER".to_string());
    }

    if can_create_db {
        options.push("CREATEDB".to_string());
    } else {
        options.push("NOCREATEDB".to_string());
    }

    if can_create_role {
        options.push("CREATEROLE".to_string());
    } else {
        options.push("NOCREATEROLE".to_string());
    }

    if let Some(limit) = connection_limit {
        options.push(format!("CONNECTION LIMIT {}", limit));
    }

    let options_str = options.join(" ");

    let create_cmd = format!(
        r#"sudo -u postgres psql -c "CREATE USER {} WITH {};" || true"#,
        username, options_str
    );

    ssh.execute(&create_cmd)
        .map_err(|e| Error::Configuration(format!("Failed to create user: {}", e)))?;

    info!("✓ User {} created successfully", username);
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_privilege_as_str() {
        assert_eq!(Privilege::Select.as_str(), "SELECT");
        assert_eq!(Privilege::All.as_str(), "ALL PRIVILEGES");
        assert_eq!(Privilege::Insert.as_str(), "INSERT");
    }

    #[test]
    fn test_privilege_equality() {
        assert_eq!(Privilege::Select, Privilege::Select);
        assert_ne!(Privilege::Select, Privilege::Insert);
    }
}