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

use crate::error::{TuskError, TuskResult};
use crate::models::connection::ConnectionConfig;
use crate::state::AppState;
use sqlx::PgPool;
use sqlx::Row;
use std::fs;
use std::sync::Arc;
use tauri::{AppHandle, Manager, State};

fn get_connections_path(app: &AppHandle) -> TuskResult<std::path::PathBuf> {
    let dir = app
        .path()
        .app_data_dir()
        .map_err(|e| TuskError::Custom(e.to_string()))?;
    fs::create_dir_all(&dir)?;
    Ok(dir.join("connections.json"))
}

#[tauri::command]
pub async fn get_connections(app: AppHandle) -> TuskResult<Vec<ConnectionConfig>> {
    let path = get_connections_path(&app)?;
    if !path.exists() {
        return Ok(vec![]);
    }
    let data = fs::read_to_string(&path)?;
    let connections: Vec<ConnectionConfig> = serde_json::from_str(&data)?;
    Ok(connections)
}

#[tauri::command]
pub async fn save_connection(app: AppHandle, config: ConnectionConfig) -> TuskResult<()> {
    let path = get_connections_path(&app)?;
    let mut connections = if path.exists() {
        let data = fs::read_to_string(&path)?;
        serde_json::from_str::<Vec<ConnectionConfig>>(&data)?
    } else {
        vec![]
    };

    if let Some(pos) = connections.iter().position(|c| c.id == config.id) {
        connections[pos] = config;
    } else {
        connections.push(config);
    }

    let data = serde_json::to_string_pretty(&connections)?;
    fs::write(&path, data)?;
    Ok(())
}

#[tauri::command]
pub async fn delete_connection(
    app: AppHandle,
    state: State<'_, Arc<AppState>>,
    id: String,
) -> TuskResult<()> {
    let path = get_connections_path(&app)?;
    if path.exists() {
        let data = fs::read_to_string(&path)?;
        let mut connections: Vec<ConnectionConfig> = serde_json::from_str(&data)?;
        connections.retain(|c| c.id != id);
        let data = serde_json::to_string_pretty(&connections)?;
        fs::write(&path, data)?;
    }

    // Close pool if connected
    let mut pools = state.pools.write().await;
    if let Some(pool) = pools.remove(&id) {
        pool.close().await;
    }

    let mut ro = state.read_only.write().await;
    ro.remove(&id);

    Ok(())
}

#[tauri::command]
pub async fn test_connection(config: ConnectionConfig) -> TuskResult<String> {
    let pool = PgPool::connect(&config.connection_url())
        .await
        .map_err(TuskError::Database)?;

    let row = sqlx::query("SELECT version()")
        .fetch_one(&pool)
        .await
        .map_err(TuskError::Database)?;

    let version: String = row.get(0);
    pool.close().await;
    Ok(version)
}

#[tauri::command]
pub async fn connect(state: State<'_, Arc<AppState>>, config: ConnectionConfig) -> TuskResult<()> {
    let pool = PgPool::connect(&config.connection_url())
        .await
        .map_err(TuskError::Database)?;

    // Verify connection
    sqlx::query("SELECT 1")
        .execute(&pool)
        .await
        .map_err(TuskError::Database)?;

    let mut pools = state.pools.write().await;
    pools.insert(config.id.clone(), pool);

    let mut ro = state.read_only.write().await;
    ro.insert(config.id.clone(), true);

    Ok(())
}

#[tauri::command]
pub async fn switch_database(
    state: State<'_, Arc<AppState>>,
    config: ConnectionConfig,
    database: String,
) -> TuskResult<()> {
    let mut switched = config.clone();
    switched.database = database;

    let pool = PgPool::connect(&switched.connection_url())
        .await
        .map_err(TuskError::Database)?;

    sqlx::query("SELECT 1")
        .execute(&pool)
        .await
        .map_err(TuskError::Database)?;

    let mut pools = state.pools.write().await;
    if let Some(old_pool) = pools.remove(&config.id) {
        old_pool.close().await;
    }
    pools.insert(config.id.clone(), pool);

    Ok(())
}

#[tauri::command]
pub async fn disconnect(state: State<'_, Arc<AppState>>, id: String) -> TuskResult<()> {
    let mut pools = state.pools.write().await;
    if let Some(pool) = pools.remove(&id) {
        pool.close().await;
    }

    let mut ro = state.read_only.write().await;
    ro.remove(&id);

    Ok(())
}

#[tauri::command]
pub async fn set_read_only(
    state: State<'_, Arc<AppState>>,
    connection_id: String,
    read_only: bool,
) -> TuskResult<()> {
    let mut map = state.read_only.write().await;
    map.insert(connection_id, read_only);
    Ok(())
}

#[tauri::command]
pub async fn get_read_only(
    state: State<'_, Arc<AppState>>,
    connection_id: String,
) -> TuskResult<bool> {
    Ok(state.is_read_only(&connection_id).await)
}