use std::collections::HashMap;
use std::future::Future;
use buffa::Message;
use bytes::Bytes;
use exoware_sdk::common::kv::v1::Entry;
use exoware_sdk::log::stream::v1::GetResponse as StreamGetResponse;
use exoware_sdk::prune_policy::PrunePolicyDocument;
use exoware_sdk::retention::RetentionPolicy;
use crate::stream::{apply_filter, CompiledMatchers};
pub type QueryExtra = HashMap<String, buffa_types::google::protobuf::Value>;
#[derive(Clone, Debug, Default)]
pub struct RangeScanBatch {
pub rows: Vec<(Bytes, Bytes)>,
pub extra: QueryExtra,
}
pub trait RangeScan: Send {
fn next_batch(
&mut self,
max_items: usize,
) -> impl Future<Output = Result<RangeScanBatch, String>> + Send;
}
pub trait Sequence: Send + Sync + 'static {
fn current_sequence(&self) -> u64;
}
#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
pub enum IngestError {
#[error("unavailable: {message}")]
Unavailable { message: String },
#[error("internal: {message}")]
Internal { message: String },
}
pub trait Ingest: Send + Sync + 'static {
fn put_batch(
&self,
kvs: Vec<(Bytes, Bytes)>,
) -> impl Future<Output = Result<u64, IngestError>> + Send;
}
pub trait Query: Sequence {
type RangeScan: RangeScan + 'static;
fn get(
&self,
key: Bytes,
) -> impl Future<Output = Result<(Option<Bytes>, QueryExtra), String>> + Send;
fn range_scan(
&self,
start: Bytes,
end: Bytes,
limit: usize,
forward: bool,
) -> impl Future<Output = Result<Self::RangeScan, String>> + Send;
fn get_many(
&self,
keys: Vec<Bytes>,
) -> impl Future<Output = Result<(Vec<(Bytes, Option<Bytes>)>, QueryExtra), String>> + Send;
}
pub trait Prune: Send + Sync + 'static {
fn apply_prune_policies(
&self,
document: PrunePolicyDocument,
) -> impl Future<Output = Result<(), String>> + Send;
}
#[derive(Clone, Debug)]
pub struct LogBatch {
sequence_number: u64,
response_bytes: Bytes,
}
impl LogBatch {
pub fn from_response_bytes(sequence_number: u64, response_bytes: impl Into<Bytes>) -> Self {
Self {
sequence_number,
response_bytes: response_bytes.into(),
}
}
pub fn from_entries(sequence_number: u64, kvs: Vec<(Bytes, Bytes)>) -> Self {
Self::from_response(StreamGetResponse {
sequence_number,
entries: kvs
.into_iter()
.map(|(key, value)| Entry {
key: key.to_vec(),
value,
..Default::default()
})
.collect(),
..Default::default()
})
}
pub fn from_response(response: StreamGetResponse) -> Self {
Self {
sequence_number: response.sequence_number,
response_bytes: Bytes::from(response.encode_to_vec()),
}
}
pub fn sequence_number(&self) -> u64 {
self.sequence_number
}
pub fn into_response_bytes(self) -> Bytes {
self.response_bytes
}
pub fn decode_response(&self) -> Result<StreamGetResponse, String> {
StreamGetResponse::decode_from_slice(&self.response_bytes)
.map_err(|err| format!("failed to decode sequence log value: {err}"))
}
}
#[derive(Clone, Debug)]
pub struct FilteredBatch {
entries: Vec<Entry>,
}
impl FilteredBatch {
pub fn from_entries(entries: Vec<Entry>) -> Self {
Self { entries }
}
pub fn from_kvs(kvs: Vec<(Bytes, Bytes)>) -> Self {
Self {
entries: kvs
.into_iter()
.map(|(key, value)| Entry {
key: key.to_vec(),
value,
..Default::default()
})
.collect(),
}
}
pub fn into_entries(self) -> Vec<Entry> {
self.entries
}
}
pub trait Log: Sequence {
fn get_batch(
&self,
sequence_number: u64,
) -> impl Future<Output = Result<Option<LogBatch>, String>> + Send;
fn get_batch_filtered(
&self,
sequence_number: u64,
matchers: &CompiledMatchers,
) -> impl Future<Output = Result<Option<FilteredBatch>, String>> + Send {
async move {
let Some(batch) = self.get_batch(sequence_number).await? else {
return Ok(None);
};
let response = batch.decode_response()?;
let entries = apply_filter(matchers, response.entries);
Ok(Some(FilteredBatch::from_entries(entries)))
}
}
fn oldest_retained_batch(&self) -> impl Future<Output = Result<Option<u64>, String>> + Send;
}
pub trait Retention: Send + Sync + 'static {
fn set_retention(
&self,
policy: Option<RetentionPolicy>,
) -> impl Future<Output = Result<Option<u64>, String>> + Send;
}
pub trait StoreEngine: Ingest + Query + Prune + Log + Retention {}
impl<T: Ingest + Query + Prune + Log + Retention> StoreEngine for T {}
#[cfg(test)]
mod tests {
use std::future::ready;
use super::*;
use exoware_sdk::kv_codec::Utf8;
use exoware_sdk::selector::Selector;
use exoware_sdk::stream_filter::StreamFilter;
#[derive(Clone)]
struct TestLog {
batch: Option<LogBatch>,
}
impl Sequence for TestLog {
fn current_sequence(&self) -> u64 {
u64::from(self.batch.is_some())
}
}
impl Log for TestLog {
fn get_batch(
&self,
sequence_number: u64,
) -> impl Future<Output = Result<Option<LogBatch>, String>> + Send {
ready(Ok((sequence_number == 1)
.then(|| self.batch.clone())
.flatten()))
}
fn oldest_retained_batch(
&self,
) -> impl Future<Output = Result<Option<u64>, String>> + Send {
ready(Ok(self.batch.as_ref().map(|_| 1)))
}
}
fn matchers() -> CompiledMatchers {
crate::stream::compile_matchers(&StreamFilter {
selectors: vec![Selector {
prefix: Bytes::from_static(&[1]),
payload_regex: Utf8::from("^keep$"),
}],
value_filters: vec![],
})
.unwrap()
}
#[test]
fn default_filtered_batch_preserves_retained_and_missing_states() {
let matching = (
Bytes::from_static(&[1, b'k', b'e', b'e', b'p']),
Bytes::from_static(b"one"),
);
let excluded = (
Bytes::from_static(&[2, b'k', b'e', b'e', b'p']),
Bytes::from_static(b"two"),
);
let log = TestLog {
batch: Some(LogBatch::from_entries(1, vec![matching, excluded])),
};
let filtered = futures::executor::block_on(log.get_batch_filtered(1, &matchers()))
.unwrap()
.unwrap();
let entries = filtered.into_entries();
assert_eq!(entries.len(), 1);
assert_eq!(entries[0].value.as_ref(), b"one");
let retained_without_matches = TestLog {
batch: Some(LogBatch::from_entries(
1,
vec![(Bytes::from_static(&[2]), Bytes::from_static(b"value"))],
)),
};
let filtered = futures::executor::block_on(
retained_without_matches.get_batch_filtered(1, &matchers()),
)
.unwrap()
.unwrap();
assert!(filtered.into_entries().is_empty());
let missing = TestLog { batch: None };
let filtered =
futures::executor::block_on(missing.get_batch_filtered(1, &matchers())).unwrap();
assert!(filtered.is_none());
}
}