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

use crate::db::sql_guard::ensure_readonly_sql;
use crate::error::{TuskError, TuskResult};
use crate::models::query_result::QueryResult;
use crate::state::{AppState, DbFlavor};
use serde_json::Value;
use sqlx::postgres::PgRow;
use sqlx::{Column, Row, TypeInfo};
use std::sync::Arc;
use std::time::Instant;
use tauri::State;

pub fn pg_value_to_json(row: &PgRow, index: usize) -> Value {
    let col = &row.columns()[index];
    let type_name = col.type_info().name();

    macro_rules! try_get {
        ($t:ty) => {
            match row.try_get::<Option<$t>, _>(index) {
                Ok(Some(v)) => return serde_json::to_value(v).unwrap_or(Value::Null),
                Ok(None) => return Value::Null,
                Err(_) => {}
            }
        };
    }

    match type_name {
        "BOOL" => try_get!(bool),
        "INT2" => try_get!(i16),
        "INT4" => try_get!(i32),
        "INT8" => try_get!(i64),
        "FLOAT4" => try_get!(f32),
        "FLOAT8" => try_get!(f64),
        "NUMERIC" => {
            try_get!(bigdecimal::BigDecimal);
        }
        "TEXT" | "VARCHAR" | "CHAR" | "BPCHAR" | "NAME" => try_get!(String),
        "JSON" | "JSONB" => try_get!(Value),
        "UUID" => try_get!(uuid::Uuid),
        "TIMESTAMP" => {
            try_get!(chrono::NaiveDateTime);
        }
        "TIMESTAMPTZ" => {
            try_get!(chrono::DateTime<chrono::Utc>);
        }
        "DATE" => try_get!(chrono::NaiveDate),
        "TIME" => try_get!(chrono::NaiveTime),
        "INTERVAL" => match row.try_get::<Option<sqlx::postgres::types::PgInterval>, _>(index) {
            Ok(Some(v)) => return Value::String(format_pg_interval(&v)),
            Ok(None) => return Value::Null,
            Err(_) => {}
        },
        "BYTEA" => match row.try_get::<Option<Vec<u8>>, _>(index) {
            Ok(Some(v)) => return Value::String(format!("\\x{}", hex::encode(&v))),
            Ok(None) => return Value::Null,
            Err(_) => {}
        },
        "OID" => match row.try_get::<Option<i32>, _>(index) {
            Ok(Some(v)) => return Value::Number(serde_json::Number::from(v)),
            Ok(None) => return Value::Null,
            Err(_) => {}
        },
        "VOID" => return Value::Null,
        // Array types (PG prefixes array type names with underscore)
        "_BOOL" => try_get!(Vec<bool>),
        "_INT2" => try_get!(Vec<i16>),
        "_INT4" => try_get!(Vec<i32>),
        "_INT8" => try_get!(Vec<i64>),
        "_FLOAT4" => try_get!(Vec<f32>),
        "_FLOAT8" => try_get!(Vec<f64>),
        "_TEXT" | "_VARCHAR" | "_CHAR" | "_BPCHAR" | "_NAME" => try_get!(Vec<String>),
        "_UUID" => try_get!(Vec<uuid::Uuid>),
        "_JSON" | "_JSONB" => try_get!(Vec<Value>),
        _ => {}
    }

    // Fallback: try as string
    match row.try_get::<Option<String>, _>(index) {
        Ok(Some(v)) => Value::String(v),
        Ok(None) => Value::Null,
        Err(_) => Value::String(format!("<unsupported type: {}>", type_name)),
    }
}

/// Render a PostgreSQL INTERVAL the same way `psql` does:
///   "1 year 2 mons 3 days 04:05:06.789012"
/// Components are emitted only when non-zero; the time component fires
/// whenever microseconds != 0 OR everything else is zero.
fn format_pg_interval(iv: &sqlx::postgres::types::PgInterval) -> String {
    let years = iv.months / 12;
    let months = iv.months % 12;
    let mut parts: Vec<String> = Vec::new();
    if years != 0 {
        parts.push(format!("{} year{}", years, if years.abs() == 1 { "" } else { "s" }));
    }
    if months != 0 {
        parts.push(format!("{} mon{}", months, if months.abs() == 1 { "" } else { "s" }));
    }
    if iv.days != 0 {
        parts.push(format!("{} day{}", iv.days, if iv.days.abs() == 1 { "" } else { "s" }));
    }
    if iv.microseconds != 0 || parts.is_empty() {
        let total_us = iv.microseconds.unsigned_abs();
        let total_seconds = total_us / 1_000_000;
        let micros = (total_us % 1_000_000) as u32;
        let h = total_seconds / 3600;
        let m = (total_seconds / 60) % 60;
        let s = total_seconds % 60;
        let sign = if iv.microseconds < 0 { "-" } else { "" };
        let time_part = if micros == 0 {
            format!("{}{:02}:{:02}:{:02}", sign, h, m, s)
        } else {
            format!(
                "{}{:02}:{:02}:{:02}.{:06}",
                sign, h, m, s, micros
            )
        };
        parts.push(time_part);
    }
    parts.join(" ")
}

pub async fn execute_query_core(
    state: &AppState,
    connection_id: &str,
    sql: &str,
) -> TuskResult<QueryResult> {
    let read_only = state.is_read_only(connection_id).await;
    let flavor = state.get_flavor(connection_id).await;

    if read_only {
        ensure_readonly_sql(sql)?;
    }

    if matches!(flavor, DbFlavor::ClickHouse) {
        let client = state.get_ch_client(connection_id).await?;
        return client.execute_query(sql, read_only).await;
    }

    let pools = state.pools.read().await;
    let pool = pools
        .get(connection_id)
        .ok_or_else(|| TuskError::NotConnected(connection_id.to_string()))?;

    let start = Instant::now();
    let rows = if read_only {
        let mut tx = pool.begin().await.map_err(TuskError::Database)?;
        sqlx::query("SET TRANSACTION READ ONLY")
            .execute(&mut *tx)
            .await
            .map_err(TuskError::Database)?;
        let result = sqlx::query(sql)
            .fetch_all(&mut *tx)
            .await
            .map_err(TuskError::Database);
        tx.rollback().await.map_err(TuskError::Database)?;
        result?
    } else {
        sqlx::query(sql)
            .fetch_all(pool)
            .await
            .map_err(TuskError::Database)?
    };
    let execution_time_ms = start.elapsed().as_millis() as u64;

    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(QueryResult {
        columns,
        types,
        rows: result_rows,
        row_count,
        execution_time_ms,
    })
}

#[tauri::command]
pub async fn execute_query(
    state: State<'_, Arc<AppState>>,
    connection_id: String,
    sql: String,
) -> TuskResult<QueryResult> {
    execute_query_core(&state, &connection_id, &sql).await
}

#[cfg(test)]
mod tests {
    use super::format_pg_interval;
    use sqlx::postgres::types::PgInterval;

    #[test]
    fn interval_zero_renders_as_zero_time() {
        let iv = PgInterval { months: 0, days: 0, microseconds: 0 };
        assert_eq!(format_pg_interval(&iv), "00:00:00");
    }

    #[test]
    fn interval_pure_time_micros() {
        // 1h 30m
        let iv = PgInterval { months: 0, days: 0, microseconds: 90 * 60 * 1_000_000 };
        assert_eq!(format_pg_interval(&iv), "01:30:00");
    }

    #[test]
    fn interval_days_only() {
        let iv = PgInterval { months: 0, days: 3, microseconds: 0 };
        assert_eq!(format_pg_interval(&iv), "3 days");
    }

    #[test]
    fn interval_one_day() {
        let iv = PgInterval { months: 0, days: 1, microseconds: 0 };
        assert_eq!(format_pg_interval(&iv), "1 day");
    }

    #[test]
    fn interval_mixed_components() {
        // 1 year 2 mons 3 days 04:05:06
        let iv = PgInterval {
            months: 14,
            days: 3,
            microseconds: ((4 * 3600) + (5 * 60) + 6) * 1_000_000,
        };
        assert_eq!(format_pg_interval(&iv), "1 year 2 mons 3 days 04:05:06");
    }

    #[test]
    fn interval_negative_time() {
        let iv = PgInterval { months: 0, days: 0, microseconds: -3_600_000_000 };
        assert_eq!(format_pg_interval(&iv), "-01:00:00");
    }

    #[test]
    fn interval_with_microseconds_fraction() {
        let iv = PgInterval { months: 0, days: 0, microseconds: 1_500_000 };
        assert_eq!(format_pg_interval(&iv), "00:00:01.500000");
    }
}