fix(sql): tokenize Oracle PL/SQL blocks with sqlparser
This commit is contained in:
t8y2
parent
8c097a3fd1
commit
b53aec5e01
|
|
@ -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<bool> = 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<String> {
|
||||
struct OraclePlSqlBlock {
|
||||
tokens: Vec<OraclePlSqlToken>,
|
||||
}
|
||||
|
||||
#[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<bool> = 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<Self> {
|
||||
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<OraclePlSqlToken> {
|
||||
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<OraclePlSqlToken> {
|
||||
let mut tokens = Vec::new();
|
||||
let mut iter = sql.char_indices().peekable();
|
||||
|
||||
|
|
@ -1660,7 +1713,7 @@ fn oracle_plsql_tokens(sql: &str) -> Vec<String> {
|
|||
}
|
||||
|
||||
if ch == ';' {
|
||||
tokens.push(";".to_string());
|
||||
tokens.push(OraclePlSqlToken::Semicolon);
|
||||
continue;
|
||||
}
|
||||
|
||||
|
|
@ -1675,7 +1728,7 @@ fn oracle_plsql_tokens(sql: &str) -> Vec<String> {
|
|||
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 = "\
|
||||
|
|
|
|||
Loading…
Reference in New Issue