#[cfg(test)]
mod memory_tracking_tests {
use vibesql_storage::Database;
use crate::{
errors::ExecutorError, limits::MAX_MEMORY_BYTES, select::executor::builder::SelectExecutor,
};
#[test]
fn test_memory_limit_exceeded() {
let db = Database::new();
let executor = SelectExecutor::new(&db);
let result = executor.track_memory_allocation(15 * 1024 * 1024 * 1024);
assert!(
matches!(result, Err(ExecutorError::MemoryLimitExceeded { .. })),
"Expected MemoryLimitExceeded error"
);
if let Err(ExecutorError::MemoryLimitExceeded { used_bytes, max_bytes }) = result {
assert_eq!(used_bytes, 15 * 1024 * 1024 * 1024);
assert_eq!(max_bytes, MAX_MEMORY_BYTES);
} else {
panic!("Expected MemoryLimitExceeded variant");
}
}
#[test]
fn test_memory_allocation_below_limit() {
let db = Database::new();
let executor = SelectExecutor::new(&db);
let result = executor.track_memory_allocation(2 * 1024 * 1024 * 1024);
assert!(result.is_ok(), "Should allow allocation below limit");
}
#[test]
fn test_memory_deallocation() {
let db = Database::new();
let executor = SelectExecutor::new(&db);
executor.track_memory_allocation(3 * 1024 * 1024 * 1024).unwrap();
executor.track_memory_deallocation(1024 * 1024 * 1024);
let result = executor.track_memory_allocation(3 * 1024 * 1024 * 1024);
assert!(result.is_ok(), "Should allow allocation when within limit");
}
#[test]
fn test_memory_deallocation_underflow() {
let db = Database::new();
let executor = SelectExecutor::new(&db);
executor.track_memory_deallocation(100 * 1024 * 1024 * 1024);
executor
.track_memory_allocation(1)
.expect("Should work with minimal allocation after underflow");
}
}
#[cfg(test)]
mod integration_tests {
use vibesql_storage::Database;
use crate::{errors::ExecutorError, select::executor::builder::SelectExecutor};
#[test]
fn test_normal_query_within_memory_limit() {
let mut db = Database::new();
let schema = vibesql_catalog::TableSchema::new(
"small_table".to_string(),
vec![
vibesql_catalog::ColumnSchema::new(
"id".to_string(),
vibesql_types::DataType::Integer,
false,
),
vibesql_catalog::ColumnSchema::new(
"name".to_string(),
vibesql_types::DataType::Varchar { max_length: Some(100) },
false,
),
],
);
db.create_table(schema).unwrap();
db.insert_row(
"small_table",
vibesql_storage::Row::new(vec![
vibesql_types::SqlValue::Integer(1),
vibesql_types::SqlValue::Varchar(arcstr::ArcStr::from("Alice")),
]),
)
.unwrap();
db.insert_row(
"small_table",
vibesql_storage::Row::new(vec![
vibesql_types::SqlValue::Integer(2),
vibesql_types::SqlValue::Varchar(arcstr::ArcStr::from("Bob")),
]),
)
.unwrap();
let executor = SelectExecutor::new(&db);
let stmt = vibesql_ast::SelectStmt {
into_table: None,
into_variables: None,
with_clause: None,
set_operation: None,
values: None,
distinct: false,
select_list: vec![vibesql_ast::SelectItem::Wildcard { alias: None }],
from: Some(vibesql_ast::FromClause::Table {
index_hint: None,
name: "small_table".to_string(),
alias: None,
column_aliases: None,
quoted: false,
}),
where_clause: None,
group_by: None,
having: None,
window_definitions: None,
order_by: None,
limit: None,
offset: None,
};
let result = executor.execute(&stmt);
assert!(result.is_ok(), "Normal query should succeed");
let rows = result.unwrap();
assert_eq!(rows.len(), 2);
}
#[test]
fn test_massive_join_caught_early() {
let mut db = Database::new();
let schema1 = vibesql_catalog::TableSchema::new(
"t1".to_string(),
vec![vibesql_catalog::ColumnSchema::new(
"id".to_string(),
vibesql_types::DataType::Integer,
false,
)],
);
db.create_table(schema1).unwrap();
let schema2 = vibesql_catalog::TableSchema::new(
"t2".to_string(),
vec![vibesql_catalog::ColumnSchema::new(
"id".to_string(),
vibesql_types::DataType::Integer,
false,
)],
);
db.create_table(schema2).unwrap();
for i in 0..10_100 {
db.insert_row(
"t1",
vibesql_storage::Row::new(vec![vibesql_types::SqlValue::Integer(i)]),
)
.unwrap();
db.insert_row(
"t2",
vibesql_storage::Row::new(vec![vibesql_types::SqlValue::Integer(i)]),
)
.unwrap();
}
let executor = SelectExecutor::new(&db);
let cross_join_stmt = vibesql_ast::SelectStmt {
into_table: None,
into_variables: None,
with_clause: None,
set_operation: None,
values: None,
distinct: false,
select_list: vec![vibesql_ast::SelectItem::Wildcard { alias: None }],
from: Some(vibesql_ast::FromClause::Join {
left: Box::new(vibesql_ast::FromClause::Table {
index_hint: None,
name: "t1".to_string(),
alias: None,
column_aliases: None,
quoted: false,
}),
right: Box::new(vibesql_ast::FromClause::Table {
index_hint: None,
name: "t2".to_string(),
alias: None,
column_aliases: None,
quoted: false,
}),
join_type: vibesql_ast::JoinType::Cross,
condition: None,
using_columns: None,
natural: false,
alias: None,
}),
where_clause: None,
group_by: None,
having: None,
window_definitions: None,
order_by: None,
limit: None,
offset: None,
};
let result = executor.execute(&cross_join_stmt);
assert!(
matches!(result, Err(ExecutorError::MemoryLimitExceeded { .. })),
"Massive CROSS JOIN should be caught before execution, got: {:?}",
result
);
let inner_join_stmt = vibesql_ast::SelectStmt {
into_table: None,
into_variables: None,
with_clause: None,
set_operation: None,
values: None,
distinct: false,
select_list: vec![vibesql_ast::SelectItem::Wildcard { alias: None }],
from: Some(vibesql_ast::FromClause::Join {
left: Box::new(vibesql_ast::FromClause::Table {
index_hint: None,
name: "t1".to_string(),
alias: None,
column_aliases: None,
quoted: false,
}),
right: Box::new(vibesql_ast::FromClause::Table {
index_hint: None,
name: "t2".to_string(),
alias: None,
column_aliases: None,
quoted: false,
}),
join_type: vibesql_ast::JoinType::Inner,
condition: Some(vibesql_ast::Expression::Literal(
vibesql_types::SqlValue::Boolean(true),
)),
using_columns: None,
natural: false,
alias: None,
}),
where_clause: None,
group_by: None,
having: None,
window_definitions: None,
order_by: None,
limit: None,
offset: None,
};
let result = executor.execute(&inner_join_stmt);
assert!(
matches!(result, Err(ExecutorError::MemoryLimitExceeded { .. })),
"Massive INNER JOIN should be caught before execution, got: {:?}",
result
);
}
}