lox/crates/lox-distributor/src/db_handler.rs

use std::sync::{Arc, Mutex};

use crate::metrics::Metrics;
use crate::{lox_context, DbConfig};
use chrono::{naive::Days, DateTime, Local, NaiveDateTime, Utc};
use lox_library::{BridgeAuth, BridgeDb};
use sled::IVec;
use std::collections::HashMap;
use thiserror::Error;
use troll_patrol::bridge_info::BridgeInfo as TPBridgeInfo;

#[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 {
    db: sled::Db,
}

impl DB {
    // Writes the Lox context to the lox database with "context_%Y-%m-%d_%H:%M:%S" as the
    // database key
    pub fn write_context(&mut self, context: lox_context::LoxServerContext) {
        let date = Local::now().format("context_%Y-%m-%d_%H:%M:%S").to_string();
        /* Uncomment to generate test file for this function after making changes to lox library
            let file = OpenOptions::new()
            .create(true)
            .write(true)
            .truncate(true)
            .open("db_test_file.json").unwrap();
        serde_json::to_writer(&file, &context).unwrap();
        */
        let json_result = serde_json::to_vec(&context).unwrap();
        println!("Writing context to the db with key: {:?}", date);
        let _new_ivec = self.db.insert(
            IVec::from(date.as_bytes().to_vec()),
            IVec::from(json_result.clone()),
        );
        assert_eq!(
            self.db
                .get(IVec::from(date.as_bytes().to_vec()))
                .unwrap()
                .unwrap(),
            IVec::from(json_result)
        );
    }

    // If roll_back_date is empty, opens the most recent entry in the lox database or if none exists, creates a
    // new database. If roll_back_date is not empty, use the specified date to roll back to a previous lox-context
    // either exactly the entry at the roll_back_date or within 24 hours from the roll_back_date.
    pub fn open_new_or_existing_db(
        db_config: DbConfig,
        roll_back_date: Option<String>,
        metrics: Metrics,
    ) -> Result<(DB, lox_context::LoxServerContext), sled::Error> {
        let mut context: lox_context::LoxServerContext;
        let (lox_db, context) = match sled::open(db_config.db_path) {
            Ok(lox_db) => {
                // Check if the lox_db already exists
                if lox_db.was_recovered() && !lox_db.is_empty() {
                    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 {
                    let new_db = BridgeDb::new();
                    let new_ba = BridgeAuth::new(new_db.pubkey);
                    context = lox_context::LoxServerContext {
                        db: Arc::new(Mutex::new(new_db)),
                        ba: Arc::new(Mutex::new(new_ba)),
                        extra_bridges: Arc::new(Mutex::new(Vec::new())),
                        to_be_replaced_bridges: Arc::new(Mutex::new(Vec::new())),
                        tp_bridge_infos: Arc::new(Mutex::new(
                            HashMap::<[u8; 20], TPBridgeInfo>::new(),
                        )),
                        metrics,
                    };
                }
                (DB { db: lox_db }, context)
            }
            Err(e) => {
                panic!("Unable to read or create lox database! {:?}", e);
            }
        };
        Ok((lox_db, context))
    }
}

// Logic for finding the correct context to open from the database
fn read_lox_context_from_db(
    lox_db: sled::Db,
    roll_back_date: Option<String>,
) -> 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
    match roll_back_date {
        // If roll back date has been specified, either the exact date or range should be set
        Some(roll_back_date) => {
            // If the date is specified and it's in the database, use that to populate the context
            if lox_db.contains_key(roll_back_date.clone()).unwrap() {
                // Find date/time in db and use the context from that date.
                let ivec_context = lox_db
                    .get(IVec::from(roll_back_date.as_bytes().to_vec()))
                    .unwrap()
                    .unwrap();
                context = serde_json::from_slice(&ivec_context).unwrap();
                println!("Successfully used exact key {:?}", roll_back_date);
            } else {
                // If the exact date is not found, compute the start of the range as 24 hours prior to the specified date
                match compute_startdate_string(roll_back_date.clone()){
                    Some(start_date) => {
                        let start_date = start_date.format("context_%Y-%m-%d_%H:%M:%S").to_string();
                        let r: sled::Iter = lox_db.range(IVec::from(start_date.as_bytes().to_vec())..IVec::from(roll_back_date.as_bytes().to_vec()), );
                        match r.last(){
                            Some(entry) => {
                                let ivec_context = entry.unwrap();
                                let key: String = String::from_utf8(ivec_context.0.to_vec()).unwrap();
                                println!("Successfully used range and key {:?}", key);
                                context = serde_json::from_slice(&ivec_context.1).unwrap();
                            },
                            //If no entries are found within the 24 hours prior to the specified date, Exit
                            None => panic!("UNEXPECTED DATE: No entries found in the 24 hours prior to the input roll_back_date"),
                        }
                    }
                    None => panic!("UNEXPECTED DATE: Tried to use input date to calculate 24 hour range but failed. Exiting."),
                }
            }
        }
        // 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)?,
    }
    Ok(context)
}

fn compute_startdate_string(date_range_end: String) -> Option<DateTime<Utc>> {
    let parsed_end =
        NaiveDateTime::parse_from_str(&date_range_end, "context_%Y-%m-%d_%H:%M:%S").unwrap();
    let dt = DateTime::<Utc>::from_naive_utc_and_offset(parsed_end, Utc);
    dt.with_timezone(&Utc).checked_sub_days(Days::new(1))
}

// Use the last context that was entered into the database
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)]
mod tests {
    use std::env;
    use std::fs;

    use super::lox_context::LoxServerContext;
    use super::DbConfig;
    use super::Metrics;
    use super::DB;

    #[test]
    fn test_write_context() {
        // TODO: Fix db_test_file.json
        env::set_var("TEST_FILE_PATH", "db_test_file.json");
        let (mut lox_db, _context) =
            DB::open_new_or_existing_db(DbConfig::default(), None, Metrics::default()).unwrap();
        assert!(
            lox_db.db.is_empty(),
            "db read from context that shouldn't exist"
        );
        let path = env::var("TEST_FILE_PATH").unwrap();
        let contents = fs::File::open(&path).unwrap();
        //let test_string = std::str::from_utf8(&contents).unwrap();
        let test_context: LoxServerContext = serde_json::from_reader(contents).unwrap();
        lox_db.write_context(test_context);
        assert!(
            lox_db.db.len() == 1,
            "db should have written only one context"
        );
    }
}