basic changes to the Database struct and updated functions to not clone &MochiFiles

This commit is contained in:
MrDulfin 2025-05-24 13:51:56 -04:00
parent 21977fe74f
commit 5763e3d6f8
2 changed files with 56 additions and 88 deletions

View file

@ -1,18 +1,12 @@
mod schema; pub mod schema;
use std::{ use std::{
collections::{hash_map::Values, HashMap, HashSet}, collections::{HashMap, HashSet}, ffi::OsStr, fs::{self}, io::{self}, path::{Path, PathBuf}, str::FromStr, sync::{Arc, Mutex, RwLock}
ffi::OsStr,
fs::{self, File},
io::{self, Write},
path::{Path, PathBuf},
sync::{Arc, Mutex, RwLock},
}; };
use blake3::Hash; use blake3::Hash;
use chrono::{DateTime, NaiveDateTime, TimeDelta, Utc}; use chrono::{DateTime, NaiveDateTime, TimeDelta, Utc};
use ciborium::{from_reader, into_writer}; use dotenvy::dotenv;
use diesel::{prelude::Queryable, Selectable};
use log::{error, info, warn}; use log::{error, info, warn};
use rand::distributions::{Alphanumeric, DistString}; use rand::distributions::{Alphanumeric, DistString};
use rocket::{ use rocket::{
@ -22,74 +16,45 @@ use rocket::{
use serde_with::serde_as; use serde_with::serde_as;
use uuid::Uuid; use uuid::Uuid;
use diesel::prelude::*;
pub struct Mochibase { pub struct Mochibase {
path: PathBuf, path: PathBuf,
/// connection to the db /// connection to the db
db: Arc<Mutex<diesel::sqlite::SqliteConnection>>, db: Arc<Mutex<SqliteConnection>>,
} }
impl Mochibase { impl Mochibase {
/// Create a new database initialized with no data, and save it to the
/// provided path
pub fn new<P: AsRef<Path>>(path: &P) -> Result<Self, io::Error> {
let output = Self {
path: path.as_ref().to_path_buf(),
entries: HashMap::new(),
hashes: HashMap::new(),
};
// Save the database initially after creating it
output.save()?;
Ok(output)
}
/// Open the database from a path
pub fn open<P: AsRef<Path>>(path: &P) -> Result<Self, io::Error> {
let mut file = File::open(path)?;
from_reader(&mut file)
.map_err(|e| io::Error::other(format!("failed to open database: {e}")))
}
/// Open the database from a path, **or create it if it does not exist** /// Open the database from a path, **or create it if it does not exist**
pub fn open_or_new<P: AsRef<Path>>(path: &P) -> Result<Self, io::Error> { pub fn open_or_new<P: AsRef<str>>(path: &P) -> Result<Self, io::Error> {
if !path.as_ref().exists() { dotenv().ok();
Self::new(path) let connection = SqliteConnection::establish(path.as_ref())
} else { .unwrap_or_else(|e| panic!("Failed to connect, error: {}", e));
Self::open(path) Ok(
Self {
path: PathBuf::from_str(path.as_ref()).unwrap(),
db: Arc::new(Mutex::new(connection))
} }
} )
/// Save the database to its file
pub fn save(&self) -> Result<(), io::Error> {
// Create a file and write the LZ4 compressed stream into it
let mut file = File::create(self.path.with_extension("bkp"))?;
into_writer(self, &mut file)
.map_err(|e| io::Error::other(format!("failed to save database: {e}")))?;
file.flush()?;
fs::rename(self.path.with_extension("bkp"), &self.path).unwrap();
Ok(())
} }
/// Insert a [`MochiFile`] into the database. /// Insert a [`MochiFile`] into the database.
/// ///
/// If the database already contained this value, then `false` is returned. /// If the database already contained this value, then `false` is returned.
pub fn insert(&mut self, mmid: &Mmid, entry: MochiFile) -> bool { pub fn insert(&mut self, mmid_: &Mmid, entry: MochiFile) -> bool {
if let Some(s) = self.hashes.get_mut(&entry.hash) { use schema::mochifiles::dsl::*;
let hash_matched_mmids: Vec<Mmid> = mochifiles
.filter(hash.eq(entry.hash()))
.select(mmid)
.load(&mut *self.db.lock().unwrap())
.expect("Error getting mmids");
// If the database already contains the hash, make sure the file is unique // If the database already contains the hash, make sure the file is unique
if !s.insert(mmid.clone()) { if hash_matched_mmids.contains(mmid_) {
return false; return false;
} }
} else { entry.insert_into(mochifiles).on_conflict_do_nothing();
// If the database does not contain the hash, create a new set for it
self.hashes
.insert(entry.hash, HashSet::from([mmid.clone()]));
}
self.entries.insert(mmid.clone(), entry.clone());
true true
} }
@ -97,19 +62,14 @@ impl Mochibase {
/// Remove an [`Mmid`] from the database entirely. /// Remove an [`Mmid`] from the database entirely.
/// ///
/// If the database did not contain this value, then `false` is returned. /// If the database did not contain this value, then `false` is returned.
pub fn remove_mmid(&mut self, mmid: &Mmid) -> bool { pub fn remove_mmid(&mut self, mmid_: &Mmid) -> bool {
let hash = if let Some(h) = self.entries.get(mmid).map(|e| e.hash) { use schema::mochifiles::dsl::*;
self.entries.remove(mmid);
h
} else {
return false;
};
if let Some(s) = self.hashes.get_mut(&hash) {
s.remove(mmid);
}
if diesel::delete(mochifiles.filter(mmid.eq(mmid_))).execute(&mut *self.db.lock().unwrap()).expect("Error deleting posts") > 0 {
true true
} else {
false
}
} }
/// Remove a hash from the database entirely. /// Remove a hash from the database entirely.
@ -134,21 +94,29 @@ impl Mochibase {
} }
/// Get an entry by its [`Mmid`]. Returns [`None`] if the value does not exist. /// Get an entry by its [`Mmid`]. Returns [`None`] if the value does not exist.
pub fn get(&self, mmid: &Mmid) -> Option<&MochiFile> { pub fn get(&self, mmid_: &Mmid) -> Option<MochiFile> {
self.entries.get(mmid) use schema::mochifiles::dsl::*;
mochifiles.filter(mmid.eq(mmid_)).select(MochiFile::as_select()).load(&mut *self.db.lock().unwrap()).unwrap().get(0).map(|f| f.clone())
} }
pub fn get_hash(&self, hash: &Hash) -> Option<&HashSet<Mmid>> { pub fn get_hash(&self, hash_: &String) -> Option<Vec<MochiFile>> {
self.hashes.get(hash) use schema::mochifiles::dsl::*;
let files = mochifiles.filter(hash.eq(hash_)).select(MochiFile::as_select()).load(&mut *self.db.lock().unwrap()).expect("failed to load mochifiles by hash");
if files.is_empty() {
None
} else {
Some(files)
}
} }
pub fn entries(&self) -> Values<'_, Mmid, MochiFile> { pub fn entries(&self) -> Vec<MochiFile> {
self.entries.values() use schema::mochifiles::dsl::*;
mochifiles.select(MochiFile::as_select()).load(&mut *self.db.lock().unwrap()).expect("failed to load all mochifiles")
} }
} }
/// An entry in the database storing metadata about a file /// An entry in the database storing metadata about a file
#[derive(Queryable, Selectable)] #[derive(Queryable, Selectable, Insertable)]
#[diesel(table_name = crate::database::schema::mochifiles)] #[diesel(table_name = crate::database::schema::mochifiles)]
#[diesel(check_for_backend(diesel::sqlite::Sqlite))] #[diesel(check_for_backend(diesel::sqlite::Sqlite))]
#[derive(Debug, Clone, Serialize, Deserialize)] #[derive(Debug, Clone, Serialize, Deserialize)]

View file

@ -35,7 +35,7 @@ pub fn server_info(settings: &State<Settings>) -> Json<ServerInfo> {
#[get("/info/<mmid>")] #[get("/info/<mmid>")]
pub async fn file_info(db: &State<Arc<RwLock<Mochibase>>>, mmid: &str) -> Option<Json<MochiFile>> { pub async fn file_info(db: &State<Arc<RwLock<Mochibase>>>, mmid: &str) -> Option<Json<MochiFile>> {
let mmid: Mmid = mmid.try_into().ok()?; let mmid: Mmid = mmid.try_into().ok()?;
let entry = db.read().unwrap().get(&mmid).cloned()?; let entry = db.read().unwrap().get(&mmid)?;
Some(Json(entry)) Some(Json(entry))
} }
@ -47,7 +47,7 @@ pub async fn file_info_opengraph(
mmid: &str, mmid: &str,
) -> Option<Markup> { ) -> Option<Markup> {
let mmid: Mmid = mmid.try_into().ok()?; let mmid: Mmid = mmid.try_into().ok()?;
let entry = db.read().unwrap().get(&mmid).cloned()?; let entry = db.read().unwrap().get(&mmid)?;
let file = File::open(settings.file_dir.join(entry.hash().to_string())) let file = File::open(settings.file_dir.join(entry.hash().to_string()))
.await .await
@ -55,7 +55,7 @@ pub async fn file_info_opengraph(
let size = to_pretty_size(file.metadata().await.ok()?.len()); let size = to_pretty_size(file.metadata().await.ok()?.len());
let seconds_till_expiry = entry.expiry().signed_duration_since(Utc::now()).num_seconds(); let seconds_till_expiry = entry.expiry().and_utc().signed_duration_since(Utc::now()).num_seconds();
let expiry = to_pretty_time(seconds_till_expiry as u32, BreakStyle::Space, TimeGranularity::Minutes); let expiry = to_pretty_time(seconds_till_expiry as u32, BreakStyle::Space, TimeGranularity::Minutes);
let title = entry.name().clone() + " - " + &size + " - " + &expiry; let title = entry.name().clone() + " - " + &size + " - " + &expiry;
@ -95,7 +95,7 @@ pub struct ServerInfo {
#[get("/f/<mmid>")] #[get("/f/<mmid>")]
pub async fn lookup_mmid(db: &State<Arc<RwLock<Mochibase>>>, mmid: &str) -> Option<Redirect> { pub async fn lookup_mmid(db: &State<Arc<RwLock<Mochibase>>>, mmid: &str) -> Option<Redirect> {
let mmid: Mmid = mmid.try_into().ok()?; let mmid: Mmid = mmid.try_into().ok()?;
let entry = db.read().unwrap().get(&mmid).cloned()?; let entry = db.read().unwrap().get(&mmid)?;
Some(Redirect::to(uri!(lookup_mmid_name( Some(Redirect::to(uri!(lookup_mmid_name(
mmid.to_string(), mmid.to_string(),
@ -111,7 +111,7 @@ pub async fn lookup_mmid_noredir(
download: bool, download: bool,
) -> Option<FileDownloader> { ) -> Option<FileDownloader> {
let mmid: Mmid = mmid.try_into().ok()?; let mmid: Mmid = mmid.try_into().ok()?;
let entry = db.read().unwrap().get(&mmid).cloned()?; let entry = db.read().unwrap().get(&mmid)?;
let file = File::open(settings.file_dir.join(entry.hash().to_string())) let file = File::open(settings.file_dir.join(entry.hash().to_string()))
.await .await
@ -162,7 +162,7 @@ pub async fn lookup_mmid_name(
name: &str, name: &str,
) -> Option<(ContentType, File)> { ) -> Option<(ContentType, File)> {
let mmid: Mmid = mmid.try_into().ok()?; let mmid: Mmid = mmid.try_into().ok()?;
let entry = db.read().unwrap().get(&mmid).cloned()?; let entry = db.read().unwrap().get(&mmid)?;
// If the name does not match, then this is invalid // If the name does not match, then this is invalid
if name != entry.name() { if name != entry.name() {