diff --git a/crates/dbx-core/src/agent_explain.rs b/crates/dbx-core/src/agent_explain.rs index 9db8d5a1ac..15028992e4 100644 --- a/crates/dbx-core/src/agent_explain.rs +++ b/crates/dbx-core/src/agent_explain.rs @@ -1,7 +1,7 @@ use serde_json::Value; use crate::connection::{AppState, PoolKind}; -use crate::query_execution_sql::{is_safe_dameng_autotrace_sql, is_safe_explain_sql}; +use crate::query_execution_sql::{is_safe_dameng_autotrace_sql, is_safe_explain_sql_for_database}; pub async fn get_agent_explain_info_core( state: &AppState, @@ -12,12 +12,12 @@ pub async fn get_agent_explain_info_core( mode: Option<&str>, ) -> Result { let mode = mode.unwrap_or("explain"); - let safe = if mode.eq_ignore_ascii_case("autotrace") { - is_safe_dameng_autotrace_sql(sql) - } else { - is_safe_explain_sql(sql) + let (database_type, timeout_secs) = { + let configs = state.configs.read().await; + let config = configs.get(connection_id).ok_or_else(|| "Connection config not found".to_string())?; + (config.db_type, config.query_timeout_secs) }; - if !safe { + if !is_safe_agent_explain_sql(sql, mode, database_type) { return Err("unsafe".to_string()); } @@ -33,11 +33,6 @@ pub async fn get_agent_explain_info_core( } }; - let timeout_secs = { - let configs = state.configs.read().await; - configs.get(connection_id).ok_or_else(|| "Connection config not found".to_string())?.query_timeout_secs - }; - let params = serde_json::json!({ "sql": sql, "database": database.unwrap_or_default(), @@ -50,6 +45,14 @@ pub async fn get_agent_explain_info_core( decode_agent_explain_result(result) } +fn is_safe_agent_explain_sql(sql: &str, mode: &str, database_type: crate::models::connection::DatabaseType) -> bool { + if mode.eq_ignore_ascii_case("autotrace") { + is_safe_dameng_autotrace_sql(sql) + } else { + is_safe_explain_sql_for_database(sql, Some(database_type)) + } +} + fn decode_agent_explain_result(result: Value) -> Result { match result { Value::String(plan) => Ok(plan), @@ -61,6 +64,17 @@ fn decode_agent_explain_result(result: Value) -> Result { #[cfg(test)] mod tests { use super::*; + use crate::models::connection::DatabaseType; + + #[test] + fn allows_oracle_dml_explain_without_relaxing_autotrace() { + let update = "UPDATE AA_PAY_VOUCHER_TEMP t SET t.AGENCY_ID = '1'"; + + assert!(is_safe_agent_explain_sql(update, "explain", DatabaseType::Oracle)); + assert!(!is_safe_agent_explain_sql(update, "autotrace", DatabaseType::Oracle)); + assert!(!is_safe_agent_explain_sql(update, "explain", DatabaseType::Dameng)); + assert!(!is_safe_agent_explain_sql("DROP TABLE AA_PAY_VOUCHER_TEMP", "explain", DatabaseType::Oracle)); + } #[test] fn decodes_string_and_object_agent_explain_results() { diff --git a/crates/dbx-core/src/query_execution_sql.rs b/crates/dbx-core/src/query_execution_sql.rs index ecfdb7a7e1..0c9f897ee2 100644 --- a/crates/dbx-core/src/query_execution_sql.rs +++ b/crates/dbx-core/src/query_execution_sql.rs @@ -54,7 +54,7 @@ pub fn build_explain_sql(options: ExplainSqlOptions) -> ExplainSqlBuildResult { if source.is_empty() { return explain_err("empty"); } - if !is_safe_explain_sql(&source) { + if !is_safe_explain_sql_for_database(&source, options.database_type) { return explain_err("unsafe"); } if options.analyze == Some(true) @@ -133,11 +133,18 @@ pub fn supports_explain_plan(database_type: Option) -> bool { } pub fn is_safe_explain_sql(sql: &str) -> bool { + is_safe_explain_sql_for_database(sql, None) +} + +pub fn is_safe_explain_sql_for_database(sql: &str, database_type: Option) -> bool { let source = strip_trailing_semicolons(sql.trim()); - !source.is_empty() - && !has_extra_statement_after_semicolon(&source) - && is_safe_explain_source(&source) - && !contains_dangerous_sql_keyword(&source) + if source.is_empty() || has_extra_statement_after_semicolon(&source) { + return false; + } + if database_type == Some(DatabaseType::Oracle) && is_safe_oracle_explain_dml_source(&source) { + return true; + } + is_safe_explain_source(&source) && !contains_dangerous_sql_keyword(&source) } /// Returns true for databases that support SQL query execution (execute_query / get_sample_data). @@ -183,6 +190,12 @@ fn is_safe_explain_source(sql: &str) -> bool { }) } +fn is_safe_oracle_explain_dml_source(sql: &str) -> bool { + let source = strip_sql_comments_and_literals(sql); + let source = source.trim_start().to_ascii_uppercase(); + ["INSERT", "UPDATE", "DELETE", "MERGE"].iter().any(|keyword| starts_with_keyword(&source, keyword)) +} + pub fn contains_dangerous_sql_keyword(sql: &str) -> bool { let source = strip_sql_comments_and_literals(sql).to_lowercase(); ["drop", "delete", "truncate", "alter", "update", "merge", "replace", "insert", "create"] @@ -984,6 +997,67 @@ mod tests { ); } + #[test] + fn builds_oracle_explain_plan_for_update() { + let result = build_explain_sql(ExplainSqlOptions { + database_type: Some(DatabaseType::Oracle), + format: None, + analyze: None, + sql: "UPDATE AA_PAY_VOUCHER_TEMP t SET t.AGENCY_ID = '1';".to_string(), + }); + + assert_eq!( + result, + ExplainSqlBuildResult { + ok: true, + sql: Some("EXPLAIN PLAN FOR UPDATE AA_PAY_VOUCHER_TEMP t SET t.AGENCY_ID = '1'".to_string()), + reason: None, + } + ); + } + + #[test] + fn allows_oracle_explain_plan_for_dml_only() { + for (sql, expected) in [ + ("INSERT INTO audit_log (id) VALUES (1)", "EXPLAIN PLAN FOR INSERT INTO audit_log (id) VALUES (1)"), + ("DELETE FROM audit_log WHERE id = 1", "EXPLAIN PLAN FOR DELETE FROM audit_log WHERE id = 1"), + ( + "MERGE INTO target t USING source s ON (t.id = s.id) WHEN MATCHED THEN UPDATE SET t.name = s.name", + "EXPLAIN PLAN FOR MERGE INTO target t USING source s ON (t.id = s.id) WHEN MATCHED THEN UPDATE SET t.name = s.name", + ), + ] { + assert_eq!( + build_explain_sql(ExplainSqlOptions { + database_type: Some(DatabaseType::Oracle), + format: None, + analyze: None, + sql: sql.to_string(), + }) + .sql, + Some(expected.to_string()), + "{sql}" + ); + } + + for sql in [ + "DROP TABLE audit_log", + "ALTER TABLE audit_log ADD created_at DATE", + "BEGIN DELETE FROM audit_log; END;", + "UPDATE audit_log SET id = 2; DROP TABLE audit_log", + ] { + assert_eq!( + build_explain_sql(ExplainSqlOptions { + database_type: Some(DatabaseType::Oracle), + format: None, + analyze: None, + sql: sql.to_string(), + }), + ExplainSqlBuildResult { ok: false, sql: None, reason: Some("unsafe".to_string()) }, + "{sql}" + ); + } + } + #[test] fn builds_sqlserver_showplan_xml_batches() { let result = build_explain_sql(ExplainSqlOptions {