use std::collections::HashSet;
use std::path::Path;
use redb::{ReadableDatabase, ReadableTableMetadata, TableError};
use uuid::Uuid;
use crate::memory_core::retrieval::PalaceHandle;
use crate::memory_core::store::kg_store::SESSIONS;
use crate::redb_open::INCOMPATIBLE_SUFFIX;
pub type LegacyProbe = dyn Fn(&Path, &HashSet<Uuid>) -> Option<String> + Send + Sync;
#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
pub enum PalaceNotEmpty {
#[error("the palace could not be confirmed empty: {0}")]
Unconfirmed(String),
#[error("the palace holds legacy data: {0}")]
LegacyData(String),
#[error("the palace has {0} drawer(s)")]
Drawers(usize),
#[error("the palace has {0} chat session(s)")]
ChatSessions(u64),
}
pub fn check_palace_empty(
handle: &PalaceHandle,
data_dir: &Path,
legacy: &LegacyProbe,
) -> Result<(), PalaceNotEmpty> {
if handle.drawer_load_degraded {
return Err(PalaceNotEmpty::Unconfirmed(
"the drawer table loaded degraded".to_string(),
));
}
let live = handle
.kg
.load_drawer_ids()
.map_err(|e| PalaceNotEmpty::Unconfirmed(format!("drawer ids could not be read: {e:#}")))?;
if let Some(reason) = legacy(data_dir, &live) {
return Err(PalaceNotEmpty::LegacyData(reason));
}
let drawers = handle.drawers.read().len().max(live.len());
if drawers > 0 {
return Err(PalaceNotEmpty::Drawers(drawers));
}
match count_chat_sessions(data_dir) {
Ok(0) => Ok(()),
Ok(n) => Err(PalaceNotEmpty::ChatSessions(n)),
Err(reason) => Err(PalaceNotEmpty::Unconfirmed(reason)),
}
}
pub fn conservative_legacy_probe(data_dir: &Path, _live: &HashSet<Uuid>) -> Option<String> {
let mut reasons = Vec::new();
match data_dir.join("kg.db").try_exists() {
Ok(false) => {}
Ok(true) => reasons.push("a legacy kg.db is present".to_string()),
Err(e) => reasons.push(format!("kg.db could not be checked: {e}")),
}
match std::fs::read_dir(data_dir) {
Ok(entries) => {
let quarantined = entries
.filter_map(|e| e.ok())
.filter(|e| {
e.file_name()
.to_string_lossy()
.contains(INCOMPATIBLE_SUFFIX)
})
.count();
if quarantined > 0 {
reasons.push(format!("{quarantined} {INCOMPATIBLE_SUFFIX} file(s)"));
}
}
Err(e) if e.kind() == std::io::ErrorKind::NotFound => {}
Err(e) => reasons.push(format!("{} could not be listed: {e}", data_dir.display())),
}
(!reasons.is_empty()).then(|| reasons.join("; "))
}
pub fn count_chat_sessions(data_dir: &Path) -> Result<u64, String> {
let path = data_dir.join("chat_sessions.redb");
match path.try_exists() {
Ok(false) => return Ok(0),
Ok(true) => {}
Err(e) => return Err(format!("{} could not be checked: {e}", path.display())),
}
let fail = |e: &dyn std::fmt::Display| format!("{} could not be read: {e}", path.display());
let db = crate::redb_cache::open_palace_db_read_only(&path).map_err(|e| fail(&e))?;
let rtx = db.begin_read().map_err(|e| fail(&e))?;
match rtx.open_table(SESSIONS) {
Ok(table) => table.len().map_err(|e| fail(&e)),
Err(TableError::TableDoesNotExist(_)) => Ok(0),
Err(e) => Err(fail(&e)),
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::memory_core::store::chat_sessions::ChatSessionStore;
#[test]
fn chat_sessions_make_a_palace_non_empty() {
let tmp = tempfile::tempdir().unwrap();
assert_eq!(count_chat_sessions(tmp.path()), Ok(0), "absent store");
let path = tmp.path().join("chat_sessions.redb");
drop(ChatSessionStore::open(&path).unwrap());
assert_eq!(
count_chat_sessions(tmp.path()),
Ok(0),
"store with no sessions"
);
let store = ChatSessionStore::open(&path).unwrap();
store.create_session(Some("t".to_string())).unwrap();
drop(store);
assert_eq!(count_chat_sessions(tmp.path()), Ok(1));
}
#[test]
fn conservative_probe_flags_kg_db_and_quarantined_stores() {
let tmp = tempfile::tempdir().unwrap();
let live = HashSet::new();
assert_eq!(conservative_legacy_probe(tmp.path(), &live), None);
std::fs::write(tmp.path().join("kg.db"), b"x").unwrap();
assert!(conservative_legacy_probe(tmp.path(), &live).is_some());
std::fs::remove_file(tmp.path().join("kg.db")).unwrap();
std::fs::write(
tmp.path().join(format!("kg.redb{INCOMPATIBLE_SUFFIX}")),
b"x",
)
.unwrap();
assert!(conservative_legacy_probe(tmp.path(), &live).is_some());
}
}