sal-postgresclient 0.1.0

// PostgreSQL installer module
//
// This module provides functionality to install and configure PostgreSQL using nerdctl.

use std::collections::HashMap;
use std::env;
use std::fs;
use std::path::Path;
use std::process::Command;
use std::thread;
use std::time::Duration;

use sal_virt::nerdctl::Container;
use std::error::Error;
use std::fmt;

// Custom error type for PostgreSQL installer
#[derive(Debug)]
pub enum PostgresInstallerError {
    IoError(std::io::Error),
    NerdctlError(String),
    PostgresError(String),
}

impl fmt::Display for PostgresInstallerError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            PostgresInstallerError::IoError(e) => write!(f, "I/O error: {}", e),
            PostgresInstallerError::NerdctlError(e) => write!(f, "Nerdctl error: {}", e),
            PostgresInstallerError::PostgresError(e) => write!(f, "PostgreSQL error: {}", e),
        }
    }
}

impl Error for PostgresInstallerError {
    fn source(&self) -> Option<&(dyn Error + 'static)> {
        match self {
            PostgresInstallerError::IoError(e) => Some(e),
            _ => None,
        }
    }
}

impl From<std::io::Error> for PostgresInstallerError {
    fn from(error: std::io::Error) -> Self {
        PostgresInstallerError::IoError(error)
    }
}

/// PostgreSQL installer configuration
pub struct PostgresInstallerConfig {
    /// Container name for PostgreSQL
    pub container_name: String,
    /// PostgreSQL version to install
    pub version: String,
    /// Port to expose PostgreSQL on
    pub port: u16,
    /// Username for PostgreSQL
    pub username: String,
    /// Password for PostgreSQL
    pub password: String,
    /// Data directory for PostgreSQL
    pub data_dir: Option<String>,
    /// Environment variables for PostgreSQL
    pub env_vars: HashMap<String, String>,
    /// Whether to use persistent storage
    pub persistent: bool,
}

impl Default for PostgresInstallerConfig {
    fn default() -> Self {
        Self {
            container_name: "postgres".to_string(),
            version: "latest".to_string(),
            port: 5432,
            username: "postgres".to_string(),
            password: "postgres".to_string(),
            data_dir: None,
            env_vars: HashMap::new(),
            persistent: true,
        }
    }
}

impl PostgresInstallerConfig {
    /// Create a new PostgreSQL installer configuration with default values
    pub fn new() -> Self {
        Self::default()
    }

    /// Set the container name
    pub fn container_name(mut self, name: &str) -> Self {
        self.container_name = name.to_string();
        self
    }

    /// Set the PostgreSQL version
    pub fn version(mut self, version: &str) -> Self {
        self.version = version.to_string();
        self
    }

    /// Set the port to expose PostgreSQL on
    pub fn port(mut self, port: u16) -> Self {
        self.port = port;
        self
    }

    /// Set the username for PostgreSQL
    pub fn username(mut self, username: &str) -> Self {
        self.username = username.to_string();
        self
    }

    /// Set the password for PostgreSQL
    pub fn password(mut self, password: &str) -> Self {
        self.password = password.to_string();
        self
    }

    /// Set the data directory for PostgreSQL
    pub fn data_dir(mut self, data_dir: &str) -> Self {
        self.data_dir = Some(data_dir.to_string());
        self
    }

    /// Add an environment variable
    pub fn env_var(mut self, key: &str, value: &str) -> Self {
        self.env_vars.insert(key.to_string(), value.to_string());
        self
    }

    /// Set whether to use persistent storage
    pub fn persistent(mut self, persistent: bool) -> Self {
        self.persistent = persistent;
        self
    }
}

/// Install PostgreSQL using nerdctl
///
/// # Arguments
///
/// * `config` - PostgreSQL installer configuration
///
/// # Returns
///
/// * `Result<Container, PostgresInstallerError>` - Container instance or error
pub fn install_postgres(
    config: PostgresInstallerConfig,
) -> Result<Container, PostgresInstallerError> {
    // Create the data directory if it doesn't exist and persistent storage is enabled
    let data_dir = if config.persistent {
        let dir = config.data_dir.unwrap_or_else(|| {
            let home_dir = env::var("HOME").unwrap_or_else(|_| "/tmp".to_string());
            format!("{}/.postgres-data", home_dir)
        });

        if !Path::new(&dir).exists() {
            fs::create_dir_all(&dir).map_err(|e| PostgresInstallerError::IoError(e))?;
        }

        Some(dir)
    } else {
        None
    };

    // Build the image name
    let image = format!("postgres:{}", config.version);

    // Pull the PostgreSQL image to ensure we have the latest version
    println!("Pulling PostgreSQL image: {}...", image);
    let pull_result = Command::new("nerdctl")
        .args(&["pull", &image])
        .output()
        .map_err(|e| PostgresInstallerError::IoError(e))?;

    if !pull_result.status.success() {
        return Err(PostgresInstallerError::NerdctlError(format!(
            "Failed to pull PostgreSQL image: {}",
            String::from_utf8_lossy(&pull_result.stderr)
        )));
    }

    // Create the container
    let mut container = Container::new(&config.container_name).map_err(|e| {
        PostgresInstallerError::NerdctlError(format!("Failed to create container: {}", e))
    })?;

    // Set the image
    container.image = Some(image);

    // Set the port
    container = container.with_port(&format!("{}:5432", config.port));

    // Set environment variables
    container = container.with_env("POSTGRES_USER", &config.username);
    container = container.with_env("POSTGRES_PASSWORD", &config.password);
    container = container.with_env("POSTGRES_DB", "postgres");

    // Add custom environment variables
    for (key, value) in &config.env_vars {
        container = container.with_env(key, value);
    }

    // Add volume for persistent storage if enabled
    if let Some(dir) = data_dir {
        container = container.with_volume(&format!("{}:/var/lib/postgresql/data", dir));
    }

    // Set restart policy
    container = container.with_restart_policy("unless-stopped");

    // Set detach mode
    container = container.with_detach(true);

    // Build and start the container
    let container = container.build().map_err(|e| {
        PostgresInstallerError::NerdctlError(format!("Failed to build container: {}", e))
    })?;

    // Wait for PostgreSQL to start
    println!("Waiting for PostgreSQL to start...");
    thread::sleep(Duration::from_secs(5));

    // Set environment variables for PostgreSQL client
    env::set_var("POSTGRES_HOST", "localhost");
    env::set_var("POSTGRES_PORT", config.port.to_string());
    env::set_var("POSTGRES_USER", config.username);
    env::set_var("POSTGRES_PASSWORD", config.password);
    env::set_var("POSTGRES_DB", "postgres");

    Ok(container)
}

/// Create a new database in PostgreSQL
///
/// # Arguments
///
/// * `container` - PostgreSQL container
/// * `db_name` - Database name
///
/// # Returns
///
/// * `Result<(), PostgresInstallerError>` - Ok if successful, Err otherwise
pub fn create_database(container: &Container, db_name: &str) -> Result<(), PostgresInstallerError> {
    // Check if container is running
    if container.container_id.is_none() {
        return Err(PostgresInstallerError::PostgresError(
            "Container is not running".to_string(),
        ));
    }

    // Execute the command to create the database
    let command = format!(
        "createdb -U {} {}",
        env::var("POSTGRES_USER").unwrap_or_else(|_| "postgres".to_string()),
        db_name
    );

    container.exec(&command).map_err(|e| {
        PostgresInstallerError::NerdctlError(format!("Failed to create database: {}", e))
    })?;

    Ok(())
}

/// Execute a SQL script in PostgreSQL
///
/// # Arguments
///
/// * `container` - PostgreSQL container
/// * `db_name` - Database name
/// * `sql` - SQL script to execute
///
/// # Returns
///
/// * `Result<String, PostgresInstallerError>` - Output of the command or error
pub fn execute_sql(
    container: &Container,
    db_name: &str,
    sql: &str,
) -> Result<String, PostgresInstallerError> {
    // Check if container is running
    if container.container_id.is_none() {
        return Err(PostgresInstallerError::PostgresError(
            "Container is not running".to_string(),
        ));
    }

    // Create a temporary file with the SQL script
    let temp_file = "/tmp/postgres_script.sql";
    fs::write(temp_file, sql).map_err(|e| PostgresInstallerError::IoError(e))?;

    // Copy the file to the container
    let container_id = container.container_id.as_ref().unwrap();
    let copy_result = Command::new("nerdctl")
        .args(&[
            "cp",
            temp_file,
            &format!("{}:/tmp/script.sql", container_id),
        ])
        .output()
        .map_err(|e| PostgresInstallerError::IoError(e))?;

    if !copy_result.status.success() {
        return Err(PostgresInstallerError::PostgresError(format!(
            "Failed to copy SQL script to container: {}",
            String::from_utf8_lossy(&copy_result.stderr)
        )));
    }

    // Execute the SQL script
    let command = format!(
        "psql -U {} -d {} -f /tmp/script.sql",
        env::var("POSTGRES_USER").unwrap_or_else(|_| "postgres".to_string()),
        db_name
    );

    let result = container.exec(&command).map_err(|e| {
        PostgresInstallerError::NerdctlError(format!("Failed to execute SQL script: {}", e))
    })?;

    // Clean up
    fs::remove_file(temp_file).ok();

    Ok(result.stdout)
}

/// Check if PostgreSQL is running
///
/// # Arguments
///
/// * `container` - PostgreSQL container
///
/// # Returns
///
/// * `Result<bool, PostgresInstallerError>` - true if running, false otherwise, or error
pub fn is_postgres_running(container: &Container) -> Result<bool, PostgresInstallerError> {
    // Check if container is running
    if container.container_id.is_none() {
        return Ok(false);
    }

    // Execute a simple query to check if PostgreSQL is running
    let command = format!(
        "psql -U {} -c 'SELECT 1'",
        env::var("POSTGRES_USER").unwrap_or_else(|_| "postgres".to_string())
    );

    match container.exec(&command) {
        Ok(_) => Ok(true),
        Err(_) => Ok(false),
    }
}