First version of ORM trait

optd-persistent/src/bin/migrate_test.rs

@@ -0,0 +1,23 @@
+use optd_persistent::{migrate, TEST_DATABASE_URL};
+use sea_orm::*;
+use sea_orm_migration::prelude::*;
+
+#[tokio::main]
+async fn main() {
+    let _ = std::fs::remove_file(TEST_DATABASE_URL);
+
+    let db = Database::connect(TEST_DATABASE_URL)
+        .await
+        .expect("Unable to connect to the database");
+
+    migrate(&db)
+        .await
+        .expect("Something went wrong during migration");
+
+    db.execute(sea_orm::Statement::from_string(
+        sea_orm::DatabaseBackend::Sqlite,
+        "PRAGMA foreign_keys = ON;".to_owned(),
+    ))
+    .await
+    .expect("Unable to enable foreign keys");
+}

optd-persistent/src/cost_model_orm_manager_impl.rs

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

optd-persistent/src/cost_model_storage_layer.rs

@@ -0,0 +1,75 @@
+#![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::{EpochId, ExprId, StatId, StorageResult};
+use sea_orm::*;
+use sea_orm_migration::prelude::*;
+use serde_json::json;
+use std::sync::Arc;
+
+pub enum CatalogSource {
+    Iceberg(),
+}
+
+pub trait CostModelStorageLayer {
+    // TODO: Change EpochId to event::Model::epoch_id
+    async fn create_new_epoch(&mut self, source: String, data: String) -> StorageResult<EpochId>;
+
+    async fn update_stats_from_catalog(
+        &self,
+        c: CatalogSource,
+        epoch_id: EpochId,
+    ) -> StorageResult<()>;
+
+    // i32 in `stats:i32` is a placeholder for the stats type
+    async fn update_stats(&self, stats: i32, epoch_id: EpochId) -> StorageResult<()>;
+
+    async fn store_cost(&self, expr_id: ExprId, cost: i32, epoch_id: EpochId) -> StorageResult<()>;
+
+    async fn store_expr_stats_mappings(
+        &self,
+        expr_id: ExprId,
+        stat_ids: Vec<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: i32,
+        stat_type: i32,
+        epoch_id: Option<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: i32,
+        stat_type: i32,
+        epoch_id: Option<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<i32>,
+        stat_type: i32,
+        epoch_id: Option<EpochId>,
+    ) -> StorageResult<Option<f32>>;
+
+    async fn get_cost_analysis(
+        &self,
+        expr_id: ExprId,
+        epoch_id: EpochId,
+    ) -> StorageResult<Option<i32>>;
+
+    async fn get_cost(&self, expr_id: ExprId) -> StorageResult<Option<i32>>;
+}

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,26 @@
 use sea_orm::*;
 use sea_orm_migration::prelude::*;
 
+mod cost_model_orm_manager_impl;
+mod cost_model_storage_layer;
+mod entities;
+mod memo_orm_manager_impl;
+mod memo_storage_layer;
 mod migrator;
+mod orm_manager;
+
 use migrator::Migrator;
 
+pub type GroupId = i32;
+pub type ExprId = i32;
+pub type EpochId = i32;
+pub type StatId = i32;
+
+pub type StorageResult<T> = Result<T, DbErr>;
+
 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 async fn migrate(db: &DatabaseConnection) -> Result<(), DbErr> {
     Migrator::refresh(db).await

optd-persistent/src/memo_orm_manager_impl.rs

@@ -0,0 +1,86 @@
+#![allow(dead_code, unused_imports, unused_variables)]
+
+use crate::memo_storage_layer::MemoStorageLayer;
+use crate::orm_manager::ORMManager;
+
+impl MemoStorageLayer for ORMManager {
+    async fn get_group_winner_from_group_id(
+        &self,
+        group_id: i32,
+    ) -> crate::StorageResult<Option<crate::entities::physical_expression::ActiveModel>> {
+        todo!()
+    }
+
+    async fn add_new_expr(
+        &mut self,
+        expr: crate::memo_storage_layer::Expression,
+    ) -> crate::StorageResult<(crate::GroupId, crate::ExprId)> {
+        todo!()
+    }
+
+    async fn add_expr_to_group(
+        &mut self,
+        expr: crate::memo_storage_layer::Expression,
+        group_id: crate::GroupId,
+    ) -> crate::StorageResult<Option<crate::ExprId>> {
+        todo!()
+    }
+
+    async fn get_group_id(&self, expr_id: crate::ExprId) -> crate::StorageResult<crate::GroupId> {
+        todo!()
+    }
+
+    async fn get_expr_memoed(
+        &self,
+        expr_id: crate::ExprId,
+    ) -> crate::StorageResult<crate::memo_storage_layer::Expression> {
+        todo!()
+    }
+
+    async fn get_all_group_ids(&self) -> crate::StorageResult<Vec<crate::GroupId>> {
+        todo!()
+    }
+
+    async fn get_group(
+        &self,
+        group_id: crate::GroupId,
+    ) -> crate::StorageResult<crate::entities::cascades_group::ActiveModel> {
+        todo!()
+    }
+
+    async fn update_group_winner(
+        &mut self,
+        group_id: crate::GroupId,
+        latest_winner: Option<crate::ExprId>,
+    ) -> crate::StorageResult<()> {
+        todo!()
+    }
+
+    async fn get_all_exprs_in_group(
+        &self,
+        group_id: crate::GroupId,
+    ) -> crate::StorageResult<Vec<crate::ExprId>> {
+        todo!()
+    }
+
+    async fn get_group_info(
+        &self,
+        group_id: crate::GroupId,
+    ) -> crate::StorageResult<&Option<crate::ExprId>> {
+        todo!()
+    }
+
+    async fn get_predicate_binding(
+        &self,
+        group_id: crate::GroupId,
+    ) -> crate::StorageResult<Option<crate::memo_storage_layer::Expression>> {
+        todo!()
+    }
+
+    async fn try_get_predicate_binding(
+        &self,
+        group_id: crate::GroupId,
+    ) -> crate::StorageResult<Option<crate::memo_storage_layer::Expression>> {
+        todo!()
+    }
+}