diff --git a/desktop/src-tauri/src/commands/query.rs b/desktop/src-tauri/src/commands/query.rs index 783b42f..b4743a6 100644 --- a/desktop/src-tauri/src/commands/query.rs +++ b/desktop/src-tauri/src/commands/query.rs @@ -10,7 +10,7 @@ //! 3. Every query is recorded in the local `query_history` table. use crate::db::pool::DbHandle; -use crate::models::db_viewer::{ColumnInfo, QueryResult}; +use crate::models::db_viewer::{ColumnInfo, MultiQueryResult, QueryResult, StatementNotice}; use serde::{Deserialize, Serialize}; use sqlx::{Column, Row}; use std::time::Instant; @@ -708,6 +708,507 @@ async fn execute_mysql_raw( }) } +// --------------------------------------------------------------------------- +// Multiple result sets (execute_query_multi) +// --------------------------------------------------------------------------- +// Per the spec: statements execute sequentially; SELECT-producing statements +// get wrapped pagination (first page per set + total count — per-set paging is +// deliberately omitted because re-running a script would re-execute DML); +// DML/DDL produce notices; the first error stops the run, prior sets survive. + +const MAX_STATEMENTS: usize = 50; + +pub(crate) enum StatementOutcome { + Set(QueryResult), + Notice { + kind: &'static str, + text: String, + affected: Option, + }, +} + +/// Classify a statement by its first significant keyword (execution-time +/// classification per spec §4 — NOT in the tokenizer). +pub(crate) fn statement_kind(stmt: &str) -> &'static str { + let k: String = stmt + .trim_start() + .chars() + .take_while(|c| c.is_ascii_alphabetic()) + .collect::() + .to_lowercase(); + match k.as_str() { + "select" | "with" | "values" | "table" => "select", + "insert" | "update" | "delete" => "dml", + _ => "ddl", + } +} + +/// Parse a PG CommandComplete tag ("INSERT 0 5" → 5; "CREATE TABLE" → None). +pub(crate) fn affected_from_tag(tag: &str) -> Option { + tag.split_whitespace().last()?.parse::().ok() +} + +fn notice_for( + idx: usize, + kind: &str, + text: String, + affected: Option, +) -> StatementNotice { + StatementNotice { + statement_index: idx, + kind: kind.to_string(), + text, + affected, + } +} + +// ── SQLite walker ───────────────────────────────────────────────── + +pub(crate) fn execute_sqlite_statement( + conn: &rusqlite::Connection, + stmt: &str, + page: i64, + page_size: i64, +) -> Result { + match conn.execute(stmt, []) { + Ok(n) => Ok(StatementOutcome::Notice { + kind: if statement_kind(stmt) == "ddl" { "ddl" } else { "dml" }, + text: format!("OK — {n} rows affected"), + affected: Some(n as i64), + }), + Err(rusqlite::Error::ExecuteReturnedResults) => { + let (columns, all_rows) = execute_sqlite_with_query(conn, stmt)?; + let total_rows = all_rows.len() as i64; + let rows: Vec> = all_rows + .into_iter() + .take(page_size.max(0) as usize) + .collect(); + Ok(StatementOutcome::Set(QueryResult { + columns, + rows, + total_rows, + page, + page_size, + execution_time_ms: None, + })) + } + Err(rusqlite::Error::SqliteFailure(e, _)) + if e.code == rusqlite::ErrorCode::OperationInterrupted => + { + Err("Query cancelled".to_string()) + } + Err(e) => Err(e.to_string()), + } +} + +pub(crate) fn run_sqlite_statements( + conn: &rusqlite::Connection, + statements: &[String], + page: i64, + page_size: i64, + result_sets: &mut Vec, + notices: &mut Vec, +) -> Option { + for (idx, stmt) in statements.iter().enumerate() { + match execute_sqlite_statement(conn, stmt, page, page_size) { + Ok(StatementOutcome::Set(qr)) => result_sets.push(qr), + Ok(StatementOutcome::Notice { kind, text, affected }) => { + notices.push(notice_for(idx, kind, text, affected)) + } + Err(e) if e == "Query cancelled" => return Some(e), + Err(e) => { + notices.push(StatementNotice { + statement_index: idx, + kind: "error".into(), + text: e.clone(), + affected: None, + }); + return Some(format!("Statement {} failed: {e}", idx + 1)); + } + } + } + None +} + +// ── PostgreSQL walker ───────────────────────────────────────────── + +pub(crate) async fn execute_pg_statement( + client: &tokio_postgres::Client, + stmt: &str, + page: i64, + page_size: i64, +) -> Result { + let trimmed = stmt.trim(); + let off = (page.saturating_sub(1).max(0)) * page_size; + + // 1. Try wrapped pagination — a parse error here executes nothing. + let wrapped_count = format!("SELECT COUNT(*) FROM ({}) AS _gridline_cnt", trimmed); + if let Ok(row) = client.query_one(&wrapped_count, &[]).await { + let total_rows: i64 = row.get(0); + let wrapped_data = format!( + "SELECT * FROM ({}) AS _gridline_data LIMIT $1 OFFSET $2", + trimmed + ); + let data_rows = client + .query(&wrapped_data, &[&page_size, &off]) + .await + .map_err(|e| { + if is_pg_cancel_error(&e) { + "Query cancelled".to_string() + } else { + format!("{e}") + } + })?; + let columns: Vec = match data_rows.first() { + Some(first) => first + .columns() + .iter() + .map(|c| ColumnInfo { + name: c.name().to_string(), + data_type: format!("{:?}", c.type_()), + is_nullable: true, + is_pk: false, + is_fk: false, + fk_ref: None, + default_value: None, + editable: true, + is_generated: false, + }) + .collect(), + None => Vec::new(), + }; + let rows: Vec> = data_rows + .iter() + .map(|row| { + (0..row.len()) + .map(|i| crate::commands::db_viewer::pg_value_to_json(row, i)) + .collect() + }) + .collect(); + return Ok(StatementOutcome::Set(QueryResult { + columns, + rows, + total_rows, + page, + page_size, + execution_time_ms: None, + })); + } + + // 2. Raw single-statement execution via simple_query (captures the + // CommandComplete tag so DML notices carry affected counts). + let messages = client + .simple_query(trimmed) + .await + .map_err(|e| { + if is_pg_cancel_error(&e) { + "Query cancelled".to_string() + } else { + crate::commands::db_viewer::pg_error_message(&e) + } + })?; + + let mut columns: Vec = Vec::new(); + let mut all_rows: Vec> = Vec::new(); + let mut last_tag: Option = None; + for msg in messages { + match msg { + tokio_postgres::SimpleQueryMessage::Row(row) => { + if columns.is_empty() { + columns = row + .columns() + .iter() + .map(|c| ColumnInfo { + name: c.name().to_string(), + data_type: "text".to_string(), + is_nullable: true, + is_pk: false, + is_fk: false, + fk_ref: None, + default_value: None, + editable: true, + is_generated: false, + }) + .collect(); + } + let values: Vec = (0..row.len()) + .map(|i| match row.try_get::(i) { + Ok(Some(s)) => serde_json::Value::String(s.to_string()), + _ => serde_json::Value::Null, + }) + .collect(); + all_rows.push(values); + } + tokio_postgres::SimpleQueryMessage::CommandComplete(tag) => { + last_tag = Some(tag.to_string()); + } + _ => {} + } + } + + if all_rows.is_empty() { + let tag = last_tag.clone().unwrap_or_else(|| "OK".to_string()); + let affected = last_tag.as_deref().and_then(affected_from_tag); + return Ok(StatementOutcome::Notice { + kind: if statement_kind(stmt) == "ddl" { "ddl" } else { "dml" }, + text: tag, + affected, + }); + } + + let total_rows = all_rows.len() as i64; + let rows: Vec> = all_rows + .into_iter() + .take(page_size.max(0) as usize) + .collect(); + Ok(StatementOutcome::Set(QueryResult { + columns, + rows, + total_rows, + page, + page_size, + execution_time_ms: None, + })) +} + +pub(crate) async fn run_pg_statements( + client: &tokio_postgres::Client, + statements: &[String], + page: i64, + page_size: i64, + result_sets: &mut Vec, + notices: &mut Vec, +) -> Option { + for (idx, stmt) in statements.iter().enumerate() { + match execute_pg_statement(client, stmt, page, page_size).await { + Ok(StatementOutcome::Set(qr)) => result_sets.push(qr), + Ok(StatementOutcome::Notice { kind, text, affected }) => { + notices.push(notice_for(idx, kind, text, affected)) + } + Err(e) if e == "Query cancelled" => return Some(e), + Err(e) => { + notices.push(StatementNotice { + statement_index: idx, + kind: "error".into(), + text: e.clone(), + affected: None, + }); + return Some(format!("Statement {} failed: {e}", idx + 1)); + } + } + } + None +} + +// ── MySQL walker ────────────────────────────────────────────────── + +async fn execute_mysql_statement_on( + conn: &mut sqlx::mysql::MySqlConnection, + stmt: &str, + page: i64, + page_size: i64, +) -> Result { + let trimmed = stmt.trim(); + let off = (page.saturating_sub(1).max(0)) * page_size; + + let total_rows: i64 = match sqlx::query_scalar::<_, i64>(&mysql_wrap_count(trimmed)) + .fetch_one(&mut *conn) + .await + { + Ok(n) => n, + Err(e) if is_mysql_cancel_error(&e) => return Err("Query cancelled".to_string()), + Err(_) => { + return execute_mysql_raw(conn, trimmed, page, page_size, off) + .await + .map(|qr| { + if qr.columns.is_empty() { + StatementOutcome::Notice { + kind: if statement_kind(stmt) == "ddl" { "ddl" } else { "dml" }, + text: "OK".to_string(), + affected: None, + } + } else { + StatementOutcome::Set(qr) + } + }); + } + }; + + let data_rows = match sqlx::query(&mysql_wrap_data(trimmed)) + .bind(page_size) + .bind(off) + .fetch_all(&mut *conn) + .await + { + Ok(rows) => rows, + Err(e) if is_mysql_cancel_error(&e) => return Err("Query cancelled".to_string()), + Err(e) => return Err(crate::commands::db_viewer::sanitize_error(&format!("{e}"))), + }; + let columns: Vec = match data_rows.first() { + Some(first) => first + .columns() + .iter() + .map(|c| ColumnInfo { + name: c.name().to_string(), + data_type: c.type_info().to_string(), + is_nullable: true, + is_pk: false, + is_fk: false, + fk_ref: None, + default_value: None, + editable: true, + is_generated: false, + }) + .collect(), + None => Vec::new(), + }; + let rows: Vec> = data_rows + .iter() + .map(|r| (0..r.len()).map(|i| mysql_cell_to_json(r, i)).collect()) + .collect(); + Ok(StatementOutcome::Set(QueryResult { + columns, + rows, + total_rows, + page, + page_size, + execution_time_ms: None, + })) +} + +async fn run_mysql_statements( + conn: &mut sqlx::mysql::MySqlConnection, + statements: &[String], + page: i64, + page_size: i64, + result_sets: &mut Vec, + notices: &mut Vec, + cancel_registry: &crate::cancel::CancelRegistry, + connection_id: &str, +) -> Option { + let conn_id: i64 = sqlx::query_scalar("SELECT CAST(CONNECTION_ID() AS SIGNED)") + .fetch_one(&mut *conn) + .await + .unwrap_or(-1); + cancel_registry.set_mysql_conn_id(connection_id, Some(conn_id)); + for (idx, stmt) in statements.iter().enumerate() { + match execute_mysql_statement_on(conn, stmt, page, page_size).await { + Ok(StatementOutcome::Set(qr)) => result_sets.push(qr), + Ok(StatementOutcome::Notice { kind, text, affected }) => { + notices.push(notice_for(idx, kind, text, affected)) + } + Err(e) if e == "Query cancelled" => { + cancel_registry.set_mysql_conn_id(connection_id, None); + return Some(e); + } + Err(e) => { + notices.push(StatementNotice { + statement_index: idx, + kind: "error".into(), + text: e.clone(), + affected: None, + }); + cancel_registry.set_mysql_conn_id(connection_id, None); + return Some(format!("Statement {} failed: {e}", idx + 1)); + } + } + } + cancel_registry.set_mysql_conn_id(connection_id, None); + None +} + +/// Execute a multi-statement script, returning every result set and notices. +/// `execute_query` is untouched; the query workspace calls this instead. +#[tauri::command] +pub async fn execute_query_multi( + connection_id: String, + query: String, + page: Option, + page_size: Option, + state: State<'_, crate::AppState>, +) -> Result { + let start = Instant::now(); + let p = page.unwrap_or(1); + let ps = page_size.unwrap_or(50); + + let statements = crate::db::sql_split::split_statements(&query); + if statements.is_empty() { + return Err("Query cannot be empty".to_string()); + } + if statements.len() > MAX_STATEMENTS { + return Err(format!( + "Too many statements: {} (max {MAX_STATEMENTS})", + statements.len() + )); + } + + let mut result_sets: Vec = Vec::new(); + let mut notices: Vec = Vec::new(); + + let err_text: Option; + let mut pm = state.pool_manager.lock().await; + match pm.get(&connection_id) { + Some(DbHandle::Postgresql(client, _)) => { + err_text = + run_pg_statements(client, &statements, p, ps, &mut result_sets, &mut notices).await; + } + Some(DbHandle::Sqlite(conn)) => { + err_text = run_sqlite_statements(conn, &statements, p, ps, &mut result_sets, &mut notices); + } + Some(DbHandle::MySql(pool)) => { + let mut conn = pool + .acquire() + .await + .map_err(|e| crate::commands::db_viewer::sanitize_error(&format!("{e}")))?; + err_text = run_mysql_statements( + &mut *conn, + &statements, + p, + ps, + &mut result_sets, + &mut notices, + &state.cancel_registry, + &connection_id, + ) + .await; + } + None => return Err("Connection not found".to_string()), + } + + let elapsed = start.elapsed().as_millis() as i64; + let row_count = result_sets.first().map(|q| q.rows.len() as i64).unwrap_or(0); + let history_id = Uuid::new_v4().to_string(); + if let Some(e) = &err_text { + insert_history( + &state.db_store, + &history_id, + &connection_id, + &query, + Some(elapsed), + Some(row_count), + "error", + Some(e.as_str()), + ); + } else { + insert_history( + &state.db_store, + &history_id, + &connection_id, + &query, + Some(elapsed), + Some(row_count), + "success", + None, + ); + } + + Ok(MultiQueryResult { + result_sets, + notices, + execution_time_ms: elapsed, + }) +} + // --------------------------------------------------------------------------- // Query history commands // --------------------------------------------------------------------------- @@ -1378,3 +1879,7 @@ mod tests { assert!(res.unwrap_err().contains("Query cancelled")); } } + +#[cfg(test)] +#[path = "query.test.rs"] +mod query_tests; diff --git a/desktop/src-tauri/src/commands/query.test.rs b/desktop/src-tauri/src/commands/query.test.rs new file mode 100644 index 0000000..54ffb59 --- /dev/null +++ b/desktop/src-tauri/src/commands/query.test.rs @@ -0,0 +1,134 @@ +use super::*; + +// ── statement_kind / affected_from_tag (pure) ────────────────────── + +#[test] +fn statement_kind_classifies_by_first_keyword() { + assert_eq!(statement_kind("SELECT 1"), "select"); + assert_eq!(statement_kind(" with c as (select 1) select * from c"), "select"); + assert_eq!(statement_kind("insert into t values (1)"), "dml"); + assert_eq!(statement_kind("UPDATE t SET a=1"), "dml"); + assert_eq!(statement_kind("DELETE FROM t"), "dml"); + assert_eq!(statement_kind("CREATE TABLE t (a int)"), "ddl"); + assert_eq!(statement_kind("-- c\nDROP TABLE t"), "ddl"); +} + +#[test] +fn affected_from_tag_parses_pg_command_tags() { + assert_eq!(affected_from_tag("INSERT 0 5"), Some(5)); + assert_eq!(affected_from_tag("UPDATE 3"), Some(3)); + assert_eq!(affected_from_tag("CREATE TABLE"), None); +} + +// ── SQLite walker — real in-memory database, no env vars, not ignored ── + +fn mem_conn() -> rusqlite::Connection { + let conn = rusqlite::Connection::open_in_memory().expect("open in-memory"); + conn.execute_batch( + "CREATE TABLE t (a INTEGER, b TEXT); + INSERT INTO t VALUES (1, 'x'), (2, 'y'), (3, 'z');", + ) + .expect("seed"); + conn +} + +#[test] +fn sqlite_walker_yields_notice_for_dml_with_affected() { + let conn = mem_conn(); + let out = execute_sqlite_statement(&conn, "INSERT INTO t VALUES (4, 'w')", 1, 50) + .expect("statement runs"); + match out { + StatementOutcome::Notice { kind, affected, .. } => { + assert_eq!(kind, "dml"); + assert_eq!(affected, Some(1)); + } + _ => panic!("expected Notice, got Set"), + } +} + +#[test] +fn sqlite_walker_yields_set_for_select() { + let conn = mem_conn(); + let out = execute_sqlite_statement(&conn, "SELECT a FROM t ORDER BY a", 1, 2) + .expect("statement runs"); + match out { + StatementOutcome::Set(qr) => { + assert_eq!(qr.columns.len(), 1); + assert_eq!(qr.rows.len(), 2, "page size caps rows"); + assert_eq!(qr.total_rows, 3, "total counts all rows"); + } + _ => panic!("expected Set, got Notice"), + } +} + +#[test] +fn sqlite_walker_notices_ddl() { + let conn = mem_conn(); + let out = execute_sqlite_statement(&conn, "CREATE TABLE u (x INT)", 1, 50).unwrap(); + match out { + StatementOutcome::Notice { kind, .. } => assert_eq!(kind, "ddl"), + _ => panic!("expected Notice"), + } +} + +#[test] +fn run_sqlite_statements_stops_at_first_error_with_prior_sets() { + let conn = mem_conn(); + let statements = crate::db::sql_split::split_statements( + "SELECT a FROM t; SELECT * FROM missing; SELECT 1", + ); + let mut sets = Vec::new(); + let mut notices = Vec::new(); + let err = run_sqlite_statements(&conn, &statements, 1, 50, &mut sets, &mut notices); + assert_eq!(sets.len(), 1, "first set is returned"); + assert!(err.is_some(), "the failed statement stops the run"); +} + +/// Live PG walker test — mirror the env/client conventions of the live PG +/// tests in `commands/objects.test.rs`. Assert: `SELECT 1; SELECT 2` (split +/// via sql_split) produces two Sets via `run_pg_statements`, and +/// `UPDATE`-on-missing-table produces a first-error stop with the error +/// notice attached. +#[tokio::test] +#[ignore] +async fn pg_multi_statement_live() { + let h = std::env::var("GRIDLINE_TEST_PG_HOST").expect("set GRIDLINE_TEST_PG_HOST"); + let p: u16 = std::env::var("GRIDLINE_TEST_PG_PORT") + .unwrap_or_else(|_| "5432".into()) + .parse() + .unwrap(); + let u = std::env::var("GRIDLINE_TEST_PG_USER").expect("set GRIDLINE_TEST_PG_USER"); + let d = std::env::var("GRIDLINE_TEST_PG_DB").expect("set GRIDLINE_TEST_PG_DB"); + let pw = std::env::var("GRIDLINE_TEST_PG_PASSWORD").unwrap_or_default(); + let (client, conn) = tokio_postgres::connect( + &format!("host={h} port={p} user={u} dbname={d} password={pw}"), + tokio_postgres::NoTls, + ) + .await + .expect("connect to test PG"); + let handle = tokio::spawn(async move { + let _ = conn.await; + }); + + // Two SELECTs → two Sets, no notices. + let stmts = crate::db::sql_split::split_statements("SELECT 1; SELECT 2"); + let mut sets = vec![]; + let mut notices = vec![]; + let err = run_pg_statements(&client, &stmts, 1, 50, &mut sets, &mut notices).await; + assert!(err.is_none()); + assert_eq!(sets.len(), 2); + assert!(notices.is_empty()); + + // UPDATE on a missing table → first-error stop with an error notice. + let bad = crate::db::sql_split::split_statements("UPDATE gridline_missing SET a=1"); + let mut sets2 = vec![]; + let mut notices2 = vec![]; + let err2 = run_pg_statements(&client, &bad, 1, 50, &mut sets2, &mut notices2).await; + assert!(err2.is_some(), "missing table must stop the run"); + assert!( + notices2.iter().any(|n| n.kind == "error"), + "an error notice must be attached" + ); + + let _ = handle; +} diff --git a/desktop/src-tauri/src/lib.rs b/desktop/src-tauri/src/lib.rs index e8f75ff..777d131 100644 --- a/desktop/src-tauri/src/lib.rs +++ b/desktop/src-tauri/src/lib.rs @@ -174,6 +174,7 @@ pub fn run() { settings::import_settings, schema_graph::get_schema_graph, query::execute_query, + query::execute_query_multi, query::cancel_query, query::get_query_history, query::clear_query_history, diff --git a/desktop/src-tauri/src/models/db_viewer.rs b/desktop/src-tauri/src/models/db_viewer.rs index 371af06..f884ea3 100644 --- a/desktop/src-tauri/src/models/db_viewer.rs +++ b/desktop/src-tauri/src/models/db_viewer.rs @@ -72,6 +72,22 @@ pub struct QueryResult { pub execution_time_ms: Option, } +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct StatementNotice { + pub statement_index: usize, + /// "dml" | "ddl" | "error" + pub kind: String, + pub text: String, + pub affected: Option, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct MultiQueryResult { + pub result_sets: Vec, + pub notices: Vec, + pub execution_time_ms: i64, +} + #[derive(Debug, Clone, Serialize, Deserialize)] pub struct Pagination { pub page: i64,