use std::io::{Read, Write};
mod error;
pub use error::StoreError;
pub mod mem;
pub use mem::MemStore;
#[cfg(all(feature = "store-fs", not(target_arch = "wasm32")))]
pub mod fs;
#[cfg(all(feature = "store-fs", not(target_arch = "wasm32")))]
pub use fs::FsStore;
#[cfg(feature = "store-callback")]
pub mod callback;
#[cfg(feature = "store-callback")]
pub use callback::{CallbackStore, CallbackStoreBuilder};
#[cfg(feature = "store-s3")]
pub mod s3;
#[cfg(feature = "store-s3")]
pub use s3::{S3Config, S3Store};
#[cfg(feature = "store-redis")]
pub mod redis;
#[cfg(feature = "store-redis")]
pub use redis::{RedisConfig, RedisStore};
#[cfg(feature = "store-pg")]
pub mod pg;
#[cfg(feature = "store-pg")]
pub use pg::{PgConfig, PgStore};
#[cfg(any(test, feature = "store-testkit"))]
pub mod contract;
pub type Result<T> = std::result::Result<T, StoreError>;
pub trait Store: Send + Sync {
fn get(&self, key: &str) -> Result<Option<Vec<u8>>>;
fn put(&self, key: &str, bytes: &[u8]) -> Result<()>;
fn exists(&self, key: &str) -> Result<bool>;
fn delete(&self, key: &str) -> Result<()>;
fn list(&self, prefix: &str) -> Result<Vec<String>>;
fn health(&self) -> Result<()>;
fn put_if_absent(&self, key: &str, bytes: &[u8]) -> Result<bool> {
if self.exists(key)? {
Ok(false)
} else {
self.put(key, bytes)?;
Ok(true)
}
}
fn list_paginated(
&self,
prefix: &str,
after: Option<&str>,
limit: usize,
) -> Result<(Vec<String>, Option<String>)> {
let mut all = self.list(prefix)?;
all.sort();
let start = match after {
Some(a) => all.partition_point(|k| k.as_str() <= a),
None => 0,
};
let page: Vec<String> = all.into_iter().skip(start).take(limit).collect();
let next = if limit > 0 && page.len() == limit {
page.last().cloned()
} else {
None
};
Ok((page, next))
}
fn reader(&self, key: &str) -> Result<Box<dyn Read + '_>> {
match self.get(key)? {
Some(bytes) => Ok(Box::new(std::io::Cursor::new(bytes))),
None => Err(StoreError::NotFound(key.to_string())),
}
}
fn writer(&self, key: &str) -> Result<Box<dyn Write + '_>> {
Ok(Box::new(BufferingWriter::new(self, key.to_string())))
}
}
struct BufferingWriter<'a, S: Store + ?Sized> {
store: &'a S,
key: String,
buf: Vec<u8>,
committed: bool,
}
impl<'a, S: Store + ?Sized> BufferingWriter<'a, S> {
fn new(store: &'a S, key: String) -> Self {
Self {
store,
key,
buf: Vec::new(),
committed: false,
}
}
fn commit(&mut self) -> std::io::Result<()> {
self.store
.put(&self.key, &self.buf)
.map_err(|e| std::io::Error::other(e.to_string()))?;
self.committed = true;
Ok(())
}
}
impl<S: Store + ?Sized> Write for BufferingWriter<'_, S> {
fn write(&mut self, data: &[u8]) -> std::io::Result<usize> {
self.buf.extend_from_slice(data);
Ok(data.len())
}
fn flush(&mut self) -> std::io::Result<()> {
self.commit()
}
}
impl<S: Store + ?Sized> Drop for BufferingWriter<'_, S> {
fn drop(&mut self) {
if !self.committed {
if let Err(e) = self.commit() {
eprintln!(
"nucleation::store: writer for key {:?} failed to commit on drop \
({e}); data was NOT persisted. Call flush() to handle this error.",
self.key
);
debug_assert!(false, "store writer failed to commit on drop: {e}");
}
}
}
}
#[cfg(any(feature = "store-s3", feature = "store-redis", feature = "store-pg"))]
pub(crate) fn block_on<F>(rt: &tokio::runtime::Runtime, fut: F) -> F::Output
where
F: std::future::Future + Send,
F::Output: Send,
{
if tokio::runtime::Handle::try_current().is_ok() {
std::thread::scope(|s| s.spawn(|| rt.block_on(fut)).join().unwrap())
} else {
rt.block_on(fut)
}
}
#[cfg(test)]
#[cfg(any(feature = "store-s3", feature = "store-redis", feature = "store-pg"))]
mod block_on_tests {
#[test]
fn block_on_from_within_a_runtime_does_not_panic() {
let outer = tokio::runtime::Runtime::new().unwrap();
let inner = tokio::runtime::Runtime::new().unwrap();
let result = outer.block_on(async { super::block_on(&inner, async { 21 + 21 }) });
assert_eq!(result, 42);
}
}
pub fn open(url: &str) -> Result<Box<dyn Store>> {
if url == "mem://" || url.starts_with("mem://") {
return Ok(Box::new(MemStore::new()));
}
#[cfg(all(feature = "store-fs", not(target_arch = "wasm32")))]
if let Some(path) = url.strip_prefix("file://") {
return Ok(Box::new(FsStore::new(path)));
}
#[cfg(feature = "store-s3")]
if let Some(rest) = url.strip_prefix("s3://") {
let (bucket, prefix) = match rest.split_once('/') {
Some((b, p)) => (b.to_string(), p.to_string()),
None => (rest.to_string(), String::new()),
};
let cfg = s3::S3Config {
prefix,
region: std::env::var("AWS_REGION").ok(),
endpoint: std::env::var("AWS_ENDPOINT_URL").ok(),
force_path_style: matches!(
std::env::var("AWS_S3_FORCE_PATH_STYLE").as_deref(),
Ok("true") | Ok("1")
),
..s3::S3Config::new(bucket)
};
return Ok(Box::new(s3::S3Store::connect(cfg)?));
}
#[cfg(feature = "store-redis")]
if url.starts_with("redis://") || url.starts_with("rediss://") {
return Ok(Box::new(redis::RedisStore::connect(
redis::RedisConfig::new(url),
)?));
}
#[cfg(feature = "store-pg")]
if url.starts_with("postgres://") || url.starts_with("postgresql://") {
let table =
std::env::var("NUC_STORE_PG_TABLE").unwrap_or_else(|_| "nucleation_store".to_string());
return Ok(Box::new(pg::PgStore::connect(pg::PgConfig::new(
url, table,
))?));
}
let scheme = url.split("://").next().unwrap_or(url);
Err(StoreError::Unsupported(format!(
"no store backend for scheme `{scheme}` (is its feature enabled?)"
)))
}
#[cfg(test)]
mod open_tests {
use super::*;
#[test]
fn mem_scheme_opens_a_working_store() {
let store = open("mem://").expect("open mem");
store.put("k", b"v").expect("put");
assert_eq!(store.get("k").expect("get"), Some(b"v".to_vec()));
}
#[test]
fn unknown_scheme_is_unsupported() {
match open("ftp://host/path") {
Err(StoreError::Unsupported(_)) => {}
Err(other) => panic!("expected Unsupported, got {other:?}"),
Ok(_) => panic!("expected Unsupported, got Ok"),
}
}
#[cfg(all(feature = "store-fs", not(target_arch = "wasm32")))]
#[test]
fn file_scheme_opens_an_fs_store() {
let dir = std::env::temp_dir().join(format!("nucleation-open-fs-{}", std::process::id()));
let _ = std::fs::remove_dir_all(&dir);
let url = format!("file://{}", dir.display());
let store = open(&url).expect("open file");
store.put("a/b", b"xy").expect("put");
assert_eq!(store.get("a/b").expect("get"), Some(b"xy".to_vec()));
let _ = std::fs::remove_dir_all(&dir);
}
}