use crate::error::{TuskError, TuskResult}; use crate::models::query_result::QueryResult; use crate::state::AppState; 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::, _>(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); } "DATE" => try_get!(chrono::NaiveDate), "TIME" => try_get!(chrono::NaiveTime), "BYTEA" => match row.try_get::>, _>(index) { Ok(Some(v)) => return Value::String(format!("\\x{}", hex::encode(&v))), Ok(None) => return Value::Null, Err(_) => {} }, "OID" => match row.try_get::, _>(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), "_INT2" => try_get!(Vec), "_INT4" => try_get!(Vec), "_INT8" => try_get!(Vec), "_FLOAT4" => try_get!(Vec), "_FLOAT8" => try_get!(Vec), "_TEXT" | "_VARCHAR" | "_CHAR" | "_BPCHAR" | "_NAME" => try_get!(Vec), "_UUID" => try_get!(Vec), "_JSON" | "_JSONB" => try_get!(Vec), _ => {} } // Fallback: try as string match row.try_get::, _>(index) { Ok(Some(v)) => Value::String(v), Ok(None) => Value::Null, Err(_) => Value::String(format!("", type_name)), } } pub async fn execute_query_core( state: &AppState, connection_id: &str, sql: &str, ) -> TuskResult { let read_only = state.is_read_only(connection_id).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(); 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> = 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>, connection_id: String, sql: String, ) -> TuskResult { execute_query_core(&state, &connection_id, &sql).await }