matdb 0.1.0

use anyhow::{bail, Context as AnyhowContext, Result};

use crate::{
    ast::Set, catalog::Index, eval::Context, kv::RangeIterKV, program::inner_insert, Db, Value,
};

use super::{execute_plan_next, PlanNode};

#[derive(Debug)]
pub struct Update {
    table: String,
    sets: Vec<Set>,
    child: Box<PlanNode>,
}

impl Update {
    pub fn new(table: String, sets: Vec<Set>, child: PlanNode) -> Self {
        Self {
            table,
            sets,
            child: Box::new(child),
        }
    }

    pub fn execute(&mut self, db: &mut Db, ctx: Context) -> Result<Option<Context>> {
        let schema = db.get_table_schema(db.this_tx_id, &self.table)?;
        let Some(ctx) = execute_plan_next(&mut self.child, db, ctx)? else {
            return Ok(None);
        };

        let mut new_ctx = ctx.clone();

        for set in &self.sets {
            if set.lhs.0 != schema.name
                || schema
                    .columns
                    .iter()
                    .find(|c| c.name == set.lhs.1)
                    .is_none()
            {
                bail!("{} is not a column of table {}", set.lhs, schema.name)
            }

            new_ctx.set(set.lhs.clone(), new_ctx.eval(&set.rhs)?);
        }

        for reference in &schema.referenced_by {
            let ref_schema = db.get_table_schema(db.this_tx_id, &reference)?;
            for foreign_key in ref_schema
                .foreign_keys
                .iter()
                .filter(|f| f.rhs_table == schema.name)
            {
                let old_sk = foreign_key
                    .rhs_exprs
                    .iter()
                    .map(|e| ctx.eval(e))
                    .collect::<Result<Vec<_>, _>>()?;

                let new_sk = foreign_key
                    .rhs_exprs
                    .iter()
                    .map(|e| new_ctx.eval(e))
                    .collect::<Result<Vec<_>, _>>()?;

                if old_sk == new_sk {
                    continue;
                }

                if RangeIterKV::new_simple(
                    if ref_schema.primary_key == foreign_key.lhs_exprs {
                        ref_schema.name.clone()
                    } else {
                        Index {
                            exprs: foreign_key.lhs_exprs.clone(),
                            unique: false,
                        }
                        .name(&ref_schema.name)
                    },
                    db.this_tx_id,
                    new_sk.clone(),
                )
                .next(db)?
                .is_some()
                {
                    bail!("row is referenced by table {}", reference)
                }
            }
        }

        let old_primary_key = schema
            .primary_key
            .iter()
            .map(|e| ctx.eval(e))
            .collect::<Result<Vec<_>, _>>()?;

        let old_secondary_keys = schema
            .indexes
            .iter()
            .map(|index| -> Result<Vec<Value>> {
                Ok(index
                    .exprs
                    .iter()
                    .map(|e| ctx.eval(e))
                    .collect::<Result<Vec<_>, _>>()?
                    .into_iter()
                    .chain(old_primary_key.iter().cloned())
                    .collect::<Vec<_>>())
            })
            .collect::<Result<Vec<_>>>()?;

        db.delete_key(&self.table, old_primary_key)?;

        for (sk, index) in old_secondary_keys.into_iter().zip(schema.indexes.iter()) {
            db.delete_key(&index.name(&schema.name), sk)?;
        }

        inner_insert(schema, &mut new_ctx, db).context("failed to update row")?;

        Ok(Some(new_ctx))
    }
}