use color_eyre::Result;
use std::fs;
use std::io::{BufRead, BufReader, Write};
use std::path::{Path, PathBuf};
const CACHE_FILES: &[&str] = &["query_history.txt"];
#[derive(Clone)]
pub struct CacheManager {
pub(crate) cache_dir: PathBuf,
}
impl CacheManager {
pub fn with_dir(cache_dir: PathBuf) -> Self {
Self { cache_dir }
}
pub fn new(app_name: &str) -> Result<Self> {
if let Some(dir) = std::env::var_os("DATUI_CACHE_DIR") {
return Ok(Self {
cache_dir: PathBuf::from(dir),
});
}
let cache_dir = dirs::cache_dir()
.ok_or_else(|| color_eyre::eyre::eyre!("Could not determine cache directory"))?
.join(app_name);
Ok(Self { cache_dir })
}
pub fn cache_dir(&self) -> &Path {
&self.cache_dir
}
pub fn cache_file(&self, filename: &str) -> PathBuf {
self.cache_dir.join(filename)
}
pub fn ensure_cache_dir(&self) -> Result<()> {
if !self.cache_dir.exists() {
fs::create_dir_all(&self.cache_dir)?;
}
Ok(())
}
pub fn clear_file(&self, filename: &str) -> Result<()> {
let file_path = self.cache_file(filename);
if file_path.exists() {
fs::remove_file(&file_path)?;
}
Ok(())
}
pub fn clear_all(&self) -> Result<()> {
for filename in CACHE_FILES {
let file_path = self.cache_file(filename);
if file_path.exists() {
if let Err(_e) = fs::remove_file(&file_path) {
}
}
}
Ok(())
}
pub fn load_history_file(&self, history_id: &str) -> Result<Vec<String>> {
let history_file = self.cache_file(&format!("{}_history.txt", history_id));
if !history_file.exists() {
return Ok(Vec::new());
}
let file = fs::File::open(&history_file)?;
let reader = BufReader::new(file);
let mut history = Vec::new();
for line in reader.lines() {
let line = line?;
if !line.trim().is_empty() {
history.push(line);
}
}
Ok(history)
}
pub fn update_history_file<F>(&self, history_id: &str, update: F) -> Result<()>
where
F: FnOnce(&mut Vec<String>),
{
use fs2::FileExt;
self.ensure_cache_dir()?;
let lock_path = self.cache_file(&format!("{}_history.lock", history_id));
let lock = fs::OpenOptions::new()
.create(true)
.write(true)
.truncate(false)
.open(&lock_path)?;
let deadline = std::time::Instant::now() + LOCK_TIMEOUT;
let mut held = false;
loop {
if lock.try_lock_exclusive().is_ok() {
held = true;
break;
}
if std::time::Instant::now() >= deadline {
break;
}
std::thread::sleep(std::time::Duration::from_millis(2));
}
if !held {
return Ok(());
}
let mut entries = self.load_history_file(history_id).unwrap_or_default();
update(&mut entries);
let result = self.save_history_file(history_id, &entries);
let _ = FileExt::unlock(&lock);
result
}
pub fn save_history_file(&self, history_id: &str, history: &[String]) -> Result<()> {
self.ensure_cache_dir()?;
let history_file = self.cache_file(&format!("{}_history.txt", history_id));
let temp_file = self.cache_file(&format!(
"{}_history.{}.tmp",
history_id,
std::process::id()
));
{
let mut file = fs::File::create(&temp_file)?;
for entry in history {
writeln!(file, "{}", entry)?;
}
file.sync_all()?;
}
fs::rename(&temp_file, &history_file).inspect_err(|_| {
let _ = fs::remove_file(&temp_file);
})?;
Ok(())
}
}
const LOCK_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(2);
pub const MAX_RECENTS: usize = 50;
impl CacheManager {
pub fn load_recents(&self) -> Vec<std::path::PathBuf> {
self.load_history_file("recents")
.unwrap_or_default()
.into_iter()
.map(std::path::PathBuf::from)
.collect()
}
pub fn forget_recent(&self, path: &std::path::Path) {
let target = path.to_string_lossy().into_owned();
let _ = self.update_history_file("recents", |recents| {
recents.retain(|p| p != &target);
});
}
pub fn clear_recents(&self) {
let _ = self.update_history_file("recents", |recents| recents.clear());
}
pub fn push_recent(&self, path: &std::path::Path) {
let looks_like_url = path.to_string_lossy().contains("://");
let stored = if looks_like_url {
path.to_path_buf()
} else {
path.canonicalize().unwrap_or_else(|_| path.to_path_buf())
};
let entry = stored.to_string_lossy().into_owned();
let _ = self.update_history_file("recents", |recents| {
recents.retain(|p| p != &entry);
recents.insert(0, entry.clone());
recents.truncate(MAX_RECENTS);
});
}
}
#[derive(Debug, Clone, Default, serde::Serialize, serde::Deserialize)]
pub struct DatasetFacts {
pub mtime: u64,
pub size: u64,
pub rows: Option<usize>,
pub cols: Option<usize>,
#[serde(default)]
pub columns: Vec<String>,
#[serde(default)]
pub kind: Option<crate::discover::EntryKind>,
#[serde(default)]
pub cost: crate::discover::Cost,
}
pub const MAX_DATASET_FACTS: usize = 4096;
impl CacheManager {
fn dataset_index_path(&self) -> PathBuf {
self.cache_file("datasets.json")
}
pub fn load_dataset_facts(&self) -> std::collections::HashMap<PathBuf, DatasetFacts> {
let Ok(text) = fs::read_to_string(self.dataset_index_path()) else {
return Default::default();
};
serde_json::from_str::<std::collections::HashMap<PathBuf, DatasetFacts>>(&text)
.unwrap_or_default()
}
pub fn record_dataset_facts(&self, facts: &[(PathBuf, DatasetFacts)]) {
if facts.is_empty() {
return;
}
let _ = self.with_cache_lock("datasets", || {
let mut index = self.load_dataset_facts();
for (path, entry) in facts {
index.insert(path.clone(), entry.clone());
}
if index.len() > MAX_DATASET_FACTS {
let mut kept: Vec<_> = index.into_iter().collect();
kept.sort_by_key(|(_, f)| std::cmp::Reverse(f.mtime));
kept.truncate(MAX_DATASET_FACTS);
index = kept.into_iter().collect();
}
let json = serde_json::to_string(&index)?;
let temp = self.cache_file(&format!("datasets.{}.tmp", std::process::id()));
fs::write(&temp, json)?;
fs::rename(&temp, self.dataset_index_path()).inspect_err(|_| {
let _ = fs::remove_file(&temp);
})?;
Ok(())
});
}
fn with_cache_lock<F>(&self, name: &str, work: F) -> Result<()>
where
F: FnOnce() -> Result<()>,
{
use fs2::FileExt;
self.ensure_cache_dir()?;
let lock = fs::OpenOptions::new()
.create(true)
.write(true)
.truncate(false)
.open(self.cache_file(&format!("{name}.lock")))?;
let deadline = std::time::Instant::now() + LOCK_TIMEOUT;
loop {
if lock.try_lock_exclusive().is_ok() {
break;
}
if std::time::Instant::now() >= deadline {
return Ok(());
}
std::thread::sleep(std::time::Duration::from_millis(2));
}
let result = work();
let _ = FileExt::unlock(&lock);
result
}
}