matdb 0.1.0

//! An experimental embedded SQL-like DBMS featuring:
//!
//! - **Core SQL**: `SELECT`, `UPDATE`, `DELETE`, `INSERT`
//! - **Graph expressions** - support for graph-style joins and navigation
//! - **ACID compliance** - fully transactional storage engine
//! - **System catalogue** - introspection via internal tables
//! - **E-graph based query optimization** - advanced rewrite and cost-based planning
//! - **Time-traveling queries** - access to full historical data versions
//!
//! More about these features in the [readme](https://gitlab.com/vizdun/matdb).
//!
//! Typical library usage:
//! ```rust
//! let mut db = Db::open("db.matdb".into())?;
//!
//! db.run("CREATE TABLE messages (id INT PRIMARY KEY, body TEXT NOT NULL);", Vec::new())?;
//! db.run("INSERT INTO messages (id, body) VALUES (?0, ?1);", vec![1.into(), "Hello, matdb!".to_string().into()])?;
//! db.finish_tx()?;
//!
//! assert_eq!(
//!     Value::String("Hello, matdb!".to_string()),
//!     db.run("SELECT body FROM messages WHERE id = 1;", Vec::new())?.unwrap().rows[0][0]
//! );
//! ```

use std::{
    collections::{BTreeMap, BTreeSet, HashMap},
    fs::{File, OpenOptions},
    path::PathBuf,
};

use anyhow::{anyhow, Context as AnyhowContext, Result};
use catalog::{
    matdb_revs_schema, matdb_schemas_schema, matdb_tables_schema, ColumnSchema, TableSchema,
    MATDB_REVS, MATDB_SCHEMAS, MATDB_TABLES,
};
use kv::RangeIterKV;
use parser::{lexer::Lexer, Parser};
use program::run_program;
use value::Type;
pub use value::Value;

// TODO: do this some other way later
const PAGE_SIZE: usize = 4096;
const MAX_KEY_SIZE: u16 = 500;
const CACHE_SIZE: usize = 0x20000;
const WRITE_SIZE: usize = 0x20000;

mod ast;
mod catalog;
mod encoding;
mod eval;
mod executor;
mod kv;
mod parser;
mod plan;
mod program;
mod value;

/// This is the main `matdb` API through which all interaction with the database is done.
#[derive(Debug)]
pub struct Db {
    main_tree_root: u64,
    this_tx_id: u64,
    next_freelist_ptr: Option<u64>,
    previous_main_root: Option<u64>,

    file: Option<File>,
    page_count: u64,

    next_time: u64,

    time_to_page: BTreeMap<u64, u64>,
    cached_pages: HashMap<u64, (u64, Vec<u8>)>,

    write_pages: BTreeMap<u64, Vec<u8>>,

    /// contains all the pages that have been freed that can be written to again
    free_list: BTreeSet<u64>,
    /// contains all the pages that can be written to beginning with the next transaction, typically freelist pages
    next_free_list: BTreeSet<u64>,

    has_recovered: bool,
    written_anything_this_tx: bool,
}

impl Default for Db {
    fn default() -> Self {
        Self {
            main_tree_root: 0,
            this_tx_id: 0,
            next_freelist_ptr: None,
            previous_main_root: None,
            file: None,
            page_count: 0,
            next_time: 0,
            time_to_page: BTreeMap::new(),
            cached_pages: HashMap::new(),
            write_pages: BTreeMap::new(),
            free_list: BTreeSet::new(),
            next_free_list: BTreeSet::new(),
            has_recovered: false,
            written_anything_this_tx: false,
        }
    }
}

impl Db {
    fn open_rest(&mut self) -> Result<()> {
        self.recover()?;

        let schemas_exists = RangeIterKV::new_simple(
            MATDB_TABLES.to_string(),
            self.this_tx_id,
            vec![Value::String(MATDB_SCHEMAS.to_string())],
        )
        .next(self)?
        .is_some();

        if !schemas_exists {
            self.create_table(MATDB_SCHEMAS, matdb_schemas_schema())?;
            self.create_table(MATDB_TABLES, matdb_tables_schema())?;
            self.create_table(MATDB_REVS, matdb_revs_schema())?;
            self.finish_tx()?;
        }

        Ok(())
    }

    /// Open a database file, creating one if it doesn't exist.
    pub fn open(path: PathBuf) -> Result<Self> {
        let mut db = Self {
            file: Some(
                OpenOptions::new()
                    .write(true)
                    .create(true)
                    .truncate(false)
                    .read(true)
                    .open(&path)
                    .with_context(|| format!("failed to open {path:?}"))?,
            ),
            ..Default::default()
        };

        if let Some(file) = &db.file {
            file.sync_all().unwrap()
        }

        db.open_rest()?;

        Ok(db)
    }

    /// Opens an in-memory database with no persistence.
    pub fn open_memory() -> Result<Self> {
        let mut db = Self {
            file: None,
            ..Default::default()
        };

        db.open_rest()?;

        Ok(db)
    }

    /// Executes an SQL snippet. If the last statement returns a set, a [`ResultSet`] is returned. Variables can be passed through `bindings` to prevent SQL injections.
    ///
    /// Example:
    /// ```rust
    /// let mut db = Db::open("db.matdb".into())?;
    ///
    /// db.run("INSERT INTO messages (id, body) VALUES (?0, ?1);", vec![1.into(), "Hello, matdb!".to_string().into()],)?;
    /// ```
    pub fn run(&mut self, query: &str, bindings: Vec<Value>) -> Result<Option<ResultSet>> {
        fn critical(
            db: &mut Db,
            program: Vec<ast::Statement>,
            bindings: Vec<Value>,
        ) -> Result<Option<ResultSet>> {
            let out = run_program(db, &program, bindings)?;

            Ok(out.map(|(column_names, rows)| ResultSet { column_names, rows }))
        }

        if !self.has_recovered {
            self.recover().context("failed to recover database")?;
        }

        let lexer = Lexer::new(query);
        let mut parser = Parser::new(lexer);
        let program = parser.parse_program()?;

        match critical(self, program, bindings) {
            Err(e) => {
                self.has_recovered = false;
                Err(e)
            }
            r => r,
        }
    }

    /// Commits the current transaction.
    pub fn finish_tx(&mut self) -> Result<()> {
        self.finish_transaction()
    }

    fn create_table(&mut self, table: &str, schema: TableSchema) -> Result<()> {
        self.create_table_without_schema(table)?;

        self.insert_key(
            MATDB_SCHEMAS,
            vec![Value::String(table.to_string())],
            vec![
                Value::String(table.to_string()),
                Value::String(schema.to_string()),
            ],
        )?;

        for index in &schema.indexes {
            let key = index.exprs.iter().chain(schema.primary_key.iter());

            let index_schema = TableSchema {
                readonly: true,
                name: index.name(table),
                columns: key
                    .clone()
                    .map(|e| ColumnSchema {
                        name: format!("{e}"),
                        ty: Type::Any,
                        nullable: true,
                    })
                    .collect(),
                checks: vec![],
                indexes: vec![],
                primary_key: key.cloned().collect(),
                foreign_keys: vec![],
                referenced_by: vec![],
            };

            self.create_table_without_schema(&index.name(table))?;

            self.insert_key(
                MATDB_SCHEMAS,
                vec![Value::String(index.name(table))],
                vec![
                    Value::String(index.name(table)),
                    Value::String(index_schema.to_string()),
                ],
            )?;
        }

        for foreign_key in &schema.foreign_keys {
            let mut rhs_table_schema =
                self.get_table_schema(self.this_tx_id, &foreign_key.rhs_table)?;

            rhs_table_schema.referenced_by.push(schema.name.clone());

            self.insert_key(
                MATDB_SCHEMAS,
                vec![Value::String(foreign_key.rhs_table.clone())],
                vec![
                    Value::String(table.to_string()),
                    Value::String(rhs_table_schema.to_string()),
                ],
            )?
        }

        Ok(())
    }

    fn create_table_without_schema(&mut self, table: &str) -> Result<()> {
        let id = self.new_tree()?;

        self.insert_key(
            MATDB_TABLES,
            vec![Value::String(table.to_string())],
            vec![Value::String(table.to_string()), Value::Int(id as i64)],
        )
    }

    fn delete_table(&mut self, table: &str) -> Result<()> {
        let schema = self.get_table_schema(self.this_tx_id, table)?;

        self.delete_key(MATDB_SCHEMAS, vec![Value::String(table.to_string())])?;
        self.delete_table_without_schema(table)?;

        for index in schema.indexes {
            self.delete_key(MATDB_SCHEMAS, vec![Value::String(index.name(table))])?;
            self.delete_table_without_schema(&index.name(table))?;
        }

        for fk in schema.foreign_keys {
            let mut ref_schema = self.get_table_schema(self.this_tx_id, &fk.rhs_table)?;
            ref_schema.referenced_by = ref_schema
                .referenced_by
                .into_iter()
                .filter(|r| r != table)
                .collect::<Vec<_>>();

            self.insert_key(
                MATDB_SCHEMAS,
                vec![Value::String(ref_schema.name.clone())],
                vec![
                    Value::String(ref_schema.name.clone()),
                    Value::String(ref_schema.to_string()),
                ],
            )?;
        }

        Ok(())
    }

    fn delete_table_without_schema(&mut self, table: &str) -> Result<()> {
        let tree = self.table_to_kv_id_or_error(self.this_tx_id, table)?;
        self.delete_tree(tree)?;

        self.delete_key(MATDB_TABLES, vec![Value::String(table.to_string())])
    }

    fn get_table_schema(&mut self, tx_id: u64, table: &str) -> Result<TableSchema> {
        if table == MATDB_TABLES {
            Ok(matdb_tables_schema())
        } else {
            let schema_string = RangeIterKV::new_simple(
                MATDB_SCHEMAS.to_string(),
                tx_id,
                vec![Value::String(table.to_string())],
            )
            .next(self)?
            .ok_or(anyhow!("couldn't find schema for table {table}"))?
            .1[1]
                .clone()
                .into_string()
                .ok_or(anyhow!("schema for table {table} isn't a string"))?;

            Parser::new(Lexer::new(&schema_string))
                .parse_create_table_statement()
                .context("failed to parse table schema in catalog")
        }
    }

    fn insert_key(&mut self, table: &str, primary_key: Vec<Value>, row: Vec<Value>) -> Result<()> {
        if table == MATDB_TABLES {
            self.set_value_root(primary_key, row)
        } else {
            let kv_id = self.table_to_kv_id_or_error(self.this_tx_id, table)?;
            let id = self.set_value(kv_id, primary_key, row)?;

            if let Some(new_id) = id {
                self.insert_key(
                    MATDB_TABLES,
                    vec![Value::String(table.to_string())],
                    vec![Value::String(table.to_string()), Value::Int(new_id as i64)],
                )?;
            }

            Ok(())
        }
    }

    fn delete_key(&mut self, table: &str, primary_key: Vec<Value>) -> Result<()> {
        if table == MATDB_TABLES {
            self.delete_pair_root(primary_key)
        } else {
            let kv_id = self.table_to_kv_id_or_error(self.this_tx_id, table)?;
            let id = self.delete_pair(kv_id, primary_key)?;

            if let Some(new_id) = id {
                self.insert_key(
                    MATDB_TABLES,
                    vec![Value::String(table.to_string())],
                    vec![Value::String(table.to_string()), Value::Int(new_id as i64)],
                )?;
            }

            Ok(())
        }
    }

    fn table_to_kv_id_or_error(&mut self, tx_id: u64, table: &str) -> Result<u64> {
        self.table_to_kv_id(tx_id, table)?
            .ok_or_else(|| anyhow!("table {table} doesn't exist"))
    }

    fn table_to_kv_id(&mut self, tx_id: u64, table: &str) -> Result<Option<u64>> {
        Ok(if table == MATDB_TABLES {
            if tx_id == self.this_tx_id {
                Some(self.tree_root())
            } else if tx_id == self.this_tx_id - 1 {
                self.previous_main_root
            } else {
                RangeIterKV::new_simple(
                    MATDB_REVS.into(),
                    self.this_tx_id,
                    vec![Value::Int(tx_id as i64)],
                )
                .next(self)?
                .map(|(_k, v)| {
                    v[1].clone()
                        .into_i64()
                        .ok_or_else(|| anyhow!("kv id of {table} is not an INT"))
                        .map(|i| i as u64)
                })
                .transpose()?
            }
        } else {
            RangeIterKV::new_simple(
                MATDB_TABLES.to_string(),
                tx_id,
                vec![Value::String(table.to_string())],
            )
            .next(self)?
            .map(|(_k, v)| {
                v[1].clone()
                    .into_i64()
                    .ok_or_else(|| anyhow!("kv id of {table} is not an INT"))
                    .map(|i| i as u64)
            })
            .transpose()?
        })
    }
}

/// The result of an SQL query.
pub struct ResultSet {
    /// the names of the columns of the resulting rows
    pub column_names: Vec<String>,
    /// the resulting rows
    pub rows: Vec<Vec<Value>>,
}