main
use sqlparser::dialect::PostgreSqlDialect;
use sqlparser::keywords::Keyword;
use sqlparser::tokenizer::{Token, Tokenizer};
pub fn requires_primary(sql: &str) -> bool {
let Ok(tokens) = Tokenizer::new(&PostgreSqlDialect {}, sql).tokenize() else {
return true;
};
let mut tokens = tokens
.into_iter()
.filter(|token| !matches!(token, Token::Whitespace(_)))
.peekable();
let mut previous = [Keyword::NoKeyword; 2];
let mut statement_start = true;
while let Some(token) = tokens.next() {
match token {
Token::SemiColon => {
previous = [Keyword::NoKeyword; 2];
statement_start = true;
continue;
}
Token::Word(word) => {
if word.quote_style.is_none() {
if matches!(
word.keyword,
Keyword::ABORT
| Keyword::ALTER
| Keyword::ANALYZE
| Keyword::BEGIN
| Keyword::CALL
| Keyword::CLUSTER
| Keyword::COMMENT
| Keyword::COMMIT
| Keyword::COPY
| Keyword::CREATE
| Keyword::DEALLOCATE
| Keyword::DELETE
| Keyword::DISCARD
| Keyword::DO
| Keyword::DROP
| Keyword::EXECUTE
| Keyword::GRANT
| Keyword::INSERT
| Keyword::INTO
| Keyword::LISTEN
| Keyword::LOAD
| Keyword::LOCK
| Keyword::MERGE
| Keyword::NOTIFY
| Keyword::PREPARE
| Keyword::REFRESH
| Keyword::REINDEX
| Keyword::RELEASE
| Keyword::RESET
| Keyword::REVOKE
| Keyword::ROLLBACK
| Keyword::SAVEPOINT
| Keyword::SET
| Keyword::START
| Keyword::TRUNCATE
| Keyword::UNLISTEN
| Keyword::UPDATE
| Keyword::VACUUM
) {
return true;
}
if statement_start
&& (word.keyword == Keyword::END
|| [
"checkpoint",
"reassign",
"declare",
"close",
"fetch",
"move",
]
.iter()
.any(|command| word.value.eq_ignore_ascii_case(command)))
{
return true;
}
if word.keyword == Keyword::SHARE
&& (previous[1] == Keyword::FOR || previous == [Keyword::FOR, Keyword::KEY])
{
return true;
}
}
if matches!(tokens.peek(), Some(Token::LParen)) && primary_function(&word.value) {
return true;
}
previous = [previous[1], word.keyword];
}
_ => previous = [Keyword::NoKeyword; 2],
}
statement_start = false;
}
false
}
fn primary_function(name: &str) -> bool {
let name = name.to_ascii_lowercase();
matches!(
name.as_str(),
"nextval"
| "setval"
| "set_config"
| "pg_notify"
| "pg_cancel_backend"
| "pg_terminate_backend"
) || name.starts_with("pg_advisory_")
|| name.starts_with("pg_try_advisory_")
}