use std::ops::ControlFlow; use serde::{Deserialize, Serialize}; use sqlparser::ast::{ visit_expressions, BinaryOperator, Expr, FromTable, OnConflictAction, OnInsert, Query, SetExpr, SqliteOnConflict, Statement, TableFactor, UnaryOperator, Value, Visit, Visitor, }; use sqlparser::dialect::{ ClickHouseDialect, DuckDbDialect, GenericDialect, MsSqlDialect, MySqlDialect, PostgreSqlDialect, SQLiteDialect, }; use sqlparser::parser::Parser; use sqlparser::tokenizer::{Token, Tokenizer}; use crate::models::connection::DatabaseType; /// SQL risk level for agent tool safety classification. #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] pub enum SqlRisk { /// SELECT, SHOW, DESCRIBE, EXPLAIN, WITH (pure read CTE) ReadOnly, /// INSERT, UPDATE, DELETE, MERGE, REPLACE, CALL/EXEC Write, /// CREATE, ALTER, DROP, TRUNCATE, GRANT, REVOKE Ddl, /// BEGIN, COMMIT, ROLLBACK should not be issued by agent Transaction, } impl std::fmt::Display for SqlRisk { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { match self { SqlRisk::ReadOnly => write!(f, "read-only"), SqlRisk::Write => write!(f, "write"), SqlRisk::Ddl => write!(f, "DDL"), SqlRisk::Transaction => write!(f, "transaction"), } } } /// Normalize database dialect string to a canonical form for sqlparser. /// 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" | "kingbase" | "highgo" | "uxdb" | "vastbase" | "kwdb" => "postgres", "mysql" | "mariadb" | "doris" | "starrocks" | "manticoresearch" | "oceanbase" => "mysql", "sqlite" => "sqlite", "sqlserver" | "mssql" => "sqlserver", "clickhouse" => "clickhouse", "duckdb" => "duckdb", _ => "generic", } } /// Resolve dialect string to a sqlparser Dialect trait object. fn resolve_dialect(dialect: &str) -> Box { match dialect { "postgres" => Box::new(PostgreSqlDialect {}), "mysql" => Box::new(MySqlDialect {}), "sqlite" => Box::new(SQLiteDialect {}), "sqlserver" => Box::new(MsSqlDialect {}), "clickhouse" => Box::new(ClickHouseDialect {}), "duckdb" => Box::new(DuckDbDialect {}), _ => Box::new(GenericDialect {}), } } /// Classify a single SQL statement into a risk level using AST analysis. fn classify_statement(stmt: &Statement, detect_select_into: bool) -> SqlRisk { match stmt { // Pure reads Statement::Query(query) => { if query_is_write_capable(query, detect_select_into) { SqlRisk::Write } else { SqlRisk::ReadOnly } } Statement::Explain { analyze, statement, .. } => { if *analyze { classify_statement(statement, detect_select_into) } else { SqlRisk::ReadOnly } } Statement::ExplainTable { .. } => SqlRisk::ReadOnly, // Show/Describe variants stmt if is_show_read_only_statement(stmt) => SqlRisk::ReadOnly, // Write operations Statement::Insert { .. } | Statement::Update { .. } | Statement::Delete { .. } | Statement::Merge { .. } => { SqlRisk::Write } // DDL operations Statement::CreateTable { .. } | Statement::CreateView { .. } | Statement::CreateIndex { .. } | Statement::CreateSchema { .. } | Statement::CreateSequence { .. } | Statement::CreateRole { .. } | Statement::CreateType { .. } | Statement::AlterTable { .. } | Statement::AlterIndex { .. } | Statement::AlterView { .. } | Statement::Drop { .. } | Statement::Truncate { .. } => SqlRisk::Ddl, // Grant/Revoke Statement::Grant { .. } | Statement::Revoke { .. } => SqlRisk::Ddl, // Transaction control Statement::StartTransaction { .. } | Statement::Commit { .. } | Statement::Rollback { .. } => { SqlRisk::Transaction } // COPY FROM mutates data; keep COPY conservative because sqlparser does // not expose enough dialect-specific direction detail here. Statement::Copy { .. } => SqlRisk::Write, // SQLite/DuckDB PRAGMA statements can mutate database/session state. Statement::Pragma { .. } => SqlRisk::Write, // Catch-all: conservative write classification _ => SqlRisk::Write, } } /// SHOW/DESCRIBE-style statement variants sqlparser models as reads. Shared by /// the risk classifier and the manual-transaction proof (#9018) so the two /// cannot drift; unparseable SHOW variants fail closed upstream in the proof /// (parse error) and to the keyword fallback in the risk classifier. fn is_show_read_only_statement(statement: &Statement) -> bool { matches!( statement, Statement::ShowTables { .. } | Statement::ShowColumns { .. } | Statement::ShowCatalogs { .. } | Statement::ShowDatabases { .. } | Statement::ShowSchemas { .. } | Statement::ShowViews { .. } | Statement::ShowFunctions { .. } | Statement::ShowCreate { .. } | Statement::ShowVariable { .. } | Statement::ShowVariables { .. } | Statement::ShowStatus { .. } | Statement::ShowProcessList { .. } | Statement::ShowCharset(_) | Statement::ShowObjects(_) | Statement::ShowCollation { .. } ) } fn statement_is_dangerous(stmt: &Statement, detect_select_into: bool) -> bool { match stmt { // `USE` only changes connection-local state. MCP separately requires // a pinned session for it, so it does not need dangerous-SQL approval. Statement::Use(_) => false, Statement::Query(query) => query_is_dangerous(query, detect_select_into), Statement::Insert(insert) => { insert.replace_into || insert.overwrite || matches!(insert.or, Some(SqliteOnConflict::Replace)) || insert.on.as_ref().is_some_and(|on| match on { OnInsert::DuplicateKeyUpdate(_) => true, OnInsert::OnConflict(conflict) => matches!(conflict.action, OnConflictAction::DoUpdate(_)), _ => true, }) || insert.source.as_ref().is_some_and(|query| !matches!(query.body.as_ref(), SetExpr::Values(_))) } Statement::Update(update) => { !update.table.joins.is_empty() || update.from.is_some() || matches!(update.or, Some(SqliteOnConflict::Replace)) || update.selection.as_ref().is_none_or(predicate_is_obviously_unbounded) } Statement::Delete(delete) => { let from = match &delete.from { FromTable::WithFromKeyword(from) | FromTable::WithoutKeyword(from) => from, }; !delete.tables.is_empty() || delete.using.is_some() || from.len() != 1 || from.first().is_none_or(|table| !table.joins.is_empty()) || delete.selection.as_ref().is_none_or(predicate_is_obviously_unbounded) } Statement::Merge { .. } => true, Statement::Explain { analyze: true, statement, .. } => statement_is_dangerous(statement, detect_select_into), _ => !matches!(classify_statement(stmt, detect_select_into), SqlRisk::ReadOnly), } } fn predicate_is_obviously_unbounded(expr: &Expr) -> bool { if expr_contains_subquery(expr) { return true; } if !expr_contains_column_reference(expr) { return true; } match expr { Expr::Nested(expr) => predicate_is_obviously_unbounded(expr), Expr::Value(value) => value_is_truthy(&value.value), Expr::UnaryOp { op: UnaryOperator::Not, expr } => predicate_is_obviously_false(expr), Expr::BinaryOp { left, op: BinaryOperator::And, right } => { predicate_is_obviously_unbounded(left) && predicate_is_obviously_unbounded(right) } Expr::BinaryOp { left, op: BinaryOperator::Or, right } => { or_contains_conjunctive_branch(expr) || or_contains_complementary_null_checks(expr) || or_contains_complementary_comparisons(expr) || or_contains_complementary_in_lists(expr) || or_contains_complementary_between_checks(expr) || predicate_is_obviously_unbounded(left) || predicate_is_obviously_unbounded(right) } Expr::BinaryOp { left, op, right } if is_comparison_operator(op) => constant_comparison_truth(left, op, right) .unwrap_or_else(|| { left == right && matches!( op, BinaryOperator::Eq | BinaryOperator::GtEq | BinaryOperator::LtEq | BinaryOperator::Spaceship ) }), Expr::IsNotDistinctFrom(left, right) => strip_nested_expr(left) == strip_nested_expr(right), Expr::IsTrue(expr) | Expr::IsNotFalse(expr) => predicate_is_obviously_unbounded(expr), Expr::Like { negated: false, any: false, expr, pattern, escape_char: None } | Expr::ILike { negated: false, any: false, expr, pattern, escape_char: None } => { expr_contains_column_reference(expr) && like_pattern_matches_all(pattern) } _ => false, } } fn or_contains_conjunctive_branch(expr: &Expr) -> bool { match strip_nested_expr(expr) { Expr::BinaryOp { left, op: BinaryOperator::Or, right } => { or_contains_conjunctive_branch(left) || or_contains_conjunctive_branch(right) } Expr::BinaryOp { op: BinaryOperator::And, .. } => true, _ => false, } } fn expr_contains_column_reference(expr: &Expr) -> bool { visit_expressions(expr, |expr| { let is_column = match expr { Expr::Identifier(identifier) => !is_sql_value_keyword(&identifier.value), Expr::CompoundIdentifier(_) => true, _ => false, }; if is_column { ControlFlow::Break(()) } else { ControlFlow::Continue(()) } }) .is_break() } fn is_sql_value_keyword(value: &str) -> bool { matches!( value.to_ascii_lowercase().as_str(), "current_catalog" | "current_date" | "current_role" | "current_schema" | "current_time" | "current_timestamp" | "current_user" | "localtime" | "localtimestamp" | "session_user" | "system_user" | "user" ) } fn expr_contains_subquery(expr: &Expr) -> bool { visit_expressions(expr, |expr| { if matches!(expr, Expr::InSubquery { .. } | Expr::Exists { .. } | Expr::Subquery(_)) { ControlFlow::Break(()) } else { ControlFlow::Continue(()) } }) .is_break() } fn or_contains_complementary_null_checks(expr: &Expr) -> bool { let mut checks = Vec::new(); collect_or_null_checks(expr, &mut checks); checks.iter().enumerate().any(|(index, (candidate, negated))| { checks[index + 1..].iter().any(|(other, other_negated)| candidate == other && negated != other_negated) }) } #[derive(Clone, Copy, PartialEq, Eq)] enum PredicateComparisonOperator { Eq, NotEq, Gt, GtEq, Lt, LtEq, } struct PredicateComparisonCheck<'a> { left: &'a Expr, right: &'a Expr, operator: PredicateComparisonOperator, } fn or_contains_complementary_comparisons(expr: &Expr) -> bool { let mut checks = Vec::new(); collect_or_comparison_checks(expr, &mut checks); checks.iter().enumerate().any(|(index, candidate)| { checks[index + 1..].iter().any(|other| { candidate.left == other.left && candidate.right == other.right && comparison_operators_are_complementary(candidate.operator, other.operator) }) }) } fn collect_or_comparison_checks<'a>(expr: &'a Expr, checks: &mut Vec>) { let expr = strip_nested_expr(expr); if let Expr::BinaryOp { left, op: BinaryOperator::Or, right } = expr { collect_or_comparison_checks(left, checks); collect_or_comparison_checks(right, checks); } else if let Some(check) = comparison_check(expr, false) { checks.push(check); } } fn comparison_check(expr: &Expr, outer_negated: bool) -> Option> { let expr = strip_nested_expr(expr); if let Expr::UnaryOp { op: UnaryOperator::Not, expr } = expr { return comparison_check(expr, !outer_negated); } let Expr::BinaryOp { left, op, right } = expr else { return None; }; if !expr_contains_column_reference(expr) { return None; } let mut operator = predicate_comparison_operator(op)?; if outer_negated { operator = complementary_comparison_operator(operator); } let (left, right) = (strip_nested_expr(left), strip_nested_expr(right)); if left <= right { Some(PredicateComparisonCheck { left, right, operator }) } else { Some(PredicateComparisonCheck { left: right, right: left, operator: reverse_comparison_operator(operator) }) } } fn predicate_comparison_operator(operator: &BinaryOperator) -> Option { match operator { BinaryOperator::Eq | BinaryOperator::Spaceship => Some(PredicateComparisonOperator::Eq), BinaryOperator::NotEq => Some(PredicateComparisonOperator::NotEq), BinaryOperator::Gt => Some(PredicateComparisonOperator::Gt), BinaryOperator::GtEq => Some(PredicateComparisonOperator::GtEq), BinaryOperator::Lt => Some(PredicateComparisonOperator::Lt), BinaryOperator::LtEq => Some(PredicateComparisonOperator::LtEq), _ => None, } } fn complementary_comparison_operator(operator: PredicateComparisonOperator) -> PredicateComparisonOperator { match operator { PredicateComparisonOperator::Eq => PredicateComparisonOperator::NotEq, PredicateComparisonOperator::NotEq => PredicateComparisonOperator::Eq, PredicateComparisonOperator::Gt => PredicateComparisonOperator::LtEq, PredicateComparisonOperator::GtEq => PredicateComparisonOperator::Lt, PredicateComparisonOperator::Lt => PredicateComparisonOperator::GtEq, PredicateComparisonOperator::LtEq => PredicateComparisonOperator::Gt, } } fn reverse_comparison_operator(operator: PredicateComparisonOperator) -> PredicateComparisonOperator { match operator { PredicateComparisonOperator::Eq | PredicateComparisonOperator::NotEq => operator, PredicateComparisonOperator::Gt => PredicateComparisonOperator::Lt, PredicateComparisonOperator::GtEq => PredicateComparisonOperator::LtEq, PredicateComparisonOperator::Lt => PredicateComparisonOperator::Gt, PredicateComparisonOperator::LtEq => PredicateComparisonOperator::GtEq, } } fn comparison_operators_are_complementary( left: PredicateComparisonOperator, right: PredicateComparisonOperator, ) -> bool { complementary_comparison_operator(left) == right } struct PredicateInListCheck<'a> { expr: &'a Expr, list: Vec<&'a Expr>, negated: bool, } fn or_contains_complementary_in_lists(expr: &Expr) -> bool { let mut checks = Vec::new(); collect_or_in_list_checks(expr, &mut checks); checks.iter().enumerate().any(|(index, candidate)| { checks[index + 1..].iter().any(|other| { candidate.expr == other.expr && candidate.list == other.list && candidate.negated != other.negated }) }) } fn collect_or_in_list_checks<'a>(expr: &'a Expr, checks: &mut Vec>) { let expr = strip_nested_expr(expr); if let Expr::BinaryOp { left, op: BinaryOperator::Or, right } = expr { collect_or_in_list_checks(left, checks); collect_or_in_list_checks(right, checks); } else if let Some(check) = in_list_check(expr, false) { checks.push(check); } } fn in_list_check(expr: &Expr, outer_negated: bool) -> Option> { match strip_nested_expr(expr) { Expr::UnaryOp { op: UnaryOperator::Not, expr } => in_list_check(expr, !outer_negated), Expr::InList { expr, list, negated } if expr_contains_column_reference(expr) && !list.is_empty() => { let mut normalized_list = list.iter().map(strip_nested_expr).collect::>(); if normalized_list.iter().all(|item| matches!(strip_nested_expr(item), Expr::Value(_))) { normalized_list.sort_unstable(); normalized_list.dedup(); } Some(PredicateInListCheck { expr: strip_nested_expr(expr), list: normalized_list, negated: *negated != outer_negated, }) } _ => None, } } struct PredicateBetweenCheck<'a> { expr: &'a Expr, low: &'a Expr, high: &'a Expr, negated: bool, } fn or_contains_complementary_between_checks(expr: &Expr) -> bool { let mut checks = Vec::new(); collect_or_between_checks(expr, &mut checks); checks.iter().enumerate().any(|(index, candidate)| { checks[index + 1..].iter().any(|other| { candidate.expr == other.expr && candidate.low == other.low && candidate.high == other.high && candidate.negated != other.negated }) }) } fn collect_or_between_checks<'a>(expr: &'a Expr, checks: &mut Vec>) { let expr = strip_nested_expr(expr); if let Expr::BinaryOp { left, op: BinaryOperator::Or, right } = expr { collect_or_between_checks(left, checks); collect_or_between_checks(right, checks); } else if let Some(check) = between_check(expr, false) { checks.push(check); } } fn between_check(expr: &Expr, outer_negated: bool) -> Option> { match strip_nested_expr(expr) { Expr::UnaryOp { op: UnaryOperator::Not, expr } => between_check(expr, !outer_negated), Expr::Between { expr, negated, low, high } if expr_contains_column_reference(expr) => { Some(PredicateBetweenCheck { expr: strip_nested_expr(expr), low: strip_nested_expr(low), high: strip_nested_expr(high), negated: *negated != outer_negated, }) } _ => None, } } fn like_pattern_matches_all(expr: &Expr) -> bool { matches!(strip_nested_expr(expr), Expr::Value(value) if matches!(&value.value, Value::SingleQuotedString(pattern) if !pattern.is_empty() && pattern.chars().all(|character| character == '%'))) } fn collect_or_null_checks<'a>(expr: &'a Expr, checks: &mut Vec<(&'a Expr, bool)>) { let expr = strip_nested_expr(expr); if let Expr::BinaryOp { left, op: BinaryOperator::Or, right } = expr { collect_or_null_checks(left, checks); collect_or_null_checks(right, checks); } else if let Some(check) = null_check(expr) { checks.push(check); } } fn null_check(expr: &Expr) -> Option<(&Expr, bool)> { null_check_with_negation(expr, false) } fn null_check_with_negation(expr: &Expr, outer_negated: bool) -> Option<(&Expr, bool)> { match strip_nested_expr(expr) { Expr::UnaryOp { op: UnaryOperator::Not, expr } => null_check_with_negation(expr, !outer_negated), Expr::IsNull(expr) => Some((strip_nested_expr(expr), outer_negated)), Expr::IsNotNull(expr) => Some((strip_nested_expr(expr), !outer_negated)), _ => None, } } fn strip_nested_expr(mut expr: &Expr) -> &Expr { while let Expr::Nested(inner) = expr { expr = inner; } expr } fn predicate_is_obviously_false(expr: &Expr) -> bool { match expr { Expr::Nested(expr) => predicate_is_obviously_false(expr), Expr::Value(value) => value_is_falsy(&value.value), Expr::UnaryOp { op: UnaryOperator::Not, expr } => predicate_is_obviously_unbounded(expr), Expr::BinaryOp { left, op: BinaryOperator::And, right } => { predicate_is_obviously_false(left) || predicate_is_obviously_false(right) } Expr::BinaryOp { left, op: BinaryOperator::Or, right } => { predicate_is_obviously_false(left) && predicate_is_obviously_false(right) } Expr::BinaryOp { left, op, right } if is_comparison_operator(op) => { constant_comparison_truth(left, op, right).is_some_and(|result| !result) || (left == right && matches!(op, BinaryOperator::NotEq | BinaryOperator::Gt | BinaryOperator::Lt)) } Expr::IsDistinctFrom(left, right) => strip_nested_expr(left) == strip_nested_expr(right), Expr::IsFalse(expr) | Expr::IsNotTrue(expr) => predicate_is_obviously_unbounded(expr), _ => false, } } fn is_comparison_operator(operator: &BinaryOperator) -> bool { matches!( operator, BinaryOperator::Eq | BinaryOperator::NotEq | BinaryOperator::Gt | BinaryOperator::Lt | BinaryOperator::GtEq | BinaryOperator::LtEq | BinaryOperator::Spaceship ) } fn constant_comparison_truth(left: &Expr, operator: &BinaryOperator, right: &Expr) -> Option { let (Expr::Value(left), Expr::Value(right)) = (left, right) else { return None; }; match operator { BinaryOperator::Eq | BinaryOperator::Spaceship => Some(left.value == right.value), BinaryOperator::NotEq => Some(left.value != right.value), BinaryOperator::Gt | BinaryOperator::Lt | BinaryOperator::GtEq | BinaryOperator::LtEq => { let (Value::Number(left, _), Value::Number(right, _)) = (&left.value, &right.value) else { return None; }; let left = left.parse::().ok()?; let right = right.parse::().ok()?; Some(match operator { BinaryOperator::Gt => left > right, BinaryOperator::Lt => left < right, BinaryOperator::GtEq => left >= right, BinaryOperator::LtEq => left <= right, _ => unreachable!(), }) } _ => None, } } fn value_is_truthy(value: &Value) -> bool { match value { Value::Boolean(value) => *value, Value::Number(value, _) => value.parse::().is_ok_and(|number| number != 0.0), _ => false, } } fn value_is_falsy(value: &Value) -> bool { match value { Value::Boolean(value) => !*value, Value::Number(value, _) => value.parse::().is_ok_and(|number| number == 0.0), Value::Null => true, _ => false, } } const SIDE_EFFECT_SELECT_FUNCTIONS: &[&str] = &[ "lo_create", "lo_import", "lo_unlink", "nextval", "pg_advisory_lock", "pg_advisory_unlock", "pg_advisory_unlock_all", "pg_advisory_xact_lock", "pg_cancel_backend", "pg_reload_conf", "pg_terminate_backend", "pg_try_advisory_lock", "pg_try_advisory_xact_lock", "setval", "sys_cancel_backend", ]; fn query_is_write_capable(query: &Query, detect_select_into: bool) -> bool { query .with .as_ref() .is_some_and(|with| with.cte_tables.iter().any(|cte| query_is_write_capable(&cte.query, detect_select_into))) || set_expr_is_write_capable(&query.body, detect_select_into) || !query.locks.is_empty() || query_calls_known_side_effect_function(query) } fn set_expr_is_write_capable(expr: &SetExpr, detect_select_into: bool) -> bool { match expr { SetExpr::Select(select) => detect_select_into && select.into.is_some(), SetExpr::Query(query) => query_is_write_capable(query, detect_select_into), SetExpr::SetOperation { left, right, .. } => { set_expr_is_write_capable(left, detect_select_into) || set_expr_is_write_capable(right, detect_select_into) } SetExpr::Insert(_) | SetExpr::Update(_) | SetExpr::Delete(_) | SetExpr::Merge(_) => true, SetExpr::Values(_) | SetExpr::Table(_) => false, } } fn query_is_dangerous(query: &Query, detect_select_into: bool) -> bool { query .with .as_ref() .is_some_and(|with| with.cte_tables.iter().any(|cte| query_is_dangerous(&cte.query, detect_select_into))) || set_expr_is_dangerous(&query.body, detect_select_into) || !query.locks.is_empty() || query_calls_known_side_effect_function(query) } fn set_expr_is_dangerous(expr: &SetExpr, detect_select_into: bool) -> bool { match expr { SetExpr::Select(select) => detect_select_into && select.into.is_some(), SetExpr::Query(query) => query_is_dangerous(query, detect_select_into), SetExpr::SetOperation { left, right, .. } => { set_expr_is_dangerous(left, detect_select_into) || set_expr_is_dangerous(right, detect_select_into) } SetExpr::Insert(statement) | SetExpr::Update(statement) | SetExpr::Delete(statement) | SetExpr::Merge(statement) => statement_is_dangerous(statement, detect_select_into), SetExpr::Values(_) | SetExpr::Table(_) => false, } } fn query_calls_known_side_effect_function(query: &Query) -> bool { visit_expressions(query, |expr| { let is_side_effect = if let Expr::Function(function) = expr { function.name.0.last().and_then(|part| part.as_ident()).is_some_and(|name| { SIDE_EFFECT_SELECT_FUNCTIONS.iter().any(|candidate| name.value.eq_ignore_ascii_case(candidate)) }) } else { false }; if is_side_effect { ControlFlow::Break(()) } else { ControlFlow::Continue(()) } }) .is_break() } /// Classify SQL risk using sqlparser AST analysis. /// /// If parsing fails (non-standard SQL, non-SQL databases), falls back to /// keyword-based `query_execution_sql::is_write_sql()`. /// /// 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); 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 { if let Some(risk) = crate::query_execution_sql::classify_search_engine_query_risk(sql, database_type) { return Ok(match risk { crate::query_execution_sql::SearchEngineQueryRisk::ReadOnly => SqlRisk::ReadOnly, crate::query_execution_sql::SearchEngineQueryRisk::Write => SqlRisk::Write, crate::query_execution_sql::SearchEngineQueryRisk::Dangerous => SqlRisk::Ddl, }); } let database_type_name = format!("{database_type:?}"); let normalized = normalize_dialect(&database_type_name); classify_sql_risk_with_database(sql, normalized, Some(database_type)) } /// Return whether MCP must require the central dangerous-operation permission. /// Parse failures fail closed for writes. Safe-write mode permits plain INSERT /// and single-table UPDATE/DELETE statements with an effective predicate; /// broader or opaque mutations require central high-risk permission. pub fn is_dangerous_sql_for_database(sql: &str, database_type: DatabaseType) -> bool { if let Some(risk) = crate::query_execution_sql::classify_search_engine_query_risk(sql, database_type) { return risk == crate::query_execution_sql::SearchEngineQueryRisk::Dangerous; } let database_type_name = format!("{database_type:?}"); let normalized = normalize_dialect(&database_type_name); let parser_dialect = resolve_dialect(normalized); let detect_select_into = supports_select_into_table_creation(database_type); let has_locking_clause = sql_contains_locking_clause(sql, parser_dialect.as_ref()); // PostgreSQL-family and SQL Server SELECT INTO forms are represented in // the AST. Their text fallback also matches ordinary INSERT INTO, so only // use it for dialect-specific writes that the selected AST cannot express // (notably MySQL INTO OUTFILE/DUMPFILE and executable comments). let has_unparsed_dialect_specific_write = !detect_select_into && crate::query_execution_sql::has_dialect_specific_write(sql, database_type); match Parser::parse_sql(parser_dialect.as_ref(), sql) { Ok(statements) if !statements.is_empty() => { has_locking_clause || has_unparsed_dialect_specific_write || statements.iter().any(|statement| statement_is_dangerous(statement, detect_select_into)) } _ => has_locking_clause || crate::query_execution_sql::is_write_sql_for_database(sql, database_type), } } /// MCP requests must select their database through the explicit request scope. /// A `USE` statement mutates pooled/session state and could redirect later SQL, /// so it is forbidden independently of read/write and high-risk permissions. pub fn mcp_sql_has_forbidden_database_switch(sql: &str, database_type: DatabaseType) -> bool { if crate::query_execution_sql::classify_search_engine_query_risk(sql, database_type).is_some() { return false; } let database_type_name = format!("{database_type:?}"); let normalized = normalize_dialect(&database_type_name); let dialect = resolve_dialect(normalized); if sql_has_use_statement(sql, dialect.as_ref()) { return true; } if matches!( database_type, DatabaseType::Mysql | DatabaseType::Doris | DatabaseType::StarRocks | DatabaseType::ManticoreSearch | DatabaseType::Goldendb ) { let executable_comments_expanded = crate::query_execution_sql::strip_sql_comments(sql); return sql_has_use_statement(&executable_comments_expanded, dialect.as_ref()); } false } fn sql_has_use_statement(sql: &str, dialect: &dyn sqlparser::dialect::Dialect) -> bool { if let Ok(statements) = Parser::parse_sql(dialect, sql) { return statements.iter().any(|statement| matches!(statement, Statement::Use(_))); } let Ok(tokens) = Tokenizer::new(dialect, sql).tokenize() else { return false; }; let mut statement_start = true; for token in tokens { match token { Token::Whitespace(_) => {} Token::SemiColon => statement_start = true, Token::Word(word) if statement_start => { if word.value.eq_ignore_ascii_case("use") { return true; } statement_start = false; } Token::EOF => {} _ if statement_start => statement_start = false, _ => {} } } false } fn supports_select_into_table_creation(database_type: DatabaseType) -> bool { matches!( database_type, DatabaseType::Postgres | DatabaseType::Redshift | DatabaseType::Gaussdb | DatabaseType::OpenGauss | DatabaseType::Kingbase | DatabaseType::Highgo | DatabaseType::Uxdb | DatabaseType::Vastbase | DatabaseType::Kwdb | DatabaseType::SqlServer ) } 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_locking_clause = sql_contains_locking_clause(sql, parser_dialect.as_ref()); let has_dialect_specific_write = match database_type { Some(database_type) => crate::query_execution_sql::has_dialect_specific_write(sql, database_type), // Preserve MySQL executable-comment and file-output detection even // when callers provide a dialect string instead of a database type. None if normalized_dialect == "mysql" => { crate::query_execution_sql::has_dialect_specific_write(sql, DatabaseType::Mysql) } None => false, }; 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, detect_select_into); if risk as u8 > max_risk as u8 { max_risk = risk; } } if max_risk == SqlRisk::ReadOnly && (has_dialect_specific_write || has_locking_clause) { Ok(SqlRisk::Write) } else { Ok(max_risk) } } _ => { // Fallback: keyword-based classification 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 || has_locking_clause { Ok(SqlRisk::Write) } else { Ok(SqlRisk::ReadOnly) } } } } fn sql_contains_locking_clause(sql: &str, dialect: &dyn sqlparser::dialect::Dialect) -> bool { let Ok(tokens) = Tokenizer::new(dialect, sql).tokenize() else { return false; }; let mut words: Vec = Vec::with_capacity(4); for token in tokens { match token { Token::Word(word) if word.quote_style.is_none() => { words.push(word.value.to_ascii_uppercase()); if words.len() > 4 { words.remove(0); } if words.windows(2).any(|window| { matches!(window, [r#for, lock] if r#for == "FOR" && matches!(lock.as_str(), "UPDATE" | "SHARE")) }) || words.windows(3).any(|window| { matches!(window, [r#for, key, share] if r#for == "FOR" && key == "KEY" && share == "SHARE") }) || words.windows(4).any(|window| { matches!(window, [r#for, no, key, update] if r#for == "FOR" && no == "NO" && key == "KEY" && update == "UPDATE") }) || words.windows(4).any(|window| { matches!(window, [lock, r#in, share, mode] if lock == "LOCK" && r#in == "IN" && share == "SHARE" && mode == "MODE") }) { return true; } } Token::Whitespace(_) => {} _ => words.clear(), } } false } // --------------------------------------------------------------------------- // Strict manual-transaction read-only proof (#9018) // // Unlike `classify_sql_risk*` (keyword fallback on parse failure may answer // ReadOnly), the proof below is fail-closed end to end: the parse must succeed // and yield exactly one statement, no write-capable construct may appear, and // every called function must be an allowlisted pure builtin. The outcome only // drives commit/rollback button visibility in manual transaction mode — it is // a heuristic, never a security boundary (PG custom operators/casts and MySQL // UDFs can hide arbitrary side effects static syntax cannot see). // --------------------------------------------------------------------------- /// Strict manual-transaction UX proof outcome. UI-only. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum ReadProof { /// DBX's strict heuristic is satisfied for this statement. ProvenReadOnly, /// Anything unprovable. Buttons stay visible; execution is unaffected. Unproven, } /// Pure built-in functions accepted by the MySQL proof. Fail-closed: a call to /// anything not listed (UDFs, session-state setters like `LAST_INSERT_ID(expr)`, /// `SLEEP`/`GET_LOCK`, ...) is Unproven. Only ever grows by review. const MYSQL_PROOF_SAFE_FUNCTIONS: &[&str] = &[ "abs", "ascii", "avg", "bin", "ceiling", "char_length", "character_length", "coalesce", "concat", "concat_ws", "conv", "count", "crc32", "curdate", "curtime", "current_date", "current_time", "current_timestamp", "date_format", "datediff", "dayname", "dayofmonth", "dayofweek", "dayofyear", "exp", "floor", "format", "greatest", "hex", "hour", "if", "ifnull", "inet_aton", "inet_ntoa", "instr", "isnull", "json_extract", "json_length", "json_unquote", "json_valid", "last_day", "lcase", "least", "left", "length", "ln", "locate", "log", "log10", "log2", "lower", "lpad", "ltrim", "max", "md5", "microsecond", "min", "minute", "mod", "month", "monthname", "now", "nullif", "oct", "ord", "position", "pow", "power", "quarter", "rand", "repeat", "replace", "reverse", "right", "round", "rpad", "rtrim", "second", "sha", "sha1", "sha2", "sign", "space", "sqrt", "str_to_date", "substring", "substr", "time_format", "timediff", "timestampadd", "timestampdiff", "truncate", "unhex", "unix_timestamp", "upper", "ucase", "utc_date", "utc_time", "utc_timestamp", "uuid", "week", "weekday", "year", ]; /// Pure built-in functions accepted by the PostgreSQL proof. Same contract as /// `MYSQL_PROOF_SAFE_FUNCTIONS`; sequence mutators/readers (`nextval`/`setval`), /// `pg_sleep`, advisory-lock and large-object functions are deliberately absent. const POSTGRES_PROOF_SAFE_FUNCTIONS: &[&str] = &[ "abs", "age", "array_agg", "array_length", "ascii", "avg", "btrim", "cardinality", "ceil", "ceiling", "char_length", "character_length", "chr", "coalesce", "concat", "concat_ws", "count", "current_catalog", "current_date", "current_schema", "current_setting", "current_time", "current_timestamp", "current_user", "date_part", "date_trunc", "decode", "div", "encode", "every", "exp", "floor", "format", "gen_random_uuid", "greatest", "left", "length", "localtime", "localtimestamp", "log", "lower", "lpad", "ltrim", "max", "md5", "min", "mod", "now", "nullif", "position", "power", "repeat", "replace", "reverse", "right", "round", "rpad", "rtrim", "sha224", "sha256", "sha384", "sha512", "sign", "split_part", "sqrt", "starts_with", "strpos", "string_agg", "substr", "substring", "to_char", "to_date", "to_hex", "to_number", "to_timestamp", "trunc", "unnest", "upper", "version", ]; /// Prove that a single SQL statement is, by DBX's strict heuristic, an /// ordinary read. Only MySQL/PostgreSQL take part in this proof; every other /// database type is Unproven so its manual-transaction toolbar keeps the /// legacy behavior. Oracle keeps its own lexical classifier /// (`is_oracle_proven_read_only_statement`) — do not reroute it here. pub fn prove_read_only_for_database(sql: &str, database_type: DatabaseType) -> ReadProof { match database_type { DatabaseType::Mysql => prove_read_only_statement(sql, "mysql", MYSQL_PROOF_SAFE_FUNCTIONS), DatabaseType::Postgres => prove_read_only_statement(sql, "postgres", POSTGRES_PROOF_SAFE_FUNCTIONS), _ => ReadProof::Unproven, } } fn prove_read_only_statement(sql: &str, dialect: &str, allowed_functions: &[&str]) -> ReadProof { let database_type = if dialect == "mysql" { DatabaseType::Mysql } else { DatabaseType::Postgres }; // Lexical rejections that must not depend on parser support: dialect-specific // write syntax (executable comments, INTO OUTFILE/DUMPFILE, PostgreSQL // SELECT INTO) and MySQL session writes the parser models inconsistently // (SELECT @a := 1, SELECT ... INTO @var — INTO may sit in several positions). // `supports_select_into_table_creation` is deliberately NOT reused here: it // answers "SELECT INTO table creation" risk and excludes MySQL on purpose, // while the proof must reject any INTO target as session state. let cleaned = crate::query_execution_sql::strip_sql_comments_and_literals(sql); if cleaned.trim().is_empty() || crate::query_execution_sql::has_dialect_specific_write(sql, database_type) { return ReadProof::Unproven; } let parser_dialect = resolve_dialect(dialect); if dialect == "mysql" && (cleaned.contains(":=") || crate::query_execution_sql::contains_unquoted_keyword(&cleaned, parser_dialect.as_ref(), "INTO")) { return ReadProof::Unproven; } // The strict proof has NO keyword fallback: parse failure is Unproven, and // the API stays fail-closed even for callers that skipped pre-splitting. let Ok(statements) = Parser::parse_sql(parser_dialect.as_ref(), sql) else { return ReadProof::Unproven; }; let [statement] = statements.as_slice() else { return ReadProof::Unproven; }; prove_statement(statement, allowed_functions) } fn prove_statement(statement: &Statement, allowed_functions: &[&str]) -> ReadProof { match statement { Statement::Query(query) => { // Reuses the risk engine's structural walk: writable CTEs, DML set // branches, locking clauses, SELECT INTO and denylisted side-effect // functions. detect_select_into is always true — a proof is about a // clean session, not table-creation risk. if query_is_write_capable(query, true) { return ReadProof::Unproven; } let mut visitor = ProofFunctionVisitor { allowed_functions, rejected: false }; let _ = query.visit(&mut visitor); if visitor.rejected { ReadProof::Unproven } else { ReadProof::ProvenReadOnly } } // Plain EXPLAIN never executes the statement; EXPLAIN ANALYZE does. Statement::Explain { analyze: false, statement, .. } => prove_statement(statement, allowed_functions), Statement::Explain { .. } => ReadProof::Unproven, // MySQL `DESC t` / `DESCRIBE t`. Statement::ExplainTable { .. } => ReadProof::ProvenReadOnly, statement if is_show_read_only_statement(statement) => ReadProof::ProvenReadOnly, _ => ReadProof::Unproven, } } /// Full-statement allowlist scan. Deliberately uses the `visitor`-feature walk /// (covers CTE bodies, FROM subqueries, LATERAL derived tables and ORDER BY /// expressions) instead of `visit_expressions`, which skips FROM subqueries. /// Window functions and table functions (`FROM generate_series(...)`) are /// rejected outright in v1 — fail-closed, relaxed only by review. struct ProofFunctionVisitor<'a> { allowed_functions: &'a [&'a str], rejected: bool, } impl Visitor for ProofFunctionVisitor<'_> { type Break = (); fn pre_visit_expr(&mut self, expr: &Expr) -> ControlFlow<()> { if let Expr::Function(function) = expr { if function.over.is_some() { self.rejected = true; return ControlFlow::Break(()); } let parts = &function.name.0; let allowed = parts.len() == 1 && parts.last().and_then(|part| part.as_ident()).is_some_and(|ident| { self.allowed_functions.iter().any(|candidate| ident.value.eq_ignore_ascii_case(candidate)) }); if !allowed { self.rejected = true; return ControlFlow::Break(()); } } ControlFlow::Continue(()) } fn pre_visit_table_factor(&mut self, table_factor: &TableFactor) -> ControlFlow<()> { if let TableFactor::Table { args: Some(_), .. } = table_factor { self.rejected = true; return ControlFlow::Break(()); } ControlFlow::Continue(()) } } #[cfg(test)] mod tests { use super::*; #[test] fn mysql_proof_accepts_plain_reads_and_allowlisted_functions() { for sql in [ "SELECT * FROM users", "SELECT COUNT(*) FROM orders", "SELECT id, name FROM users WHERE id IN (1, 2, 3)", "SELECT id, name FROM users WHERE id IN (SELECT user_id FROM roles) ORDER BY created_at LIMIT 10", "SELECT 'update' FROM t", "SELECT a FROM t UNION ALL SELECT b FROM u", "WITH totals AS (SELECT COUNT(*) AS n FROM orders) SELECT n FROM totals", "SELECT CONCAT(first_name, ' ', last_name) AS full_name FROM users", "SELECT id FROM users ORDER BY rand() LIMIT 1", "SHOW TABLES", "SHOW CREATE TABLE users", "EXPLAIN SELECT * FROM users", "DESC users", "DESCRIBE users", "-- only a comment header\nSELECT id FROM users", "SELECT 1", ] { assert_eq!( prove_read_only_for_database(sql, DatabaseType::Mysql), ReadProof::ProvenReadOnly, "expected proven: {sql}" ); } } #[test] fn mysql_proof_rejects_session_writes_locking_and_unknown_functions() { for sql in [ "SELECT @a := 1", "SELECT 1 INTO @current_id", "SELECT id FROM users INTO @current_id LIMIT 1", "SELECT * FROM users INTO OUTFILE '/tmp/x'", "SELECT * FROM users FOR UPDATE", "SELECT * FROM users FOR SHARE", "SELECT * FROM users LOCK IN SHARE MODE", "SELECT SLEEP(10)", "SELECT GET_LOCK('x', 1)", "SELECT LAST_INSERT_ID()", "SELECT my_custom_udf(1)", "SELECT secret_schema.fn(1) FROM t", "SELECT id FROM users WHERE EXISTS (SELECT SLEEP(1))", "WITH x AS (SELECT SLEEP(1)) SELECT 1", "SELECT ROW_NUMBER() OVER (ORDER BY id) FROM users", "SELECT id FROM users; SELECT 1", "EXPLAIN ANALYZE SELECT * FROM users", "SELEC * FORM users", "SET autocommit = 1", ] { assert_eq!( prove_read_only_for_database(sql, DatabaseType::Mysql), ReadProof::Unproven, "expected unproven: {sql}" ); } } #[test] fn postgres_proof_accepts_plain_reads_and_rejects_side_effects() { for sql in [ "SELECT * FROM users", "WITH x AS (SELECT 1) SELECT * FROM x", "SELECT * FROM t WHERE id IN (SELECT id FROM u)", "SELECT coalesce(a, b) FROM t", "SELECT split_part(email, '@', 2) FROM users", "EXPLAIN SELECT 1", "SELECT 1", ] { assert_eq!( prove_read_only_for_database(sql, DatabaseType::Postgres), ReadProof::ProvenReadOnly, "expected proven: {sql}" ); } for sql in [ "WITH w AS (INSERT INTO t VALUES (1) RETURNING *) SELECT * FROM w", "SELECT * INTO new_t FROM t", "SELECT * FROM t FOR UPDATE", // FOR SHARE parses into query.locks like FOR UPDATE; the other two // PG lock strengths (FOR NO KEY UPDATE / FOR KEY SHARE) fail the // parse itself in sqlparser 0.62 and stay Unproven fail-closed. "SELECT * FROM t FOR SHARE", "SELECT * FROM t FOR NO KEY UPDATE", "SELECT * FROM t FOR KEY SHARE", "SELECT nextval('seq')", "SELECT setval('seq', 1)", "SELECT pg_sleep(1)", "SELECT lo_import('/etc/passwd')", "EXPLAIN ANALYZE SELECT 1", "SELECT my_udf()", "SELECT id, row_number() OVER () FROM t", "SELECT 1; SELECT 2", ] { assert_eq!( prove_read_only_for_database(sql, DatabaseType::Postgres), ReadProof::Unproven, "expected unproven: {sql}" ); } } #[test] fn proof_visitor_covers_from_subqueries_ctes_lateral_and_order_by() { // visit_expressions skips FROM subqueries; the proof's visitor must not. for sql in [ "SELECT * FROM (SELECT SLEEP(1)) AS x", "SELECT * FROM t ORDER BY SLEEP(1)", "SELECT a FROM t, LATERAL (SELECT my_udf() FROM u) d", "SELECT * FROM generate_series(1, 10)", ] { assert_eq!( prove_read_only_for_database(sql, DatabaseType::Postgres), ReadProof::Unproven, "expected unproven: {sql}" ); } assert_eq!( prove_read_only_for_database("SELECT * FROM (SELECT 1) AS x ORDER BY md5(id)", DatabaseType::Postgres), ReadProof::ProvenReadOnly ); } #[test] fn proof_only_applies_to_enabled_dialects() { // Oracle keeps its own lexical classifier; the gate in query.rs must // never route it through the generic proof. assert_eq!(prove_read_only_for_database("SELECT 1", DatabaseType::Oracle), ReadProof::Unproven); // Family members without manual-transaction UI stay unproven in v1. assert_eq!(prove_read_only_for_database("SELECT 1", DatabaseType::Doris), ReadProof::Unproven); } #[test] fn classify_select_statements() { assert_eq!(classify_sql_risk("SELECT * FROM users", "postgres").unwrap(), SqlRisk::ReadOnly); assert_eq!( classify_sql_risk("SELECT id, name FROM users WHERE active = true", "mysql").unwrap(), SqlRisk::ReadOnly ); assert_eq!(classify_sql_risk("SHOW TABLES", "mysql").unwrap(), SqlRisk::ReadOnly); assert_eq!(classify_sql_risk("DESCRIBE users", "mysql").unwrap(), SqlRisk::ReadOnly); assert_eq!(classify_sql_risk("EXPLAIN SELECT * FROM users", "postgres").unwrap(), SqlRisk::ReadOnly); } #[test] fn classify_mysql_show_triggers_as_read_only() { for sql in [ "SHOW TRIGGERS;", "SHOW TRIGGERS FROM `rs_main` LIKE 'trg_order_items_after_%';", "show triggers in `rs_main` where `Event` = 'INSERT';", ] { assert_eq!(classify_sql_risk(sql, "mysql").unwrap(), SqlRisk::ReadOnly, "expected read-only: {sql}"); assert_eq!( classify_sql_risk_for_database(sql, DatabaseType::Mysql).unwrap(), SqlRisk::ReadOnly, "expected read-only: {sql}" ); assert!(!is_dangerous_sql_for_database(sql, DatabaseType::Mysql), "expected safe SQL: {sql}"); } } #[test] fn classify_supported_show_statements_as_read_only() { for (sql, database_type) in [ ("SHOW COLLATION", DatabaseType::Mysql), ("SHOW CHARACTER SET", DatabaseType::Mysql), ("SHOW search_path", DatabaseType::Postgres), ] { assert_eq!( classify_sql_risk_for_database(sql, database_type).unwrap(), SqlRisk::ReadOnly, "expected read-only: {sql}" ); } } #[test] fn mysql_show_triggers_preserves_write_detection() { for (sql, expected_risk) in [ ("SHOW TRIGGERS; DELETE FROM order_items", SqlRisk::Write), ("SHOW TRIGGERS; DROP TABLE order_items", SqlRisk::Ddl), ("SHOW TRIGGERS; /*!50000 DELETE FROM order_items */", SqlRisk::Write), ] { assert_eq!(classify_sql_risk(sql, "mysql").unwrap(), expected_risk, "expected write detection: {sql}"); assert_eq!( classify_sql_risk_for_database(sql, DatabaseType::Mysql).unwrap(), expected_risk, "expected write detection: {sql}" ); assert!(is_dangerous_sql_for_database(sql, DatabaseType::Mysql), "expected dangerous SQL: {sql}"); } } #[test] fn classify_mysql_legacy_shared_lock_reads_like_for_share() { // MySQL's pre-8.0.1 spelling of `FOR SHARE` is still valid in 8.0.x and // takes the same shared row locks, so it must not classify as a plain read. for sql in [ "SELECT * FROM users LOCK IN SHARE MODE", "SELECT * FROM users lock in share mode", "SELECT id FROM users WHERE id = 1 LOCK IN SHARE MODE", "WITH c AS (SELECT 1) SELECT * FROM users LOCK IN SHARE MODE", "SELECT * FROM (SELECT * FROM users LOCK IN SHARE MODE) AS locked_users", "WITH locked AS (SELECT * FROM users LOCK IN SHARE MODE) SELECT * FROM locked", "SELECT (SELECT id FROM users LOCK IN SHARE MODE)", ] { assert_eq!( classify_sql_risk_for_database(sql, DatabaseType::Mysql).unwrap(), SqlRisk::Write, "expected locking read to be write-capable: {sql}" ); assert!(is_dangerous_sql_for_database(sql, DatabaseType::Mysql), "expected high-risk SQL: {sql}"); } // A column or alias literally named "mode" must stay a plain read. for sql in [ "SELECT lock_in_share_mode FROM settings", "SELECT mode FROM share WHERE id = 1", "SELECT 'LOCK IN SHARE MODE'", "SELECT lock, `in`, share, mode FROM settings", ] { assert_eq!( classify_sql_risk_for_database(sql, DatabaseType::Mysql).unwrap(), SqlRisk::ReadOnly, "expected plain read: {sql}" ); } } #[test] fn classify_cte_read() { assert_eq!( classify_sql_risk("WITH cte AS (SELECT 1) SELECT * FROM cte", "postgres").unwrap(), SqlRisk::ReadOnly ); } #[test] fn classify_writable_ctes_recursively() { for sql in [ "WITH inserted AS (INSERT INTO users (id) VALUES (1) RETURNING id) SELECT * FROM inserted", "WITH updated AS (UPDATE users SET active = true WHERE id = 1 RETURNING id) SELECT * FROM updated", "WITH deleted AS (DELETE FROM users WHERE id = 1 RETURNING id) SELECT * FROM deleted", "WITH merged AS (MERGE INTO users USING staged_users ON users.id = staged_users.id WHEN MATCHED THEN UPDATE SET active = true RETURNING users.id) SELECT * FROM merged", "WITH outer_cte AS (WITH deleted AS (DELETE FROM users WHERE id = 1 RETURNING id) SELECT * FROM deleted) SELECT * FROM outer_cte", ] { assert_eq!(classify_sql_risk(sql, "postgres").unwrap(), SqlRisk::Write, "expected writable CTE: {sql}"); assert!( crate::query_execution_sql::is_write_sql_for_database(sql, DatabaseType::Postgres), "expected writable CTE to trip read-only enforcement: {sql}" ); } } #[test] fn writable_cte_danger_tracks_the_nested_mutation() { for sql in [ "WITH updated AS (UPDATE users SET active = true WHERE id = 1 RETURNING id) SELECT * FROM updated", "WITH deleted AS (DELETE FROM users WHERE id = 1 RETURNING id) SELECT * FROM deleted", "WITH outer_cte AS (WITH deleted AS (DELETE FROM users WHERE id = 1 RETURNING id) SELECT * FROM deleted) SELECT * FROM outer_cte", ] { assert!(!is_dangerous_sql_for_database(sql, DatabaseType::Postgres), "expected guarded write: {sql}"); } for sql in [ "WITH updated AS (UPDATE users SET active = true RETURNING id) SELECT * FROM updated", "WITH deleted AS (DELETE FROM users RETURNING id) SELECT * FROM deleted", "WITH outer_cte AS (WITH deleted AS (DELETE FROM users RETURNING id) SELECT * FROM deleted) SELECT * FROM outer_cte", ] { assert!(is_dangerous_sql_for_database(sql, DatabaseType::Postgres), "expected dangerous write: {sql}"); } } #[test] fn classify_write_statements() { assert_eq!(classify_sql_risk("INSERT INTO users VALUES (1)", "postgres").unwrap(), SqlRisk::Write); assert_eq!(classify_sql_risk("UPDATE users SET name = 'x'", "postgres").unwrap(), SqlRisk::Write); assert_eq!(classify_sql_risk("DELETE FROM users", "postgres").unwrap(), SqlRisk::Write); assert_eq!(classify_sql_risk("EXPLAIN ANALYZE DELETE FROM users", "postgres").unwrap(), SqlRisk::Write); assert_eq!(classify_sql_risk("SELECT * INTO backup_users FROM users", "postgres").unwrap(), SqlRisk::Write); assert_eq!( classify_sql_risk("SELECT * FROM users INTO OUTFILE '/tmp/users.csv'", "mysql").unwrap(), SqlRisk::Write ); assert_eq!(classify_sql_risk("/*! DELETE FROM users */", "mysql").unwrap(), SqlRisk::Write); } #[test] fn mcp_forbids_persistent_database_switching_in_every_permission_mode() { for (sql, database_type) in [ ("USE reporting", DatabaseType::Mysql), ("-- target database\nUSE reporting", DatabaseType::Mysql), ("SELECT 1; USE reporting", DatabaseType::Mysql), ("USE [reporting]", DatabaseType::SqlServer), ("USE DATABASE reporting", DatabaseType::Snowflake), ("/*!50000 USE reporting */", DatabaseType::Mysql), ] { assert!(mcp_sql_has_forbidden_database_switch(sql, database_type), "expected blocked SQL: {sql}"); } for sql in ["SELECT use FROM feature_flags", "SELECT 'USE reporting'", "SELECT 1"] { assert!(!mcp_sql_has_forbidden_database_switch(sql, DatabaseType::Mysql), "expected allowed SQL: {sql}"); } assert!(!mcp_sql_has_forbidden_database_switch("/*!50000 USE reporting */", DatabaseType::Postgres)); } #[test] fn classify_known_side_effect_selects_and_copy_as_writes() { for sql in [ "SELECT setval('user_id_seq', 42)", "SELECT nextval('user_id_seq')", "SELECT pg_terminate_backend(42)", "SELECT * FROM users FOR UPDATE", "SELECT * FROM users FOR KEY SHARE", "COPY users TO '/tmp/users.csv'", "COPY (SELECT * FROM users) TO PROGRAM 'cat > /tmp/users.csv'", ] { assert_eq!( classify_sql_risk_for_database(sql, DatabaseType::Postgres).unwrap(), SqlRisk::Write, "expected write-capable SQL: {sql}" ); assert!(is_dangerous_sql_for_database(sql, DatabaseType::Postgres), "expected high-risk SQL: {sql}"); } } #[test] fn classify_backend_query_cancellation_as_a_side_effect() { for (sql, database_type) in [ ("SELECT pg_cancel_backend(42)", DatabaseType::Postgres), ("SELECT sys_cancel_backend(42)", DatabaseType::Kingbase), ] { assert_eq!( classify_sql_risk_for_database(sql, database_type).unwrap(), SqlRisk::Write, "expected query cancellation to be write-capable: {sql}" ); assert!( is_dangerous_sql_for_database(sql, database_type), "expected query cancellation to require production protection: {sql}" ); } } #[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); assert_eq!(classify_sql_risk("DROP TABLE users", "postgres").unwrap(), SqlRisk::Ddl); assert_eq!(classify_sql_risk("ALTER TABLE users ADD COLUMN age INT", "postgres").unwrap(), SqlRisk::Ddl); assert_eq!(classify_sql_risk("TRUNCATE TABLE users", "postgres").unwrap(), SqlRisk::Ddl); } #[test] fn high_risk_sql_requires_central_permission_for_unbounded_changes() { assert!(is_dangerous_sql_for_database("TRUNCATE TABLE users", DatabaseType::Mysql)); assert!(is_dangerous_sql_for_database("DELETE FROM users", DatabaseType::Mysql)); assert!(is_dangerous_sql_for_database("UPDATE users SET active = 0", DatabaseType::Mysql)); assert!(is_dangerous_sql_for_database("DELETE FROM users WHERE 1 = 1", DatabaseType::Mysql)); assert!(is_dangerous_sql_for_database("UPDATE users SET active = 0 WHERE TRUE", DatabaseType::Mysql)); assert!(is_dangerous_sql_for_database("DELETE FROM users WHERE id = id", DatabaseType::Mysql)); assert!(is_dangerous_sql_for_database( "DELETE FROM users WHERE lower(email) = lower(email)", DatabaseType::Mysql )); assert!(is_dangerous_sql_for_database("DELETE FROM users WHERE NOT (1 = 0)", DatabaseType::Mysql)); assert!(is_dangerous_sql_for_database("UPDATE users SET active = 0 WHERE 2 > 1", DatabaseType::Mysql)); assert!(is_dangerous_sql_for_database( "DELETE FROM users WHERE id IS NULL OR id IS NOT NULL", DatabaseType::Mysql )); assert!(is_dangerous_sql_for_database( "DELETE FROM users WHERE id IS NULL OR NOT (((id IS NULL)))", DatabaseType::Mysql )); assert!(is_dangerous_sql_for_database( "DELETE FROM users WHERE id IS NOT NULL OR NOT (((id IS NOT NULL)))", DatabaseType::Mysql )); assert!(is_dangerous_sql_for_database("DELETE FROM users WHERE id = 1 OR id <> 1", DatabaseType::Mysql)); assert!(is_dangerous_sql_for_database("DELETE FROM users WHERE id != 1 OR 1 = id", DatabaseType::Mysql)); assert!(is_dangerous_sql_for_database( "DELETE FROM users WHERE status = 'disabled' OR status != 'disabled'", DatabaseType::Mysql )); assert!(is_dangerous_sql_for_database( "DELETE FROM users WHERE (id = 1 OR status = 'disabled') OR 1 != id OR id IS NULL", DatabaseType::Mysql )); assert!(is_dangerous_sql_for_database("DELETE FROM users WHERE id > 1 OR id <= 1", DatabaseType::Mysql)); assert!(is_dangerous_sql_for_database("DELETE FROM users WHERE 1 >= id OR id > 1", DatabaseType::Mysql)); assert!(is_dangerous_sql_for_database("DELETE FROM users WHERE id >= 1 OR 1 > id", DatabaseType::Mysql)); assert!(is_dangerous_sql_for_database( "DELETE FROM users WHERE id IN (1) OR id NOT IN (1) OR id IS NULL", DatabaseType::Mysql )); assert!(is_dangerous_sql_for_database( "DELETE FROM users WHERE id IN (1, 2) OR id NOT IN (2, 1) OR id IS NULL", DatabaseType::Mysql )); assert!(is_dangerous_sql_for_database( "DELETE FROM users WHERE id BETWEEN 1 AND 2 OR id NOT BETWEEN 1 AND 2 OR id IS NULL", DatabaseType::Mysql )); assert!(is_dangerous_sql_for_database( "DELETE FROM users WHERE id = 1 OR (id <> 1 AND TRUE) OR id IS NULL", DatabaseType::Mysql )); assert!(is_dangerous_sql_for_database( "DELETE FROM users WHERE id IS NOT DISTINCT FROM id", DatabaseType::Postgres )); assert!(is_dangerous_sql_for_database("DELETE FROM users WHERE id <=> id", DatabaseType::Mysql)); assert!(is_dangerous_sql_for_database( "UPDATE users SET active = 0 WHERE name LIKE '%' OR name IS NULL", DatabaseType::Mysql )); assert!(is_dangerous_sql_for_database( "UPDATE users SET active = 0 WHERE name LIKE '%%' OR name IS NULL", DatabaseType::Mysql )); assert!(is_dangerous_sql_for_database( "DELETE FROM users WHERE (id IS NULL OR status = 'disabled') OR id IS NOT NULL", DatabaseType::Mysql )); assert!(is_dangerous_sql_for_database("UPDATE users SET active = 0 WHERE abs(1) = 1", DatabaseType::Mysql)); assert!(is_dangerous_sql_for_database("DELETE FROM users WHERE lower('A') = 'a'", DatabaseType::Mysql)); assert!(is_dangerous_sql_for_database("DELETE FROM users WHERE coalesce(NULL, 1) = 1", DatabaseType::Mysql)); assert!(is_dangerous_sql_for_database("DELETE FROM users WHERE LOWER(_utf8mb4'A') = 'a'", DatabaseType::Mysql)); assert!(is_dangerous_sql_for_database( "DELETE FROM users WHERE EXTRACT(YEAR FROM DATE '2026-01-01') = 2026", DatabaseType::Postgres )); assert!(is_dangerous_sql_for_database( "DELETE FROM users WHERE DATE '2026-01-01' < CURRENT_DATE", DatabaseType::Postgres )); assert!(is_dangerous_sql_for_database("DELETE FROM users WHERE USER = CURRENT_USER", DatabaseType::Postgres)); assert!(is_dangerous_sql_for_database( "DELETE FROM users WHERE id IN (SELECT id FROM archived_users)", DatabaseType::Postgres )); assert!(is_dangerous_sql_for_database( "DELETE FROM users WHERE EXISTS (SELECT 1 FROM archived_users)", DatabaseType::Postgres )); assert!(!is_dangerous_sql_for_database("DELETE FROM users WHERE id = 1", DatabaseType::Mysql)); assert!(!is_dangerous_sql_for_database("UPDATE users SET active = 0 WHERE id = 1", DatabaseType::Mysql)); assert!(!is_dangerous_sql_for_database("UPDATE users SET active = 0 WHERE abs(id) = 1", DatabaseType::Mysql)); assert!(!is_dangerous_sql_for_database( "DELETE FROM users WHERE lower(email) = 'disabled@example.com'", DatabaseType::Mysql )); assert!(!is_dangerous_sql_for_database( "DELETE FROM users WHERE EXTRACT(YEAR FROM created_at) = 2026", DatabaseType::Postgres )); assert!(!is_dangerous_sql_for_database( "DELETE FROM users WHERE id IS NULL OR status IS NOT NULL", DatabaseType::Mysql )); assert!(!is_dangerous_sql_for_database( "DELETE FROM users WHERE (id IS NULL OR NOT (((id IS NULL)))) AND tenant_id = 1", DatabaseType::Mysql )); assert!(!is_dangerous_sql_for_database( "DELETE FROM users WHERE (id = 1 OR id <> 1) AND tenant_id = 1", DatabaseType::Mysql )); assert!(!is_dangerous_sql_for_database( "DELETE FROM users WHERE status = 'pending' OR status <> 'disabled'", DatabaseType::Mysql )); assert!(!is_dangerous_sql_for_database("DELETE FROM users WHERE id IN (1) OR id IN (2)", DatabaseType::Mysql)); assert!(!is_dangerous_sql_for_database( "DELETE FROM users WHERE id BETWEEN 1 AND 2 OR id BETWEEN 4 AND 5", DatabaseType::Mysql )); assert!(!is_dangerous_sql_for_database( "UPDATE users SET active = 0 WHERE name LIKE 'admin%'", DatabaseType::Mysql )); assert!(!is_dangerous_sql_for_database( "UPDATE users SET active = 0 WHERE status = 'inactive' AND tenant_id = 1", DatabaseType::Mysql )); assert!(is_dangerous_sql_for_database( "UPDATE users JOIN accounts ON accounts.id = users.account_id SET users.active = 0 WHERE users.id = 1", DatabaseType::Mysql )); assert!(is_dangerous_sql_for_database( "UPDATE users SET active = false FROM accounts WHERE users.account_id = accounts.id", DatabaseType::Postgres )); assert!(is_dangerous_sql_for_database( "DELETE users FROM users JOIN accounts ON accounts.id = users.account_id WHERE users.id = 1", DatabaseType::Mysql )); assert!(is_dangerous_sql_for_database("REPLACE INTO users (id) VALUES (1)", DatabaseType::Mysql)); assert!(is_dangerous_sql_for_database( "INSERT INTO users (id) VALUES (1) ON DUPLICATE KEY UPDATE active = 1", DatabaseType::Mysql )); assert!(is_dangerous_sql_for_database( "INSERT INTO users (id) VALUES (1) ON CONFLICT (id) DO UPDATE SET active = true", DatabaseType::Postgres )); assert!(!is_dangerous_sql_for_database("INSERT INTO users (id) VALUES (1)", DatabaseType::Mysql)); assert!(is_dangerous_sql_for_database( "INSERT INTO users (id) SELECT id FROM staged_users", DatabaseType::Mysql )); assert!(!is_dangerous_sql_for_database( "INSERT INTO users (id) VALUES (1) ON CONFLICT (id) DO NOTHING", DatabaseType::Postgres )); } #[test] fn select_into_requires_central_high_risk_permission_for_supported_databases() { for database_type in [ DatabaseType::Postgres, DatabaseType::Redshift, DatabaseType::Gaussdb, DatabaseType::OpenGauss, DatabaseType::Kingbase, DatabaseType::Highgo, DatabaseType::Vastbase, DatabaseType::Kwdb, ] { assert!( is_dangerous_sql_for_database("SELECT * INTO copied_users FROM users", database_type), "expected PostgreSQL-family SELECT INTO to require high-risk permission for {database_type:?}" ); } assert!(is_dangerous_sql_for_database("SELECT * INTO #copied_users FROM users", DatabaseType::SqlServer)); assert!(is_dangerous_sql_for_database("SELECT 1 INTO OUTFILE '/tmp/dbx-probe.txt'", DatabaseType::Mysql)); assert!(!is_dangerous_sql_for_database("SELECT 1", DatabaseType::Postgres)); assert!(!is_dangerous_sql_for_database("SELECT 1 INTO unsupported", DatabaseType::Sqlite)); } #[test] fn classify_transaction_statements() { assert_eq!(classify_sql_risk("BEGIN", "postgres").unwrap(), SqlRisk::Transaction); assert_eq!(classify_sql_risk("COMMIT", "postgres").unwrap(), SqlRisk::Transaction); assert_eq!(classify_sql_risk("ROLLBACK", "postgres").unwrap(), SqlRisk::Transaction); } #[test] fn classify_multi_statement_returns_highest_risk() { // SELECT + INSERT = Write assert_eq!(classify_sql_risk("SELECT 1; INSERT INTO users VALUES (1)", "postgres").unwrap(), SqlRisk::Write); } #[test] fn classify_fallback_on_parse_error() { // Non-standard SQL should fall back to keyword matching assert_eq!(classify_sql_risk("SELECT * FROM users", "generic").unwrap(), SqlRisk::ReadOnly); } #[test] fn classify_unknown_statement_is_write() { // Statements not explicitly handled should be conservative (Write) // This depends on sqlparser's coverage, but we can test the catch-all assert_eq!(classify_sql_risk("GRANT SELECT ON users TO admin", "postgres").unwrap(), SqlRisk::Ddl); } #[test] fn classifies_search_engine_rest_risk_by_method_and_path() { for database_type in [DatabaseType::Elasticsearch, DatabaseType::Easysearch] { assert_eq!( classify_sql_risk_for_database("GET /_cluster/health", database_type).unwrap(), SqlRisk::ReadOnly ); assert_eq!( classify_sql_risk_for_database("POST /products/_search\n{}", database_type).unwrap(), SqlRisk::ReadOnly ); assert_eq!( classify_sql_risk_for_database("PUT /products/_doc/1\n{}", database_type).unwrap(), SqlRisk::Write ); assert!(!is_dangerous_sql_for_database("PUT /products/_doc/1\n{}", database_type)); assert_eq!(classify_sql_risk_for_database("DELETE /products", database_type).unwrap(), SqlRisk::Ddl); assert!(is_dangerous_sql_for_database("DELETE /products", database_type)); } } }