From 7d94b3c8619d261c1adc7ee1e424d0c61878301f Mon Sep 17 00:00:00 2001
From: Cecylia Bocovich <cohosh@torproject.org>
Date: Mon, 18 Dec 2023 14:17:13 -0500
Subject: [PATCH] Add LoxDBError and refactor to return Result

Add a new error type to capture multiple possible database reading
errors and refactor both read_lox_context_from_db and use_last_context
to return a Result.
---
 Cargo.lock                               |  1 +
 crates/lox-distributor/Cargo.toml        |  1 +
 crates/lox-distributor/src/db_handler.rs | 37 ++++++++++++++++++------
 3 files changed, 30 insertions(+), 9 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 9cdf39f..9f414ee 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1066,6 +1066,7 @@ dependencies = [
  "serde_json",
  "serde_with",
  "sled",
+ "thiserror",
  "time",
  "tokio",
 ]
diff --git a/crates/lox-distributor/Cargo.toml b/crates/lox-distributor/Cargo.toml
index 0a3196d..d71cbb5 100644
--- a/crates/lox-distributor/Cargo.toml
+++ b/crates/lox-distributor/Cargo.toml
@@ -35,6 +35,7 @@ serde_json = "1.0.108"
 prometheus = "0.13.3"
 sled = "0.34.7"
 prometheus-client = "0.22.0"
+thiserror = "1"
 
 [dependencies.chrono]
 version = "0.4.31"
diff --git a/crates/lox-distributor/src/db_handler.rs b/crates/lox-distributor/src/db_handler.rs
index c5a63f4..7493aba 100644
--- a/crates/lox-distributor/src/db_handler.rs
+++ b/crates/lox-distributor/src/db_handler.rs
@@ -5,6 +5,18 @@ use crate::{lox_context, DbConfig};
 use chrono::{naive::Days, DateTime, Local, NaiveDateTime, Utc};
 use lox_library::{BridgeAuth, BridgeDb};
 use sled::IVec;
+use thiserror::Error;
+
+#[derive(Error, Debug)]
+pub enum LoxDBError {
+    //failed to get last db entry
+    #[error("Failed to read last database entry")]
+    ReadFailure(#[from] sled::Error),
+
+    //no last db entries
+    #[error("No database entries stored")]
+    DatabaseEmpty,
+}
 
 // Database of Lox Distributor State
 pub struct DB {
@@ -52,7 +64,10 @@ impl DB {
             Ok(lox_db) => {
                 // Check if the lox_db already exists
                 if lox_db.was_recovered() && !lox_db.is_empty() {
-                    context = read_lox_context_from_db(lox_db.clone(), roll_back_date);
+                    context = match read_lox_context_from_db(lox_db.clone(), roll_back_date) {
+                        Ok(ctx) => ctx,
+                        Err(e) => panic!("Unable to read lox database {:?}", e)
+                    };
                     context.metrics = metrics;
                 //Otherwise, create a new Lox context
                 } else {
@@ -80,7 +95,7 @@ impl DB {
 fn read_lox_context_from_db(
     lox_db: sled::Db,
     roll_back_date: Option<String>,
-) -> lox_context::LoxServerContext {
+) -> Result<lox_context::LoxServerContext, LoxDBError> {
     let context: lox_context::LoxServerContext;
     // Check if there is a roll back date and try to choose the appropriate context
     // to rollback to, otherwise, take the last saved context
@@ -119,10 +134,10 @@ fn read_lox_context_from_db(
         }
         // Use the last entry to populate the Lox context if no rollback date is set (which should be most common)
         None => {
-            context = use_last_context(lox_db);
+            context = use_last_context(lox_db)?
         }
     }
-    context
+    Ok(context)
 }
 
 fn compute_startdate_string(date_range_end: String) -> Option<DateTime<Utc>> {
@@ -133,11 +148,15 @@ fn compute_startdate_string(date_range_end: String) -> Option<DateTime<Utc>> {
 }
 
 // Use the last context that was entered into the database
-fn use_last_context(lox_db: sled::Db) -> lox_context::LoxServerContext {
-    let ivec_context = lox_db.last().unwrap().unwrap();
-    let ivec_date: String = String::from_utf8(ivec_context.0.to_vec()).unwrap();
-    println!("Using last context with date: {:?}", ivec_date);
-    serde_json::from_slice(&ivec_context.1).unwrap()
+fn use_last_context(lox_db: sled::Db) -> Result<lox_context::LoxServerContext, LoxDBError> {
+        match lox_db.last()? {
+            Some(ivec_context) => {
+                let ivec_date: String = String::from_utf8(ivec_context.0.to_vec()).unwrap();
+                println!("Using last context with date: {:?}", ivec_date);
+                Ok(serde_json::from_slice(&ivec_context.1).unwrap())
+            },
+            None => Err(LoxDBError::DatabaseEmpty)
+        }
 }
 
 #[cfg(test)]