diff --git a/agents/drivers/duckdb/src/query.rs b/agents/drivers/duckdb/src/query.rs index 48147455bb..813f61b55d 100644 --- a/agents/drivers/duckdb/src/query.rs +++ b/agents/drivers/duckdb/src/query.rs @@ -45,6 +45,22 @@ mod tests { assert_eq!(result.rows[2][0], serde_json::json!(99.999)); } + #[test] + fn duckdb_execute_returns_dml_returning_rows() { + let con = duckdb::Connection::open_in_memory().expect("connect in-memory DuckDB"); + con.execute_batch("CREATE TABLE users (id INTEGER, name VARCHAR)").expect("create table"); + + let inserted = duckdb_execute(&con, "INSERT INTO users VALUES (1, 'Ada') RETURNING id, name") + .expect("insert returning"); + let updated = duckdb_execute(&con, "UPDATE users SET name = 'Ada Lovelace' RETURNING id, name") + .expect("update returning"); + let deleted = duckdb_execute(&con, "DELETE FROM users RETURNING id, name").expect("delete returning"); + + assert_eq!(inserted.rows, vec![vec![serde_json::json!(1), serde_json::json!("Ada")]]); + assert_eq!(updated.rows, vec![vec![serde_json::json!(1), serde_json::json!("Ada Lovelace")]]); + assert_eq!(deleted.rows, vec![vec![serde_json::json!(1), serde_json::json!("Ada Lovelace")]]); + } + #[test] fn duckdb_execute_returns_rows_for_from_first_query() { let con = duckdb::Connection::open_in_memory().expect("connect in-memory DuckDB"); diff --git a/agents/drivers/duckdb/src/sql.rs b/agents/drivers/duckdb/src/sql.rs index 2354ad7b70..b69c0c8c81 100644 --- a/agents/drivers/duckdb/src/sql.rs +++ b/agents/drivers/duckdb/src/sql.rs @@ -202,6 +202,62 @@ pub fn starts_with_duckdb_result_sql_keyword(sql: &str) -> bool { .any(|keyword| { token.eq_ignore_ascii_case(keyword) || (*keyword == "DESCRIBE" && token.eq_ignore_ascii_case("DESC")) }) + || matches!(token.to_ascii_uppercase().as_str(), "INSERT" | "UPDATE" | "DELETE" | "MERGE") + && contains_unquoted_sql_keyword(sql, "RETURNING") +} + +fn contains_unquoted_sql_keyword(sql: &str, keyword: &str) -> bool { + let bytes = sql.as_bytes(); + let mut index = 0; + let mut in_single = false; + let mut in_double = false; + + while index < bytes.len() { + match bytes[index] { + b'\'' if !in_double => { + if in_single && bytes.get(index + 1) == Some(&b'\'') { + index += 2; + continue; + } + in_single = !in_single; + index += 1; + } + b'"' if !in_single => { + if in_double && bytes.get(index + 1) == Some(&b'"') { + index += 2; + continue; + } + in_double = !in_double; + index += 1; + } + b'-' if !in_single && !in_double && bytes.get(index + 1) == Some(&b'-') => { + index += 2; + while index < bytes.len() && bytes[index] != b'\n' { + index += 1; + } + } + b'/' if !in_single && !in_double && bytes.get(index + 1) == Some(&b'*') => { + index += 2; + while index + 1 < bytes.len() && !(bytes[index] == b'*' && bytes[index + 1] == b'/') { + index += 1; + } + index = (index + 2).min(bytes.len()); + } + byte if !in_single && !in_double && (byte.is_ascii_alphabetic() || byte == b'_') => { + let start = index; + index += 1; + while index < bytes.len() && (bytes[index].is_ascii_alphanumeric() || bytes[index] == b'_') { + index += 1; + } + if sql[start..index].eq_ignore_ascii_case(keyword) { + return true; + } + } + _ => index += 1, + } + } + + false } fn first_executable_sql_token(sql: &str) -> Option<&str> { @@ -269,6 +325,8 @@ mod tests { fn detects_result_statements_after_comments() { assert!(starts_with_duckdb_result_sql_keyword("/* note */ WITH rows AS (SELECT 1) SELECT * FROM rows")); assert!(starts_with_duckdb_result_sql_keyword("DESC SELECT 1")); + assert!(starts_with_duckdb_result_sql_keyword("INSERT INTO items VALUES (1) RETURNING id")); + assert!(!starts_with_duckdb_result_sql_keyword("INSERT INTO items(note) VALUES ('RETURNING')")); assert!(!starts_with_duckdb_result_sql_keyword("INSERT INTO items VALUES (1)")); } } diff --git a/crates/dbx-core/src/agent_tools.rs b/crates/dbx-core/src/agent_tools.rs index 721154721e..cffa8089c4 100644 --- a/crates/dbx-core/src/agent_tools.rs +++ b/crates/dbx-core/src/agent_tools.rs @@ -616,8 +616,10 @@ async fn execute_execute_query( /// Format a QueryResult as a Markdown table for LLM consumption. fn format_query_result_as_text(result: &QueryResult, limit: usize) -> Result { - if result.rows.is_empty() { - return Ok("Query returned 0 rows.".to_string()); + // A result without columns is a command result, not an empty result set. + // This is how drivers represent DML that does not use RETURNING. + if result.columns.is_empty() { + return Ok(format!("Query executed. {} row(s) affected.", result.affected_rows)); } let mut lines = Vec::new(); @@ -1135,6 +1137,50 @@ for line in sys.stdin: assert!(names.contains(&"explain_query")); } + fn query_result(columns: Vec<&str>, rows: Vec>, affected_rows: u64) -> QueryResult { + QueryResult { + columns: columns.into_iter().map(str::to_string).collect(), + column_types: Vec::new(), + column_sortables: Vec::new(), + spatial_columns: Vec::new(), + spatial_values: Vec::new(), + rows, + affected_rows, + execution_time_ms: 1, + truncated: false, + session_id: None, + has_more: false, + elasticsearch_raw_body: None, + } + } + + #[test] + fn query_result_formatter_reports_dml_affected_rows() { + let result = query_result(vec![], vec![], 2); + + assert_eq!(format_query_result_as_text(&result, 50).unwrap(), "Query executed. 2 row(s) affected."); + } + + #[test] + fn query_result_formatter_distinguishes_zero_row_dml_from_an_empty_result_set() { + let dml = query_result(vec![], vec![], 0); + let returning = query_result(vec!["id", "name"], vec![], 0); + + assert_eq!(format_query_result_as_text(&dml, 50).unwrap(), "Query executed. 0 row(s) affected."); + assert_eq!(format_query_result_as_text(&returning, 50).unwrap(), "| id | name |\n|---|---|\n(0 rows, 1ms)"); + } + + #[test] + fn query_result_formatter_renders_returning_rows() { + let result = + query_result(vec!["id", "name"], vec![vec![serde_json::json!(5), serde_json::json!("returning")]], 0); + + assert_eq!( + format_query_result_as_text(&result, 50).unwrap(), + "| id | name |\n|---|---|\n| 5 | returning |\n(1 rows, 1ms)" + ); + } + #[test] fn sample_data_sql_uses_database_identifier_and_limit_syntax() { assert_eq!( diff --git a/crates/dbx-core/src/db/mysql.rs b/crates/dbx-core/src/db/mysql.rs index 5feffe0191..68415c8f5e 100644 --- a/crates/dbx-core/src/db/mysql.rs +++ b/crates/dbx-core/src/db/mysql.rs @@ -4,6 +4,9 @@ use mysql_async::consts::ColumnType; use mysql_async::prelude::*; use percent_encoding::percent_decode_str; use rust_decimal::Decimal; +use sqlparser::ast::Statement; +use sqlparser::dialect::MySqlDialect; +use sqlparser::parser::Parser; use std::borrow::Cow; use std::collections::{HashMap, HashSet}; use std::path::PathBuf; @@ -4243,12 +4246,31 @@ fn prefers_text_protocol_query(sql: &str, dialect: MySqlQueryDialect) -> bool { is_result_set_query(sql, dialect) || requires_text_protocol_query(sql, dialect) } -fn is_result_set_query(sql: &str, dialect: MySqlQueryDialect) -> bool { +pub(crate) fn is_result_set_query(sql: &str, dialect: MySqlQueryDialect) -> bool { starts_with_executable_sql_keyword_for_database( sql, &["SELECT", "SHOW", "DESCRIBE", "EXPLAIN", "WITH", "CALL"], DatabaseType::Mysql, - ) || dialect.supports_admin_show_results && is_admin_show_query(sql) + ) || mysql_statement_returns_rows(sql) + || dialect.supports_admin_show_results && is_admin_show_query(sql) +} + +/// MariaDB 10.5+ returns a result set for INSERT/DELETE/REPLACE ... RETURNING. +/// Route it through the existing query path; MySQL servers that do not support +/// the syntax still return their normal SQL syntax error. +fn mysql_statement_returns_rows(sql: &str) -> bool { + let Ok(statements) = Parser::parse_sql(&MySqlDialect {}, sql) else { + return false; + }; + let [statement] = statements.as_slice() else { + return false; + }; + + match statement { + Statement::Insert(insert) => insert.returning.is_some(), + Statement::Delete(delete) => delete.returning.is_some(), + _ => false, + } } fn requires_text_protocol_query(sql: &str, dialect: MySqlQueryDialect) -> bool { @@ -4677,6 +4699,15 @@ mod tests { assert!(is_result_set_query(sql, MySqlQueryDialect::default())); } + #[test] + fn mariadb_returning_dml_is_treated_as_a_result_set() { + let dialect = MySqlQueryDialect::default(); + + assert!(is_result_set_query("INSERT INTO users (id) VALUES (1) RETURNING id", dialect)); + assert!(is_result_set_query("DELETE FROM users WHERE id = 1 RETURNING id", dialect)); + assert!(!is_result_set_query("UPDATE users SET name = 'Ada'", dialect)); + } + #[test] fn mysql_hash_comments_before_queries_preserve_result_sets_per_issue_3830() { let dialect = MySqlQueryDialect::default(); diff --git a/crates/dbx-core/src/db/postgres.rs b/crates/dbx-core/src/db/postgres.rs index 7b48c3bc23..2add896ece 100644 --- a/crates/dbx-core/src/db/postgres.rs +++ b/crates/dbx-core/src/db/postgres.rs @@ -8,6 +8,9 @@ use rustls::client::verify_server_cert_signed_by_trust_anchor; use rustls::crypto::{verify_tls12_signature, verify_tls13_signature, CryptoProvider}; use rustls::pki_types::{CertificateDer, PrivateKeyDer, ServerName, UnixTime}; use rustls::server::ParsedCertificate; +use sqlparser::ast::Statement; +use sqlparser::dialect::PostgreSqlDialect; +use sqlparser::parser::Parser; use std::collections::{BTreeSet, HashMap}; use std::fs::File; use std::future::Future; @@ -3219,6 +3222,30 @@ fn query_result_row_limit(max_rows: Option) -> usize { max_rows.unwrap_or(crate::query::MAX_ROWS).max(1) } +/// Returns whether PostgreSQL should execute this statement through the row +/// retrieval path. DML without `RETURNING` needs `execute` for its command +/// tag/affected-row count, while DML with `RETURNING` produces a result set. +pub(crate) fn postgres_statement_returns_rows(sql: &str) -> bool { + if starts_with_executable_sql_keyword(sql, &["SELECT", "SHOW", "EXPLAIN", "WITH", "TABLE"]) { + return true; + } + + let Ok(statements) = Parser::parse_sql(&PostgreSqlDialect {}, sql) else { + return false; + }; + let [statement] = statements.as_slice() else { + return false; + }; + + match statement { + Statement::Insert(insert) => insert.returning.is_some(), + Statement::Update(update) => update.returning.is_some(), + Statement::Delete(delete) => delete.returning.is_some(), + Statement::Merge(merge) => merge.output.is_some(), + _ => false, + } +} + pub async fn execute_query(pool: &Pool, sql: &str) -> Result { execute_query_with_max_rows(pool, sql, None).await } @@ -3231,7 +3258,7 @@ pub async fn execute_query_with_max_rows( let start = Instant::now(); let row_limit = query_result_row_limit(max_rows); - if starts_with_executable_sql_keyword(sql, &["SELECT", "SHOW", "EXPLAIN", "WITH", "TABLE"]) { + if postgres_statement_returns_rows(sql) { let client = checkout_postgres_client(pool, None, super::connection_timeout()).await?; execute_select_query(&client, sql, start, row_limit).await } else { @@ -3729,7 +3756,7 @@ async fn execute_query_with_max_rows_inner( let start = Instant::now(); let row_limit = query_result_row_limit(max_rows); - if starts_with_executable_sql_keyword(sql, &["SELECT", "SHOW", "EXPLAIN", "WITH", "TABLE"]) { + if postgres_statement_returns_rows(sql) { if prefer_text_protocol { execute_select_text(client, sql, start, row_limit, None).await } else { @@ -4678,6 +4705,29 @@ mod tests { assert!(!postgres_requires_single_schema_search_path("ERROR: permission denied for schema admaterial")); } + #[test] + fn postgres_statement_returns_rows_for_returning_dml_only() { + for sql in [ + "INSERT INTO users (id) VALUES (1) RETURNING id", + "UPDATE users SET name = 'Ada' RETURNING id, name", + "DELETE FROM users WHERE id = 1 RETURNING id", + "MERGE INTO users AS target USING updates AS source ON target.id = source.id WHEN MATCHED THEN UPDATE SET name = source.name RETURNING target.id", + "WITH removed AS (DELETE FROM users WHERE id = 1 RETURNING id) SELECT * FROM removed", + ] { + assert!(postgres_statement_returns_rows(sql), "expected result rows for: {sql}"); + } + + for sql in [ + "INSERT INTO users (id) VALUES (1)", + "UPDATE users SET name = 'Ada'", + "DELETE FROM users WHERE id = 1", + "MERGE INTO users AS target USING updates AS source ON target.id = source.id WHEN MATCHED THEN UPDATE SET name = source.name", + "INSERT INTO users (note) VALUES ('RETURNING is text')", + ] { + assert!(!postgres_statement_returns_rows(sql), "expected command result for: {sql}"); + } + } + #[test] fn database_list_does_not_collect_storage_usage() { assert!(list_databases_sql().contains("pg_database")); @@ -4846,6 +4896,58 @@ mod tests { } } + #[tokio::test] + async fn postgres_dml_returning_preserves_result_rows() { + let Some(container) = start_docker_postgres().await else { + return; + }; + let pool = connect(&container.url(), Duration::from_secs(5)).await.expect("connect postgres"); + + execute_query(&pool, "CREATE TABLE dml_returning (id integer PRIMARY KEY, name text NOT NULL)") + .await + .expect("create table"); + + let inserted = execute_query(&pool, "INSERT INTO dml_returning VALUES (1, 'alice'), (2, 'bob')") + .await + .expect("insert rows"); + let updated = execute_query(&pool, "UPDATE dml_returning SET name = 'bob-updated' WHERE id = 2") + .await + .expect("update rows"); + let deleted = execute_query(&pool, "DELETE FROM dml_returning WHERE id = 1").await.expect("delete row"); + let unmatched = + execute_query(&pool, "UPDATE dml_returning SET name = name WHERE id = 999").await.expect("update no rows"); + + assert_eq!(inserted.affected_rows, 2); + assert_eq!(updated.affected_rows, 1); + assert_eq!(deleted.affected_rows, 1); + assert_eq!(unmatched.affected_rows, 0); + + let insert_returning = execute_query(&pool, "INSERT INTO dml_returning VALUES (3, 'carol') RETURNING id, name") + .await + .expect("insert returning"); + let update_returning = + execute_query(&pool, "UPDATE dml_returning SET name = 'carol-updated' WHERE id = 3 RETURNING id, name") + .await + .expect("update returning"); + let delete_returning = execute_query(&pool, "DELETE FROM dml_returning WHERE id = 3 RETURNING id, name") + .await + .expect("delete returning"); + let empty_returning = + execute_query(&pool, "UPDATE dml_returning SET name = name WHERE id = 999 RETURNING id, name") + .await + .expect("empty returning"); + + for result in [&insert_returning, &update_returning, &delete_returning] { + assert_eq!(result.columns, vec!["id", "name"]); + assert_eq!(result.rows.len(), 1); + } + assert_eq!(insert_returning.rows[0], vec![serde_json::json!(3), serde_json::json!("carol")]); + assert_eq!(update_returning.rows[0], vec![serde_json::json!(3), serde_json::json!("carol-updated")]); + assert_eq!(delete_returning.rows[0], vec![serde_json::json!(3), serde_json::json!("carol-updated")]); + assert_eq!(empty_returning.columns, vec!["id", "name"]); + assert!(empty_returning.rows.is_empty()); + } + async fn assert_postgres_18(pool: &Pool) { let version = execute_query(pool, "SHOW server_version_num").await.expect("query PostgreSQL version"); let version_num = version.rows[0][0] diff --git a/crates/dbx-core/src/db/sqlite.rs b/crates/dbx-core/src/db/sqlite.rs index b17445e981..035361bcd0 100644 --- a/crates/dbx-core/src/db/sqlite.rs +++ b/crates/dbx-core/src/db/sqlite.rs @@ -2,6 +2,9 @@ use percent_encoding::percent_decode_str; use rusqlite::functions::{Context, FunctionFlags}; use rusqlite::types::{Value, ValueRef}; use rusqlite::{Connection, LoadExtensionGuard, OpenFlags}; +use sqlparser::ast::Statement; +use sqlparser::dialect::SQLiteDialect; +use sqlparser::parser::Parser; use std::collections::HashSet; use std::io::Read; use std::path::Path; @@ -581,6 +584,45 @@ mod tests { assert_eq!(result.rows[0][0], serde_json::json!("Ada")); } + #[tokio::test] + async fn sqlite_dml_returning_preserves_result_rows() { + let pool = connect_path(":memory:").await.expect("connect in-memory SQLite"); + execute_query(&pool, "CREATE TABLE dml_returning (id INTEGER PRIMARY KEY, name TEXT NOT NULL)") + .await + .expect("create table"); + + let inserted = execute_query(&pool, "INSERT INTO dml_returning VALUES (1, 'alice')").await.expect("insert row"); + let unmatched = + execute_query(&pool, "UPDATE dml_returning SET name = name WHERE id = 999").await.expect("update no rows"); + assert_eq!(inserted.affected_rows, 1); + assert_eq!(unmatched.affected_rows, 0); + + let insert_returning = execute_query(&pool, "INSERT INTO dml_returning VALUES (2, 'bob') RETURNING id, name") + .await + .expect("insert returning"); + let update_returning = + execute_query(&pool, "UPDATE dml_returning SET name = 'bob-updated' WHERE id = 2 RETURNING id, name") + .await + .expect("update returning"); + let delete_returning = execute_query(&pool, "DELETE FROM dml_returning WHERE id = 2 RETURNING id, name") + .await + .expect("delete returning"); + let empty_returning = + execute_query(&pool, "UPDATE dml_returning SET name = name WHERE id = 999 RETURNING id, name") + .await + .expect("empty returning"); + + for result in [&insert_returning, &update_returning, &delete_returning] { + assert_eq!(result.columns, vec!["id", "name"]); + assert_eq!(result.rows.len(), 1); + } + assert_eq!(insert_returning.rows[0], vec![serde_json::json!(2), serde_json::json!("bob")]); + assert_eq!(update_returning.rows[0], vec![serde_json::json!(2), serde_json::json!("bob-updated")]); + assert_eq!(delete_returning.rows[0], vec![serde_json::json!(2), serde_json::json!("bob-updated")]); + assert_eq!(empty_returning.columns, vec!["id", "name"]); + assert!(empty_returning.rows.is_empty()); + } + #[tokio::test] async fn text_affinity_blob_bytes_display_as_utf8_text() { let pool = connect_path(":memory:").await.expect("connect in-memory SQLite"); @@ -2322,6 +2364,28 @@ fn query_result_row_limit(max_rows: Option) -> usize { max_rows.unwrap_or(crate::query::MAX_ROWS).max(1) } +/// DML with `RETURNING` produces a result set in SQLite and must use the +/// statement query API instead of `execute_batch`. +fn sqlite_statement_returns_rows(sql: &str) -> bool { + if starts_with_executable_sql_keyword(sql, &["SELECT", "PRAGMA", "EXPLAIN", "WITH"]) { + return true; + } + + let Ok(statements) = Parser::parse_sql(&SQLiteDialect {}, sql) else { + return false; + }; + let [statement] = statements.as_slice() else { + return false; + }; + + match statement { + Statement::Insert(insert) => insert.returning.is_some(), + Statement::Update(update) => update.returning.is_some(), + Statement::Delete(delete) => delete.returning.is_some(), + _ => false, + } +} + const SQLITE_FUNCTION_ALIASES: &[(&str, &str)] = &[("if", "IIF"), ("substring", "substr")]; fn normalize_sqlite_sql(sql: &str) -> String { @@ -2426,7 +2490,7 @@ fn execute_query_blocking(pool: &SqliteHandle, sql: &str, max_rows: Option>(); let column_decl_types = diff --git a/crates/dbx-core/src/db/sqlserver.rs b/crates/dbx-core/src/db/sqlserver.rs index 52bcfa908e..de7f89d200 100644 --- a/crates/dbx-core/src/db/sqlserver.rs +++ b/crates/dbx-core/src/db/sqlserver.rs @@ -2479,7 +2479,9 @@ pub async fn execute_query_with_max_rows( ) -> Result { let start = Instant::now(); - if starts_with_executable_sql_keyword(sql, &["SELECT", "EXEC", "WITH", "TABLE"]) { + if starts_with_executable_sql_keyword(sql, &["SELECT", "EXEC", "WITH", "TABLE"]) + || sqlserver_dml_output_returns_rows(sql) + { let query = match sqlserver_unsafe_type_query(client, sql).await { Ok(Some(query)) => query, Ok(None) => SqlServerUnsafeTypeQuery::plain(sql), @@ -3156,6 +3158,32 @@ mod tests { assert!(sqlserver_dml_output_returns_rows("DELETE FROM dbo.users OUTPUT deleted.id WHERE id = 1;")); } + #[test] + fn sqlserver_dml_output_detection_distinguishes_client_rows_from_output_into() { + assert!(sqlserver_dml_output_returns_rows( + "INSERT INTO dbo.users(name) OUTPUT inserted.id, inserted.name VALUES (N'Ada')" + )); + assert!(sqlserver_dml_output_returns_rows("UPDATE dbo.users SET active = 1 OUTPUT inserted.id WHERE id = 1")); + assert!(sqlserver_dml_output_returns_rows("DELETE FROM dbo.users OUTPUT deleted.id WHERE id = 1")); + assert!(!sqlserver_dml_output_returns_rows( + "INSERT INTO dbo.audit OUTPUT inserted.id INTO dbo.audit_ids VALUES (1)" + )); + assert!(!sqlserver_dml_output_returns_rows( + "UPDATE dbo.users SET note = N'OUTPUT inserted.id' WHERE id = 1 -- OUTPUT deleted.id" + )); + assert!(!sqlserver_dml_output_returns_rows("SELECT N'OUTPUT inserted.id'")); + } + + #[test] + fn sqlserver_dml_output_uses_the_result_set_execution_path() { + let source = include_str!("sqlserver.rs"); + let execute_query = source.split("pub async fn execute_query_with_max_rows").nth(1).unwrap(); + let execute_query = execute_query.split("pub async fn execute_batch").next().unwrap(); + + assert!(execute_query.contains("sqlserver_dml_output_returns_rows(sql)")); + assert!(execute_query.contains("client.query(query.sql.as_str(), &[])")); + } + #[test] fn sqlserver_user_query_paths_do_not_collect_full_results_before_limiting() { let source = include_str!("sqlserver.rs"); diff --git a/crates/dbx-core/src/query.rs b/crates/dbx-core/src/query.rs index 8506a2bc2e..78cd941fc3 100644 --- a/crates/dbx-core/src/query.rs +++ b/crates/dbx-core/src/query.rs @@ -22,7 +22,7 @@ use crate::database_capabilities; use crate::db; use crate::models::connection::{ConnectionConfig, DatabaseType}; use crate::query_execution_sql::is_write_sql; -use crate::sql::{split_sql_batches, split_sql_statements, starts_with_executable_sql_keyword}; +use crate::sql::{split_sql_batches, split_sql_statements}; use crate::sql_dialect::{resolve_for_db, CAP_TRANSACTIONAL_DDL}; use crate::sql_risk::{classify_sql_risk_for_database, SqlRisk}; @@ -3522,7 +3522,7 @@ async fn execute_manual_txn_postgres_statement( sql: &str, row_limit: usize, ) -> Result { - if starts_with_executable_sql_keyword(sql, &["SELECT", "SHOW", "EXPLAIN", "WITH", "TABLE"]) { + if db::postgres::postgres_statement_returns_rows(sql) { db::postgres::execute_select_query(conn, sql, std::time::Instant::now(), row_limit).await } else { let affected = conn.execute(sql, &[]).await.map_err(|e| format!("Query failed: {e}"))?; @@ -3548,7 +3548,7 @@ async fn execute_manual_txn_mysql_statement( sql: &str, row_limit: usize, ) -> Result { - if starts_with_executable_sql_keyword(sql, &["SELECT", "SHOW", "DESCRIBE", "EXPLAIN", "WITH"]) { + if db::mysql::is_result_set_query(sql, db::mysql::MySqlQueryDialect::default()) { let start = std::time::Instant::now(); let mut result = conn.query_iter(sql).await.map_err(|e| format!("Query failed: {e}"))?; let columns: Vec = result.columns_ref().iter().map(|c| c.name_str().to_string()).collect();