From 0e779f9edc3b42c04a1ef695aad12a647c8cefac Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E4=BA=8C=E4=B8=AB=E8=AE=B2=E6=A2=B5?= Date: Sun, 12 Jul 2026 23:05:32 +0800 Subject: [PATCH] fix(sql): improve read-only statement detection --- crates/dbx-core/src/agent_tools.rs | 8 +- crates/dbx-core/src/query.rs | 30 +- crates/dbx-core/src/query_execution_sql.rs | 393 ++++++++++++++++++++- crates/dbx-core/src/query_result_export.rs | 8 +- crates/dbx-core/src/sql_risk.rs | 101 +++++- 5 files changed, 495 insertions(+), 45 deletions(-) diff --git a/crates/dbx-core/src/agent_tools.rs b/crates/dbx-core/src/agent_tools.rs index 0ad3125ce..2ddda1e17 100644 --- a/crates/dbx-core/src/agent_tools.rs +++ b/crates/dbx-core/src/agent_tools.rs @@ -453,9 +453,8 @@ async fn execute_execute_query( .map(|l| (l as usize).min(MAX_ALLOWED_ROWS)) .unwrap_or(EXECUTE_QUERY_LIMIT); - // Classify SQL risk using sqlparser AST - let db_type_str = format!("{:?}", db_type).to_lowercase(); - let risk = crate::sql_risk::classify_sql_risk(sql, &db_type_str)?; + // Classify SQL risk using the concrete database dialect. + let risk = crate::sql_risk::classify_sql_risk_for_database(sql, *db_type)?; let connection_config = state.configs.read().await.get(connection_id).cloned(); if let Some(config) = connection_config { if risk != SqlRisk::ReadOnly && crate::production_safety::targets_production_database(&config, database, sql) { @@ -586,8 +585,7 @@ async fn execute_explain_query( } // Classify SQL risk – only ReadOnly queries can be explained - let db_type_str = format!("{:?}", db_type).to_lowercase(); - let risk = match crate::sql_risk::classify_sql_risk(sql, &db_type_str) { + let risk = match crate::sql_risk::classify_sql_risk_for_database(sql, *db_type) { Ok(r) => r, Err(e) => return (Err(e), None), }; diff --git a/crates/dbx-core/src/query.rs b/crates/dbx-core/src/query.rs index c82b57cd7..f33c0dbf7 100644 --- a/crates/dbx-core/src/query.rs +++ b/crates/dbx-core/src/query.rs @@ -130,12 +130,14 @@ impl DbOperationBudget { /// Uses config_for_pool_key to correctly resolve configs when pool_key includes /// a database suffix (e.g., "prod:app" → config stored under "prod"). pub async fn check_read_only_for_connection(state: &AppState, pool_key: &str, sql: &str) -> Result<(), String> { - let conn_name = { + let connection = { let configs = state.configs.read().await; - crate::connection::config_for_pool_key(pool_key, &configs).filter(|c| c.read_only).map(|c| c.name.clone()) + crate::connection::config_for_pool_key(pool_key, &configs) + .filter(|config| config.read_only) + .map(|config| (config.name.clone(), config.db_type)) }; - if let Some(name) = conn_name { - crate::query_execution_sql::check_read_only(sql, &name)?; + if let Some((name, database_type)) = connection { + crate::query_execution_sql::check_read_only(sql, &name, database_type)?; } Ok(()) } @@ -146,13 +148,15 @@ pub async fn check_read_only_for_connection_multi( pool_key: &str, statements: &[impl AsRef], ) -> Result<(), String> { - let conn_name = { + let connection = { let configs = state.configs.read().await; - crate::connection::config_for_pool_key(pool_key, &configs).filter(|c| c.read_only).map(|c| c.name.clone()) + crate::connection::config_for_pool_key(pool_key, &configs) + .filter(|config| config.read_only) + .map(|config| (config.name.clone(), config.db_type)) }; - if let Some(name) = conn_name { + if let Some((name, database_type)) = connection { for sql in statements { - crate::query_execution_sql::check_read_only(sql.as_ref(), &name)?; + crate::query_execution_sql::check_read_only(sql.as_ref(), &name, database_type)?; } } Ok(()) @@ -1100,18 +1104,18 @@ pub async fn do_execute( let _activity_touch = state.pool_activity_touch(pool_key); let query_timeout = resolve_query_timeout(options.timeout_secs); - let (_duckdb_attached_names, conn_name_if_readonly) = { + let (_duckdb_attached_names, read_only_connection) = { let configs = state.configs.read().await; let config = crate::connection::config_for_pool_key(pool_key, &configs); let attached = config .map(|c| c.attached_databases.iter().map(|db| db.name.clone()).collect::>()) .unwrap_or_default(); - let conn_name = config.filter(|c| c.read_only).map(|c| c.name.clone()); - (attached, conn_name) + let connection = config.filter(|config| config.read_only).map(|config| (config.name.clone(), config.db_type)); + (attached, connection) }; let operation_budget = operation_budget_for_pool_key(state, pool_key, query_timeout).await; - if let Some(name) = conn_name_if_readonly { - crate::query_execution_sql::check_read_only(sql, &name)?; + if let Some((name, database_type)) = read_only_connection { + crate::query_execution_sql::check_read_only(sql, &name, database_type)?; } let pool_db_type = connection_database_type_for_pool_key(state, pool_key).await; let connections = state.connections.read().await; diff --git a/crates/dbx-core/src/query_execution_sql.rs b/crates/dbx-core/src/query_execution_sql.rs index e6c3f48fb..5a3350f88 100644 --- a/crates/dbx-core/src/query_execution_sql.rs +++ b/crates/dbx-core/src/query_execution_sql.rs @@ -1,4 +1,6 @@ use serde::{Deserialize, Serialize}; +use sqlparser::dialect::{MsSqlDialect, PostgreSqlDialect}; +use sqlparser::tokenizer::{Token, Tokenizer}; use crate::models::connection::DatabaseType; @@ -182,9 +184,124 @@ const SAFE_READ_PRAGMA_NAMES: &[&str] = &[ ]; /// Returns true if the SQL statement is a write operation (not a pure read). +/// +/// Callers that know the connection database type should use +/// [`is_write_sql_for_database`] so executable comments are interpreted using +/// the correct dialect. This untyped helper deliberately remains conservative +/// for executable comments. pub fn is_write_sql(sql: &str) -> bool { + is_write_sql_with_database_type(sql, None) +} + +/// Returns true if the SQL statement is a write operation for a database +/// dialect. In addition to ordinary write statements, this recognizes MySQL +/// executable comments and file exports, plus PostgreSQL-family/SQL Server +/// `SELECT ... INTO` table creation. +pub fn is_write_sql_for_database(sql: &str, database_type: DatabaseType) -> bool { + is_write_sql_with_database_type(sql, Some(database_type)) +} + +fn is_write_sql_with_database_type(sql: &str, database_type: Option) -> bool { + if database_type.is_some_and(|database_type| has_dialect_specific_write(sql, database_type)) { + return true; + } + + // The untyped helper remains conservative for MySQL executable comments. + // Typed callers handle those comments in has_dialect_specific_write above. + let detect_mysql_executable_comments = database_type.is_none(); + let detect_select_into = database_type.is_none(); + let statements = match database_type { + Some(database_type) => crate::sql::split_sql_statements_for_database(sql, database_type), + None => crate::sql::split_sql_statements(sql), + }; + + statements + .iter() + .any(|statement| is_write_sql_statement(statement, detect_mysql_executable_comments, detect_select_into)) +} + +fn is_mysql_compatible_database(database_type: DatabaseType) -> bool { + matches!( + database_type, + DatabaseType::Mysql + | DatabaseType::Doris + | DatabaseType::StarRocks + | DatabaseType::ManticoreSearch + | DatabaseType::Goldendb + ) +} + +fn is_postgresql_family_database(database_type: DatabaseType) -> bool { + matches!( + database_type, + DatabaseType::Postgres + | DatabaseType::Redshift + | DatabaseType::Gaussdb + | DatabaseType::OpenGauss + | DatabaseType::Kingbase + | DatabaseType::Highgo + | DatabaseType::Vastbase + | DatabaseType::Kwdb + ) +} + +/// Detects write-capable syntax that otherwise looks like a read query and is +/// interpreted differently depending on the database dialect. +pub(crate) fn has_dialect_specific_write(sql: &str, database_type: DatabaseType) -> bool { + let statements = crate::sql::split_sql_statements_for_database(sql, database_type); + statements.iter().any(|statement| has_dialect_specific_write_statement(statement, database_type)) +} + +fn has_dialect_specific_write_statement(sql: &str, database_type: DatabaseType) -> bool { + if is_mysql_compatible_database(database_type) { + let (cleaned, has_executable_comment) = strip_sql_comments_and_literals_with_metadata(sql, true); + return has_executable_comment + || contains_keyword_sequence(&cleaned, "INTO", "OUTFILE") + || contains_keyword_sequence(&cleaned, "INTO", "DUMPFILE"); + } + + if is_postgresql_family_database(database_type) { + contains_unquoted_keyword(sql, &PostgreSqlDialect {}, "INTO") + } else { + match database_type { + DatabaseType::SqlServer => contains_unquoted_keyword(sql, &MsSqlDialect {}, "INTO"), + _ => false, + } + } +} + +fn contains_keyword_sequence(sql: &str, first: &str, second: &str) -> bool { + let words = sql.split(|ch: char| !ch.is_ascii_alphanumeric() && ch != '_').filter(|word| !word.is_empty()); + + let mut previous_matches = false; + for word in words { + if previous_matches && word.eq_ignore_ascii_case(second) { + return true; + } + previous_matches = word.eq_ignore_ascii_case(first); + } + false +} + +fn contains_unquoted_keyword(sql: &str, dialect: &dyn sqlparser::dialect::Dialect, keyword: &str) -> bool { + Tokenizer::new(dialect, sql).tokenize().is_ok_and(|tokens| { + tokens.into_iter().any(|token| { + matches!(token, Token::Word(word) if word.quote_style.is_none() && word.value.eq_ignore_ascii_case(keyword)) + }) + }) +} + +fn is_write_sql_statement(sql: &str, detect_mysql_executable_comments: bool, detect_select_into: bool) -> bool { // 1. Strip comments and string literals - let cleaned = strip_sql_comments_and_literals(sql); + let (cleaned, has_mysql_executable_comment) = + strip_sql_comments_and_literals_with_metadata(sql, detect_mysql_executable_comments); + // MySQL/MariaDB executable comments may contain arbitrary SQL, including + // writes that are not represented by the outer statement (for example, + // INTO OUTFILE inside a SELECT). Treat them as writes rather than + // attempting to parse every supported MySQL dialect extension here. + if has_mysql_executable_comment { + return true; + } let trimmed = cleaned.trim_start(); if trimmed.is_empty() { return false; @@ -201,7 +318,14 @@ pub fn is_write_sql(sql: &str) -> bool { return !is_safe_read_pragma(&upper); } - if starts_with_keyword(&upper, "SELECT") && select_contains_top_level_into(&upper) { + // SHOW CREATE returns object metadata; CREATE is part of its read-only syntax. + if starts_with_show_create(&upper) { + return false; + } + + // Untyped callers stay conservative. Typed callers handle SELECT ... INTO + // only for database families where the syntax performs a write. + if detect_select_into && starts_with_keyword(&upper, "SELECT") && select_contains_top_level_into(&upper) { return true; } @@ -210,6 +334,16 @@ pub fn is_write_sql(sql: &str) -> bool { !starts_with_read || contains_dangerous_sql_keyword(sql) } +fn starts_with_show_create(upper: &str) -> bool { + let Some(after_show) = upper.strip_prefix("SHOW") else { + return false; + }; + if !after_show.is_empty() && after_show.as_bytes()[0].is_ascii_alphanumeric() { + return false; + } + starts_with_keyword(after_show.trim_start(), "CREATE") +} + fn select_contains_top_level_into(upper: &str) -> bool { let mut token = String::new(); let mut depth = 0usize; @@ -283,8 +417,8 @@ fn starts_with_keyword(upper: &str, keyword: &str) -> bool { /// Check whether a SQL statement is allowed under read-only mode. /// Returns Err with a descriptive message if the statement is a write operation. -pub fn check_read_only(sql: &str, connection_name: &str) -> Result<(), String> { - if is_write_sql(sql) { +pub fn check_read_only(sql: &str, connection_name: &str, database_type: DatabaseType) -> Result<(), String> { + if is_write_sql_for_database(sql, database_type) { return Err(format!( "Read-only mode: connection '{}' has read-only protection enabled. Write operation (including stored procedure calls) blocked.", connection_name @@ -374,12 +508,18 @@ fn strip_sql_comments(sql: &str) -> String { } pub fn strip_sql_comments_and_literals(sql: &str) -> String { + strip_sql_comments_and_literals_with_metadata(sql, false).0 +} + +fn strip_sql_comments_and_literals_with_metadata(sql: &str, detect_mysql_executable_comments: bool) -> (String, bool) { let mut output = String::with_capacity(sql.len()); let mut chars = sql.chars().peekable(); let mut in_line_comment = false; let mut in_block_comment = false; let mut in_single_quote = false; let mut in_double_quote = false; + let mut in_backtick_quote = false; + let mut has_mysql_executable_comment = false; while let Some(ch) = chars.next() { if in_line_comment { @@ -423,6 +563,18 @@ pub fn strip_sql_comments_and_literals(sql: &str) -> String { continue; } + if in_backtick_quote { + if ch == '`' { + if chars.peek() == Some(&'`') { + chars.next(); + } else { + in_backtick_quote = false; + } + } + output.push(' '); + continue; + } + if ch == '-' && chars.peek() == Some(&'-') { chars.next(); in_line_comment = true; @@ -434,13 +586,10 @@ pub fn strip_sql_comments_and_literals(sql: &str) -> String { } if ch == '/' && chars.peek() == Some(&'*') { chars.next(); - if let Some(body) = read_mysql_executable_comment_body(&mut chars) { - output.push(' '); - output.push_str(&strip_sql_comments_and_literals(&body)); - output.push(' '); - } else { - in_block_comment = true; + if detect_mysql_executable_comments && is_mysql_executable_comment_start(&chars) { + has_mysql_executable_comment = true; } + in_block_comment = true; continue; } if ch == '\'' { @@ -453,11 +602,27 @@ pub fn strip_sql_comments_and_literals(sql: &str) -> String { output.push(' '); continue; } + if ch == '`' { + in_backtick_quote = true; + output.push(' '); + continue; + } output.push(ch); } - output + (output, has_mysql_executable_comment) +} + +/// `/*! ... */` is executable in MySQL, while `/*M! ... */` (optionally +/// followed by a version number) is executable in MariaDB. +fn is_mysql_executable_comment_start(chars: &std::iter::Peekable>) -> bool { + let mut marker = chars.clone(); + match marker.next() { + Some('!') => true, + Some('M') => marker.next() == Some('!'), + _ => false, + } } fn read_mysql_executable_comment_body(chars: &mut std::iter::Peekable) -> Option @@ -681,6 +846,15 @@ mod tests { assert!(!contains_dangerous_sql_keyword("SELECT \"CREATE TABLE\" FROM t")); } + #[test] + fn contains_dangerous_sql_keyword_ignores_backtick_identifiers() { + for keyword in ["drop", "delete", "truncate", "alter", "update", "merge", "replace", "insert", "create"] { + let sql = format!("SELECT 1 AS `{keyword}`"); + assert!(!contains_dangerous_sql_keyword(&sql), "expected safe SQL: {sql}"); + } + assert!(!contains_dangerous_sql_keyword("SELECT 1 AS `before``delete`")); + } + #[test] fn is_write_sql_detects_simple_writes() { assert!(is_write_sql("INSERT INTO users VALUES (1)")); @@ -714,6 +888,58 @@ mod tests { assert!(!is_write_sql("FROM users SELECT *")); } + #[test] + fn is_write_sql_allows_backtick_identifiers_with_dangerous_keywords() { + for keyword in ["drop", "delete", "truncate", "alter", "update", "merge", "replace", "insert", "create"] { + let sql = format!("SELECT 1 AS `{keyword}`"); + assert!(!is_write_sql(&sql), "expected read-only SQL: {sql}"); + } + assert!(!is_write_sql("SHOW COLUMNS FROM `delete`")); + assert!(!is_write_sql("DESC `delete`")); + assert!(!is_write_sql("SELECT 1 AS `before``delete`")); + assert!(!is_write_sql("SELECT 1 AS `semi;delete`; SELECT 2")); + } + + #[test] + fn is_write_sql_still_blocks_writes_with_backtick_identifiers() { + assert!(is_write_sql("DELETE FROM `users`")); + assert!(is_write_sql("DROP TABLE `users`")); + assert!(is_write_sql("UPDATE `users` SET name = 'Ada'")); + } + + #[test] + fn is_write_sql_allows_show_create_statements() { + for sql in [ + "SHOW CREATE TABLE users", + "show create view active_users", + "SHOW CREATE PROCEDURE refresh_users", + "SHOW CREATE FUNCTION user_count", + " /* metadata */\nSHOW\nCREATE TABLE users;", + "SHOW CREATE TABLE users; SELECT ';' AS separator", + "SHOW CREATE TABLE users /* ; DELETE FROM users */", + ] { + assert!(!is_write_sql(sql), "expected read-only SQL: {sql}"); + } + } + + #[test] + fn is_write_sql_blocks_writes_after_show_create() { + for write_sql in [ + "DROP TABLE users", + "DELETE FROM users", + "TRUNCATE TABLE users", + "ALTER TABLE users ADD COLUMN active BOOLEAN", + "UPDATE users SET active = true", + "MERGE INTO users USING source ON users.id = source.id WHEN MATCHED THEN UPDATE SET active = true", + "REPLACE INTO users VALUES (1)", + "INSERT INTO users VALUES (1)", + "CREATE TABLE audit (id INT)", + ] { + let sql = format!("SHOW CREATE TABLE users; {write_sql}"); + assert!(is_write_sql(&sql), "expected write SQL: {sql}"); + } + } + #[test] fn is_write_sql_ignores_leading_whitespace_and_comments() { assert!(!is_write_sql(" /* comment */ SELECT * FROM users")); @@ -721,6 +947,103 @@ mod tests { assert!(is_write_sql(" /* comment */ INSERT INTO users VALUES (1)")); } + #[test] + fn is_write_sql_blocks_mysql_and_mariadb_executable_comments() { + for sql in [ + "SELECT 3156 /*! INTO OUTFILE '/var/lib/mysql-files/dbx_ro_probe.txt' */", + "SELECT 3156 /*!50000 INTO OUTFILE '/var/lib/mysql-files/dbx_ro_probe.txt' */", + "SELECT 3156 /*M! INTO OUTFILE '/var/lib/mysql-files/dbx_ro_probe.txt' */", + "SELECT 3156 /*M!100100 INTO OUTFILE '/var/lib/mysql-files/dbx_ro_probe.txt' */", + ] { + assert!(is_write_sql_for_database(sql, DatabaseType::Mysql), "expected write SQL: {sql}"); + } + + assert!(!is_write_sql_for_database( + "SELECT 3156 /* INTO OUTFILE '/var/lib/mysql-files/dbx_ro_probe.txt' */", + DatabaseType::Mysql + )); + assert!(!is_write_sql_for_database( + "SELECT '/*!50000 INTO OUTFILE \'/tmp/probe\' */' AS note", + DatabaseType::Mysql + )); + } + + #[test] + fn is_write_sql_treats_mysql_executable_comment_syntax_as_plain_for_other_dialects() { + for database_type in [DatabaseType::Postgres, DatabaseType::Sqlite] { + for sql in [ + "SELECT 3156 /*!50000 INTO OUTFILE '/var/lib/mysql-files/dbx_ro_probe.txt' */", + "SELECT 3156 /*M!100100 INTO OUTFILE '/var/lib/mysql-files/dbx_ro_probe.txt' */", + "SELECT 3156 /* ordinary block comment */", + ] { + assert!( + !is_write_sql_for_database(sql, database_type), + "expected read-only SQL for {database_type:?}: {sql}" + ); + } + } + } + + #[test] + fn is_write_sql_blocks_dialect_specific_select_into_writes() { + for sql in [ + "SELECT 3156 INTO OUTFILE '/var/lib/mysql-files/dbx_ro_probe.txt'", + "SELECT 3156 INTO DUMPFILE '/var/lib/mysql-files/dbx_ro_probe.bin'", + "WITH probe AS (SELECT 3156) SELECT * FROM probe INTO OUTFILE '/var/lib/mysql-files/dbx_ro_probe.txt'", + ] { + assert!(is_write_sql_for_database(sql, DatabaseType::Mysql), "expected write SQL: {sql}"); + } + + for database_type in [ + DatabaseType::Postgres, + DatabaseType::Redshift, + DatabaseType::Gaussdb, + DatabaseType::OpenGauss, + DatabaseType::Kingbase, + DatabaseType::Highgo, + DatabaseType::Vastbase, + DatabaseType::Kwdb, + ] { + for sql in [ + "SELECT * INTO copied_users FROM users", + "WITH active AS (SELECT * FROM users WHERE active) SELECT * INTO active_users FROM active", + ] { + assert!( + is_write_sql_for_database(sql, database_type), + "expected write SQL for {database_type:?}: {sql}" + ); + } + } + + for sql in [ + "SELECT * INTO dbo.copied_users FROM dbo.users", + "WITH active AS (SELECT * FROM users WHERE active = 1) SELECT * INTO #active_users FROM active", + ] { + assert!(is_write_sql_for_database(sql, DatabaseType::SqlServer), "expected write SQL: {sql}"); + } + } + + #[test] + fn is_write_sql_does_not_globally_block_into() { + for sql in [ + "SELECT 3156 INTO @probe", + "SELECT 'INTO OUTFILE /tmp/probe' AS note", + "SELECT 3156 /* INTO OUTFILE '/tmp/probe' */", + "SELECT 1 AS `into`, 2 AS `outfile`", + ] { + assert!(!is_write_sql_for_database(sql, DatabaseType::Mysql), "expected read-only SQL: {sql}"); + } + + for sql in + ["SELECT 'INTO copied_users' AS note", "SELECT $$INTO copied_users$$ AS note", "SELECT 1 AS \"into\""] + { + assert!(!is_write_sql_for_database(sql, DatabaseType::Postgres), "expected read-only SQL: {sql}"); + } + + assert!(!is_write_sql_for_database("SELECT 1 AS [into]", DatabaseType::SqlServer)); + assert!(!is_write_sql_for_database("SELECT 1 INTO unsupported", DatabaseType::Sqlite)); + } + #[test] fn is_write_sql_cte_with_nested_write() { // CTE starting with WITH but containing a write operation inside @@ -799,19 +1122,59 @@ mod tests { #[test] fn check_read_only_success_and_error() { - assert_eq!(check_read_only("SELECT * FROM users", "prod-db"), Ok(())); - assert_eq!(check_read_only("WITH cte AS (SELECT 1) SELECT * FROM cte", "prod-db"), Ok(())); + assert_eq!(check_read_only("SELECT * FROM users", "prod-db", DatabaseType::Mysql), Ok(())); + assert_eq!(check_read_only("WITH cte AS (SELECT 1) SELECT * FROM cte", "prod-db", DatabaseType::Mysql), Ok(())); + assert_eq!(check_read_only("SHOW CREATE TABLE users", "prod-db", DatabaseType::Mysql), Ok(())); + assert_eq!(check_read_only("SELECT 1 AS `delete`", "prod-db", DatabaseType::Mysql), Ok(())); - let err = check_read_only("DELETE FROM users", "prod-db"); + let err = check_read_only("DELETE FROM users", "prod-db", DatabaseType::Mysql); assert!(err.is_err()); assert_eq!( err.unwrap_err(), "Read-only mode: connection 'prod-db' has read-only protection enabled. Write operation (including stored procedure calls) blocked." ); - let err2 = check_read_only("UPDATE users SET name = 'x'", "reporting-db"); + let err2 = check_read_only("UPDATE users SET name = 'x'", "reporting-db", DatabaseType::Mysql); assert!(err2.is_err()); assert!(err2.unwrap_err().contains("reporting-db")); + + let show_create_err = + check_read_only("SHOW CREATE TABLE users; DELETE FROM users", "prod-db", DatabaseType::Mysql); + assert!(show_create_err.is_err()); + assert!(show_create_err.unwrap_err().contains("Write operation")); + } + + #[test] + fn check_read_only_only_treats_executable_comments_as_writes_for_mysql_compatible_connections() { + let mysql_executable_comment = "SELECT 3156 /*!50000 INTO OUTFILE '/var/lib/mysql-files/dbx_ro_probe.txt' */"; + let mariadb_executable_comment = + "SELECT 3156 /*M!100100 INTO OUTFILE '/var/lib/mysql-files/dbx_ro_probe.txt' */"; + + assert!(check_read_only(mysql_executable_comment, "mysql", DatabaseType::Mysql).is_err()); + assert!(check_read_only(mariadb_executable_comment, "mariadb", DatabaseType::Mysql).is_err()); + + for database_type in [DatabaseType::Postgres, DatabaseType::Sqlite] { + assert_eq!(check_read_only(mysql_executable_comment, "readonly", database_type), Ok(())); + assert_eq!(check_read_only(mariadb_executable_comment, "readonly", database_type), Ok(())); + } + } + + #[test] + fn check_read_only_blocks_dialect_specific_select_into_writes() { + for (sql, database_type) in [ + ("SELECT 3156 INTO OUTFILE '/var/lib/mysql-files/dbx_ro_probe.txt'", DatabaseType::Mysql), + ("SELECT * INTO copied_users FROM users", DatabaseType::Postgres), + ("SELECT * INTO copied_users FROM users", DatabaseType::Redshift), + ("SELECT * INTO copied_users FROM users", DatabaseType::Gaussdb), + ("SELECT * INTO copied_users FROM users", DatabaseType::OpenGauss), + ("SELECT * INTO copied_users FROM users", DatabaseType::Kingbase), + ("SELECT * INTO copied_users FROM users", DatabaseType::Highgo), + ("SELECT * INTO copied_users FROM users", DatabaseType::Vastbase), + ("SELECT * INTO copied_users FROM users", DatabaseType::Kwdb), + ("SELECT * INTO #copied_users FROM users", DatabaseType::SqlServer), + ] { + assert!(check_read_only(sql, "readonly", database_type).is_err(), "expected blocked SQL: {sql}"); + } } #[test] diff --git a/crates/dbx-core/src/query_result_export.rs b/crates/dbx-core/src/query_result_export.rs index 799f2a61f..d9d510bc4 100644 --- a/crates/dbx-core/src/query_result_export.rs +++ b/crates/dbx-core/src/query_result_export.rs @@ -785,7 +785,7 @@ async fn try_export_mysql_query_result_stream( state.touch_pool_activity(&pool_key).await; let _activity_touch = state.pool_activity_touch(&pool_key); - let (mysql_dialect, read_only_connection_name) = { + let (mysql_dialect, read_only_connection) = { let configs = state.configs.read().await; let config = configs.get(&request.connection_id); ( @@ -797,11 +797,11 @@ async fn try_export_mysql_query_result_stream( ) }) .unwrap_or_default(), - config.filter(|config| config.read_only).map(|config| config.name.clone()), + config.filter(|config| config.read_only).map(|config| (config.name.clone(), config.db_type)), ) }; - if let Some(name) = read_only_connection_name { - crate::query_execution_sql::check_read_only(&request.sql, &name)?; + if let Some((name, database_type)) = read_only_connection { + crate::query_execution_sql::check_read_only(&request.sql, &name, database_type)?; } let xlsx_hard_limit_active = xlsx_hard_limit_active(format, request); diff --git a/crates/dbx-core/src/sql_risk.rs b/crates/dbx-core/src/sql_risk.rs index 7d1d3d249..629f37cb7 100644 --- a/crates/dbx-core/src/sql_risk.rs +++ b/crates/dbx-core/src/sql_risk.rs @@ -5,6 +5,8 @@ use sqlparser::dialect::{ }; use sqlparser::parser::Parser; +use crate::models::connection::DatabaseType; + /// SQL risk level for agent tool safety classification. #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] pub enum SqlRisk { @@ -33,7 +35,8 @@ impl std::fmt::Display for SqlRisk { /// Mirrors the logic in `sql_analysis::normalize_dialect`. fn normalize_dialect(dialect: &str) -> &'static str { match dialect.to_ascii_lowercase().as_str() { - "postgres" | "postgresql" | "redshift" | "opengauss" | "gaussdb" | "highgo" => "postgres", + "postgres" | "postgresql" | "redshift" | "opengauss" | "gaussdb" | "kingbase" | "highgo" | "vastbase" + | "kwdb" => "postgres", "mysql" | "mariadb" | "doris" | "starrocks" | "manticoresearch" | "oceanbase" => "mysql", "sqlite" => "sqlite", "sqlserver" | "mssql" => "sqlserver", @@ -57,11 +60,11 @@ fn resolve_dialect(dialect: &str) -> Box { } /// Classify a single SQL statement into a risk level using AST analysis. -fn classify_statement(stmt: &Statement) -> SqlRisk { +fn classify_statement(stmt: &Statement, detect_select_into: bool) -> SqlRisk { match stmt { // Pure reads Statement::Query(query) => { - if query_contains_select_into(query) { + if detect_select_into && query_contains_select_into(query) { SqlRisk::Write } else { SqlRisk::ReadOnly @@ -69,7 +72,7 @@ fn classify_statement(stmt: &Statement) -> SqlRisk { } Statement::Explain { analyze, statement, .. } => { if *analyze { - classify_statement(statement) + classify_statement(statement, detect_select_into) } else { SqlRisk::ReadOnly } @@ -148,22 +151,49 @@ fn set_expr_contains_select_into(expr: &SetExpr) -> bool { /// Multi-statement input: returns the highest risk level across all statements. pub fn classify_sql_risk(sql: &str, dialect: &str) -> Result { let normalized = normalize_dialect(dialect); - let parser_dialect = resolve_dialect(normalized); + classify_sql_risk_with_database(sql, normalized, None) +} + +/// Classify SQL risk using both the parser dialect and the concrete database +/// type so dialect-specific write forms cannot be mistaken for read queries. +pub fn classify_sql_risk_for_database(sql: &str, database_type: DatabaseType) -> Result { + let database_type_name = format!("{database_type:?}"); + let normalized = normalize_dialect(&database_type_name); + classify_sql_risk_with_database(sql, normalized, Some(database_type)) +} + +fn classify_sql_risk_with_database( + sql: &str, + normalized_dialect: &str, + database_type: Option, +) -> Result { + let parser_dialect = resolve_dialect(normalized_dialect); + let detect_select_into = database_type.is_none(); + let has_dialect_specific_write = database_type + .is_some_and(|database_type| crate::query_execution_sql::has_dialect_specific_write(sql, database_type)); match Parser::parse_sql(parser_dialect.as_ref(), sql) { Ok(stmts) if !stmts.is_empty() => { let mut max_risk = SqlRisk::ReadOnly; for stmt in &stmts { - let risk = classify_statement(stmt); + let risk = classify_statement(stmt, detect_select_into); if risk as u8 > max_risk as u8 { max_risk = risk; } } - Ok(max_risk) + if max_risk == SqlRisk::ReadOnly && has_dialect_specific_write { + Ok(SqlRisk::Write) + } else { + Ok(max_risk) + } } _ => { // Fallback: keyword-based classification - if crate::query_execution_sql::is_write_sql(sql) { + let is_write = database_type.map_or_else( + || crate::query_execution_sql::is_write_sql(sql), + |database_type| crate::query_execution_sql::is_write_sql_for_database(sql, database_type), + ); + if is_write { Ok(SqlRisk::Write) } else { Ok(SqlRisk::ReadOnly) @@ -210,6 +240,61 @@ mod tests { assert_eq!(classify_sql_risk("/*! DELETE FROM users */", "mysql").unwrap(), SqlRisk::Write); } + #[test] + fn classify_dialect_specific_select_into_as_write() { + for sql in [ + "SELECT 3156 INTO OUTFILE '/var/lib/mysql-files/dbx_ro_probe.txt'", + "SELECT 3156 INTO DUMPFILE '/var/lib/mysql-files/dbx_ro_probe.bin'", + ] { + assert_eq!(classify_sql_risk_for_database(sql, DatabaseType::Mysql).unwrap(), SqlRisk::Write); + } + + for database_type in [ + DatabaseType::Postgres, + DatabaseType::Redshift, + DatabaseType::Gaussdb, + DatabaseType::OpenGauss, + DatabaseType::Kingbase, + DatabaseType::Highgo, + DatabaseType::Vastbase, + DatabaseType::Kwdb, + ] { + assert_eq!( + classify_sql_risk_for_database("SELECT * INTO copied_users FROM users", database_type).unwrap(), + SqlRisk::Write, + "expected PostgreSQL-family SELECT INTO to be a write for {database_type:?}" + ); + } + assert_eq!( + classify_sql_risk_for_database("SELECT * INTO #copied_users FROM users", DatabaseType::SqlServer).unwrap(), + SqlRisk::Write + ); + assert_eq!( + classify_sql_risk_for_database( + "SELECT 3156 /*!50000 INTO OUTFILE '/var/lib/mysql-files/dbx_ro_probe.txt' */", + DatabaseType::Mysql, + ) + .unwrap(), + SqlRisk::Write + ); + } + + #[test] + fn typed_classification_preserves_existing_risk_levels() { + assert_eq!( + classify_sql_risk_for_database("SELECT * FROM users", DatabaseType::Postgres).unwrap(), + SqlRisk::ReadOnly + ); + assert_eq!( + classify_sql_risk_for_database("CREATE TABLE users (id INT)", DatabaseType::Postgres).unwrap(), + SqlRisk::Ddl + ); + assert_eq!( + classify_sql_risk_for_database("SELECT 1 INTO unsupported", DatabaseType::Sqlite).unwrap(), + SqlRisk::ReadOnly + ); + } + #[test] fn classify_ddl_statements() { assert_eq!(classify_sql_risk("CREATE TABLE users (id INT)", "postgres").unwrap(), SqlRisk::Ddl);