local.rs

use super::{Catalog, Transaction as _};
use crate::encoding::{self, Key as _, Value as _};
use crate::error::Result;
use crate::sql::types::schema::Table;
use crate::sql::types::{Expression, Row, Rows, Value};
use crate::storage::{self, mvcc};
use crate::{errdata, errinput};

use serde::{Deserialize, Serialize};
use std::borrow::Cow;
use std::collections::{HashMap, HashSet};

/// A SQL engine using local storage. This provides the main SQL storage logic,
/// also for the Raft SQL engine which dispatches to this engine for node-local
/// SQL storage.
pub struct Local<E: storage::Engine> {
    /// The local MVCC storage engine.
    pub(super) mvcc: mvcc::MVCC<E>,
}

impl<E: storage::Engine> Local<E> {
    /// Creates a new local SQL engine using the given storage engine.
    pub fn new(engine: E) -> Self {
        Self { mvcc: mvcc::MVCC::new(engine) }
    }

    /// Resumes a transaction from the given state. This is usually encapsulated
    /// in `mvcc::Transaction`, but the Raft-based engine can't retain the MVCC
    /// transaction between each request since it may be executed across
    /// multiple leader nodes, so it instead keeps the state in the session.
    pub fn resume(&self, state: mvcc::TransactionState) -> Result<Transaction<E>> {
        Ok(Transaction::new(self.mvcc.resume(state)?))
    }

    /// Fetches an unversioned key.
    pub fn get_unversioned(&self, key: &[u8]) -> Result<Option<Vec<u8>>> {
        self.mvcc.get_unversioned(key)
    }

    /// Sets an unversioned key.
    pub fn set_unversioned(&self, key: &[u8], value: Vec<u8>) -> Result<()> {
        self.mvcc.set_unversioned(key, value)
    }
}

impl<'a, E: storage::Engine + 'a> super::Engine<'a> for Local<E> {
    type Transaction = Transaction<E>;

    fn begin(&self) -> Result<Self::Transaction> {
        Ok(Self::Transaction::new(self.mvcc.begin()?))
    }

    fn begin_read_only(&self) -> Result<Self::Transaction> {
        Ok(Self::Transaction::new(self.mvcc.begin_read_only()?))
    }

    fn begin_as_of(&self, version: mvcc::Version) -> Result<Self::Transaction> {
        Ok(Self::Transaction::new(self.mvcc.begin_as_of(version)?))
    }
}

/// A SQL transaction, wrapping an MVCC transaction.
pub struct Transaction<E: storage::Engine> {
    txn: mvcc::Transaction<E>,
}

impl<E: storage::Engine> Transaction<E> {
    /// Creates a new SQL transaction using the given MVCC transaction.
    fn new(txn: mvcc::Transaction<E>) -> Self {
        Self { txn }
    }

    /// Returns the transaction's internal state.
    pub(super) fn state(&self) -> &mvcc::TransactionState {
        self.txn.state()
    }

    /// Fetches the matching primary keys for the given indexed column value, or
    /// an empty set if there is none.
    fn get_index(&self, table: &str, column: &str, value: &Value) -> Result<HashSet<Value>> {
        Ok(self
            .txn
            .get(&Key::Index(table.into(), column.into(), value.into()).encode())?
            .map(|v| HashSet::decode(&v))
            .transpose()?
            .unwrap_or_default())
    }

    /// Stores an index entry for the given column value.
    fn set_index(
        &self,
        table: &str,
        column: &str,
        value: &Value,
        ids: HashSet<Value>,
    ) -> Result<()> {
        let key = Key::Index(table.into(), column.into(), value.into()).encode();
        if ids.is_empty() {
            self.txn.delete(&key)
        } else {
            self.txn.set(&key, ids.encode())
        }
    }

    /// Returns all foreign key references to a table, as table -> columns.
    /// This includes references from the table itself.
    fn references(&self, table: &str) -> Result<Vec<(String, Vec<String>)>> {
        Ok(self
            .list_tables()?
            .into_iter()
            .map(|t| {
                let references: Vec<String> = t
                    .columns
                    .into_iter()
                    .filter(|c| c.references.as_deref() == Some(table))
                    .map(|c| c.name)
                    .collect();
                (t.name, references)
            })
            .filter(|(_, references)| !references.is_empty())
            .collect())
    }
}

impl<E: storage::Engine> super::Transaction for Transaction<E> {
    fn version(&self) -> u64 {
        self.txn.version()
    }

    fn read_only(&self) -> bool {
        self.txn.read_only()
    }

    fn commit(self) -> Result<()> {
        self.txn.commit()
    }

    fn rollback(self) -> Result<()> {
        self.txn.rollback()
    }

    fn delete(&self, table: &str, ids: &[Value]) -> Result<()> {
        // Check for foreign key referenes.

        // TODO: try to be more clever than simply iterating over each ID.
        for id in ids {
            let table = self.must_get_table(table)?;
            for (t, cs) in self.references(&table.name)? {
                let t = self.must_get_table(&t)?;
                let cs = cs
                    .into_iter()
                    .map(|c| Ok((t.get_column_index(&c)?, c)))
                    .collect::<Result<Vec<_>>>()?;
                let mut scan = self.scan(&t.name, None)?;
                while let Some(row) = scan.next().transpose()? {
                    for (i, c) in &cs {
                        if &row[*i] == id
                            && (table.name != t.name || id != table.get_row_key(&row)?)
                        {
                            return errinput!(
                                "primary key {id} referenced by table {} column {c}",
                                t.name
                            );
                        }
                    }
                }
            }

            let indexes: Vec<_> =
                table.columns.iter().enumerate().filter(|(_, c)| c.index).collect();
            if !indexes.is_empty() {
                for row in self.get(&table.name, &[id.clone()])? {
                    for (i, column) in &indexes {
                        let mut index = self.get_index(&table.name, &column.name, &row[*i])?;
                        index.remove(id);
                        self.set_index(&table.name, &column.name, &row[*i], index)?;
                    }
                }
            }
            self.txn.delete(&Key::Row(table.name.into(), id.into()).encode())?;
        }
        Ok(())
    }

    fn get(&self, table: &str, ids: &[Value]) -> Result<Vec<Row>> {
        ids.iter()
            .filter_map(|id| {
                self.txn
                    .get(&Key::Row(table.into(), id.into()).encode())
                    .transpose()
                    .map(|r| r.and_then(|v| Row::decode(&v)))
            })
            .collect()
    }

    fn insert(&self, table: &str, rows: Vec<Row>) -> Result<()> {
        let table = self.must_get_table(table)?;
        for row in rows {
            table.validate_row(&row, self)?;
            let id = table.get_row_key(&row)?;
            if !self.get(&table.name, &[id.clone()])?.is_empty() {
                return errinput!("primary key {id} already exists for table {}", table.name);
            }
            self.txn.set(&Key::Row((&table.name).into(), id.into()).encode(), row.encode())?;

            // Update indexes
            for (i, column) in table.columns.iter().enumerate().filter(|(_, c)| c.index) {
                let mut index = self.get_index(&table.name, &column.name, &row[i])?;
                index.insert(id.clone());
                self.set_index(&table.name, &column.name, &row[i], index)?;
            }
        }
        Ok(())
    }

    fn lookup_index(&self, table: &str, column: &str, value: &[Value]) -> Result<HashSet<Value>> {
        if !self.must_get_table(table)?.get_column(column)?.index {
            return errinput!("no index on {table}.{column}");
        }
        let mut pks = HashSet::new();
        for v in value {
            pks.extend(self.get_index(table, column, v)?)
        }
        Ok(pks)
    }

    fn scan(&self, table: &str, filter: Option<Expression>) -> Result<Rows> {
        let table = self.must_get_table(table)?;
        Ok(Box::new(
            self.txn
                .scan_prefix(&KeyPrefix::Row((&table.name).into()).encode())
                .iter()
                .map(|r| r.and_then(|(_, v)| Row::decode(&v)))
                .filter_map(move |r| match r {
                    Ok(row) => match &filter {
                        Some(filter) => match filter.evaluate(Some(&row)) {
                            Ok(Value::Boolean(b)) if b => Some(Ok(row)),
                            Ok(Value::Boolean(_)) | Ok(Value::Null) => None,
                            Ok(v) => Some(errdata!("filter returned {v}, expected boolean")),
                            Err(err) => Some(Err(err)),
                        },
                        None => Some(Ok(row)),
                    },
                    err => Some(err),
                })
                .collect::<Vec<_>>()
                .into_iter(),
        ))
    }

    fn scan_index(&self, table: &str, column: &str) -> Result<super::IndexScan> {
        let table = self.must_get_table(table)?;
        let column = table.get_column(column)?;
        if !column.index {
            return errinput!("no index for {}.{}", table.name, column.name);
        }
        Ok(Box::new(
            self.txn
                .scan_prefix(
                    &KeyPrefix::Index((&table.name).into(), (&column.name).into()).encode(),
                )
                .iter()
                .map(|r| -> Result<(Value, HashSet<Value>)> {
                    let (k, v) = r?;
                    let value = match Key::decode(&k)? {
                        Key::Index(_, _, pk) => pk.into_owned(),
                        _ => return errdata!("invalid index key"),
                    };
                    Ok((value, HashSet::<Value>::decode(&v)?))
                })
                .collect::<Vec<_>>()
                .into_iter(),
        ))
    }

    fn update(&self, table: &str, rows: HashMap<Value, Row>) -> Result<()> {
        let table = self.must_get_table(table)?;
        // TODO: be more clever than just iterating here.
        for (id, row) in rows {
            // If the primary key changes we do a delete and create, otherwise we replace the row
            if id != *table.get_row_key(&row)? {
                self.delete(&table.name, &[id.clone()])?;
                self.insert(&table.name, vec![row])?;
                return Ok(());
            }

            // Update indexes, knowing that the primary key has not changed
            let indexes: Vec<_> =
                table.columns.iter().enumerate().filter(|(_, c)| c.index).collect();
            if !indexes.is_empty() {
                let old = self.get(&table.name, &[id.clone()])?.remove(0);
                for (i, column) in indexes {
                    if old[i] == row[i] {
                        continue;
                    }
                    let mut index = self.get_index(&table.name, &column.name, &old[i])?;
                    index.remove(&id);
                    self.set_index(&table.name, &column.name, &old[i], index)?;

                    let mut index = self.get_index(&table.name, &column.name, &row[i])?;
                    index.insert(id.clone());
                    self.set_index(&table.name, &column.name, &row[i], index)?;
                }
            }

            table.validate_row(&row, self)?;
            self.txn.set(&Key::Row((&table.name).into(), id.into()).encode(), row.encode())?;
        }
        Ok(())
    }
}

impl<E: storage::Engine> Catalog for Transaction<E> {
    fn create_table(&self, table: Table) -> Result<()> {
        if self.get_table(&table.name)?.is_some() {
            return errinput!("table {} already exists", table.name);
        }
        table.validate(self)?;
        self.txn.set(&Key::Table((&table.name).into()).encode(), table.encode())
    }

    fn drop_table(&self, table: &str, if_exists: bool) -> Result<bool> {
        let table = if !if_exists {
            self.must_get_table(table)?
        } else if let Some(table) = self.get_table(table)? {
            table
        } else {
            return Ok(false);
        };
        if let Some((t, cs)) = self.references(&table.name)?.iter().find(|(t, _)| *t != table.name)
        {
            return errinput!("table {} is referenced by table {} column {}", table.name, t, cs[0]);
        }
        let mut scan = self.scan(&table.name, None)?;
        while let Some(row) = scan.next().transpose()? {
            self.delete(&table.name, &[table.get_row_key(&row)?.clone()])?
        }
        self.txn.delete(&Key::Table(table.name.into()).encode())?;
        Ok(true)
    }

    fn get_table(&self, table: &str) -> Result<Option<Table>> {
        self.txn.get(&Key::Table(table.into()).encode())?.map(|v| Table::decode(&v)).transpose()
    }

    fn list_tables(&self) -> Result<Vec<Table>> {
        self.txn
            .scan_prefix(&KeyPrefix::Table.encode())
            .iter()
            .map(|r| r.and_then(|(_, v)| Table::decode(&v)))
            .collect()
    }
}

/// SQL engine keys, using the KeyCode order-preserving encoding. For
/// simplicity, table and column names are used directly as identifiers in
/// keys, instead of e.g. numberic IDs. It is not possible to change
/// table/column names, so this is fine.
///
/// Uses Cow to allow encoding borrowed values but decoding owned values.
#[derive(Deserialize, Serialize)]
enum Key<'a> {
    /// A table schema by table name.
    Table(Cow<'a, str>),
    /// An index entry, by table name, index name, and index value.
    Index(Cow<'a, str>, Cow<'a, str>, Cow<'a, Value>),
    /// A table row, by table name and primary key value.
    Row(Cow<'a, str>, Cow<'a, Value>),
}

impl<'a> encoding::Key<'a> for Key<'a> {}

/// Key prefixes, allowing prefix scans of specific parts of the keyspace. These
/// must match the keys -- in particular, the enum variant indexes must match.
#[derive(Deserialize, Serialize)]
enum KeyPrefix<'a> {
    /// All table schemas.
    Table,
    /// An entire table index, by table and index name.
    Index(Cow<'a, str>, Cow<'a, str>),
    /// An entire table's rows, by table name.
    Row(Cow<'a, str>),
}

impl<'a> encoding::Key<'a> for KeyPrefix<'a> {}