diff --git a/crates/dbx-core/src/sql.rs b/crates/dbx-core/src/sql.rs index 76b15ea55..64951ca7a 100644 --- a/crates/dbx-core/src/sql.rs +++ b/crates/dbx-core/src/sql.rs @@ -1,4 +1,6 @@ use serde::{Deserialize, Serialize}; +use sqlparser::dialect::OracleDialect; +use sqlparser::tokenizer::{Token, Tokenizer}; use crate::models::connection::DatabaseType; @@ -1534,79 +1536,130 @@ fn first_executable_sql_token_with_options(sql: &str, options: SqlParsingOptions } fn starts_with_oracle_plsql_block(sql: &str) -> bool { - let tokens = oracle_plsql_tokens(sql); - match tokens.as_slice() { - [first, ..] if first.eq_ignore_ascii_case("DECLARE") => true, - [first, second, ..] - if first.eq_ignore_ascii_case("BEGIN") && !matches!(second.as_str(), ";" | "TRANSACTION" | "WORK") => - { - true - } - [first, ..] if first.eq_ignore_ascii_case("CREATE") => { - let rest = &tokens[1..]; - // Skip optional OR REPLACE - let rest = skip_or_replace(rest); - is_oracle_plsql_object_type(rest) - } - _ => false, - } -} - -/// Skip the optional `OR REPLACE` token pair. -fn skip_or_replace(tokens: &[String]) -> &[String] { - match tokens { - [or, replace, rest @ ..] if or.eq_ignore_ascii_case("OR") && replace.eq_ignore_ascii_case("REPLACE") => rest, - _ => tokens, - } -} - -/// Check whether the first token is an Oracle PL/SQL object type -/// (FUNCTION, PROCEDURE, TRIGGER, PACKAGE, or TYPE). -fn is_oracle_plsql_object_type(tokens: &[String]) -> bool { - tokens.first().is_some_and(|t| matches!(t.as_str(), "FUNCTION" | "PROCEDURE" | "TRIGGER" | "PACKAGE" | "TYPE")) + OraclePlSqlBlock::parse(sql).starts_block() } fn oracle_plsql_block_is_complete(sql: &str) -> bool { - if !starts_with_oracle_plsql_block(sql) { - return false; - } - - let mut depth = 0usize; - let mut saw_begin = false; - let mut complete = false; - let mut pending_end: Option = None; - - for token in oracle_plsql_tokens(sql) { - if token == ";" { - if let Some(is_block_end) = pending_end.take() { - if is_block_end && depth > 0 { - depth -= 1; - complete = depth == 0; - } - } - continue; - } - - if let Some(is_block_end) = pending_end.as_mut() { - if matches!(token.as_str(), "IF" | "LOOP" | "CASE") { - *is_block_end = false; - } - continue; - } - - if token == "BEGIN" { - depth += 1; - saw_begin = true; - complete = false; - } else if token == "END" { - pending_end = Some(true); - } - } - - saw_begin && complete + OraclePlSqlBlock::parse(sql).is_complete() } -fn oracle_plsql_tokens(sql: &str) -> Vec { +struct OraclePlSqlBlock { + tokens: Vec, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +enum OraclePlSqlToken { + Word(String), + Semicolon, +} + +impl OraclePlSqlBlock { + fn parse(sql: &str) -> Self { + Self { tokens: oracle_plsql_tokens(sql) } + } + + fn starts_block(&self) -> bool { + match self.tokens.as_slice() { + [first, ..] if first.is_word("DECLARE") => true, + [first, second, ..] if first.is_word("BEGIN") && !Self::is_transaction_begin_tail(second) => true, + [first, rest @ ..] if first.is_word("CREATE") => Self::starts_create_plsql_object(rest), + _ => false, + } + } + + fn is_complete(&self) -> bool { + if !self.starts_block() { + return false; + } + + let mut depth = 0usize; + let mut saw_begin = false; + let mut complete = false; + let mut pending_end: Option = None; + + for token in &self.tokens { + if token.is_semicolon() { + if let Some(is_block_end) = pending_end.take() { + if is_block_end && depth > 0 { + depth -= 1; + complete = depth == 0; + } + } + continue; + } + + if let Some(is_block_end) = pending_end.as_mut() { + if token.is_any_word(&["IF", "LOOP", "CASE"]) { + *is_block_end = false; + } + continue; + } + + if token.is_word("BEGIN") { + depth += 1; + saw_begin = true; + complete = false; + } else if token.is_word("END") { + pending_end = Some(true); + } + } + + saw_begin && complete + } + + fn starts_create_plsql_object(tokens: &[OraclePlSqlToken]) -> bool { + let tokens = Self::skip_or_replace(tokens); + tokens.first().is_some_and(|token| token.is_any_word(&["FUNCTION", "PROCEDURE", "TRIGGER", "PACKAGE", "TYPE"])) + } + + fn skip_or_replace(tokens: &[OraclePlSqlToken]) -> &[OraclePlSqlToken] { + match tokens { + [or, replace, rest @ ..] if or.is_word("OR") && replace.is_word("REPLACE") => rest, + _ => tokens, + } + } + + fn is_transaction_begin_tail(token: &OraclePlSqlToken) -> bool { + token.is_semicolon() || token.is_any_word(&["TRANSACTION", "WORK"]) + } +} + +impl OraclePlSqlToken { + fn word(value: String) -> Self { + Self::Word(value) + } + + fn from_sqlparser_token(token: Token) -> Option { + match token { + Token::Word(word) if word.quote_style.is_none() => Some(Self::Word(word.value.to_ascii_uppercase())), + Token::SemiColon => Some(Self::Semicolon), + _ => None, + } + } + + fn is_word(&self, expected: &str) -> bool { + matches!(self, Self::Word(value) if value == expected) + } + + fn is_any_word(&self, expected: &[&str]) -> bool { + expected.iter().any(|word| self.is_word(word)) + } + + fn is_semicolon(&self) -> bool { + matches!(self, Self::Semicolon) + } +} + +fn oracle_plsql_tokens(sql: &str) -> Vec { + let dialect = OracleDialect {}; + if let Ok(tokens) = Tokenizer::new(&dialect, sql).tokenize() { + return tokens.into_iter().filter_map(OraclePlSqlToken::from_sqlparser_token).collect(); + } + + oracle_plsql_tokens_fallback(sql) +} + +fn oracle_plsql_tokens_fallback(sql: &str) -> Vec { let mut tokens = Vec::new(); let mut iter = sql.char_indices().peekable(); @@ -1660,7 +1713,7 @@ fn oracle_plsql_tokens(sql: &str) -> Vec { } if ch == ';' { - tokens.push(";".to_string()); + tokens.push(OraclePlSqlToken::Semicolon); continue; } @@ -1675,7 +1728,7 @@ fn oracle_plsql_tokens(sql: &str) -> Vec { break; } } - tokens.push(token); + tokens.push(OraclePlSqlToken::word(token)); } } @@ -2395,6 +2448,33 @@ SELECT 1;"; ); } + #[test] + fn oracle_like_split_ignores_end_tokens_inside_q_quoted_strings() { + let sql = "\ +BEGIN + v_text := q'[not really END;]'; + NULL; +END; +SELECT 1;"; + + assert_eq!( + split_sql_statements_for_database(sql, DatabaseType::Oracle), + vec!["BEGIN\n v_text := q'[not really END;]';\n NULL;\nEND;", "SELECT 1"] + ); + } + + #[test] + fn oracle_plsql_tokenizer_falls_back_when_sqlparser_rejects_partial_literals() { + let sql = "\ +BEGIN + NULL; + v_text := 'partial"; + let tokens = super::oracle_plsql_tokens(sql); + + assert!(tokens.iter().any(|token| token.is_word("BEGIN"))); + assert!(tokens.iter().any(super::OraclePlSqlToken::is_semicolon)); + } + #[test] fn oracle_like_split_does_not_affect_create_table() { let sql = "\