AkurAI-Framework / crates/sqlexec/src/exec.rs

Latest change a26d2091b5fd875b4928c72190c7dbf9f37b9252 - feat(sql): add UPDATE and DELETE to parser and executor by Ólafur Búi Ólafsson

//! The top-level [`Executor`]: owns a [`BTree`] and runs SQL text against it.

use std::path::Path;

use akurai_sql::{
    parse, CreateTableStmt, DataType, DeleteStmt, Direction, InsertStmt, Literal, SelectColumn,
    SelectStmt, Statement, UpdateStmt,
};
use akurai_storage::BTree;

use crate::catalog::Schema;
use crate::error::ExecError;
use crate::row::{decode_row, encode_row};
use crate::value::Value;
use crate::{eval, ResultSet};

/// What a successful statement produced.
#[derive(Debug, PartialEq)]
pub enum Outcome {
    /// A DDL statement (`CREATE TABLE`) succeeded.
    Created(String),
    /// A DML statement wrote rows; carries the count affected.
    Affected(usize),
    /// A query produced a result set.
    Rows(ResultSet),
}

/// A SQL engine over a single B+tree file.
pub struct Executor {
    tree: BTree,
}

impl Executor {
    /// Open (or create) the database file backing this engine.
    pub fn open(path: impl AsRef<Path>) -> Result<Executor, ExecError> {
        Ok(Executor {
            tree: BTree::open(path)?,
        })
    }

    /// Parse and run one SQL statement.
    pub fn execute(&mut self, sql: &str) -> Result<Outcome, ExecError> {
        match parse(sql)? {
            Statement::CreateTable(stmt) => self.create_table(stmt),
            Statement::Insert(stmt) => self.insert(stmt),
            Statement::Update(stmt) => self.update(stmt),
            Statement::Delete(stmt) => self.delete(stmt),
            Statement::Select(stmt) => self.select(stmt),
        }
    }

    // ---- DDL --------------------------------------------------------------

    fn create_table(&mut self, stmt: CreateTableStmt) -> Result<Outcome, ExecError> {
        let key = Schema::catalog_key(&stmt.table);
        if self.tree.get(&key)?.is_some() {
            return Err(ExecError::TableExists(stmt.table));
        }
        let schema = Schema::from_defs(&stmt.table, &stmt.columns);
        self.tree.insert(&key, &schema.encode())?;
        // Seed the rowid counter at 0 so INSERT can read it unconditionally.
        self.tree
            .insert(&Schema::seq_key(&stmt.table), &0u64.to_be_bytes())?;
        self.tree.commit()?;
        Ok(Outcome::Created(stmt.table))
    }

    // ---- INSERT -----------------------------------------------------------

    fn insert(&mut self, stmt: InsertStmt) -> Result<Outcome, ExecError> {
        let schema = self.load_schema(&stmt.table)?;

        if stmt.columns.len() != stmt.values.len() {
            return Err(ExecError::ColumnCountMismatch {
                expected: stmt.columns.len(),
                got: stmt.values.len(),
            });
        }

        // Place each supplied value into its schema slot; unmentioned columns
        // default to NULL.
        let mut cells: Vec<Value> = vec![Value::Null; schema.columns.len()];
        for (name, lit) in stmt.columns.iter().zip(stmt.values.iter()) {
            let idx = schema
                .index_of(name)
                .ok_or_else(|| ExecError::UnknownColumn(name.clone()))?;
            cells[idx] = coerce(lit, &schema.columns[idx])?;
        }

        // Enforce NOT NULL after placement.
        for (i, col) in schema.columns.iter().enumerate() {
            if col.not_null && cells[i].is_null() {
                return Err(ExecError::NotNull(col.name.clone()));
            }
        }

        let rowid = self.next_rowid(&stmt.table)?;
        let key = data_key(&stmt.table, rowid);
        self.tree.insert(&key, &encode_row(&cells))?;
        self.tree.commit()?;
        Ok(Outcome::Affected(1))
    }

    // ---- UPDATE -----------------------------------------------------------

    fn update(&mut self, stmt: UpdateStmt) -> Result<Outcome, ExecError> {
        let schema = self.load_schema(&stmt.table)?;

        // Resolve every assignment up front: validates each target column exists
        // and that its value coerces to the declared type before touching rows.
        let mut sets: Vec<(usize, Value)> = Vec::with_capacity(stmt.assignments.len());
        for (name, lit) in &stmt.assignments {
            let idx = schema
                .index_of(name)
                .ok_or_else(|| ExecError::UnknownColumn(name.clone()))?;
            let value = coerce(lit, &schema.columns[idx])?;
            sets.push((idx, value));
        }

        // Scan, filter by WHERE, and stage the rewritten rows (same key/rowid).
        let (start, end) = data_bounds(&stmt.table);
        let mut rewrites: Vec<(Vec<u8>, Vec<Value>)> = Vec::new();
        for (key, bytes) in self.tree.range(&start, &end)? {
            let mut row = decode_row(&bytes)?;
            if !matches_where(&stmt.where_clause, &schema, &row)? {
                continue;
            }
            for (idx, value) in &sets {
                row[*idx] = value.clone();
            }
            // Re-enforce NOT NULL after applying the assignments.
            for (i, col) in schema.columns.iter().enumerate() {
                if col.not_null && row[i].is_null() {
                    return Err(ExecError::NotNull(col.name.clone()));
                }
            }
            rewrites.push((key, row));
        }

        let n = rewrites.len();
        for (key, row) in rewrites {
            self.tree.insert(&key, &encode_row(&row))?;
        }
        self.tree.commit()?;
        Ok(Outcome::Affected(n))
    }

    // ---- DELETE -----------------------------------------------------------

    fn delete(&mut self, stmt: DeleteStmt) -> Result<Outcome, ExecError> {
        let schema = self.load_schema(&stmt.table)?;

        let (start, end) = data_bounds(&stmt.table);
        let mut victims: Vec<Vec<u8>> = Vec::new();
        for (key, bytes) in self.tree.range(&start, &end)? {
            let row = decode_row(&bytes)?;
            if matches_where(&stmt.where_clause, &schema, &row)? {
                victims.push(key);
            }
        }

        let n = victims.len();
        for key in victims {
            self.tree.delete(&key)?;
        }
        self.tree.commit()?;
        Ok(Outcome::Affected(n))
    }

    // ---- SELECT -----------------------------------------------------------

    fn select(&mut self, stmt: SelectStmt) -> Result<Outcome, ExecError> {
        let schema = self.load_schema(&stmt.table)?;

        // 1. Scan the table's data range and decode rows.
        let (start, end) = data_bounds(&stmt.table);
        let mut rows: Vec<Vec<Value>> = Vec::new();
        for (_k, v) in self.tree.range(&start, &end)? {
            rows.push(decode_row(&v)?);
        }

        // 2. WHERE filter.
        if let Some(pred) = &stmt.where_clause {
            let mut kept = Vec::with_capacity(rows.len());
            for row in rows {
                if eval::eval_predicate(pred, &schema, &row)? {
                    kept.push(row);
                }
            }
            rows = kept;
        }

        // 3. ORDER BY.
        if let Some(order) = &stmt.order_by {
            let idx = schema
                .index_of(&order.column)
                .ok_or_else(|| ExecError::UnknownColumn(order.column.clone()))?;
            rows.sort_by(|a, b| {
                let ord = a[idx].order_key(&b[idx]);
                match order.direction {
                    Direction::Asc => ord,
                    Direction::Desc => ord.reverse(),
                }
            });
        }

        // 4. LIMIT.
        if let Some(limit) = stmt.limit {
            let n = limit.max(0) as usize;
            rows.truncate(n);
        }

        // 5. Projection.
        let proj = projection(&stmt.columns, &schema)?;
        let columns: Vec<String> = proj
            .iter()
            .map(|i| schema.columns[*i].name.clone())
            .collect();
        let out_rows: Vec<Vec<Value>> = rows
            .into_iter()
            .map(|row| proj.iter().map(|i| row[*i].clone()).collect())
            .collect();

        Ok(Outcome::Rows(ResultSet {
            columns,
            rows: out_rows,
        }))
    }

    // ---- helpers ----------------------------------------------------------

    fn load_schema(&mut self, table: &str) -> Result<Schema, ExecError> {
        let bytes = self
            .tree
            .get(&Schema::catalog_key(table))?
            .ok_or_else(|| ExecError::NoSuchTable(table.to_string()))?;
        Schema::decode(table, &bytes)
    }

    /// Read, bump, and persist the per-table rowid counter; returns the rowid
    /// to use for the new row (1-based).
    fn next_rowid(&mut self, table: &str) -> Result<u64, ExecError> {
        let key = Schema::seq_key(table);
        let current = match self.tree.get(&key)? {
            Some(b) if b.len() == 8 => {
                let mut a = [0u8; 8];
                a.copy_from_slice(&b);
                u64::from_be_bytes(a)
            }
            Some(_) => return Err(ExecError::Corrupt("rowid counter wrong size")),
            None => 0,
        };
        let next = current + 1;
        self.tree.insert(&key, &next.to_be_bytes())?;
        Ok(next)
    }
}

/// Whether `row` satisfies an optional `WHERE` predicate. A missing predicate
/// matches every row (the bare `UPDATE`/`DELETE` whole-table case).
fn matches_where(
    where_clause: &Option<akurai_sql::Expr>,
    schema: &Schema,
    row: &[Value],
) -> Result<bool, ExecError> {
    match where_clause {
        Some(pred) => eval::eval_predicate(pred, schema, row),
        None => Ok(true),
    }
}

/// Coerce a parser literal into the column's declared type, or reject it.
fn coerce(lit: &Literal, col: &crate::catalog::Column) -> Result<Value, ExecError> {
    match (lit, col.ty) {
        (Literal::Null, _) => Ok(Value::Null),
        (Literal::Int(i), DataType::Integer) => Ok(Value::Int(*i)),
        // Widen an integer literal into a REAL column.
        (Literal::Int(i), DataType::Real) => Ok(Value::Float(*i as f64)),
        (Literal::Float(f), DataType::Real) => Ok(Value::Float(*f)),
        (Literal::Str(s), DataType::Text) => Ok(Value::Str(s.clone())),
        _ => Err(ExecError::TypeMismatch {
            column: col.name.clone(),
        }),
    }
}

/// Resolve the projection list into schema column indices.
fn projection(cols: &[SelectColumn], schema: &Schema) -> Result<Vec<usize>, ExecError> {
    // `SELECT *` (or a list containing `*`) projects every column in order.
    if cols.iter().any(|c| matches!(c, SelectColumn::All)) {
        return Ok((0..schema.columns.len()).collect());
    }
    cols.iter()
        .map(|c| match c {
            SelectColumn::All => unreachable!(),
            SelectColumn::Named(name) => schema
                .index_of(name)
                .ok_or_else(|| ExecError::UnknownColumn(name.clone())),
        })
        .collect()
}

/// The data key for `table`/`rowid`: `<table>:<rowid_be>`.
fn data_key(table: &str, rowid: u64) -> Vec<u8> {
    let mut k = Vec::with_capacity(table.len() + 9);
    k.extend_from_slice(table.as_bytes());
    k.push(b':');
    k.extend_from_slice(&rowid.to_be_bytes());
    k
}

/// The half-open `[start, end)` key window covering every row of `table`.
fn data_bounds(table: &str) -> (Vec<u8>, Vec<u8>) {
    let mut prefix = Vec::with_capacity(table.len() + 1);
    prefix.extend_from_slice(table.as_bytes());
    prefix.push(b':');
    let end = prefix_upper(&prefix);
    (prefix, end)
}

/// The smallest key strictly greater than every key beginning with `prefix`.
fn prefix_upper(prefix: &[u8]) -> Vec<u8> {
    let mut end = prefix.to_vec();
    while let Some(last) = end.last_mut() {
        if *last != 0xFF {
            *last += 1;
            return end;
        }
        end.pop();
    }
    // All-0xFF prefix: no upper bound shorter than infinity. Fall back to a
    // sentinel that the storage layer treats as "scan to the end".
    vec![0xFF; prefix.len() + 1]
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn prefix_upper_increments_last_byte() {
        assert_eq!(prefix_upper(b"users:"), b"users;".to_vec());
    }

    #[test]
    fn data_key_orders_by_rowid() {
        assert!(data_key("t", 1) < data_key("t", 2));
        assert!(data_key("t", 2) < data_key("t", 256));
    }
}