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),
        "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)),
    }
}

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
}