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

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

use crate::{lox_context, DbConfig};
use chrono::{naive::Days, DateTime, Local, NaiveDateTime, Utc};
use lox_library::{BridgeAuth, BridgeDb};
use sled::IVec;

pub fn write_context_to_db(db: sled::Db, context: lox_context::LoxServerContext) {
    let date = Local::now().format("context_%Y-%m-%d_%H:%M:%S").to_string();
    let json_result = serde_json::to_vec(&context).unwrap();
    println!("Date: {:?}", date);
    let new_ivec = db.insert(
        IVec::from(date.as_bytes().to_vec()),
        IVec::from(json_result.clone()),
    );
    assert_eq!(
        db.get(IVec::from(date.as_bytes().to_vec()))
            .unwrap()
            .unwrap(),
        IVec::from(json_result)
    );
    println!("New entry key: {:?}", new_ivec);
}

pub fn open_new_or_existing_db(
    db_config: DbConfig,
    roll_back_date: Option<String>,
) -> Result<(sled::Db, lox_context::LoxServerContext), sled::Error> {
    let 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() {
                context = read_lox_context_from_db(lox_db.clone(), roll_back_date);
            //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())),
                };
            }
            (lox_db, context)
        }
        Err(e) => {
            panic!("Unable to read or create lox database! {:?}", e);
        }
    };
    Ok((lox_db, context))
}

fn read_lox_context_from_db(
    lox_db: sled::Db,
    roll_back_date: Option<String>,
) -> lox_context::LoxServerContext {
    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);
        }
    }
    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_utc(parsed_end, Utc);
    println!("WHat is datetime here: {:?}", dt);
    println!(
        "WHat is datetime - 1 subday here: {:?}",
        dt.with_timezone(&Utc).checked_sub_days(Days::new(1))
    );

    dt.with_timezone(&Utc).checked_sub_days(Days::new(1))
}

fn use_last_context(db: sled::Db) -> lox_context::LoxServerContext {
    let ivec_context = db.last().unwrap().unwrap();
    let ivec_date: String = serde_json::from_slice(&ivec_context.0).unwrap();
    println!("Using last context with date: {:?}", ivec_date);
    serde_json::from_slice(&ivec_context.1).unwrap()
}