use hawdb_core::{HawDBError, Result};
use hawdb_sql_syntax::{self as syntax, LexerConfig, SyntaxErrorCode, Token, TokenKind as T};
use std::sync::OnceLock;
fn keywords() -> &'static [&'static str] {
static KEYWORDS: OnceLock<Vec<&'static str>> = OnceLock::new();
KEYWORDS.get_or_init(|| include_str!("admission_keywords_v1.txt").lines().collect())
}
fn keyword(word: &str) -> bool {
if !word.is_ascii() {
return keywords()
.binary_search(&word.to_uppercase().as_str())
.is_ok();
}
if word.len() > 40 {
return false;
}
let mut uppercase = [0_u8; 40];
for (dst, src) in uppercase.iter_mut().zip(word.bytes()) {
*dst = src.to_ascii_uppercase();
}
let word = &uppercase[..word.len()];
keywords()
.binary_search_by(|candidate| candidate.as_bytes().cmp(word))
.is_ok()
}
#[doc(hidden)]
pub fn admit_postgres_sql_syntax(
input: &str,
max_input_bytes: usize,
max_tokens: usize,
max_recursive_tokens: usize,
) -> Result<()> {
let refusal = || HawDBError::Execution("recursive SQL syntax exceeding work admission".into());
if input.len() > max_input_bytes {
return Err(refusal());
}
let tokens = syntax::tokenize_with_config(
input,
LexerConfig {
max_input_bytes,
max_tokens: input.len(),
..LexerConfig::default()
},
)
.map_err(|error| {
if error.code == SyntaxErrorCode::InputTooLarge
|| error.code == SyntaxErrorCode::TokenLimitExceeded
{
refusal()
} else {
HawDBError::Parse(error.to_string())
}
})?;
let (mut count, mut recursive, mut at) = (0_usize, 0_usize, 0_usize);
while let Some(token) = tokens.get(at) {
if token.kind == T::End {
break;
}
count = count.saturating_add(1);
let grows = match token.kind {
T::Word => keyword(token.text(input)),
T::QuotedIdentifier
| T::Number
| T::Parameter(_)
| T::Comma
| T::Dot
| T::Semicolon
| T::RightParen
| T::RightBracket
| T::RightBrace
| T::String
| T::Question => false,
_ => true,
};
recursive = recursive.saturating_add(usize::from(grows));
if tokens
.get(at + 1)
.is_some_and(|next| paired_operator(*token, *next))
{
at += 1;
}
at += 1;
if count > max_tokens || recursive > max_recursive_tokens {
return Err(refusal());
}
}
Ok(())
}
fn paired_operator(left: Token, right: Token) -> bool {
left.span.end == right.span.start
&& matches!(
(left.kind, right.kind),
(T::Equal, T::Greater) | (T::Colon, T::Equal)
)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn keyword_work_inventory_is_sorted_unique_and_case_insensitive() {
let words: Vec<_> = include_str!("admission_keywords_v1.txt").lines().collect();
assert_eq!(words.len(), 1102);
assert!(words.windows(2).all(|pair| pair[0] < pair[1]));
assert_eq!(words.iter().map(|word| word.len()).max(), Some(40));
for word in words {
assert!(keyword(word));
assert!(keyword(&word.to_ascii_lowercase()));
}
assert!(keyword("ıN"));
assert!(keyword("ſELECT"));
for word in [
"row_id",
"列",
"select_value",
"aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa",
] {
assert!(!keyword(word));
}
}
#[test]
fn declared_keyword_and_operator_work_boundaries_are_preserved() {
let input = "SELECT value FROM items WHERE id = $1";
assert!(admit_postgres_sql_syntax(input, input.len(), 8, 7).is_ok());
assert!(matches!(
admit_postgres_sql_syntax(input, input.len(), 7, 7),
Err(HawDBError::Execution(_))
));
assert!(matches!(
admit_postgres_sql_syntax(input, input.len(), 8, 6),
Err(HawDBError::Execution(_))
));
for input in ["top_k => 1", "top_k := 1"] {
assert!(admit_postgres_sql_syntax(input, input.len(), 3, 1).is_ok());
assert!(admit_postgres_sql_syntax(input, input.len(), 2, 1).is_err());
}
assert!(admit_postgres_sql_syntax("top_k = > 1", 32, 3, 1).is_err());
assert!(admit_postgres_sql_syntax("\"SELECT\"", 32, 1, 0).is_ok());
}
}