xiaoxue_admin
/
dbx
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix(sql): detect result queries after comments

This commit is contained in:
t8y2 2026-05-10 00:48:24 +08:00
parent b013a2ac29
commit f4f3193994
10 changed files with 88 additions and 70 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
crates/dbx-core/src/db/clickhouse_driver.rs
View File

@ -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<QueryResult, String> {
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<String> = result.meta.iter().map(|c| c.name.clone()).collect();
Ok(QueryResult {

9
crates/dbx-core/src/db/dm_driver.rs
View File

@ -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<Vec
pub fn execute_query_sync(client: &DmClient, sql: &str) -> Result<QueryResult, String> {
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;

8
crates/dbx-core/src/db/gaussdb_driver.rs
View File

@ -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<QueryResult, String> {
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<String> = if let Some(first) = rows.first() {

8
crates/dbx-core/src/db/mysql.rs
View File

@ -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<QueryResult, String> {
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<String> = vec![];

10
crates/dbx-core/src/db/oracle_driver.rs
View File

@ -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<QueryResult
let sql = rewrite_fetch_first(sql);
let sql = sql.as_ref();
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"]) {
let result = conn.query(sql, &[]).await.map_err(|e| {
log::error!("[oracle] execute_query SELECT failed: {e}");
e.to_string()

17
crates/dbx-core/src/db/postgres.rs
View File

@ -7,6 +7,7 @@ use sqlx::{Column, Executor, Row, TypeInfo, ValueRef};
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};
fn pg_temporal_to_json_value(row: &PgRow, idx: usize) -> Option<serde_json::Value> {
@ -257,14 +258,8 @@ pub async fn get_columns(pool: &PgPool, schema: &str, table: &str) -> Result<Vec
pub async fn execute_query(pool: &PgPool, sql: &str) -> Result<QueryResult, 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(pool);
let mut columns: Vec<String> = vec![];
let mut column_types: Vec<String> = 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<String> = vec![];
let mut column_types: Vec<String> = vec![];

8
crates/dbx-core/src/db/sqlite.rs
View File

@ -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<SqlitePool, String> {
@ -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<QueryResult, String> {
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<String> = desc.columns().iter().map(|c| c.name().to_string()).collect();

8
crates/dbx-core/src/db/sqlserver.rs
View File

@ -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<Compat<TcpStream>>;
@ -317,13 +318,8 @@ pub async fn list_triggers(
pub async fn execute_query(client: &mut SqlServerClient, sql: &str) -> Result<QueryResult, String> {
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()

11
crates/dbx-core/src/query.rs
View File

@ -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<db::QueryResult, String> {
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")?;

70
crates/dbx-core/src/sql.rs
View File

@ -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<Vec<String>, 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"]));
}
}