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

Latest change d971623579ac6cb425eea3ca33f24af5983c18b2 - feat(sqlexec): SQL executor over the storage B+tree (CREATE/INSERT/SELECT) by Ólafur Búi Ólafsson

//! Table schemas and their persistence in the B+tree.
//!
//! The catalog is just more rows in the same tree, stored under a reserved key
//! prefix that ordinary table data can never collide with (table names are SQL
//! identifiers and so never begin with a `\x00` byte). Each table's schema is
//! encoded once at `CREATE TABLE` time and read back on demand, so the catalog
//! survives reopening the database.

use akurai_sql::{ColumnDef, DataType};

use crate::error::ExecError;

/// Reserved prefix for catalog entries: `\x00cat:<table>`.
pub const CATALOG_PREFIX: &[u8] = b"\x00cat:";
/// Reserved prefix for per-table rowid counters: `\x00seq:<table>`.
pub const SEQ_PREFIX: &[u8] = b"\x00seq:";

/// A table's columns, in declaration order.
#[derive(Debug, Clone, PartialEq)]
pub struct Schema {
    pub table: String,
    pub columns: Vec<Column>,
}

/// One column: name, type, and nullability.
#[derive(Debug, Clone, PartialEq)]
pub struct Column {
    pub name: String,
    pub ty: DataType,
    pub not_null: bool,
}

impl Schema {
    /// Build a schema from a parsed `CREATE TABLE` column list.
    pub fn from_defs(table: &str, defs: &[ColumnDef]) -> Schema {
        Schema {
            table: table.to_string(),
            columns: defs
                .iter()
                .map(|d| Column {
                    name: d.name.clone(),
                    ty: d.ty,
                    not_null: d.not_null,
                })
                .collect(),
        }
    }

    /// Find a column's position by name (case-sensitive, as the parser emits).
    pub fn index_of(&self, name: &str) -> Option<usize> {
        self.columns.iter().position(|c| c.name == name)
    }

    /// The catalog key for this table.
    pub fn catalog_key(table: &str) -> Vec<u8> {
        let mut k = CATALOG_PREFIX.to_vec();
        k.extend_from_slice(table.as_bytes());
        k
    }

    /// The rowid-counter key for this table.
    pub fn seq_key(table: &str) -> Vec<u8> {
        let mut k = SEQ_PREFIX.to_vec();
        k.extend_from_slice(table.as_bytes());
        k
    }

    /// Serialize the schema to bytes for storage.
    pub fn encode(&self) -> Vec<u8> {
        let mut out = Vec::new();
        out.extend_from_slice(&(self.columns.len() as u32).to_be_bytes());
        for c in &self.columns {
            out.extend_from_slice(&(c.name.len() as u32).to_be_bytes());
            out.extend_from_slice(c.name.as_bytes());
            out.push(type_tag(c.ty));
            out.push(c.not_null as u8);
        }
        out
    }

    /// Decode a schema previously written by [`Schema::encode`].
    pub fn decode(table: &str, bytes: &[u8]) -> Result<Schema, ExecError> {
        let mut r = Reader { bytes, pos: 0 };
        let n = r.u32()? as usize;
        let mut columns = Vec::with_capacity(n);
        for _ in 0..n {
            let len = r.u32()? as usize;
            let name = r.string(len)?;
            let ty = tag_type(r.byte()?)?;
            let not_null = r.byte()? != 0;
            columns.push(Column { name, ty, not_null });
        }
        Ok(Schema {
            table: table.to_string(),
            columns,
        })
    }
}

/// Map a [`DataType`] to its on-disk tag byte.
pub fn type_tag(ty: DataType) -> u8 {
    match ty {
        DataType::Integer => 1,
        DataType::Text => 2,
        DataType::Real => 3,
    }
}

/// Inverse of [`type_tag`].
pub fn tag_type(tag: u8) -> Result<DataType, ExecError> {
    match tag {
        1 => Ok(DataType::Integer),
        2 => Ok(DataType::Text),
        3 => Ok(DataType::Real),
        _ => Err(ExecError::Corrupt("unknown column type tag")),
    }
}

/// A tiny cursor over a byte slice, returning [`ExecError::Corrupt`] on
/// underrun rather than panicking.
pub struct Reader<'a> {
    pub bytes: &'a [u8],
    pub pos: usize,
}

impl Reader<'_> {
    pub fn byte(&mut self) -> Result<u8, ExecError> {
        let b = *self
            .bytes
            .get(self.pos)
            .ok_or(ExecError::Corrupt("unexpected end of bytes"))?;
        self.pos += 1;
        Ok(b)
    }

    pub fn take(&mut self, n: usize) -> Result<&[u8], ExecError> {
        let end = self
            .pos
            .checked_add(n)
            .filter(|e| *e <= self.bytes.len())
            .ok_or(ExecError::Corrupt("unexpected end of bytes"))?;
        let slice = &self.bytes[self.pos..end];
        self.pos = end;
        Ok(slice)
    }

    pub fn u32(&mut self) -> Result<u32, ExecError> {
        let b = self.take(4)?;
        Ok(u32::from_be_bytes([b[0], b[1], b[2], b[3]]))
    }

    pub fn u64(&mut self) -> Result<u64, ExecError> {
        let b = self.take(8)?;
        let mut a = [0u8; 8];
        a.copy_from_slice(b);
        Ok(u64::from_be_bytes(a))
    }

    pub fn string(&mut self, len: usize) -> Result<String, ExecError> {
        let b = self.take(len)?;
        String::from_utf8(b.to_vec()).map_err(|_| ExecError::Corrupt("invalid utf-8 in text"))
    }
}

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

    #[test]
    fn schema_round_trips() {
        let defs = vec![
            ColumnDef {
                name: "id".into(),
                ty: DataType::Integer,
                not_null: true,
            },
            ColumnDef {
                name: "name".into(),
                ty: DataType::Text,
                not_null: false,
            },
            ColumnDef {
                name: "score".into(),
                ty: DataType::Real,
                not_null: true,
            },
        ];
        let schema = Schema::from_defs("users", &defs);
        let bytes = schema.encode();
        let back = Schema::decode("users", &bytes).unwrap();
        assert_eq!(schema, back);
        assert_eq!(back.index_of("name"), Some(1));
        assert_eq!(back.index_of("nope"), None);
    }
}