use std::{
num::NonZeroUsize,
sync::{
atomic::{AtomicUsize, Ordering},
Arc, Mutex,
},
};
use lru::LruCache;
use vibesql_ast::Statement;
use vibesql_types::SqlValue;
use super::{extract_tables_from_statement, QuerySignature};
pub mod arena_prepared;
mod bind;
pub mod plan;
pub use plan::{
CachedPlan, ColumnProjection, PkDeletePlan, PkPointLookupPlan, ProjectionPlan,
ResolvedProjection, SimpleFastPathPlan,
};
#[derive(Debug, Clone)]
pub struct PreparedStatement {
sql: String,
statement: Statement,
signature: QuerySignature,
param_count: usize,
tables: std::collections::HashSet<String>,
cached_plan: CachedPlan,
}
impl PreparedStatement {
pub fn new(sql: String, statement: Statement) -> Self {
let signature = QuerySignature::from_ast(&statement);
let param_count = bind::count_placeholders(&statement);
let tables = extract_tables_from_statement(&statement);
let cached_plan = plan::analyze_for_plan(&statement);
Self { sql, statement, signature, param_count, tables, cached_plan }
}
pub fn sql(&self) -> &str {
&self.sql
}
pub fn statement(&self) -> &Statement {
&self.statement
}
pub fn signature(&self) -> &QuerySignature {
&self.signature
}
pub fn param_count(&self) -> usize {
self.param_count
}
pub fn tables(&self) -> &std::collections::HashSet<String> {
&self.tables
}
pub fn cached_plan(&self) -> &CachedPlan {
&self.cached_plan
}
pub fn bind(&self, params: &[SqlValue]) -> Result<Statement, PreparedStatementError> {
if params.len() != self.param_count {
return Err(PreparedStatementError::ParameterCountMismatch {
expected: self.param_count,
actual: params.len(),
});
}
if self.param_count == 0 {
return Ok(self.statement.clone());
}
Ok(bind::bind_parameters(&self.statement, params))
}
}
#[derive(Debug, Clone)]
pub enum PreparedStatementError {
ParameterCountMismatch { expected: usize, actual: usize },
ParseError(String),
}
impl std::fmt::Display for PreparedStatementError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
PreparedStatementError::ParameterCountMismatch { expected, actual } => {
write!(f, "Parameter count mismatch: expected {}, got {}", expected, actual)
}
PreparedStatementError::ParseError(msg) => write!(f, "Parse error: {}", msg),
}
}
}
impl std::error::Error for PreparedStatementError {}
#[derive(Debug, Clone)]
pub struct PreparedStatementCacheStats {
pub hits: usize,
pub misses: usize,
pub evictions: usize,
pub size: usize,
pub hit_rate: f64,
}
pub struct PreparedStatementCache {
cache: Mutex<LruCache<String, Arc<PreparedStatement>>>,
arena_cache: Mutex<LruCache<String, Arc<arena_prepared::ArenaPreparedStatement>>>,
max_size: usize,
hits: AtomicUsize,
misses: AtomicUsize,
evictions: AtomicUsize,
arena_hits: AtomicUsize,
arena_misses: AtomicUsize,
}
impl PreparedStatementCache {
pub fn new(max_size: usize) -> Self {
let cap = NonZeroUsize::new(max_size).unwrap_or(NonZeroUsize::new(1).unwrap());
Self {
cache: Mutex::new(LruCache::new(cap)),
arena_cache: Mutex::new(LruCache::new(cap)),
max_size,
hits: AtomicUsize::new(0),
misses: AtomicUsize::new(0),
evictions: AtomicUsize::new(0),
arena_hits: AtomicUsize::new(0),
arena_misses: AtomicUsize::new(0),
}
}
pub fn default_cache() -> Self {
Self::new(1000)
}
pub fn get(&self, sql: &str) -> Option<Arc<PreparedStatement>> {
let mut cache = self.cache.lock().unwrap();
if let Some(stmt) = cache.get(sql) {
self.hits.fetch_add(1, Ordering::Relaxed);
Some(Arc::clone(stmt))
} else {
self.misses.fetch_add(1, Ordering::Relaxed);
None
}
}
pub fn get_or_prepare(
&self,
sql: &str,
) -> Result<Arc<PreparedStatement>, PreparedStatementError> {
let mut cache = self.cache.lock().unwrap();
if let Some(stmt) = cache.get(sql) {
self.hits.fetch_add(1, Ordering::Relaxed);
return Ok(Arc::clone(stmt));
}
self.misses.fetch_add(1, Ordering::Relaxed);
let statement = parse_with_arena_fallback(sql)
.map_err(|e| PreparedStatementError::ParseError(e.to_string()))?;
let prepared = Arc::new(PreparedStatement::new(sql.to_string(), statement));
if cache.len() >= self.max_size {
self.evictions.fetch_add(1, Ordering::Relaxed);
}
cache.put(sql.to_string(), Arc::clone(&prepared));
Ok(prepared)
}
pub fn get_or_prepare_arena(
&self,
sql: &str,
) -> Result<Arc<arena_prepared::ArenaPreparedStatement>, arena_prepared::ArenaParseError> {
let mut cache = self.arena_cache.lock().unwrap();
if let Some(stmt) = cache.get(sql) {
self.arena_hits.fetch_add(1, Ordering::Relaxed);
return Ok(Arc::clone(stmt));
}
self.arena_misses.fetch_add(1, Ordering::Relaxed);
let prepared = Arc::new(arena_prepared::ArenaPreparedStatement::new(sql.to_string())?);
if cache.len() >= self.max_size {
self.evictions.fetch_add(1, Ordering::Relaxed);
}
cache.put(sql.to_string(), Arc::clone(&prepared));
Ok(prepared)
}
pub fn get_arena(&self, sql: &str) -> Option<Arc<arena_prepared::ArenaPreparedStatement>> {
let mut cache = self.arena_cache.lock().unwrap();
if let Some(stmt) = cache.get(sql) {
self.arena_hits.fetch_add(1, Ordering::Relaxed);
Some(Arc::clone(stmt))
} else {
None
}
}
pub fn clear(&self) {
self.cache.lock().unwrap().clear();
self.arena_cache.lock().unwrap().clear();
}
pub fn invalidate_table(&self, table: &str) {
{
let mut cache = self.cache.lock().unwrap();
let keys_to_remove: Vec<String> = cache
.iter()
.filter(|(_, stmt)| stmt.tables.iter().any(|t| t.eq_ignore_ascii_case(table)))
.map(|(k, _)| k.clone())
.collect();
for key in keys_to_remove {
cache.pop(&key);
}
}
{
let mut arena_cache = self.arena_cache.lock().unwrap();
let keys_to_remove: Vec<String> = arena_cache
.iter()
.filter(|(_, stmt)| stmt.tables().iter().any(|t| t.eq_ignore_ascii_case(table)))
.map(|(k, _)| k.clone())
.collect();
for key in keys_to_remove {
arena_cache.pop(&key);
}
}
}
pub fn stats(&self) -> PreparedStatementCacheStats {
let cache = self.cache.lock().unwrap();
let hits = self.hits.load(Ordering::Relaxed);
let misses = self.misses.load(Ordering::Relaxed);
let total = hits + misses;
let hit_rate = if total > 0 { hits as f64 / total as f64 } else { 0.0 };
PreparedStatementCacheStats {
hits,
misses,
evictions: self.evictions.load(Ordering::Relaxed),
size: cache.len(),
hit_rate,
}
}
pub fn max_size(&self) -> usize {
self.max_size
}
}
fn parse_with_arena_fallback(sql: &str) -> Result<Statement, vibesql_parser::ParseError> {
let trimmed = sql.trim_start();
let first_word = trimmed.split_whitespace().next().unwrap_or("");
if first_word.eq_ignore_ascii_case("SELECT") || first_word.eq_ignore_ascii_case("WITH") {
if let Ok(select_stmt) = vibesql_parser::arena_parser::parse_select_to_owned(sql) {
return Ok(Statement::Select(Box::new(select_stmt)));
}
}
if first_word.eq_ignore_ascii_case("INSERT") {
if let Ok(insert_stmt) = vibesql_parser::arena_parser::parse_insert_to_owned(sql) {
return Ok(Statement::Insert(insert_stmt));
}
}
if first_word.eq_ignore_ascii_case("REPLACE") {
if let Ok(insert_stmt) = vibesql_parser::arena_parser::parse_insert_to_owned(sql) {
return Ok(Statement::Insert(insert_stmt));
}
}
if first_word.eq_ignore_ascii_case("UPDATE") {
if let Ok(update_stmt) = vibesql_parser::arena_parser::parse_update_to_owned(sql) {
return Ok(Statement::Update(update_stmt));
}
}
if first_word.eq_ignore_ascii_case("DELETE") {
if let Ok(delete_stmt) = vibesql_parser::arena_parser::parse_delete_to_owned(sql) {
return Ok(Statement::Delete(delete_stmt));
}
}
vibesql_parser::Parser::parse_sql(sql)
}
#[cfg(test)]
mod tests {
use vibesql_ast::Expression;
use super::*;
#[test]
fn test_prepared_statement_no_params() {
let sql = "SELECT * FROM users";
let statement = vibesql_parser::Parser::parse_sql(sql).unwrap();
let prepared = PreparedStatement::new(sql.to_string(), statement);
assert_eq!(prepared.param_count(), 0);
assert!(prepared.bind(&[]).is_ok());
}
#[test]
fn test_prepared_statement_with_placeholder() {
let sql = "SELECT * FROM users WHERE id = ?";
let statement = vibesql_parser::Parser::parse_sql(sql).unwrap();
let prepared = PreparedStatement::new(sql.to_string(), statement);
assert_eq!(prepared.param_count(), 1);
let bound = prepared.bind(&[SqlValue::Integer(42)]).unwrap();
assert!(matches!(bound, Statement::Select(_)));
if let Statement::Select(select) = bound {
if let Some(Expression::BinaryOp { right, .. }) = &select.where_clause {
assert_eq!(**right, Expression::Literal(SqlValue::Integer(42)));
} else {
panic!("Expected BinaryOp in WHERE clause");
}
}
}
#[test]
fn test_prepared_statement_multiple_placeholders() {
let sql = "SELECT * FROM users WHERE id = ? AND name = ?";
let statement = vibesql_parser::Parser::parse_sql(sql).unwrap();
let prepared = PreparedStatement::new(sql.to_string(), statement);
assert_eq!(prepared.param_count(), 2);
let params = vec![SqlValue::Integer(42), SqlValue::Varchar(arcstr::ArcStr::from("John"))];
let bound = prepared.bind(¶ms).unwrap();
assert!(matches!(bound, Statement::Select(_)));
}
#[test]
fn test_prepared_statement_bind_param_mismatch() {
let sql = "SELECT * FROM users WHERE id = ?";
let statement = vibesql_parser::Parser::parse_sql(sql).unwrap();
let prepared = PreparedStatement::new(sql.to_string(), statement);
let result = prepared.bind(&[]);
assert!(matches!(
result,
Err(PreparedStatementError::ParameterCountMismatch { expected: 1, actual: 0 })
));
let result = prepared.bind(&[SqlValue::Integer(1), SqlValue::Integer(2)]);
assert!(matches!(
result,
Err(PreparedStatementError::ParameterCountMismatch { expected: 1, actual: 2 })
));
}
#[test]
fn test_prepared_statement_reuse() {
let sql = "SELECT * FROM users WHERE id = ?";
let statement = vibesql_parser::Parser::parse_sql(sql).unwrap();
let prepared = PreparedStatement::new(sql.to_string(), statement);
let bound1 = prepared.bind(&[SqlValue::Integer(1)]).unwrap();
let bound2 = prepared.bind(&[SqlValue::Integer(2)]).unwrap();
let bound3 = prepared.bind(&[SqlValue::Integer(3)]).unwrap();
for (bound, expected_id) in [(bound1, 1), (bound2, 2), (bound3, 3)] {
if let Statement::Select(select) = bound {
if let Some(Expression::BinaryOp { right, .. }) = &select.where_clause {
assert_eq!(**right, Expression::Literal(SqlValue::Integer(expected_id)));
}
}
}
}
#[test]
fn test_cache_get_or_prepare() {
let cache = PreparedStatementCache::new(10);
let sql = "SELECT * FROM users WHERE id = 1";
let stmt1 = cache.get_or_prepare(sql).unwrap();
let stats = cache.stats();
assert_eq!(stats.misses, 1);
assert_eq!(stats.hits, 0);
let stmt2 = cache.get_or_prepare(sql).unwrap();
let stats = cache.stats();
assert_eq!(stats.misses, 1);
assert_eq!(stats.hits, 1);
assert!(Arc::ptr_eq(&stmt1, &stmt2));
}
#[test]
fn test_cache_placeholder_reuse() {
let cache = PreparedStatementCache::new(10);
let sql = "SELECT * FROM users WHERE id = ?";
let stmt1 = cache.get_or_prepare(sql).unwrap();
assert_eq!(cache.stats().misses, 1);
assert_eq!(cache.stats().hits, 0);
let stmt2 = cache.get_or_prepare(sql).unwrap();
assert_eq!(cache.stats().misses, 1);
assert_eq!(cache.stats().hits, 1);
assert!(Arc::ptr_eq(&stmt1, &stmt2));
let bound1 = stmt1.bind(&[SqlValue::Integer(1)]).unwrap();
let bound2 = stmt2.bind(&[SqlValue::Integer(999)]).unwrap();
if let (Statement::Select(s1), Statement::Select(s2)) = (&bound1, &bound2) {
if let (
Some(Expression::BinaryOp { right: r1, .. }),
Some(Expression::BinaryOp { right: r2, .. }),
) = (&s1.where_clause, &s2.where_clause)
{
assert_eq!(**r1, Expression::Literal(SqlValue::Integer(1)));
assert_eq!(**r2, Expression::Literal(SqlValue::Integer(999)));
}
}
}
#[test]
fn test_cache_lru_eviction() {
let cache = PreparedStatementCache::new(2);
cache.get_or_prepare("SELECT * FROM users").unwrap();
cache.get_or_prepare("SELECT * FROM orders").unwrap();
assert_eq!(cache.stats().size, 2);
assert_eq!(cache.stats().evictions, 0);
cache.get_or_prepare("SELECT * FROM products").unwrap();
assert_eq!(cache.stats().size, 2);
assert_eq!(cache.stats().evictions, 1);
assert!(cache.get("SELECT * FROM users").is_none());
assert!(cache.get("SELECT * FROM orders").is_some());
assert!(cache.get("SELECT * FROM products").is_some());
}
#[test]
fn test_cache_table_invalidation() {
let cache = PreparedStatementCache::new(10);
cache.get_or_prepare("SELECT * FROM users WHERE id = ?").unwrap();
cache.get_or_prepare("SELECT * FROM orders WHERE id = ?").unwrap();
assert_eq!(cache.stats().size, 2);
cache.invalidate_table("users");
assert_eq!(cache.stats().size, 1);
assert!(cache.get("SELECT * FROM orders WHERE id = ?").is_some());
}
#[test]
fn test_arena_parse_insert() {
let sql = "INSERT INTO users (name, email) VALUES ('Alice', 'alice@example.com')";
let result = parse_with_arena_fallback(sql);
assert!(result.is_ok());
assert!(matches!(result.unwrap(), Statement::Insert(_)));
}
#[test]
fn test_arena_parse_insert_with_placeholders() {
let cache = PreparedStatementCache::new(10);
let sql = "INSERT INTO users (name, email) VALUES (?, ?)";
let stmt = cache.get_or_prepare(sql).unwrap();
assert_eq!(stmt.param_count(), 2);
let bound = stmt
.bind(&[
SqlValue::Varchar(arcstr::ArcStr::from("Bob")),
SqlValue::Varchar(arcstr::ArcStr::from("bob@example.com")),
])
.unwrap();
assert!(matches!(bound, Statement::Insert(_)));
}
#[test]
fn test_arena_parse_update() {
let sql = "UPDATE users SET name = 'Bob' WHERE id = 1";
let result = parse_with_arena_fallback(sql);
assert!(result.is_ok());
assert!(matches!(result.unwrap(), Statement::Update(_)));
}
#[test]
fn test_arena_parse_update_with_placeholders() {
let cache = PreparedStatementCache::new(10);
let sql = "UPDATE users SET name = ? WHERE id = ?";
let stmt = cache.get_or_prepare(sql).unwrap();
assert_eq!(stmt.param_count(), 2);
let bound = stmt
.bind(&[SqlValue::Varchar(arcstr::ArcStr::from("Charlie")), SqlValue::Integer(42)])
.unwrap();
assert!(matches!(bound, Statement::Update(_)));
}
#[test]
fn test_arena_parse_delete() {
let sql = "DELETE FROM users WHERE id = 1";
let result = parse_with_arena_fallback(sql);
assert!(result.is_ok());
assert!(matches!(result.unwrap(), Statement::Delete(_)));
}
#[test]
fn test_arena_parse_delete_with_placeholders() {
let cache = PreparedStatementCache::new(10);
let sql = "DELETE FROM users WHERE id = ?";
let stmt = cache.get_or_prepare(sql).unwrap();
assert_eq!(stmt.param_count(), 1);
let bound = stmt.bind(&[SqlValue::Integer(99)]).unwrap();
assert!(matches!(bound, Statement::Delete(_)));
}
#[test]
fn test_arena_parse_sysbench_insert() {
let cache = PreparedStatementCache::new(10);
let sql3 = "INSERT INTO test (a, b, c) VALUES (?, ?, ?)";
let stmt3 = cache.get_or_prepare(sql3);
assert!(stmt3.is_ok(), "3-column INSERT failed: {:?}", stmt3.err());
let sql4_gen = "INSERT INTO test (a, b, c, d) VALUES (?, ?, ?, ?)";
let stmt4_gen = cache.get_or_prepare(sql4_gen);
assert!(stmt4_gen.is_ok(), "4-column INSERT (generic) failed: {:?}", stmt4_gen.err());
let sql4 = "INSERT INTO sbtest1 (id, k, c, padding) VALUES (?, ?, ?, ?)";
let stmt4 = cache.get_or_prepare(sql4);
assert!(stmt4.is_ok(), "4-column INSERT (sysbench) failed: {:?}", stmt4.err());
assert_eq!(stmt4.unwrap().param_count(), 4);
}
}