use vibesql_ast::{IndexColumn, OrderDirection};
use vibesql_catalog::{ColumnSchema, TableSchema};
use vibesql_parser::Parser;
use vibesql_storage::{Database, Row};
use vibesql_types::{DataType, SqlValue};
use crate::select::SelectExecutor;
fn create_stock_db() -> Database {
let mut db = Database::new();
db.catalog.set_case_sensitive_identifiers(false);
let stock_schema = TableSchema::new(
"stock".to_string(),
vec![
ColumnSchema::new("s_w_id".to_string(), DataType::Integer, false),
ColumnSchema::new("s_i_id".to_string(), DataType::Integer, false),
ColumnSchema::new("s_quantity".to_string(), DataType::Integer, false),
ColumnSchema::new(
"s_data".to_string(),
DataType::Varchar { max_length: Some(50) },
true,
),
],
);
db.create_table(stock_schema).unwrap();
for i in 1..=10 {
db.insert_row(
"stock",
Row::new(vec![
SqlValue::Integer(1), SqlValue::Integer(i), SqlValue::Integer(5 + (i % 20)), SqlValue::Varchar(arcstr::ArcStr::from(format!("Stock data {}", i))), ]),
)
.unwrap();
}
for i in 1..=5 {
db.insert_row(
"stock",
Row::new(vec![
SqlValue::Integer(2), SqlValue::Integer(i), SqlValue::Integer(10 + i), SqlValue::Varchar(arcstr::ArcStr::from(format!("Stock data W2-{}", i))),
]),
)
.unwrap();
}
db
}
#[test]
fn test_covering_index_single_column_select() {
let mut db = create_stock_db();
db.create_index(
"idx_stock_quantity".to_string(),
"stock".to_string(),
false,
vec![
IndexColumn::Column {
column_name: "s_w_id".to_string(),
prefix_length: None,
direction: OrderDirection::Asc,
},
IndexColumn::Column {
column_name: "s_quantity".to_string(),
prefix_length: None,
direction: OrderDirection::Asc,
},
IndexColumn::Column {
column_name: "s_i_id".to_string(),
prefix_length: None,
direction: OrderDirection::Asc,
},
],
)
.unwrap();
let executor = SelectExecutor::new(&db);
let query = "SELECT s_i_id FROM stock WHERE s_w_id = 1 AND s_quantity < 15";
let stmt = Parser::parse_sql(query).unwrap();
if let vibesql_ast::Statement::Select(select_stmt) = stmt {
let result = executor.execute(&select_stmt).unwrap();
assert!(!result.is_empty());
let item_ids: Vec<i64> = result
.iter()
.map(|r| match &r.values[0] {
SqlValue::Integer(v) => *v,
_ => panic!("Expected integer"),
})
.collect();
for id in &item_ids {
assert!(*id >= 1 && *id <= 9, "Unexpected item id: {}", id);
}
assert_eq!(item_ids.len(), 9);
} else {
panic!("Expected SELECT statement");
}
}
#[test]
fn test_covering_index_multiple_columns_select() {
let mut db = create_stock_db();
db.create_index(
"idx_stock_covering".to_string(),
"stock".to_string(),
false,
vec![
IndexColumn::Column {
column_name: "s_w_id".to_string(),
prefix_length: None,
direction: OrderDirection::Asc,
},
IndexColumn::Column {
column_name: "s_quantity".to_string(),
prefix_length: None,
direction: OrderDirection::Asc,
},
IndexColumn::Column {
column_name: "s_i_id".to_string(),
prefix_length: None,
direction: OrderDirection::Asc,
},
],
)
.unwrap();
let executor = SelectExecutor::new(&db);
let query = "SELECT s_i_id, s_quantity FROM stock WHERE s_w_id = 1 AND s_quantity < 10";
let stmt = Parser::parse_sql(query).unwrap();
if let vibesql_ast::Statement::Select(select_stmt) = stmt {
let result = executor.execute(&select_stmt).unwrap();
assert_eq!(result.len(), 4);
for row in &result {
let s_i_id = match &row.values[0] {
SqlValue::Integer(v) => *v,
_ => panic!("Expected integer for s_i_id"),
};
let s_quantity = match &row.values[1] {
SqlValue::Integer(v) => *v,
_ => panic!("Expected integer for s_quantity"),
};
assert!(s_quantity < 10);
assert!((1..=4).contains(&s_i_id));
}
} else {
panic!("Expected SELECT statement");
}
}
#[test]
fn test_non_covering_index_falls_back_to_table_fetch() {
let mut db = create_stock_db();
db.create_index(
"idx_stock_partial".to_string(),
"stock".to_string(),
false,
vec![
IndexColumn::Column {
column_name: "s_w_id".to_string(),
prefix_length: None,
direction: OrderDirection::Asc,
},
IndexColumn::Column {
column_name: "s_i_id".to_string(),
prefix_length: None,
direction: OrderDirection::Asc,
},
],
)
.unwrap();
let executor = SelectExecutor::new(&db);
let query = "SELECT s_data FROM stock WHERE s_w_id = 1";
let stmt = Parser::parse_sql(query).unwrap();
if let vibesql_ast::Statement::Select(select_stmt) = stmt {
let result = executor.execute(&select_stmt).unwrap();
assert_eq!(result.len(), 10);
for row in &result {
match &row.values[0] {
SqlValue::Varchar(s) => assert!(s.starts_with("Stock data")),
_ => panic!("Expected varchar for s_data"),
};
}
} else {
panic!("Expected SELECT statement");
}
}
#[test]
fn test_covering_index_prefix_equality_range() {
let mut db = create_stock_db();
db.create_index(
"idx_stock_wid_qty_iid".to_string(),
"stock".to_string(),
false,
vec![
IndexColumn::Column {
column_name: "s_w_id".to_string(),
prefix_length: None,
direction: OrderDirection::Asc,
},
IndexColumn::Column {
column_name: "s_quantity".to_string(),
prefix_length: None,
direction: OrderDirection::Asc,
},
IndexColumn::Column {
column_name: "s_i_id".to_string(),
prefix_length: None,
direction: OrderDirection::Asc,
},
],
)
.unwrap();
let executor = SelectExecutor::new(&db);
let query = "SELECT s_i_id FROM stock WHERE s_w_id = 2 AND s_quantity > 12";
let stmt = Parser::parse_sql(query).unwrap();
if let vibesql_ast::Statement::Select(select_stmt) = stmt {
let result = executor.execute(&select_stmt).unwrap();
assert_eq!(result.len(), 3);
} else {
panic!("Expected SELECT statement");
}
}
#[test]
fn test_covering_index_check_function() {
use crate::select::scan::index_scan::covering::check_covering_index;
let index_cols = vec!["s_w_id", "s_quantity", "s_i_id"];
let needed = vec!["s_i_id".to_string()];
let result = check_covering_index(&index_cols, &needed);
assert!(result.is_some());
let index_cols = vec!["s_w_id", "s_quantity"];
let needed = vec!["s_i_id".to_string(), "s_data".to_string()];
let result = check_covering_index(&index_cols, &needed);
assert!(result.is_none());
let index_cols = vec!["S_W_ID", "S_QUANTITY", "S_I_ID"];
let needed = vec!["s_i_id".to_string()];
let result = check_covering_index(&index_cols, &needed);
assert!(result.is_some());
}