fix(sql): improve read-only statement detection
This commit is contained in:
二丫讲梵
GitHub
parent
1ecd8cf406
commit
0e779f9edc
|
|
@ -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),
|
||||
};
|
||||
|
|
|
|||
|
|
@ -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<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 {
|
||||
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::<Vec<_>>())
|
||||
.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;
|
||||
|
|
|
|||
|
|
@ -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<DatabaseType>) -> 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<std::str::Chars<'_>>) -> bool {
|
||||
let mut marker = chars.clone();
|
||||
match marker.next() {
|
||||
Some('!') => true,
|
||||
Some('M') => marker.next() == Some('!'),
|
||||
_ => false,
|
||||
}
|
||||
}
|
||||
|
||||
fn read_mysql_executable_comment_body<I>(chars: &mut std::iter::Peekable<I>) -> Option<String>
|
||||
|
|
@ -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]
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
|
|
|
|||
|
|
@ -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<dyn sqlparser::dialect::Dialect> {
|
|||
}
|
||||
|
||||
/// 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<SqlRisk, String> {
|
||||
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<SqlRisk, String> {
|
||||
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<DatabaseType>,
|
||||
) -> Result<SqlRisk, String> {
|
||||
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);
|
||||
|
|
|
|||
Loading…
Reference in New Issue