use std::collections::BTreeMap;
use anyhow::{anyhow, bail, Result};
use crate::{
ast::{BinOp, Column, Expr, InsertStatement, Statement},
catalog::Index,
eval::Context,
executor::execute_plan_next,
kv::RangeIterKV,
plan::{optimize_delete, optimize_select, optimize_update},
value::Type,
Db, Value,
};
pub fn run_program(
db: &mut Db,
program: &[Statement],
bindings: Vec<Value>,
) -> Result<Option<(Vec<String>, Vec<Vec<Value>>)>> {
let mut last_expr = None;
for stmnt in program {
match stmnt {
Statement::CreateTable(body) => {
let mut body = body.clone();
if db.table_to_kv_id(db.this_tx_id, &body.name)?.is_some() {
bail!("table {} already exists", body.name)
}
if body.readonly {
bail!("READ ONLY clause of CREATE TABLE is for internal use only")
}
if !body.referenced_by.is_empty() {
bail!("REFERENCED BY clause of CREATE TABLE is for internal use only")
}
for fk in &body.foreign_keys {
if db.table_to_kv_id(db.this_tx_id, &fk.rhs_table)?.is_none() {
bail!(
"table {} in foreign key constraint doesn't exist",
fk.rhs_table
)
}
let rhs_schema = db.get_table_schema(db.this_tx_id, &fk.rhs_table)?;
if rhs_schema
.indexes
.iter()
.find(|i| i.exprs == fk.rhs_exprs && i.unique)
.is_none()
&& rhs_schema.primary_key != fk.rhs_exprs
{
bail!(
"table {} does not have index required for foreign key constraint",
fk.rhs_table
)
}
if body
.indexes
.iter()
.find(|i| i.exprs == fk.lhs_exprs)
.is_none()
&& body.primary_key != fk.lhs_exprs
{
body.indexes.push(Index {
exprs: fk.lhs_exprs.clone(),
unique: false,
})
}
}
db.create_table(&body.name, body.clone())?;
last_expr = None;
}
Statement::DropTable(name) => {
let schema = db.get_table_schema(db.this_tx_id, &name)?;
if schema.readonly {
bail!("table {name} is read-only and cannot be droped")
}
if !schema.referenced_by.is_empty() {
bail!("table is referenced by another table in a foreign key constraint")
}
db.delete_table(&name)?;
last_expr = None;
}
Statement::Insert(body) => {
do_insert(db, body, bindings.clone())?;
last_expr = None;
}
Statement::Delete(body) => {
let schema = db.get_table_schema(db.this_tx_id, &body.table)?;
if schema.readonly {
bail!("table {} is read-only", body.table)
}
let mut plan = optimize_delete(db, body)?;
while let Some(_) =
execute_plan_next(&mut plan, db, Context::new(bindings.clone()))?
{
}
last_expr = None;
}
Statement::Update(body) => {
let schema = db.get_table_schema(db.this_tx_id, &body.table)?;
if schema.readonly {
bail!("table {} is read-only", body.table)
}
let mut plan = optimize_update(db, body)?;
while let Some(_) =
execute_plan_next(&mut plan, db, Context::new(bindings.clone()))?
{
}
last_expr = None;
}
Statement::Select(body) => {
fn edges(
map: &mut BTreeMap<String, (String, Vec<Expr>, Option<Vec<Expr>>)>,
e: &Expr,
) -> Result<(), String> {
match &e {
Expr::Binding(_) => Ok(()),
Expr::Literal(_) => Ok(()),
Expr::Unary(_op, expr) => edges(map, expr),
Expr::Bin(lhs, _, rhs) => {
edges(map, &lhs)?;
edges(map, &rhs)?;
Ok(())
}
Expr::Column(_) => Ok(()),
Expr::Edge(lhs, mapping, rhs, _e) => {
for e in lhs {
edges(map, e)?;
}
for e in rhs.iter().flatten() {
edges(map, e)?;
}
if let Some((g_table, g_lhs, g_rhs)) = map.get(&mapping.0) {
if g_lhs != lhs || g_rhs != rhs || &mapping.1 != g_table {
Err(mapping.0.to_string())
} else {
Ok(())
}
} else {
map.insert(
mapping.0.to_string(),
(mapping.1.to_string(), lhs.clone(), rhs.clone()),
);
Ok(())
}
}
}
}
let mut body = body.clone();
let mut joins = BTreeMap::new();
for e in &body.expr {
edges(&mut joins, e).map_err(|e| {
anyhow!("duplicate alias in edge expressions for table {e}")
})?;
}
for join in joins {
if body
.table_mappings
.iter()
.find(|(a, _t)| a == &join.0)
.is_some()
{
bail!("edge expressions and explicit joins conflict on {}", join.0)
}
let lhs_exprs = join.1 .1;
let rhs_exprs = if let Some(e) = join.1 .2 {
e
} else {
let schema = db.get_table_schema(db.this_tx_id, &join.0)?;
schema.primary_key
};
if lhs_exprs.len() != rhs_exprs.len() {
bail!("left hand side and right hand side expression lists of the edge expression do not match")
}
body.table_mappings
.push((join.0.clone(), join.1 .0.clone()));
body.cond = Some(Expr::Bin(
Box::new(body.cond.take().unwrap_or(Expr::Literal(Value::Bool(true)))),
BinOp::And,
Box::new(
lhs_exprs
.into_iter()
.zip(rhs_exprs)
.map(|(l, r)| Expr::Bin(Box::new(l), BinOp::Eq, Box::new(r)))
.reduce(|l, r| Expr::Bin(Box::new(l), BinOp::And, Box::new(r)))
.unwrap(),
),
));
}
let mut plan = optimize_select(db, &body)?;
if body.explain {
todo!()
} else {
let mut results = Vec::new();
while let Some(ctx) =
execute_plan_next(&mut plan, db, Context::new(bindings.clone()))?
{
results.push(ctx.as_slice().to_vec());
}
last_expr = Some((
match plan {
crate::executor::PlanNode::Select(select) => select.headers(),
_ => todo!(),
},
results,
))
}
}
}
}
Ok(last_expr)
}
fn do_insert(db: &mut Db, body: &InsertStatement, bindings: Vec<Value>) -> Result<()> {
let schema = db.get_table_schema(db.this_tx_id, &body.table)?;
if schema.readonly {
bail!("table {} is read-only", body.table)
}
let mut ctx = Context::new(bindings);
for (short_col_name, expr) in body.values.iter() {
let Some(col_schema) = schema.columns.iter().find(|c| c.name == *short_col_name) else {
bail!(
"table {} does not have a column named {}",
schema.name,
short_col_name
)
};
let full = Column(schema.name.clone(), col_schema.name.clone());
if ctx.get(&full).is_some() {
bail!("column {} is inputed multiple times", short_col_name)
}
let v = ctx.eval(expr)?;
ctx.set(full, v)
}
inner_insert(schema, &mut ctx, db)
}
pub fn inner_insert(
schema: crate::catalog::TableSchema,
ctx: &mut Context,
db: &mut Db,
) -> Result<()> {
for col_schema in &schema.columns {
let full = Column(schema.name.clone(), col_schema.name.clone());
if ctx.get(&full).is_some() {
continue;
}
if !col_schema.nullable {
bail!("column {} must not be null", col_schema.name)
}
ctx.set(full, Value::Null)
}
for check in &schema.checks {
match ctx.eval(check)? {
Value::Bool(true) => {}
Value::Bool(false) => bail!("check failed: {}", check),
_ => bail!("check did not return a boolean value: {}", check),
}
}
let canonical_row = schema
.columns
.iter()
.map(|n| {
ctx.get(&Column(schema.name.clone(), n.name.to_string()))
.unwrap()
})
.collect::<Vec<_>>();
for (v, col_schema) in canonical_row.iter().zip(schema.columns.iter()) {
match (v, col_schema.ty) {
(Value::Null, _) => {
if !col_schema.nullable {
bail!("column {} must not be null", col_schema.name)
}
}
(_, Type::Any)
| (Value::Bool(_), Type::Bool)
| (Value::Int(_), Type::Int)
| (Value::String(_), Type::String) => {}
(_v, _t) => {
bail!(
"column {} must be of type {}",
col_schema.name,
col_schema.ty
)
}
}
}
let primary_key_row = schema
.primary_key
.iter()
.map(|e| ctx.eval(e))
.collect::<Result<Vec<Value>, _>>()?;
if RangeIterKV::new_simple(schema.name.clone(), db.this_tx_id, primary_key_row.clone())
.next(db)?
.is_some()
{
bail!("inputed row has duplicate primary key")
}
let secondary_key_rows = schema
.indexes
.iter()
.map(|index| -> Result<Vec<Value>> {
Ok(index
.exprs
.iter()
.map(|e| ctx.eval(e))
.collect::<Result<Vec<_>, _>>()?
.into_iter()
.chain(primary_key_row.iter().cloned())
.collect::<Vec<_>>())
})
.collect::<Result<Vec<_>>>()?;
for (sk, unique_index) in secondary_key_rows
.iter()
.zip(schema.indexes.iter())
.filter(|(_, index)| index.unique)
{
if RangeIterKV::new_simple(unique_index.name(&schema.name), db.this_tx_id, sk.clone())
.next(db)?
.is_some()
{
bail!(
"unique constraint violated: {}",
unique_index
.exprs
.iter()
.map(|e| e.to_string())
.collect::<Vec<_>>()
.join(",")
)
}
}
for foreign_key_schema in &schema.foreign_keys {
let rhs_schema = db.get_table_schema(db.this_tx_id, &foreign_key_schema.rhs_table)?;
let sk = foreign_key_schema
.lhs_exprs
.iter()
.map(|e| ctx.eval(e))
.collect::<Result<Vec<_>, _>>()?;
if RangeIterKV::new_simple(
if rhs_schema.primary_key == foreign_key_schema.rhs_exprs {
rhs_schema.name
} else {
Index {
exprs: foreign_key_schema.rhs_exprs.clone(),
unique: false,
}
.name(&foreign_key_schema.rhs_table)
},
db.this_tx_id,
sk.clone(),
)
.next(db)?
.is_none()
{
bail!("foreign key constraint violated: {}", foreign_key_schema)
}
}
db.insert_key(&schema.name, primary_key_row.clone(), canonical_row.clone())?;
for (sk, index) in secondary_key_rows.into_iter().zip(schema.indexes.iter()) {
db.insert_key(&index.name(&schema.name), sk.clone(), sk)?;
}
Ok(())
}