Merge pull request #18 from cmu-db/cost-model-ORM-trait

Add cost model trait

Author: lanlou1554

optd-persistent/Cargo.toml

@@ -14,5 +14,5 @@ sea-orm = { version = "1.1", features = [
 ] }
 sea-orm-migration = "1.1"
 serde_json = "1.0"
-
 tokio = { version = "1.41", features = ["full"] }
+trait-variant = "0.1"

optd-persistent/src/bin/migrate_test.rs

@@ -0,0 +1,16 @@
+use optd_persistent::{migrate, TEST_DATABASE_FILE, TEST_DATABASE_URL};
+use sea_orm::*;
+use sea_orm_migration::prelude::*;
+
+#[tokio::main]
+async fn main() {
+    let _ = std::fs::remove_file(TEST_DATABASE_FILE);
+
+    let db = Database::connect(TEST_DATABASE_URL)
+        .await
+        .expect("Unable to connect to the database");
+
+    migrate(&db)
+        .await
+        .expect("Something went wrong during migration");
+}

optd-persistent/src/cost_model/interface.rs

@@ -0,0 +1,93 @@
+#![allow(dead_code, unused_imports)]
+
+use crate::entities::cascades_group;
+use crate::entities::event::Model as event_model;
+use crate::entities::logical_expression;
+use crate::entities::physical_expression;
+use crate::StorageResult;
+use sea_orm::*;
+use sea_orm_migration::prelude::*;
+use serde_json::json;
+use std::sync::Arc;
+
+pub enum CatalogSource {
+    Iceberg(),
+}
+
+#[trait_variant::make(Send)]
+pub trait CostModelStorageLayer {
+    type GroupId;
+    type TableId;
+    type AttrId;
+    type ExprId;
+    type EpochId;
+    type StatId;
+
+    // TODO: Change EpochId to event::Model::epoch_id
+    async fn create_new_epoch(
+        &mut self,
+        source: String,
+        data: String,
+    ) -> StorageResult<Self::EpochId>;
+
+    async fn update_stats_from_catalog(
+        &self,
+        c: CatalogSource,
+        epoch_id: Self::EpochId,
+    ) -> StorageResult<()>;
+
+    // i32 in `stats:i32` is a placeholder for the stats type
+    async fn update_stats(&self, stats: i32, epoch_id: Self::EpochId) -> StorageResult<()>;
+
+    async fn store_cost(
+        &self,
+        expr_id: Self::ExprId,
+        cost: i32,
+        epoch_id: Self::EpochId,
+    ) -> StorageResult<()>;
+
+    async fn store_expr_stats_mappings(
+        &self,
+        expr_id: Self::ExprId,
+        stat_ids: Vec<Self::StatId>,
+    ) -> StorageResult<()>;
+
+    /// Get the statistics for a given table.
+    ///
+    /// If `epoch_id` is None, it will return the latest statistics.
+    async fn get_stats_for_table(
+        &self,
+        table_id: Self::TableId,
+        // TODO: Add enum for stat_type
+        stat_type: i32,
+        epoch_id: Option<Self::EpochId>,
+    ) -> StorageResult<Option<f32>>;
+
+    /// Get the statistics for a given attribute.
+    ///
+    /// If `epoch_id` is None, it will return the latest statistics.
+    async fn get_stats_for_attr(
+        &self,
+        attr_id: Self::AttrId,
+        stat_type: i32,
+        epoch_id: Option<Self::EpochId>,
+    ) -> StorageResult<Option<f32>>;
+
+    /// Get the joint statistics for a list of attributes.
+    ///
+    /// If `epoch_id` is None, it will return the latest statistics.
+    async fn get_stats_for_attrs(
+        &self,
+        attr_ids: Vec<Self::AttrId>,
+        stat_type: i32,
+        epoch_id: Option<Self::EpochId>,
+    ) -> StorageResult<Option<f32>>;
+
+    async fn get_cost_analysis(
+        &self,
+        expr_id: Self::ExprId,
+        epoch_id: Self::EpochId,
+    ) -> StorageResult<Option<i32>>;
+
+    async fn get_cost(&self, expr_id: Self::ExprId) -> StorageResult<Option<i32>>;
+}

optd-persistent/src/cost_model/mod.rs

@@ -0,0 +1,2 @@
+pub mod interface;
+pub mod orm;

optd-persistent/src/cost_model/orm.rs

@@ -0,0 +1,90 @@
+#![allow(dead_code, unused_imports, unused_variables)]
+
+use crate::{BackendManager, CostModelStorageLayer, StorageResult};
+
+use super::interface::CatalogSource;
+
+impl CostModelStorageLayer for BackendManager {
+    type GroupId = i32;
+    type TableId = i32;
+    type AttrId = i32;
+    type ExprId = i32;
+    type EpochId = i32;
+    type StatId = i32;
+
+    async fn create_new_epoch(
+        &mut self,
+        source: String,
+        data: String,
+    ) -> StorageResult<Self::EpochId> {
+        todo!()
+    }
+
+    async fn update_stats_from_catalog(
+        &self,
+        c: CatalogSource,
+        epoch_id: Self::EpochId,
+    ) -> StorageResult<()> {
+        todo!()
+    }
+
+    async fn update_stats(&self, stats: i32, epoch_id: Self::EpochId) -> StorageResult<()> {
+        todo!()
+    }
+
+    async fn store_cost(
+        &self,
+        expr_id: Self::ExprId,
+        cost: i32,
+        epoch_id: Self::EpochId,
+    ) -> StorageResult<()> {
+        todo!()
+    }
+
+    async fn store_expr_stats_mappings(
+        &self,
+        expr_id: Self::ExprId,
+        stat_ids: Vec<Self::StatId>,
+    ) -> StorageResult<()> {
+        todo!()
+    }
+
+    async fn get_stats_for_table(
+        &self,
+        table_id: i32,
+        stat_type: i32,
+        epoch_id: Option<Self::EpochId>,
+    ) -> StorageResult<Option<f32>> {
+        todo!()
+    }
+
+    async fn get_stats_for_attr(
+        &self,
+        attr_id: i32,
+        stat_type: i32,
+        epoch_id: Option<Self::EpochId>,
+    ) -> StorageResult<Option<f32>> {
+        todo!()
+    }
+
+    async fn get_stats_for_attrs(
+        &self,
+        attr_ids: Vec<i32>,
+        stat_type: i32,
+        epoch_id: Option<Self::EpochId>,
+    ) -> StorageResult<Option<f32>> {
+        todo!()
+    }
+
+    async fn get_cost_analysis(
+        &self,
+        expr_id: Self::ExprId,
+        epoch_id: Self::EpochId,
+    ) -> StorageResult<Option<i32>> {
+        todo!()
+    }
+
+    async fn get_cost(&self, expr_id: Self::ExprId) -> StorageResult<Option<i32>> {
+        todo!()
+    }
+}

optd-persistent/src/entities/prelude.rs

@@ -1,4 +1,5 @@
 //! `SeaORM` Entity, @generated by sea-orm-codegen 1.1.0
+#![allow(dead_code, unused_imports, unused_variables)]
 
 pub use super::attribute::Entity as Attribute;
 pub use super::attribute_constraint_junction::Entity as AttributeConstraintJunction;

optd-persistent/src/lib.rs

@@ -1,11 +1,62 @@
+#![allow(dead_code)]
+
+use std::sync::atomic::AtomicUsize;
+
 use sea_orm::*;
 use sea_orm_migration::prelude::*;
 
-mod migrator;
 use migrator::Migrator;
 
+mod entities;
+mod migrator;
+
+mod cost_model;
+pub use cost_model::interface::CostModelStorageLayer;
+
+pub type StorageResult<T> = Result<T, BackendError>;
+
+pub enum CostModelError {
+    // TODO: Add more error types
+    UnknownStatisticType,
+}
+
+pub enum BackendError {
+    CostModel(CostModelError),
+    Database(DbErr),
+    // Add other variants as needed for different error types
+}
+
+impl From<CostModelError> for BackendError {
+    fn from(value: CostModelError) -> Self {
+        BackendError::CostModel(value)
+    }
+}
+
+impl From<DbErr> for BackendError {
+    fn from(value: DbErr) -> Self {
+        BackendError::Database(value)
+    }
+}
+
+pub struct BackendManager {
+    db: DatabaseConnection,
+    latest_epoch_id: AtomicUsize,
+}
+
+impl BackendManager {
+    /// Creates a new `BackendManager`.
+    pub async fn new() -> StorageResult<Self> {
+        Ok(Self {
+            db: Database::connect(DATABASE_URL).await?,
+            latest_epoch_id: AtomicUsize::new(0),
+        })
+    }
+}
+
 pub const DATABASE_URL: &str = "sqlite:./sqlite.db?mode=rwc";
 pub const DATABASE_FILE: &str = "./sqlite.db";
+pub const TEST_DATABASE_URL: &str = "sqlite:./test.db?mode=rwc";
+pub const TEST_DATABASE_FILE: &str = "./test.db";
 
 pub async fn migrate(db: &DatabaseConnection) -> Result<(), DbErr> {
     Migrator::refresh(db).await