bustubx 0.3.0

A simple relational database written in Rust
Documentation
use log::debug;
use std::sync::Arc;
use tempfile::TempDir;

use crate::buffer::BUFFER_POOL_SIZE;
use crate::catalog::load_catalog_data;
use crate::common::util::{pretty_format_logical_plan, pretty_format_physical_plan};
use crate::error::{BustubxError, BustubxResult};
use crate::optimizer::LogicalOptimizer;
use crate::planner::logical_plan::LogicalPlan;
use crate::planner::PhysicalPlanner;
use crate::{
    buffer::BufferPoolManager,
    catalog::Catalog,
    execution::{ExecutionContext, ExecutionEngine},
    planner::{LogicalPlanner, PlannerContext},
    storage::{DiskManager, Tuple},
};

pub struct Database {
    disk_manager: Arc<DiskManager>,
    pub(crate) buffer_pool: Arc<BufferPoolManager>,
    pub(crate) catalog: Catalog,
    temp_dir: Option<TempDir>,
}
impl Database {
    pub fn new_on_disk(db_path: &str) -> BustubxResult<Self> {
        let disk_manager = Arc::new(DiskManager::try_new(db_path)?);
        let buffer_pool = Arc::new(BufferPoolManager::new(
            BUFFER_POOL_SIZE,
            disk_manager.clone(),
        ));

        let catalog = Catalog::new(buffer_pool.clone());

        let mut db = Self {
            disk_manager,
            buffer_pool,
            catalog,
            temp_dir: None,
        };
        load_catalog_data(&mut db)?;
        Ok(db)
    }

    pub fn new_temp() -> BustubxResult<Self> {
        let temp_dir = TempDir::new()?;
        let temp_path = temp_dir.path().join("test.db");
        let disk_manager =
            Arc::new(DiskManager::try_new(temp_path.to_str().ok_or(
                BustubxError::Internal("Invalid temp path".to_string()),
            )?)?);
        let buffer_pool = Arc::new(BufferPoolManager::new(
            BUFFER_POOL_SIZE,
            disk_manager.clone(),
        ));

        let catalog = Catalog::new(buffer_pool.clone());

        let mut db = Self {
            disk_manager,
            buffer_pool,
            catalog,
            temp_dir: Some(temp_dir),
        };
        load_catalog_data(&mut db)?;
        Ok(db)
    }

    pub fn run(&mut self, sql: &str) -> BustubxResult<Vec<Tuple>> {
        let logical_plan = self.create_logical_plan(sql)?;
        debug!(
            "Logical Plan: \n{}",
            pretty_format_logical_plan(&logical_plan)
        );

        let optimized_logical_plan = LogicalOptimizer::new().optimize(&logical_plan)?;
        debug!(
            "Optimized Logical Plan: \n{}",
            pretty_format_logical_plan(&logical_plan)
        );

        // logical plan -> physical plan
        let physical_planner = PhysicalPlanner {
            catalog: &self.catalog,
        };
        let physical_plan = physical_planner.create_physical_plan(optimized_logical_plan);
        debug!(
            "Physical Plan: \n{}",
            pretty_format_physical_plan(&physical_plan)
        );

        let execution_ctx = ExecutionContext::new(&mut self.catalog);
        let mut execution_engine = ExecutionEngine {
            context: execution_ctx,
        };
        let tuples = execution_engine.execute(Arc::new(physical_plan))?;
        Ok(tuples)
    }

    pub fn create_logical_plan(&mut self, sql: &str) -> BustubxResult<LogicalPlan> {
        // sql -> ast
        let stmts = crate::parser::parse_sql(sql)?;
        if stmts.len() != 1 {
            return Err(BustubxError::NotSupport(
                "only support one sql statement".to_string(),
            ));
        }
        let stmt = &stmts[0];
        let mut planner = LogicalPlanner {
            context: PlannerContext {
                catalog: &self.catalog,
            },
        };
        // ast -> logical plan
        planner.plan(stmt)
    }

    pub fn flush(&self) -> BustubxResult<()> {
        self.buffer_pool.flush_all_pages()
    }
}