From f4f31939942d8698df8430531a0947e707cbf32f Mon Sep 17 00:00:00 2001 From: t8y2 <1156263951@qq.com> Date: Sun, 10 May 2026 00:48:24 +0800 Subject: [PATCH] fix(sql): detect result queries after comments --- crates/dbx-core/src/db/clickhouse_driver.rs | 9 +-- crates/dbx-core/src/db/dm_driver.rs | 9 +-- crates/dbx-core/src/db/gaussdb_driver.rs | 8 +-- crates/dbx-core/src/db/mysql.rs | 8 +-- crates/dbx-core/src/db/oracle_driver.rs | 10 +-- crates/dbx-core/src/db/postgres.rs | 17 +---- crates/dbx-core/src/db/sqlite.rs | 8 +-- crates/dbx-core/src/db/sqlserver.rs | 8 +-- crates/dbx-core/src/query.rs | 11 +--- crates/dbx-core/src/sql.rs | 70 ++++++++++++++++++++- 10 files changed, 88 insertions(+), 70 deletions(-) diff --git a/crates/dbx-core/src/db/clickhouse_driver.rs b/crates/dbx-core/src/db/clickhouse_driver.rs index 7052a15f4..005f62240 100644 --- a/crates/dbx-core/src/db/clickhouse_driver.rs +++ b/crates/dbx-core/src/db/clickhouse_driver.rs @@ -3,6 +3,7 @@ use serde::Deserialize; use std::time::Instant; use super::{connection_timeout, with_connection_timeout}; +use crate::sql::starts_with_executable_sql_keyword; use crate::types::{ColumnInfo, DatabaseInfo, QueryResult, TableInfo}; pub struct ChClient { @@ -148,14 +149,8 @@ pub async fn get_columns(client: &ChClient, database: &str, table: &str) -> Resu pub async fn execute_query(client: &ChClient, database: &str, sql: &str) -> Result { let start = Instant::now(); - let trimmed = sql.trim().to_uppercase(); - if trimmed.starts_with("SELECT") - || trimmed.starts_with("SHOW") - || trimmed.starts_with("DESCRIBE") - || trimmed.starts_with("EXPLAIN") - || trimmed.starts_with("WITH") - { + if starts_with_executable_sql_keyword(sql, &["SELECT", "SHOW", "DESCRIBE", "EXPLAIN", "WITH"]) { let result = ch_query(client, sql, Some(database)).await?; let columns: Vec = result.meta.iter().map(|c| c.name.clone()).collect(); Ok(QueryResult { diff --git a/crates/dbx-core/src/db/dm_driver.rs b/crates/dbx-core/src/db/dm_driver.rs index fc3aba154..8df20fd92 100644 --- a/crates/dbx-core/src/db/dm_driver.rs +++ b/crates/dbx-core/src/db/dm_driver.rs @@ -2,6 +2,7 @@ use std::time::Instant; use odbc_api::{buffers::TextRowSet, ConnectionOptions, Cursor, ResultSetMetadata}; +use crate::sql::starts_with_executable_sql_keyword; use crate::types::{ColumnInfo, DatabaseInfo, ForeignKeyInfo, IndexInfo, QueryResult, TableInfo, TriggerInfo}; use super::CONNECTION_TIMEOUT_SECS; @@ -260,14 +261,8 @@ pub fn list_triggers(client: &DmClient, schema: &str, table: &str) -> Result Result { let start = Instant::now(); let sql = sql.trim().trim_end_matches(';'); - let trimmed = sql.to_uppercase(); - if trimmed.starts_with("SELECT") - || trimmed.starts_with("WITH") - || trimmed.starts_with("SHOW") - || trimmed.starts_with("DESCRIBE") - || trimmed.starts_with("EXPLAIN") - { + if starts_with_executable_sql_keyword(sql, &["SELECT", "WITH", "SHOW", "DESCRIBE", "EXPLAIN"]) { match client.conn.execute(sql, (), None).map_err(|e| e.to_string())? { Some(mut cursor) => { let col_count = cursor.num_result_cols().map_err(|e| e.to_string())? as u16; diff --git a/crates/dbx-core/src/db/gaussdb_driver.rs b/crates/dbx-core/src/db/gaussdb_driver.rs index 13e0112b1..b44004995 100644 --- a/crates/dbx-core/src/db/gaussdb_driver.rs +++ b/crates/dbx-core/src/db/gaussdb_driver.rs @@ -1,5 +1,6 @@ use std::time::Instant; +use crate::sql::starts_with_executable_sql_keyword; use crate::types::{ColumnInfo, DatabaseInfo, ForeignKeyInfo, IndexInfo, QueryResult, TableInfo, TriggerInfo}; use super::CONNECTION_TIMEOUT_SECS; @@ -245,13 +246,8 @@ pub async fn list_triggers(client: &mut GaussdbClient, schema: &str, table: &str pub async fn execute_query(client: &mut GaussdbClient, sql: &str) -> Result { let start = Instant::now(); let sql = sql.trim().trim_end_matches(';'); - let trimmed = sql.to_uppercase(); - if trimmed.starts_with("SELECT") - || trimmed.starts_with("WITH") - || trimmed.starts_with("SHOW") - || trimmed.starts_with("EXPLAIN") - { + if starts_with_executable_sql_keyword(sql, &["SELECT", "WITH", "SHOW", "EXPLAIN"]) { let rows = client.client.query(sql, &[]).await.map_err(|e| e.to_string())?; let columns: Vec = if let Some(first) = rows.first() { diff --git a/crates/dbx-core/src/db/mysql.rs b/crates/dbx-core/src/db/mysql.rs index e58a80e60..5946f0436 100644 --- a/crates/dbx-core/src/db/mysql.rs +++ b/crates/dbx-core/src/db/mysql.rs @@ -5,6 +5,7 @@ use sqlx::mysql::{MySqlPool, MySqlPoolOptions, MySqlRow}; use sqlx::{Column, Executor, Row, TypeInfo, ValueRef}; use std::time::{Duration, Instant}; +use crate::sql::starts_with_executable_sql_keyword; use crate::types::{ColumnInfo, DatabaseInfo, ForeignKeyInfo, IndexInfo, QueryResult, TableInfo, TriggerInfo}; fn quote_value(s: &str) -> String { @@ -249,13 +250,8 @@ pub async fn get_columns(pool: &MySqlPool, database: &str, table: &str) -> Resul pub async fn execute_query(pool: &MySqlPool, sql: &str, bare: bool) -> Result { let start = Instant::now(); - let trimmed = sql.trim().to_uppercase(); - if trimmed.starts_with("SELECT") - || trimmed.starts_with("SHOW") - || trimmed.starts_with("DESCRIBE") - || trimmed.starts_with("EXPLAIN") - { + if starts_with_executable_sql_keyword(sql, &["SELECT", "SHOW", "DESCRIBE", "EXPLAIN"]) { if bare { let mut stream = sqlx::raw_sql(sql).fetch(&*pool); let mut columns: Vec = vec![]; diff --git a/crates/dbx-core/src/db/oracle_driver.rs b/crates/dbx-core/src/db/oracle_driver.rs index 357e671b1..ccc59235e 100644 --- a/crates/dbx-core/src/db/oracle_driver.rs +++ b/crates/dbx-core/src/db/oracle_driver.rs @@ -3,6 +3,7 @@ use rust_oracle::{Config, Connection}; use std::time::Instant; use super::{connection_timeout, CONNECTION_TIMEOUT_SECS}; +use crate::sql::starts_with_executable_sql_keyword; use crate::types::{ColumnInfo, DatabaseInfo, ForeignKeyInfo, IndexInfo, QueryResult, TableInfo, TriggerInfo}; pub type OracleClient = Connection; @@ -265,14 +266,7 @@ pub async fn execute_query(conn: &OracleClient, sql: &str) -> Result Option { @@ -257,14 +258,8 @@ pub async fn get_columns(pool: &PgPool, schema: &str, table: &str) -> Result Result { let start = Instant::now(); - let trimmed = sql.trim().to_uppercase(); - if trimmed.starts_with("SELECT") - || trimmed.starts_with("SHOW") - || trimmed.starts_with("EXPLAIN") - || trimmed.starts_with("WITH") - || trimmed.starts_with("TABLE") - { + if starts_with_executable_sql_keyword(sql, &["SELECT", "SHOW", "EXPLAIN", "WITH", "TABLE"]) { let mut stream = sqlx::query(sql).persistent(false).fetch(pool); let mut columns: Vec = vec![]; let mut column_types: Vec = vec![]; @@ -323,14 +318,8 @@ pub async fn execute_query_with_schema(pool: &PgPool, schema: &str, sql: &str) - sqlx::query(&set_path).execute(&mut *conn).await.map_err(|e| e.to_string())?; let start = Instant::now(); - let trimmed = sql.trim().to_uppercase(); - if trimmed.starts_with("SELECT") - || trimmed.starts_with("SHOW") - || trimmed.starts_with("EXPLAIN") - || trimmed.starts_with("WITH") - || trimmed.starts_with("TABLE") - { + if starts_with_executable_sql_keyword(sql, &["SELECT", "SHOW", "EXPLAIN", "WITH", "TABLE"]) { let mut stream = sqlx::query(sql).persistent(false).fetch(&mut *conn); let mut columns: Vec = vec![]; let mut column_types: Vec = vec![]; diff --git a/crates/dbx-core/src/db/sqlite.rs b/crates/dbx-core/src/db/sqlite.rs index 622c55624..a5b9ef618 100644 --- a/crates/dbx-core/src/db/sqlite.rs +++ b/crates/dbx-core/src/db/sqlite.rs @@ -4,6 +4,7 @@ use sqlx::{Column, Executor, Row}; use std::time::{Duration, Instant}; use super::file_validator::validate_file_path; +use crate::sql::starts_with_executable_sql_keyword; use crate::types::{ColumnInfo, DatabaseInfo, ForeignKeyInfo, IndexInfo, QueryResult, TableInfo, TriggerInfo}; pub async fn connect_path(path: &str) -> Result { @@ -162,13 +163,8 @@ pub async fn list_triggers(pool: &SqlitePool, _schema: &str, table: &str) -> Res pub async fn execute_query(pool: &SqlitePool, sql: &str) -> Result { let start = Instant::now(); - let trimmed = sql.trim().to_uppercase(); - if trimmed.starts_with("SELECT") - || trimmed.starts_with("PRAGMA") - || trimmed.starts_with("EXPLAIN") - || trimmed.starts_with("WITH") - { + if starts_with_executable_sql_keyword(sql, &["SELECT", "PRAGMA", "EXPLAIN", "WITH"]) { let desc = pool.describe(sql).await.map_err(|e| e.to_string())?; let columns: Vec = desc.columns().iter().map(|c| c.name().to_string()).collect(); diff --git a/crates/dbx-core/src/db/sqlserver.rs b/crates/dbx-core/src/db/sqlserver.rs index f46a5725d..0e82780ba 100644 --- a/crates/dbx-core/src/db/sqlserver.rs +++ b/crates/dbx-core/src/db/sqlserver.rs @@ -5,6 +5,7 @@ use tokio::net::TcpStream; use tokio_util::compat::{Compat, TokioAsyncWriteCompatExt}; use super::{connection_timeout, CONNECTION_TIMEOUT_SECS}; +use crate::sql::starts_with_executable_sql_keyword; use crate::types::{ColumnInfo, DatabaseInfo, ForeignKeyInfo, IndexInfo, QueryResult, TableInfo, TriggerInfo}; pub type SqlServerClient = Client>; @@ -317,13 +318,8 @@ pub async fn list_triggers( pub async fn execute_query(client: &mut SqlServerClient, sql: &str) -> Result { let start = Instant::now(); - let trimmed = sql.trim().to_uppercase(); - if trimmed.starts_with("SELECT") - || trimmed.starts_with("EXEC") - || trimmed.starts_with("WITH") - || trimmed.starts_with("TABLE") - { + if starts_with_executable_sql_keyword(sql, &["SELECT", "EXEC", "WITH", "TABLE"]) { let mut stream = client.query(sql, &[]).await.map_err(|e| e.to_string())?; let columns_meta = stream .columns() diff --git a/crates/dbx-core/src/query.rs b/crates/dbx-core/src/query.rs index e1931767e..f8c21a52d 100644 --- a/crates/dbx-core/src/query.rs +++ b/crates/dbx-core/src/query.rs @@ -5,7 +5,7 @@ use tokio_util::sync::CancellationToken; use crate::connection::{AppState, PoolKind}; use crate::db; -use crate::sql::split_sql_statements; +use crate::sql::{split_sql_statements, starts_with_executable_sql_keyword}; pub const QUERY_TIMEOUT: Duration = Duration::from_secs(30); pub const MAX_ROWS: usize = 10000; @@ -13,15 +13,8 @@ pub const QUERY_CANCELED: &str = "Query canceled"; pub fn duckdb_execute(con: &duckdb::Connection, sql: &str) -> Result { let start = std::time::Instant::now(); - let trimmed = sql.trim().to_uppercase(); - if trimmed.starts_with("SELECT") - || trimmed.starts_with("SHOW") - || trimmed.starts_with("DESCRIBE") - || trimmed.starts_with("EXPLAIN") - || trimmed.starts_with("WITH") - || trimmed.starts_with("PRAGMA") - { + if starts_with_executable_sql_keyword(sql, &["SELECT", "SHOW", "DESCRIBE", "EXPLAIN", "WITH", "PRAGMA"]) { let mut stmt = con.prepare(sql).map_err(|e| e.to_string())?; let mut rows = stmt.query([]).map_err(|e| e.to_string())?; let stmt_ref = rows.as_ref().ok_or("DuckDB statement unavailable")?; diff --git a/crates/dbx-core/src/sql.rs b/crates/dbx-core/src/sql.rs index 19551e1f6..8cdd222e4 100644 --- a/crates/dbx-core/src/sql.rs +++ b/crates/dbx-core/src/sql.rs @@ -206,6 +206,58 @@ pub fn statement_summary(statement: &str) -> String { collapsed.chars().take(MAX_LEN).collect() } +pub fn starts_with_executable_sql_keyword(sql: &str, keywords: &[&str]) -> bool { + let Some(token) = first_executable_sql_token(sql) else { + return false; + }; + keywords.iter().any(|keyword| token.eq_ignore_ascii_case(keyword)) +} + +fn first_executable_sql_token(sql: &str) -> Option<&str> { + let bytes = sql.as_bytes(); + let mut i = 0; + + while i < bytes.len() { + while i < bytes.len() && bytes[i].is_ascii_whitespace() { + i += 1; + } + + if i + 1 < bytes.len() && bytes[i] == b'-' && bytes[i + 1] == b'-' { + i += 2; + while i < bytes.len() && bytes[i] != b'\n' { + i += 1; + } + continue; + } + + if i + 1 < bytes.len() && bytes[i] == b'/' && bytes[i + 1] == b'*' { + if i + 2 < bytes.len() && (bytes[i + 2] == b'!' || (i + 3 < bytes.len() && &bytes[i + 2..i + 4] == b"M!")) { + i += if bytes[i + 2] == b'!' { 3 } else { 4 }; + while i < bytes.len() && (bytes[i].is_ascii_digit() || bytes[i].is_ascii_whitespace()) { + i += 1; + } + break; + } + + i += 2; + while i + 1 < bytes.len() && !(bytes[i] == b'*' && bytes[i + 1] == b'/') { + i += 1; + } + i = (i + 2).min(bytes.len()); + continue; + } + + break; + } + + let start = i; + while i < bytes.len() && (bytes[i].is_ascii_alphabetic() || bytes[i] == b'_') { + i += 1; + } + + (i > start).then_some(&sql[start..i]) +} + fn starts_with_chars(chars: &[char], start: usize, needle: &[char]) -> bool { start + needle.len() <= chars.len() && chars[start..start + needle.len()] == *needle } @@ -306,7 +358,7 @@ fn split_sql_script(sql: &str) -> Result, String> { #[cfg(test)] mod tests { - use super::{split_sql_script, SqlStatementSplitter}; + use super::{split_sql_script, starts_with_executable_sql_keyword, SqlStatementSplitter}; #[test] fn splits_semicolon_delimited_statements() { @@ -394,4 +446,20 @@ mod tests { vec!["/*!40101 SET @OLD_CHARACTER_SET_CLIENT=@@CHARACTER_SET_CLIENT */", "SELECT 1",] ); } + + #[test] + fn detects_result_set_keyword_after_comments() { + assert!(starts_with_executable_sql_keyword("-- comment\nselect * from users;", &["SELECT"])); + assert!(starts_with_executable_sql_keyword( + "/* comment */\nWITH rows AS (SELECT 1) SELECT * FROM rows;", + &["WITH"] + )); + assert!(!starts_with_executable_sql_keyword("-- comment only\n", &["SELECT"])); + } + + #[test] + fn detects_mysql_executable_comment_keyword() { + assert!(starts_with_executable_sql_keyword("/*!40101 SELECT 1 */", &["SELECT"])); + assert!(starts_with_executable_sql_keyword("/*M! SELECT 1 */", &["SELECT"])); + } }