From 4b1676d5bb418247323cb0ee09d44e74c42037b9 Mon Sep 17 00:00:00 2001
From: SarveshOO7 <sarveshtandonus@gmail.com>
Date: Sun, 17 Nov 2024 17:10:21 -0500
Subject: [PATCH] Fix CostModel's error system to match memo trait

---
 optd-persistent/src/cost_model/orm.rs | 23 +++++++++++++----------
 optd-persistent/src/lib.rs            |  7 +++++++
 2 files changed, 20 insertions(+), 10 deletions(-)

diff --git a/optd-persistent/src/cost_model/orm.rs b/optd-persistent/src/cost_model/orm.rs
index d172c14..f2fde13 100644
--- a/optd-persistent/src/cost_model/orm.rs
+++ b/optd-persistent/src/cost_model/orm.rs
@@ -238,10 +238,13 @@ impl CostModelStorageLayer for BackendManager {
                         match res {
                             Ok(insert_res) => insert_res.last_insert_id,
                             Err(_) => {
-                                return Err(BackendError::BackendError(format!(
-                                    "failed to insert statistic {:?} into statistic table",
-                                    stat
-                                )))
+                                return Err(BackendError::CostModel(
+                                    format!(
+                                        "failed to insert statistic {:?} into statistic table",
+                                        stat
+                                    )
+                                    .into(),
+                                ))
                             }
                         }
                     }
@@ -450,10 +453,10 @@ impl CostModelStorageLayer for BackendManager {
             .collect::<Vec<_>>();
 
         if attr_ids.len() != attr_base_indices.len() {
-            return Err(BackendError::BackendError(format!(
+            return Err(BackendError::CostModel(format!(
                 "Not all attributes found for table_id {} and base indices {:?}",
                 table_id, attr_base_indices
-            )));
+            ).into()));
         }
 
         self.get_stats_for_attr(attr_ids, stat_type, epoch_id).await
@@ -505,10 +508,10 @@ impl CostModelStorageLayer for BackendManager {
             .one(&self.db)
             .await?;
         if expr_exists.is_none() {
-            return Err(BackendError::BackendError(format!(
+            return Err(BackendError::CostModel(format!(
                 "physical expression id {} not found when storing cost",
                 physical_expression_id
-            )));
+            ).into()));
         }
 
         // Check if epoch_id exists in Event table
@@ -518,10 +521,10 @@ impl CostModelStorageLayer for BackendManager {
             .await
             .unwrap();
         if epoch_exists.is_none() {
-            return Err(BackendError::BackendError(format!(
+            return Err(BackendError::CostModel(format!(
                 "epoch id {} not found when storing cost",
                 epoch_id
-            )));
+            ).into()));
         }
 
         let new_cost = plan_cost::ActiveModel {
diff --git a/optd-persistent/src/lib.rs b/optd-persistent/src/lib.rs
index 340a07a..a548b28 100644
--- a/optd-persistent/src/lib.rs
+++ b/optd-persistent/src/lib.rs
@@ -47,6 +47,7 @@ pub enum CostModelError {
     // TODO: Add more error types
     UnknownStatisticType,
     VersionedStatisticNotFound,
+    CustomError(String),
 }
 
 /// TODO convert this to `thiserror`
@@ -69,6 +70,12 @@ pub enum BackendError {
     // TODO: Add other variants as needed for different error types
 }
 
+impl From<String> for CostModelError {
+    fn from(value: String) -> Self {
+        CostModelError::CustomError(value)
+    }
+}
+
 impl From<CostModelError> for BackendError {
     fn from(value: CostModelError) -> Self {
         BackendError::CostModel(value)