add db cache

1 files changed, 75 insertions, 0 deletions

diff --git a/src/cache.rs b/src/cache.rs
new file mode 100644
index 0000000..114ab09
--- /dev/null
+++ b/src/cache.rs
@@ -0,0 +1,75 @@
+use anyhow::Result;
+use redb::{Database, TableDefinition};
+use std::{
+    fs::{File, create_dir_all},
+    io::{Read, Write},
+    path::{Path, PathBuf},
+};
+
+pub enum Cache {
+    Directory(PathBuf),
+    Redb(Database),
+}
+const T_DOWNLOAD: TableDefinition<&str, &[u8]> = TableDefinition::new("dl");
+impl Cache {
+    pub fn new_directory() -> Result<Self> {
+        let cachedir = xdg::BaseDirectories::with_prefix("weareearth")
+            .unwrap()
+            .create_cache_directory("download")
+            .unwrap();
+        create_dir_all(cachedir.join("BulkMetadata"))?;
+        create_dir_all(cachedir.join("NodeData"))?;
+        Ok(Self::Directory(cachedir))
+    }
+    pub fn new_db(path: &Path) -> Result<Self> {
+        let db = Database::open(path)?;
+        {
+            let txn = db.begin_write()?;
+            txn.open_table(T_DOWNLOAD)?;
+            txn.commit()?;
+        }
+        Ok(Self::Redb(db))
+    }
+
+    pub fn get(&self, path: &str) -> Result<Option<Vec<u8>>> {
+        match self {
+            Cache::Directory(cachedir) => {
+                let cachepath = cachedir.join(path);
+                if cachepath.exists() {
+                    let mut buf = Vec::new();
+                    File::open(cachepath)?.read_to_end(&mut buf)?;
+                    Ok(Some(buf.into()))
+                } else {
+                    Ok(None)
+                }
+            }
+            Cache::Redb(database) => {
+                let txn = database.begin_read()?;
+                let table = txn.open_table(T_DOWNLOAD)?;
+                let res = table.get(path)?;
+                if let Some(res) = res {
+                    Ok(Some(res.value().to_vec()))
+                } else {
+                    Ok(None)
+                }
+            }
+        }
+    }
+    pub fn insert(&self, path: &str, data: &[u8]) -> Result<()> {
+        match self {
+            Cache::Directory(cachedir) => {
+                let cachepath = cachedir.join(path);
+                File::create(cachepath)?.write_all(data)?;
+                Ok(())
+            }
+            Cache::Redb(database) => {
+                let txn = database.begin_write()?;
+                let mut table = txn.open_table(T_DOWNLOAD)?;
+                table.insert(path, data)?;
+                drop(table);
+                txn.commit()?;
+                Ok(())
+            }
+        }
+    }
+}