use anyhow::{Context, Result}; use rocksdb::{ColumnFamilyDescriptor, Options, WriteBatch, DB}; use serde::{Deserialize, Serialize}; use std::collections::HashMap; use std::str::FromStr; use uuid::Uuid; #[derive(Debug, Clone, Serialize, Deserialize)] pub struct FileNode { pub node_id: String, pub label: String, pub aliases: Aliases, pub file_uuid: Option, pub sha256: Option, pub parent_id: Option, pub children: Vec, pub node_type: NodeType, pub icon: Option, pub color: Option, pub bg_color: Option, pub file_size: Option, pub registered_at: Option, pub created_at: String, pub updated_at: String, pub sort_order: i32, } #[derive(Debug, Clone, Serialize, Deserialize)] pub struct Aliases { #[serde(flatten)] pub map: HashMap, } impl Aliases { pub fn empty() -> Self { Aliases { map: HashMap::new(), } } pub fn to_json(&self) -> String { serde_json::to_string(&self.map).unwrap_or_else(|_| "{}".to_string()) } pub fn from_json(s: &str) -> Self { let map: HashMap = serde_json::from_str(s).unwrap_or_default(); Aliases { map } } } #[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq)] #[serde(rename_all = "snake_case")] pub enum NodeType { Folder, File, DynamicLayer, } impl NodeType { pub fn as_str(&self) -> &'static str { match self { NodeType::Folder => "folder", NodeType::File => "file", NodeType::DynamicLayer => "dynamic_layer", } } } impl FromStr for NodeType { type Err = String; fn from_str(s: &str) -> Result { match s { "folder" => Ok(NodeType::Folder), "file" => Ok(NodeType::File), "dynamic_layer" => Ok(NodeType::DynamicLayer), _ => Ok(NodeType::Folder), } } } pub struct FileTreeRocksDB { pub user_id: String, pub db: DB, } impl FileTreeRocksDB { pub fn user_db_path(user_id: &str) -> String { format!("data/users_rocksdb/{}.rocksdb", user_id) } pub fn init_user_db(user_id: &str) -> Result { let db_path = Self::user_db_path(user_id); let parent = std::path::Path::new(&db_path).parent().unwrap(); std::fs::create_dir_all(parent)?; let mut opts = Options::default(); opts.create_if_missing(true); opts.create_missing_column_families(true); let cfs = vec![ ColumnFamilyDescriptor::new("file_nodes", Options::default()), ColumnFamilyDescriptor::new("file_registry", Options::default()), ColumnFamilyDescriptor::new("file_locations", Options::default()), ColumnFamilyDescriptor::new("parent_index", Options::default()), ]; let db = DB::open_cf_descriptors(&opts, &db_path, cfs) .with_context(|| format!("Failed to open RocksDB at {}", db_path))?; Ok(FileTreeRocksDB { user_id: user_id.to_string(), db, }) } pub fn open_user_db(user_id: &str) -> Result { let db_path = Self::user_db_path(user_id); let mut opts = Options::default(); opts.create_if_missing(true); opts.create_missing_column_families(true); let cfs = vec![ ColumnFamilyDescriptor::new("file_nodes", Options::default()), ColumnFamilyDescriptor::new("file_registry", Options::default()), ColumnFamilyDescriptor::new("file_locations", Options::default()), ColumnFamilyDescriptor::new("parent_index", Options::default()), ]; let db = DB::open_cf_descriptors(&opts, &db_path, cfs) .with_context(|| format!("Failed to open RocksDB at {}", db_path))?; Ok(FileTreeRocksDB { user_id: user_id.to_string(), db, }) } pub fn insert_node(&self, node: &FileNode) -> Result<()> { let cf = self.db.cf_handle("file_nodes").unwrap(); let node_data = serde_json::to_vec(node)?; self.db.put_cf(cf, node.node_id.as_bytes(), node_data)?; if let Some(parent_id) = &node.parent_id { let cf_parent = self.db.cf_handle("parent_index").unwrap(); let key = format!("children:{}", parent_id); let existing = self.db.get_cf(cf_parent, key.as_bytes())?; let mut children: Vec = match existing { Some(data) => serde_json::from_slice(&data)?, None => Vec::new(), }; if !children.contains(&node.node_id) { children.push(node.node_id.clone()); self.db .put_cf(cf_parent, key.as_bytes(), serde_json::to_vec(&children)?)?; } } Ok(()) } pub fn insert_node_batch(&self, nodes: &[FileNode]) -> Result<()> { let mut batch = WriteBatch::default(); let cf = self.db.cf_handle("file_nodes").unwrap(); let cf_parent = self.db.cf_handle("parent_index").unwrap(); for node in nodes { let node_data = serde_json::to_vec(node)?; batch.put_cf(cf, node.node_id.as_bytes(), node_data); } self.db.write(batch)?; for node in nodes { if let Some(parent_id) = &node.parent_id { let key = format!("children:{}", parent_id); let existing = self.db.get_cf(cf_parent, key.as_bytes())?; let mut children: Vec = match existing { Some(data) => serde_json::from_slice(&data)?, None => Vec::new(), }; if !children.contains(&node.node_id) { children.push(node.node_id.clone()); self.db .put_cf(cf_parent, key.as_bytes(), serde_json::to_vec(&children)?)?; } } } Ok(()) } pub fn get_node(&self, node_id: &str) -> Result> { let cf = self.db.cf_handle("file_nodes").unwrap(); let value = self.db.get_cf(cf, node_id.as_bytes())?; match value { Some(data) => { let node: FileNode = serde_json::from_slice(&data)?; Ok(Some(node)) } None => Ok(None), } } pub fn get_children(&self, parent_id: &str) -> Result> { let cf = self.db.cf_handle("parent_index").unwrap(); let key = format!("children:{}", parent_id); let value = self.db.get_cf(cf, key.as_bytes())?; match value { Some(data) => { let children: Vec = serde_json::from_slice(&data)?; Ok(children) } None => Ok(Vec::new()), } } pub fn load_all(&self) -> Result> { let cf = self.db.cf_handle("file_nodes").unwrap(); let mut nodes = Vec::new(); let iter = self.db.iterator_cf(cf, rocksdb::IteratorMode::Start); for item in iter { let (_, value) = item?; let node: FileNode = serde_json::from_slice(&value)?; nodes.push(node); } nodes.sort_by(|a, b| { a.sort_order .cmp(&b.sort_order) .then_with(|| a.created_at.cmp(&b.created_at)) }); Ok(nodes) } pub fn update_node(&self, node_id: &str, updates: &FileNode) -> Result<()> { let cf = self.db.cf_handle("file_nodes").unwrap(); let node_data = serde_json::to_vec(updates)?; self.db.put_cf(cf, node_id.as_bytes(), node_data)?; Ok(()) } pub fn delete_node(&self, node_id: &str) -> Result<()> { let node = self.get_node(node_id)?; if let Some(n) = node { if let Some(parent_id) = &n.parent_id { let cf_parent = self.db.cf_handle("parent_index").unwrap(); let key = format!("children:{}", parent_id); let existing = self.db.get_cf(cf_parent, key.as_bytes())?; let mut children: Vec = match existing { Some(data) => serde_json::from_slice(&data)?, None => Vec::new(), }; children.retain(|id| id != node_id); self.db .put_cf(cf_parent, key.as_bytes(), serde_json::to_vec(&children)?)?; } } let cf = self.db.cf_handle("file_nodes").unwrap(); self.db.delete_cf(cf, node_id.as_bytes())?; Ok(()) } pub fn count_nodes(&self) -> Result { let cf = self.db.cf_handle("file_nodes").unwrap(); let mut count = 0; let iter = self.db.iterator_cf(cf, rocksdb::IteratorMode::Start); for item in iter { let _ = item?; count += 1; } Ok(count) } pub fn new_folder(label: &str, parent_id: Option<&str>) -> FileNode { FileNode { node_id: Uuid::new_v4().to_string().replace("-", ""), label: label.to_string(), aliases: Aliases::empty(), file_uuid: None, sha256: None, parent_id: parent_id.map(|s| s.to_string()), children: Vec::new(), node_type: NodeType::Folder, icon: None, color: None, bg_color: None, file_size: None, registered_at: None, created_at: chrono::Utc::now().to_rfc3339(), updated_at: chrono::Utc::now().to_rfc3339(), sort_order: 0, } } pub fn new_file_node( label: &str, file_uuid: &str, sha256: Option<&str>, original_name: &str, file_size: Option, mime_type: Option<&str>, parent_id: Option<&str>, ) -> FileNode { FileNode { node_id: Uuid::new_v4().to_string().replace("-", ""), label: label.to_string(), aliases: Aliases::empty(), file_uuid: Some(file_uuid.to_string()), sha256: sha256.map(|s| s.to_string()), parent_id: parent_id.map(|s| s.to_string()), children: Vec::new(), node_type: NodeType::File, icon: None, color: None, bg_color: None, file_size, registered_at: Some(chrono::Utc::now().to_rfc3339()), created_at: chrono::Utc::now().to_rfc3339(), updated_at: chrono::Utc::now().to_rfc3339(), sort_order: 0, } } }