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

use crate::commands::queries::pg_value_to_json;
use crate::error::{TuskError, TuskResult};
use crate::models::query_result::PaginatedQueryResult;
use crate::state::AppState;
use serde_json::Value;
use sqlx::{Column, Row, TypeInfo};
use std::time::Instant;
use tauri::State;

fn escape_ident(name: &str) -> String {
    format!("\"{}\"", name.replace('"', "\"\""))
}

#[tauri::command]
pub async fn get_table_data(
    state: State<'_, AppState>,
    connection_id: String,
    schema: String,
    table: String,
    page: u32,
    page_size: u32,
    sort_column: Option<String>,
    sort_direction: Option<String>,
    filter: Option<String>,
) -> TuskResult<PaginatedQueryResult> {
    let pools = state.pools.read().await;
    let pool = pools
        .get(&connection_id)
        .ok_or(TuskError::NotConnected(connection_id))?;

    let qualified = format!("{}.{}", escape_ident(&schema), escape_ident(&table));

    let mut where_clause = String::new();
    if let Some(ref f) = filter {
        if !f.trim().is_empty() {
            where_clause = format!(" WHERE {}", f);
        }
    }

    let mut order_clause = String::new();
    if let Some(ref col) = sort_column {
        let dir = sort_direction.as_deref().unwrap_or("ASC");
        let dir = if dir.eq_ignore_ascii_case("desc") {
            "DESC"
        } else {
            "ASC"
        };
        order_clause = format!(" ORDER BY {} {}", escape_ident(col), dir);
    }

    let offset = (page.saturating_sub(1)) * page_size;

    let data_sql = format!(
        "SELECT * FROM {}{}{} LIMIT {} OFFSET {}",
        qualified, where_clause, order_clause, page_size, offset
    );
    let count_sql = format!("SELECT COUNT(*) FROM {}{}", qualified, where_clause);

    let start = Instant::now();

    let (rows, count_row) = tokio::try_join!(
        sqlx::query(&data_sql).fetch_all(pool),
        sqlx::query(&count_sql).fetch_one(pool),
    )
    .map_err(TuskError::Database)?;

    let execution_time_ms = start.elapsed().as_millis();
    let total_rows: i64 = count_row.get(0);

    let mut columns = Vec::new();
    let mut types = Vec::new();

    if let Some(first_row) = rows.first() {
        for col in first_row.columns() {
            columns.push(col.name().to_string());
            types.push(col.type_info().name().to_string());
        }
    }

    let result_rows: Vec<Vec<Value>> = rows
        .iter()
        .map(|row| (0..columns.len()).map(|i| pg_value_to_json(row, i)).collect())
        .collect();

    let row_count = result_rows.len();

    Ok(PaginatedQueryResult {
        columns,
        types,
        rows: result_rows,
        row_count,
        execution_time_ms,
        total_rows,
        page,
        page_size,
    })
}

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

    let qualified = format!("{}.{}", escape_ident(&schema), escape_ident(&table));

    let set_clause = format!("{} = $1", escape_ident(&column));

    let where_parts: Vec<String> = pk_columns
        .iter()
        .enumerate()
        .map(|(i, col)| format!("{} = ${}", escape_ident(col), i + 2))
        .collect();
    let where_clause = where_parts.join(" AND ");

    let sql = format!(
        "UPDATE {} SET {} WHERE {}",
        qualified, set_clause, where_clause
    );

    let mut query = sqlx::query(&sql);
    query = bind_json_value(query, &value);
    for pk_val in &pk_values {
        query = bind_json_value(query, pk_val);
    }

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

    Ok(())
}

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

    let qualified = format!("{}.{}", escape_ident(&schema), escape_ident(&table));

    let col_list: Vec<String> = columns.iter().map(|c| escape_ident(c)).collect();
    let placeholders: Vec<String> = (1..=columns.len()).map(|i| format!("${}", i)).collect();

    let sql = format!(
        "INSERT INTO {} ({}) VALUES ({})",
        qualified,
        col_list.join(", "),
        placeholders.join(", ")
    );

    let mut query = sqlx::query(&sql);
    for val in &values {
        query = bind_json_value(query, val);
    }

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

    Ok(())
}

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

    let qualified = format!("{}.{}", escape_ident(&schema), escape_ident(&table));
    let mut total_affected: u64 = 0;

    for pk_values in &pk_values_list {
        let where_parts: Vec<String> = pk_columns
            .iter()
            .enumerate()
            .map(|(i, col)| format!("{} = ${}", escape_ident(col), i + 1))
            .collect();
        let where_clause = where_parts.join(" AND ");

        let sql = format!("DELETE FROM {} WHERE {}", qualified, where_clause);

        let mut query = sqlx::query(&sql);
        for val in pk_values {
            query = bind_json_value(query, val);
        }

        let result = query.execute(pool).await.map_err(TuskError::Database)?;
        total_affected += result.rows_affected();
    }

    Ok(total_affected)
}

fn bind_json_value<'q>(
    query: sqlx::query::Query<'q, sqlx::Postgres, sqlx::postgres::PgArguments>,
    value: &'q Value,
) -> sqlx::query::Query<'q, sqlx::Postgres, sqlx::postgres::PgArguments> {
    match value {
        Value::Null => query.bind(None::<String>),
        Value::Bool(b) => query.bind(*b),
        Value::Number(n) => {
            if let Some(i) = n.as_i64() {
                query.bind(i)
            } else if let Some(f) = n.as_f64() {
                query.bind(f)
            } else {
                query.bind(n.to_string())
            }
        }
        Value::String(s) => query.bind(s.as_str()),
        _ => query.bind(value.to_string()),
    }
}