use std::collections::HashSet; use std::sync::Mutex; use sqlx::sqlite::SqliteRow; use sqlx::{Column, Row, ValueRef}; use crate::{DbError, DbPool}; /// sqlx requires the SQL string to be `'static` interning stops repeated runs /// of the same text from leaking fn intern(sql: &str) -> &'static str { static CACHE: Mutex>> = Mutex::new(None); let mut guard = CACHE.lock().unwrap(); let cache = guard.get_or_insert_with(HashSet::new); if let Some(existing) = cache.get(sql) { return existing; } let leaked: &'static str = Box::leak(sql.to_owned().into_boxed_str()); cache.insert(leaked); leaked } #[derive(Debug, Clone, Default)] pub struct ConsoleQueryResult { pub columns: Vec, pub rows: Vec>, pub rows_affected: u64, pub is_select: bool, } /// Statements starting with `SELECT` `WITH` `PRAGMA` or `EXPLAIN` are fetched /// into a table everything else is executed and reports `rows_affected` pub async fn run_console_query(pool: &DbPool, sql: &str) -> Result { let sql = intern(sql.trim().trim_end_matches(';').trim()); if returns_rows(sql) { let fetched = sqlx::query(sql).persistent(false).fetch_all(pool).await?; Ok(rows_to_result(&fetched)) } else { let done = sqlx::query(sql).persistent(false).execute(pool).await?; Ok(ConsoleQueryResult { rows_affected: done.rows_affected(), is_select: false, ..Default::default() }) } } fn returns_rows(sql: &str) -> bool { let head = sql .trim_start() .split(|c: char| c.is_whitespace() || c == '(') .next() .unwrap_or("") .to_ascii_uppercase(); matches!(head.as_str(), "SELECT" | "WITH" | "PRAGMA" | "EXPLAIN") } fn rows_to_result(fetched: &[SqliteRow]) -> ConsoleQueryResult { let columns = fetched .first() .map(|row| { row.columns() .iter() .map(|c| c.name().to_string()) .collect::>() }) .unwrap_or_default(); let rows = fetched .iter() .map(|row| { (0..row.columns().len()) .map(|idx| cell_to_string(row, idx)) .collect() }) .collect(); ConsoleQueryResult { columns, rows, rows_affected: 0, is_select: true, } } /// SQLite is dynamically typed so probe the raw value and decode accordingly fn cell_to_string(row: &SqliteRow, idx: usize) -> String { match row.try_get_raw(idx) { Ok(raw) if raw.is_null() => return "NULL".to_string(), Ok(_) => {} Err(_) => return "".to_string(), } if let Ok(v) = row.try_get::(idx) { return v; } if let Ok(v) = row.try_get::(idx) { return v.to_string(); } if let Ok(v) = row.try_get::(idx) { return v.to_string(); } if let Ok(v) = row.try_get::, _>(idx) { return format!("", v.len()); } "".to_string() } #[cfg(test)] mod tests { use super::*; use sqlx::sqlite::SqlitePoolOptions; async fn pool() -> DbPool { let pool = SqlitePoolOptions::new() .connect("sqlite::memory:") .await .unwrap(); sqlx::query("CREATE TABLE t (id INTEGER PRIMARY KEY, name TEXT, ratio REAL)") .execute(&pool) .await .unwrap(); pool } #[tokio::test] async fn select_returns_columns_and_rows() { let pool = pool().await; sqlx::query("INSERT INTO t (name, ratio) VALUES ('a', 1.5), (NULL, 2.0)") .execute(&pool) .await .unwrap(); let res = run_console_query(&pool, "SELECT id, name, ratio FROM t ORDER BY id;") .await .unwrap(); assert!(res.is_select); assert_eq!(res.columns, vec!["id", "name", "ratio"]); assert_eq!(res.rows.len(), 2); assert_eq!(res.rows[0], vec!["1", "a", "1.5"]); assert_eq!(res.rows[1][1], "NULL"); } #[tokio::test] async fn non_select_reports_rows_affected() { let pool = pool().await; let res = run_console_query(&pool, "INSERT INTO t (name) VALUES ('x'), ('y')") .await .unwrap(); assert!(!res.is_select); assert_eq!(res.rows_affected, 2); assert!(res.rows.is_empty()); } #[tokio::test] async fn invalid_sql_errors() { let pool = pool().await; assert!( run_console_query(&pool, "SELECT * FROM nope") .await .is_err() ); } }