Import upstream v0.16.22, stripped
Upstream commit: 474dd0229cb20cf513036619781ed97bd8073c3f Enterprise-only files removed or emptied: 63 Enterprise-only snippets removed: 117 in 50 files Dangling module declarations removed: 5 Cargo edits turning enterprise off: 14 Verification: clean Enterprise feature gates left for rebuilt features: 19 in 18 files Produced by tools/fork/strip.py. The full report is in docs/fork/strip-reports/ on main.
This commit is contained in:
@@ -0,0 +1,195 @@
|
||||
/*
|
||||
* SPDX-FileCopyrightText: 2020 Stalwart Labs LLC <[email protected]>
|
||||
*
|
||||
* SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL
|
||||
*/
|
||||
|
||||
use crate::{BlobStore, CompressionAlgo, Store, U32_LEN};
|
||||
use std::{ops::Range, time::Instant};
|
||||
use trc::{AddContext, StoreEvent};
|
||||
|
||||
const MAGIC_MARKER: u8 = 0xa0;
|
||||
const LZ4_MARKER: u8 = MAGIC_MARKER | 0x01;
|
||||
//const ZSTD_MARKER: u8 = MAGIC_MARKER | 0x02;
|
||||
const NONE_MARKER: u8 = 0x00;
|
||||
|
||||
impl BlobStore {
|
||||
pub async fn get_blob(&self, key: &[u8], range: Range<usize>) -> trc::Result<Option<Vec<u8>>> {
|
||||
let start_time = Instant::now();
|
||||
let result = match &self {
|
||||
BlobStore::Store(store) => match store {
|
||||
#[cfg(feature = "sqlite")]
|
||||
Store::SQLite(store) => store.get_blob(key, 0..usize::MAX).await,
|
||||
#[cfg(feature = "foundation")]
|
||||
Store::FoundationDb(store) => store.get_blob(key, 0..usize::MAX).await,
|
||||
#[cfg(feature = "postgres")]
|
||||
Store::PostgreSQL(store) => store.get_blob(key, 0..usize::MAX).await,
|
||||
#[cfg(feature = "mysql")]
|
||||
Store::MySQL(store) => store.get_blob(key, 0..usize::MAX).await,
|
||||
#[cfg(feature = "rocks")]
|
||||
Store::RocksDb(store) => store.get_blob(key, 0..usize::MAX).await,
|
||||
Store::Ephemeral(store) => store.get_blob(key, 0..usize::MAX).await,
|
||||
Store::None => Err(trc::StoreEvent::NotConfigured.into()),
|
||||
},
|
||||
BlobStore::Fs(store) => store.get_blob(key, 0..usize::MAX).await,
|
||||
#[cfg(feature = "s3")]
|
||||
BlobStore::S3(store) => store.get_blob(key, 0..usize::MAX).await,
|
||||
#[cfg(feature = "azure")]
|
||||
BlobStore::Azure(store) => store.get_blob(key, 0..usize::MAX).await,
|
||||
}
|
||||
.caused_by(trc::location!())?;
|
||||
|
||||
trc::event!(
|
||||
Store(StoreEvent::BlobRead),
|
||||
Key = key,
|
||||
Elapsed = start_time.elapsed(),
|
||||
Size = result.as_ref().map_or(0, |data| data.len()),
|
||||
);
|
||||
|
||||
let Some(mut data) = result else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
let mut data = match data.last().copied() {
|
||||
Some(LZ4_MARKER) => {
|
||||
lz4_flex::decompress_size_prepended(data.get(..data.len() - 1).unwrap_or_default())
|
||||
.map_err(|err| {
|
||||
trc::StoreEvent::DecompressError
|
||||
.reason(err)
|
||||
.ctx(trc::Key::Key, key)
|
||||
.ctx(trc::Key::CausedBy, trc::location!())
|
||||
})?
|
||||
}
|
||||
Some(NONE_MARKER) => {
|
||||
if !data.is_empty() {
|
||||
data.truncate(data.len() - 1);
|
||||
}
|
||||
data
|
||||
}
|
||||
Some(_) => {
|
||||
trc::event!(Store(StoreEvent::BlobMissingMarker), Key = key);
|
||||
|
||||
data
|
||||
}
|
||||
None => {
|
||||
return Ok(Some(data));
|
||||
}
|
||||
};
|
||||
|
||||
if range.start == 0 {
|
||||
if range.end > data.len() {
|
||||
Ok(Some(data))
|
||||
} else {
|
||||
data.truncate(range.end);
|
||||
Ok(Some(data))
|
||||
}
|
||||
} else {
|
||||
Ok(Some(
|
||||
data.get(range.start..range.end)
|
||||
.unwrap_or_default()
|
||||
.to_vec(),
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn put_blob(
|
||||
&self,
|
||||
key: &[u8],
|
||||
data: &[u8],
|
||||
compression: CompressionAlgo,
|
||||
) -> trc::Result<()> {
|
||||
let data = match compression {
|
||||
CompressionAlgo::None => {
|
||||
let mut uncompressed = Vec::with_capacity(data.len() + 1);
|
||||
uncompressed.extend_from_slice(data);
|
||||
uncompressed.push(NONE_MARKER);
|
||||
uncompressed
|
||||
}
|
||||
CompressionAlgo::Lz4 => {
|
||||
let mut compressed =
|
||||
vec![
|
||||
LZ4_MARKER;
|
||||
lz4_flex::block::get_maximum_output_size(data.len()) + U32_LEN + 1
|
||||
];
|
||||
|
||||
// Compress the data
|
||||
let compressed_len =
|
||||
lz4_flex::compress_into(data, &mut compressed[U32_LEN..]).unwrap();
|
||||
|
||||
// Prepend the length of the uncompressed data
|
||||
compressed[..U32_LEN].copy_from_slice(&(data.len() as u32).to_le_bytes());
|
||||
|
||||
// Truncate to the actual size
|
||||
compressed.truncate(compressed_len + U32_LEN + 1);
|
||||
compressed
|
||||
}
|
||||
};
|
||||
|
||||
let start_time = Instant::now();
|
||||
let result = match &self {
|
||||
BlobStore::Store(store) => match store {
|
||||
#[cfg(feature = "sqlite")]
|
||||
Store::SQLite(store) => store.put_blob(key, &data).await,
|
||||
#[cfg(feature = "foundation")]
|
||||
Store::FoundationDb(store) => store.put_blob(key, &data).await,
|
||||
#[cfg(feature = "postgres")]
|
||||
Store::PostgreSQL(store) => store.put_blob(key, &data).await,
|
||||
#[cfg(feature = "mysql")]
|
||||
Store::MySQL(store) => store.put_blob(key, &data).await,
|
||||
#[cfg(feature = "rocks")]
|
||||
Store::RocksDb(store) => store.put_blob(key, &data).await,
|
||||
Store::Ephemeral(store) => store.put_blob(key, &data).await,
|
||||
Store::None => Err(trc::StoreEvent::NotConfigured.into()),
|
||||
},
|
||||
BlobStore::Fs(store) => store.put_blob(key, &data).await,
|
||||
#[cfg(feature = "s3")]
|
||||
BlobStore::S3(store) => store.put_blob(key, &data).await,
|
||||
#[cfg(feature = "azure")]
|
||||
BlobStore::Azure(store) => store.put_blob(key, &data).await,
|
||||
}
|
||||
.caused_by(trc::location!());
|
||||
|
||||
trc::event!(
|
||||
Store(StoreEvent::BlobWrite),
|
||||
Key = key,
|
||||
Elapsed = start_time.elapsed(),
|
||||
Size = data.len(),
|
||||
);
|
||||
|
||||
result
|
||||
}
|
||||
|
||||
pub async fn delete_blob(&self, key: &[u8]) -> trc::Result<bool> {
|
||||
let start_time = Instant::now();
|
||||
let result = match &self {
|
||||
BlobStore::Store(store) => match store {
|
||||
#[cfg(feature = "sqlite")]
|
||||
Store::SQLite(store) => store.delete_blob(key).await,
|
||||
#[cfg(feature = "foundation")]
|
||||
Store::FoundationDb(store) => store.delete_blob(key).await,
|
||||
#[cfg(feature = "postgres")]
|
||||
Store::PostgreSQL(store) => store.delete_blob(key).await,
|
||||
#[cfg(feature = "mysql")]
|
||||
Store::MySQL(store) => store.delete_blob(key).await,
|
||||
#[cfg(feature = "rocks")]
|
||||
Store::RocksDb(store) => store.delete_blob(key).await,
|
||||
Store::Ephemeral(store) => store.delete_blob(key).await,
|
||||
Store::None => Err(trc::StoreEvent::NotConfigured.into()),
|
||||
},
|
||||
BlobStore::Fs(store) => store.delete_blob(key).await,
|
||||
#[cfg(feature = "s3")]
|
||||
BlobStore::S3(store) => store.delete_blob(key).await,
|
||||
#[cfg(feature = "azure")]
|
||||
BlobStore::Azure(store) => store.delete_blob(key).await,
|
||||
}
|
||||
.caused_by(trc::location!());
|
||||
|
||||
trc::event!(
|
||||
Store(StoreEvent::BlobWrite),
|
||||
Key = key,
|
||||
Elapsed = start_time.elapsed(),
|
||||
);
|
||||
|
||||
result
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,625 @@
|
||||
/*
|
||||
* SPDX-FileCopyrightText: 2020 Stalwart Labs LLC <[email protected]>
|
||||
*
|
||||
* SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL
|
||||
*/
|
||||
|
||||
use registry::schema::structs::Rate;
|
||||
use std::borrow::Cow;
|
||||
use trc::AddContext;
|
||||
|
||||
#[allow(unused_imports)]
|
||||
use crate::{
|
||||
Deserialize, InMemoryStore, IterateParams, QueryResult, Store, U64_LEN, Value, ValueKey,
|
||||
write::{
|
||||
BatchBuilder, Operation, ValueClass, ValueOp,
|
||||
key::{DeserializeBigEndian, KeySerializer},
|
||||
now,
|
||||
},
|
||||
};
|
||||
use crate::{
|
||||
SerializeInfallible,
|
||||
backend::{http::lookup::HttpStoreGet, memory::StaticMemoryStore},
|
||||
write::{InMemoryClass, assert::AssertValue},
|
||||
};
|
||||
|
||||
pub struct KeyValue<T> {
|
||||
pub key: Vec<u8>,
|
||||
pub value: T,
|
||||
pub expires: Option<u64>,
|
||||
}
|
||||
|
||||
impl InMemoryStore {
|
||||
pub async fn key_set(&self, kv: KeyValue<Vec<u8>>) -> trc::Result<()> {
|
||||
match self {
|
||||
InMemoryStore::Store(store) => {
|
||||
let mut batch = BatchBuilder::new();
|
||||
batch.any_op(Operation::Value {
|
||||
class: ValueClass::InMemory(InMemoryClass::Key(kv.key)),
|
||||
op: ValueOp::Set(
|
||||
KeySerializer::new(kv.value.len() + U64_LEN)
|
||||
.write(kv.expires.map_or(u64::MAX, |expires| now() + expires))
|
||||
.write(kv.value.as_slice())
|
||||
.finalize(),
|
||||
),
|
||||
});
|
||||
store.write(batch.build_all()).await.map(|_| ())
|
||||
}
|
||||
#[cfg(feature = "redis")]
|
||||
InMemoryStore::Redis(store) => store.key_set(&kv.key, &kv.value, kv.expires).await,
|
||||
InMemoryStore::Static(_) | InMemoryStore::Http(_) => {
|
||||
Err(trc::StoreEvent::NotSupported.into_err())
|
||||
}
|
||||
}
|
||||
.caused_by(trc::location!())
|
||||
}
|
||||
|
||||
pub async fn counter_incr(&self, kv: KeyValue<i64>, return_value: bool) -> trc::Result<i64> {
|
||||
match self {
|
||||
InMemoryStore::Store(store) => {
|
||||
let mut batch = BatchBuilder::new();
|
||||
|
||||
if let Some(expires) = kv.expires {
|
||||
batch.any_op(Operation::Value {
|
||||
class: ValueClass::InMemory(InMemoryClass::Key(kv.key.clone())),
|
||||
op: ValueOp::Set(
|
||||
KeySerializer::new(U64_LEN * 2)
|
||||
.write(0u64)
|
||||
.write(now() + expires)
|
||||
.finalize(),
|
||||
),
|
||||
});
|
||||
}
|
||||
|
||||
if return_value {
|
||||
batch.any_op(Operation::Value {
|
||||
class: ValueClass::InMemory(InMemoryClass::Counter(kv.key)),
|
||||
op: ValueOp::AddAndGet(kv.value),
|
||||
});
|
||||
|
||||
store
|
||||
.write(batch.build_all())
|
||||
.await
|
||||
.and_then(|r| r.last_counter_id())
|
||||
} else {
|
||||
batch.any_op(Operation::Value {
|
||||
class: ValueClass::InMemory(InMemoryClass::Counter(kv.key)),
|
||||
op: ValueOp::AtomicAdd(kv.value),
|
||||
});
|
||||
|
||||
store.write(batch.build_all()).await.map(|_| 0)
|
||||
}
|
||||
}
|
||||
#[cfg(feature = "redis")]
|
||||
InMemoryStore::Redis(store) => store.key_incr(&kv.key, kv.value, kv.expires).await,
|
||||
InMemoryStore::Static(_) | InMemoryStore::Http(_) => {
|
||||
Err(trc::StoreEvent::NotSupported.into_err())
|
||||
}
|
||||
}
|
||||
.caused_by(trc::location!())
|
||||
}
|
||||
|
||||
pub async fn key_delete(&self, key: impl Into<LookupKey<'_>>) -> trc::Result<()> {
|
||||
match self {
|
||||
InMemoryStore::Store(store) => {
|
||||
let mut batch = BatchBuilder::new();
|
||||
batch.any_op(Operation::Value {
|
||||
class: ValueClass::InMemory(InMemoryClass::Key(key.into().into_bytes())),
|
||||
op: ValueOp::Clear,
|
||||
});
|
||||
store.write(batch.build_all()).await.map(|_| ())
|
||||
}
|
||||
#[cfg(feature = "redis")]
|
||||
InMemoryStore::Redis(store) => store.key_delete(key.into().as_bytes()).await,
|
||||
InMemoryStore::Static(_) | InMemoryStore::Http(_) => {
|
||||
Err(trc::StoreEvent::NotSupported.into_err())
|
||||
}
|
||||
}
|
||||
.caused_by(trc::location!())
|
||||
}
|
||||
|
||||
pub async fn counter_delete(&self, key: impl Into<LookupKey<'_>>) -> trc::Result<()> {
|
||||
match self {
|
||||
InMemoryStore::Store(store) => {
|
||||
let mut batch = BatchBuilder::new();
|
||||
batch.any_op(Operation::Value {
|
||||
class: ValueClass::InMemory(InMemoryClass::Counter(key.into().into_bytes())),
|
||||
op: ValueOp::Clear,
|
||||
});
|
||||
store.write(batch.build_all()).await.map(|_| ())
|
||||
}
|
||||
#[cfg(feature = "redis")]
|
||||
InMemoryStore::Redis(store) => store.key_delete(key.into().as_bytes()).await,
|
||||
InMemoryStore::Static(_) | InMemoryStore::Http(_) => {
|
||||
Err(trc::StoreEvent::NotSupported.into_err())
|
||||
}
|
||||
}
|
||||
.caused_by(trc::location!())
|
||||
}
|
||||
|
||||
pub async fn key_delete_prefix(&self, prefix: &[u8]) -> trc::Result<()> {
|
||||
match self {
|
||||
InMemoryStore::Store(store) => {
|
||||
if prefix.is_empty() {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let from_range = prefix.to_vec();
|
||||
let mut to_range = Vec::with_capacity(prefix.len() + 3);
|
||||
to_range.extend_from_slice(prefix);
|
||||
to_range.extend_from_slice([u8::MAX, u8::MAX, u8::MAX].as_ref());
|
||||
|
||||
store
|
||||
.delete_range(
|
||||
ValueKey::from(ValueClass::InMemory(InMemoryClass::Counter(
|
||||
from_range.clone(),
|
||||
))),
|
||||
ValueKey::from(ValueClass::InMemory(InMemoryClass::Counter(
|
||||
to_range.clone(),
|
||||
))),
|
||||
)
|
||||
.await?;
|
||||
|
||||
store
|
||||
.delete_range(
|
||||
ValueKey::from(ValueClass::InMemory(InMemoryClass::Key(from_range))),
|
||||
ValueKey::from(ValueClass::InMemory(InMemoryClass::Key(to_range))),
|
||||
)
|
||||
.await
|
||||
}
|
||||
#[cfg(feature = "redis")]
|
||||
InMemoryStore::Redis(store) => store.key_delete_prefix(prefix).await,
|
||||
InMemoryStore::Static(_) | InMemoryStore::Http(_) => {
|
||||
Err(trc::StoreEvent::NotSupported.into_err())
|
||||
}
|
||||
}
|
||||
.caused_by(trc::location!())
|
||||
}
|
||||
|
||||
pub async fn key_get<T: Deserialize + From<Value<'static>> + std::fmt::Debug + 'static>(
|
||||
&self,
|
||||
key: impl Into<LookupKey<'_>>,
|
||||
) -> trc::Result<Option<T>> {
|
||||
match self {
|
||||
InMemoryStore::Store(store) => store
|
||||
.get_value::<LookupValue<T>>(ValueKey::from(ValueClass::InMemory(
|
||||
InMemoryClass::Key(key.into().into_bytes()),
|
||||
)))
|
||||
.await
|
||||
.map(|value| value.and_then(|v| v.into())),
|
||||
#[cfg(feature = "redis")]
|
||||
InMemoryStore::Redis(store) => store.key_get(key.into().as_bytes()).await,
|
||||
InMemoryStore::Static(store) => Ok(match store.as_ref() {
|
||||
StaticMemoryStore::Map(map) => map
|
||||
.get(key.into().as_str())
|
||||
.map(|value| T::from(value.clone())),
|
||||
StaticMemoryStore::Set(set) => {
|
||||
if set.contains(key.into().as_str()) {
|
||||
Some(T::from(Value::Bool(true)))
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
}),
|
||||
InMemoryStore::Http(store) => {
|
||||
Ok(store.get(key.into().as_str()).map(|value| T::from(value)))
|
||||
}
|
||||
}
|
||||
.caused_by(trc::location!())
|
||||
}
|
||||
|
||||
pub async fn counter_get(&self, key: impl Into<LookupKey<'_>>) -> trc::Result<i64> {
|
||||
match self {
|
||||
InMemoryStore::Store(store) => {
|
||||
store
|
||||
.get_counter(ValueKey::from(ValueClass::InMemory(
|
||||
InMemoryClass::Counter(key.into().into_bytes()),
|
||||
)))
|
||||
.await
|
||||
}
|
||||
#[cfg(feature = "redis")]
|
||||
InMemoryStore::Redis(store) => store.counter_get(key.into().as_bytes()).await,
|
||||
InMemoryStore::Static(_) | InMemoryStore::Http(_) => {
|
||||
Err(trc::StoreEvent::NotSupported.into_err())
|
||||
}
|
||||
}
|
||||
.caused_by(trc::location!())
|
||||
}
|
||||
|
||||
pub async fn key_exists(&self, key: impl Into<LookupKey<'_>>) -> trc::Result<bool> {
|
||||
match self {
|
||||
InMemoryStore::Store(store) => store
|
||||
.get_value::<LookupValue<Empty>>(ValueKey::from(ValueClass::InMemory(
|
||||
InMemoryClass::Key(key.into().into_bytes()),
|
||||
)))
|
||||
.await
|
||||
.map(|value| matches!(value, Some(LookupValue::Value(Empty)))),
|
||||
#[cfg(feature = "redis")]
|
||||
InMemoryStore::Redis(store) => store.key_exists(key.into().as_bytes()).await,
|
||||
InMemoryStore::Static(store) => Ok(match store.as_ref() {
|
||||
StaticMemoryStore::Map(map) => map.get(key.into().as_str()).is_some(),
|
||||
StaticMemoryStore::Set(set) => set.contains(key.into().as_str()),
|
||||
}),
|
||||
InMemoryStore::Http(store) => Ok(store.contains(key.into().as_str())),
|
||||
}
|
||||
.caused_by(trc::location!())
|
||||
}
|
||||
|
||||
pub async fn is_rate_allowed(
|
||||
&self,
|
||||
prefix: u8,
|
||||
key: &[u8],
|
||||
rate: &Rate,
|
||||
soft_check: bool,
|
||||
) -> trc::Result<Option<u64>> {
|
||||
let now = now();
|
||||
let period = rate.period.as_secs().max(1);
|
||||
let range_start = now / period;
|
||||
let range_end = (range_start * period) + period;
|
||||
let expires_in = range_end - now;
|
||||
|
||||
let mut bucket = Vec::with_capacity(key.len() + U64_LEN + 1);
|
||||
bucket.push(prefix);
|
||||
bucket.extend_from_slice(key);
|
||||
bucket.extend_from_slice(range_start.to_be_bytes().as_slice());
|
||||
|
||||
let requests = if !soft_check {
|
||||
self.counter_incr(KeyValue::new(bucket, 1).expires(expires_in), true)
|
||||
.await
|
||||
.caused_by(trc::location!())?
|
||||
} else {
|
||||
self.counter_get(bucket).await.caused_by(trc::location!())? + 1
|
||||
};
|
||||
|
||||
if requests <= rate.count as i64 {
|
||||
Ok(None)
|
||||
} else {
|
||||
Ok(Some(expires_in))
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn try_lock(&self, prefix: u8, key: &[u8], duration: u64) -> trc::Result<bool> {
|
||||
match self {
|
||||
InMemoryStore::Store(store) => {
|
||||
let key = KeyValue::<()>::build_key(prefix, key);
|
||||
let lock_expiry = match store
|
||||
.get_value::<u64>(ValueKey::from(ValueClass::InMemory(InMemoryClass::Key(
|
||||
key.clone(),
|
||||
))))
|
||||
.await
|
||||
{
|
||||
Ok(lock_expiry) => lock_expiry,
|
||||
Err(err)
|
||||
if err.matches(trc::EventType::Store(trc::StoreEvent::DataCorruption)) =>
|
||||
{
|
||||
// TODO remove in 1.0
|
||||
let mut batch = BatchBuilder::new();
|
||||
batch.any_op(Operation::Value {
|
||||
class: ValueClass::InMemory(InMemoryClass::Key(key.clone())),
|
||||
op: ValueOp::Clear,
|
||||
});
|
||||
store
|
||||
.write(batch.build_all())
|
||||
.await
|
||||
.caused_by(trc::location!())?;
|
||||
None
|
||||
}
|
||||
Err(err) => {
|
||||
return Err(err
|
||||
.details("Failed to read lock.")
|
||||
.caused_by(trc::location!()));
|
||||
}
|
||||
};
|
||||
|
||||
let now = now();
|
||||
if lock_expiry.is_some_and(|expiry| expiry > now) {
|
||||
return Ok(false);
|
||||
}
|
||||
|
||||
let key: ValueClass = ValueClass::InMemory(InMemoryClass::Key(key));
|
||||
let mut batch = BatchBuilder::new();
|
||||
batch.assert_value(
|
||||
key.clone(),
|
||||
match lock_expiry {
|
||||
Some(value) => AssertValue::U64(value),
|
||||
None => AssertValue::None,
|
||||
},
|
||||
);
|
||||
batch.set(key.clone(), (now + duration).serialize());
|
||||
match store.write(batch.build_all()).await {
|
||||
Ok(_) => Ok(true),
|
||||
Err(err) if err.is_assertion_failure() => Ok(false),
|
||||
Err(err) => Err(err
|
||||
.details("Failed to lock event.")
|
||||
.caused_by(trc::location!())),
|
||||
}
|
||||
}
|
||||
#[cfg(feature = "redis")]
|
||||
InMemoryStore::Redis(store) => {
|
||||
store
|
||||
.try_lock(&KeyValue::<()>::build_key(prefix, key), duration)
|
||||
.await
|
||||
}
|
||||
InMemoryStore::Static(_) | InMemoryStore::Http(_) => {
|
||||
Err(trc::StoreEvent::NotSupported.into_err())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn remove_lock(&self, prefix: u8, key: &[u8]) -> trc::Result<()> {
|
||||
self.key_delete(KeyValue::<()>::build_key(prefix, key))
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn purge_in_memory_store(&self) -> trc::Result<()> {
|
||||
match self {
|
||||
InMemoryStore::Store(store) => {
|
||||
// Delete expired keys and counters
|
||||
let from_key = ValueKey::from(ValueClass::InMemory(InMemoryClass::Key(vec![0u8])));
|
||||
let to_key =
|
||||
ValueKey::from(ValueClass::InMemory(InMemoryClass::Key(vec![u8::MAX; 10])));
|
||||
|
||||
let current_time = now();
|
||||
let mut expired_keys = Vec::new();
|
||||
let mut expired_counters = Vec::new();
|
||||
store
|
||||
.iterate(IterateParams::new(from_key, to_key), |key, value| {
|
||||
let expiry = value.deserialize_be_u64(0).caused_by(trc::location!())?;
|
||||
if expiry == 0 {
|
||||
if value
|
||||
.deserialize_be_u64(U64_LEN)
|
||||
.caused_by(trc::location!())?
|
||||
<= current_time
|
||||
{
|
||||
expired_counters.push(key.to_vec());
|
||||
}
|
||||
} else if expiry <= current_time {
|
||||
expired_keys.push(key.to_vec());
|
||||
}
|
||||
Ok(true)
|
||||
})
|
||||
.await
|
||||
.caused_by(trc::location!())?;
|
||||
|
||||
if !expired_keys.is_empty() {
|
||||
let mut batch = BatchBuilder::new();
|
||||
for key in expired_keys {
|
||||
batch.any_op(Operation::Value {
|
||||
class: ValueClass::InMemory(InMemoryClass::Key(key)),
|
||||
op: ValueOp::Clear,
|
||||
});
|
||||
if batch.is_large_batch() {
|
||||
store
|
||||
.write(batch.build_all())
|
||||
.await
|
||||
.caused_by(trc::location!())?;
|
||||
batch = BatchBuilder::new();
|
||||
}
|
||||
}
|
||||
if !batch.is_empty() {
|
||||
store
|
||||
.write(batch.build_all())
|
||||
.await
|
||||
.caused_by(trc::location!())?;
|
||||
}
|
||||
}
|
||||
|
||||
if !expired_counters.is_empty() {
|
||||
let mut batch = BatchBuilder::new();
|
||||
for key in expired_counters {
|
||||
batch.any_op(Operation::Value {
|
||||
class: ValueClass::InMemory(InMemoryClass::Counter(key.clone())),
|
||||
op: ValueOp::Clear,
|
||||
});
|
||||
batch.any_op(Operation::Value {
|
||||
class: ValueClass::InMemory(InMemoryClass::Key(key)),
|
||||
op: ValueOp::Clear,
|
||||
});
|
||||
if batch.is_large_batch() {
|
||||
store
|
||||
.write(batch.build_all())
|
||||
.await
|
||||
.caused_by(trc::location!())?;
|
||||
batch = BatchBuilder::new();
|
||||
}
|
||||
}
|
||||
if !batch.is_empty() {
|
||||
store
|
||||
.write(batch.build_all())
|
||||
.await
|
||||
.caused_by(trc::location!())?;
|
||||
}
|
||||
}
|
||||
}
|
||||
#[cfg(feature = "redis")]
|
||||
InMemoryStore::Redis(_) => {}
|
||||
InMemoryStore::Static(_) | InMemoryStore::Http(_) => {}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn is_sql(&self) -> bool {
|
||||
match self {
|
||||
InMemoryStore::Store(store) => store.is_sql(),
|
||||
_ => false,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn is_redis(&self) -> bool {
|
||||
match self {
|
||||
#[cfg(feature = "redis")]
|
||||
InMemoryStore::Redis(_) => true,
|
||||
InMemoryStore::Static(_) => false,
|
||||
_ => false,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn into_store(self) -> Option<Store> {
|
||||
match self {
|
||||
InMemoryStore::Store(store) => Some(store),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub enum LookupKey<'x> {
|
||||
String(String),
|
||||
StringRef(&'x str),
|
||||
Bytes(Vec<u8>),
|
||||
BytesRef(&'x [u8]),
|
||||
}
|
||||
|
||||
impl<'x> From<&'x str> for LookupKey<'x> {
|
||||
fn from(key: &'x str) -> Self {
|
||||
LookupKey::StringRef(key)
|
||||
}
|
||||
}
|
||||
|
||||
impl<'x> From<&'x String> for LookupKey<'x> {
|
||||
fn from(key: &'x String) -> Self {
|
||||
LookupKey::StringRef(key.as_str())
|
||||
}
|
||||
}
|
||||
|
||||
impl<'x> From<&'x [u8]> for LookupKey<'x> {
|
||||
fn from(key: &'x [u8]) -> Self {
|
||||
LookupKey::BytesRef(key)
|
||||
}
|
||||
}
|
||||
|
||||
impl<'x> From<Cow<'x, str>> for LookupKey<'x> {
|
||||
fn from(key: Cow<'x, str>) -> Self {
|
||||
match key {
|
||||
Cow::Borrowed(key) => LookupKey::StringRef(key),
|
||||
Cow::Owned(key) => LookupKey::String(key),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl From<String> for LookupKey<'static> {
|
||||
fn from(key: String) -> Self {
|
||||
LookupKey::String(key)
|
||||
}
|
||||
}
|
||||
|
||||
impl From<Vec<u8>> for LookupKey<'static> {
|
||||
fn from(key: Vec<u8>) -> Self {
|
||||
LookupKey::Bytes(key)
|
||||
}
|
||||
}
|
||||
|
||||
impl LookupKey<'_> {
|
||||
pub fn as_str(&self) -> &str {
|
||||
match self {
|
||||
LookupKey::String(string) => string,
|
||||
LookupKey::StringRef(string) => string,
|
||||
LookupKey::Bytes(bytes) => std::str::from_utf8(bytes).unwrap_or_default(),
|
||||
LookupKey::BytesRef(bytes) => std::str::from_utf8(bytes).unwrap_or_default(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn into_bytes(self) -> Vec<u8> {
|
||||
match self {
|
||||
LookupKey::String(string) => string.into_bytes(),
|
||||
LookupKey::StringRef(string) => string.as_bytes().to_vec(),
|
||||
LookupKey::Bytes(bytes) => bytes,
|
||||
LookupKey::BytesRef(bytes) => bytes.to_vec(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn as_bytes(&self) -> &[u8] {
|
||||
match self {
|
||||
LookupKey::String(string) => string.as_bytes(),
|
||||
LookupKey::StringRef(string) => string.as_bytes(),
|
||||
LookupKey::Bytes(bytes) => bytes.as_slice(),
|
||||
LookupKey::BytesRef(bytes) => bytes,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<T> KeyValue<T> {
|
||||
pub fn build_key(prefix: u8, key: impl AsRef<[u8]>) -> Vec<u8> {
|
||||
let key_ = key.as_ref();
|
||||
let mut key = Vec::with_capacity(key_.len() + 1);
|
||||
key.push(prefix);
|
||||
key.extend_from_slice(key_);
|
||||
key
|
||||
}
|
||||
|
||||
pub fn with_prefix(prefix: u8, key: impl AsRef<[u8]>, value: T) -> Self {
|
||||
Self {
|
||||
key: Self::build_key(prefix, key),
|
||||
value,
|
||||
expires: None,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn new(key: impl Into<Vec<u8>>, value: T) -> Self {
|
||||
Self {
|
||||
key: key.into(),
|
||||
value,
|
||||
expires: None,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn expires(mut self, expires: u64) -> Self {
|
||||
self.expires = expires.into();
|
||||
self
|
||||
}
|
||||
|
||||
pub fn expires_opt(mut self, expires: Option<u64>) -> Self {
|
||||
self.expires = expires;
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
struct Empty;
|
||||
|
||||
enum LookupValue<T> {
|
||||
Value(T),
|
||||
None,
|
||||
}
|
||||
|
||||
impl<T: Deserialize> Deserialize for LookupValue<T> {
|
||||
fn deserialize(bytes: &[u8]) -> trc::Result<Self> {
|
||||
bytes.deserialize_be_u64(0).and_then(|expires| {
|
||||
Ok(if expires > now() {
|
||||
LookupValue::Value(
|
||||
T::deserialize(bytes.get(U64_LEN..).unwrap_or_default())
|
||||
.caused_by(trc::location!())?,
|
||||
)
|
||||
} else {
|
||||
LookupValue::None
|
||||
})
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl Deserialize for Empty {
|
||||
fn deserialize(_bytes: &[u8]) -> trc::Result<Self> {
|
||||
Ok(Empty)
|
||||
}
|
||||
}
|
||||
|
||||
impl<T> From<LookupValue<T>> for Option<T> {
|
||||
fn from(value: LookupValue<T>) -> Self {
|
||||
match value {
|
||||
LookupValue::Value(value) => Some(value),
|
||||
LookupValue::None => None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl From<Value<'static>> for String {
|
||||
fn from(value: Value<'static>) -> Self {
|
||||
match value {
|
||||
Value::Text(string) => string.into_owned(),
|
||||
Value::Blob(bytes) => String::from_utf8_lossy(bytes.as_ref()).into_owned(),
|
||||
Value::Bool(boolean) => boolean.to_string(),
|
||||
Value::Null => String::new(),
|
||||
Value::Integer(num) => num.to_string(),
|
||||
Value::Float(num) => num.to_string(),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,107 @@
|
||||
/*
|
||||
* SPDX-FileCopyrightText: 2020 Stalwart Labs LLC <[email protected]>
|
||||
*
|
||||
* SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL
|
||||
*/
|
||||
|
||||
use crate::Store;
|
||||
use roaring::RoaringBitmap;
|
||||
|
||||
pub mod blob;
|
||||
pub mod lookup;
|
||||
pub mod search;
|
||||
pub mod store;
|
||||
|
||||
impl Store {
|
||||
pub fn id(&self) -> &'static str {
|
||||
match self {
|
||||
#[cfg(feature = "sqlite")]
|
||||
Self::SQLite(_) => "sqlite",
|
||||
#[cfg(feature = "foundation")]
|
||||
Self::FoundationDb(_) => "foundationdb",
|
||||
#[cfg(feature = "postgres")]
|
||||
Self::PostgreSQL(_) => "postgresql",
|
||||
#[cfg(feature = "mysql")]
|
||||
Self::MySQL(_) => "mysql",
|
||||
#[cfg(feature = "rocks")]
|
||||
Self::RocksDb(_) => "rocksdb",
|
||||
Self::Ephemeral(_) => "ephemeral",
|
||||
Self::None => "none",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[allow(clippy::len_without_is_empty)]
|
||||
pub trait DocumentSet: Sync + Send {
|
||||
fn min(&self) -> u32;
|
||||
fn max(&self) -> u32;
|
||||
fn contains(&self, id: u32) -> bool;
|
||||
fn len(&self) -> usize;
|
||||
fn iterate(&self) -> impl Iterator<Item = u32>;
|
||||
}
|
||||
|
||||
impl DocumentSet for RoaringBitmap {
|
||||
fn min(&self) -> u32 {
|
||||
self.min().unwrap_or(0)
|
||||
}
|
||||
|
||||
fn max(&self) -> u32 {
|
||||
self.max().map(|m| m + 1).unwrap_or(0)
|
||||
}
|
||||
|
||||
fn contains(&self, id: u32) -> bool {
|
||||
self.contains(id)
|
||||
}
|
||||
|
||||
fn len(&self) -> usize {
|
||||
self.len() as usize
|
||||
}
|
||||
|
||||
fn iterate(&self) -> impl Iterator<Item = u32> {
|
||||
self.iter()
|
||||
}
|
||||
}
|
||||
|
||||
impl DocumentSet for Vec<u32> {
|
||||
fn contains(&self, id: u32) -> bool {
|
||||
self.binary_search(&id).is_ok()
|
||||
}
|
||||
|
||||
fn min(&self) -> u32 {
|
||||
self.first().copied().unwrap_or(0)
|
||||
}
|
||||
|
||||
fn max(&self) -> u32 {
|
||||
self.last().copied().map(|m| m + 1).unwrap_or(0)
|
||||
}
|
||||
|
||||
fn len(&self) -> usize {
|
||||
self.len()
|
||||
}
|
||||
|
||||
fn iterate(&self) -> impl Iterator<Item = u32> {
|
||||
self.iter().copied()
|
||||
}
|
||||
}
|
||||
|
||||
impl DocumentSet for () {
|
||||
fn min(&self) -> u32 {
|
||||
0
|
||||
}
|
||||
|
||||
fn max(&self) -> u32 {
|
||||
u32::MAX
|
||||
}
|
||||
|
||||
fn contains(&self, _: u32) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
fn len(&self) -> usize {
|
||||
0
|
||||
}
|
||||
|
||||
fn iterate(&self) -> impl Iterator<Item = u32> {
|
||||
std::iter::empty()
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,337 @@
|
||||
/*
|
||||
* SPDX-FileCopyrightText: 2020 Stalwart Labs LLC <[email protected]>
|
||||
*
|
||||
* SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL
|
||||
*/
|
||||
|
||||
use crate::{
|
||||
SearchStore, Store,
|
||||
search::{
|
||||
IndexDocument, SearchComparator, SearchField, SearchFilter, SearchOperator, SearchQuery,
|
||||
SearchValue,
|
||||
split::{SplitFilter, split_filters},
|
||||
},
|
||||
write::SearchIndex,
|
||||
};
|
||||
use std::cmp::Ordering;
|
||||
use trc::AddContext;
|
||||
|
||||
impl SearchStore {
|
||||
pub async fn query_account(&self, query: SearchQuery) -> trc::Result<Vec<u32>> {
|
||||
// Pre-filter by mask
|
||||
if query.mask.is_empty() {
|
||||
return Ok(vec![]);
|
||||
}
|
||||
|
||||
// If the store does not support FTS, use the internal FTS store
|
||||
if let Some(store) = self.internal_fts() {
|
||||
return store.query_account(query).await;
|
||||
}
|
||||
|
||||
// If all filters and comparators are external, delegate to the underlying store
|
||||
let mut account_id = u32::MAX;
|
||||
let mut has_local_filters = false;
|
||||
let mut has_external_filters = false;
|
||||
for filter in &query.filters {
|
||||
match filter {
|
||||
SearchFilter::Operator {
|
||||
field: SearchField::AccountId,
|
||||
op: SearchOperator::Equal,
|
||||
value: SearchValue::Uint(id),
|
||||
} => {
|
||||
account_id = *id as u32;
|
||||
}
|
||||
SearchFilter::DocumentSet(_) => {
|
||||
has_local_filters = true;
|
||||
}
|
||||
SearchFilter::Operator { .. } => {
|
||||
has_external_filters = true;
|
||||
}
|
||||
_ => (),
|
||||
}
|
||||
}
|
||||
|
||||
if account_id == u32::MAX {
|
||||
return Err(trc::StoreEvent::UnexpectedError
|
||||
.reason("Account ID filter is required for account queries")
|
||||
.caused_by(trc::location!()));
|
||||
}
|
||||
|
||||
if !has_local_filters && !has_external_filters && query.comparators.is_empty() {
|
||||
return Ok(query.mask.iter().collect());
|
||||
}
|
||||
|
||||
if !has_local_filters && query.comparators.iter().all(|c| c.is_external()) {
|
||||
return self
|
||||
.sub_query(query.index, &query.filters, &query.comparators)
|
||||
.await
|
||||
.map(|results| {
|
||||
if !results.is_empty() || has_external_filters {
|
||||
results
|
||||
.into_iter()
|
||||
.filter(|id| query.mask.contains(*id))
|
||||
.collect()
|
||||
} else {
|
||||
// Database sort is broken, return masked results
|
||||
query.mask.iter().collect()
|
||||
}
|
||||
})
|
||||
.caused_by(trc::location!());
|
||||
}
|
||||
|
||||
let filters = if has_external_filters {
|
||||
// Split filters
|
||||
let split_filters = split_filters(query.filters).ok_or_else(|| {
|
||||
trc::StoreEvent::UnexpectedError
|
||||
.reason("Invalid filter query")
|
||||
.caused_by(trc::location!())
|
||||
})?;
|
||||
|
||||
let mut filters = Vec::with_capacity(split_filters.len());
|
||||
for split_filter in split_filters {
|
||||
match split_filter {
|
||||
SplitFilter::External(external) => {
|
||||
// Execute sub-query
|
||||
filters.push(SearchFilter::DocumentSet(
|
||||
self.sub_query(query.index, &external, &[])
|
||||
.await?
|
||||
.into_iter()
|
||||
.collect(),
|
||||
));
|
||||
}
|
||||
SplitFilter::Internal(filter) => {
|
||||
filters.push(filter);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
filters
|
||||
} else {
|
||||
query.filters
|
||||
};
|
||||
|
||||
// Merge results locally
|
||||
let results = SearchQuery::new(query.index)
|
||||
.with_filters(filters)
|
||||
.with_mask(query.mask)
|
||||
.filter();
|
||||
|
||||
let total_results = results.results().len();
|
||||
match total_results.cmp(&1) {
|
||||
Ordering::Equal => Ok(vec![results.results().min().unwrap()]),
|
||||
Ordering::Less => Ok(vec![]),
|
||||
Ordering::Greater => {
|
||||
if !query.comparators.is_empty() {
|
||||
let mut local = Vec::with_capacity(query.comparators.len());
|
||||
let mut external = Vec::with_capacity(query.comparators.len());
|
||||
let mut external_first = false;
|
||||
for (pos, comparator) in query.comparators.into_iter().enumerate() {
|
||||
if comparator.is_external() {
|
||||
external.push(comparator);
|
||||
if pos == 0 {
|
||||
external_first = true;
|
||||
}
|
||||
} else {
|
||||
local.push(comparator);
|
||||
}
|
||||
}
|
||||
|
||||
if !external.is_empty() {
|
||||
let mut results = results.results().clone();
|
||||
let filters = vec![
|
||||
SearchFilter::Operator {
|
||||
field: SearchField::AccountId,
|
||||
op: SearchOperator::Equal,
|
||||
value: SearchValue::Uint(account_id as u64),
|
||||
},
|
||||
SearchFilter::Operator {
|
||||
field: SearchField::DocumentId,
|
||||
op: SearchOperator::GreaterEqualThan,
|
||||
value: SearchValue::Uint(results.min().unwrap() as u64),
|
||||
},
|
||||
SearchFilter::Operator {
|
||||
field: SearchField::DocumentId,
|
||||
op: SearchOperator::LowerEqualThan,
|
||||
value: SearchValue::Uint(results.max().unwrap() as u64),
|
||||
},
|
||||
];
|
||||
|
||||
let mut ordered_results = Vec::with_capacity(total_results as usize);
|
||||
for ordered_result in
|
||||
self.sub_query(query.index, &filters, &external).await?
|
||||
{
|
||||
if results.remove(ordered_result) {
|
||||
ordered_results.push(ordered_result);
|
||||
}
|
||||
}
|
||||
// Add any remaining results not yet in the index
|
||||
ordered_results.extend(results);
|
||||
|
||||
if local.is_empty() {
|
||||
return Ok(ordered_results);
|
||||
}
|
||||
|
||||
let comparator = SearchComparator::SortedSet {
|
||||
set: ordered_results
|
||||
.into_iter()
|
||||
.enumerate()
|
||||
.map(|(pos, id)| (id, pos as u32))
|
||||
.collect(),
|
||||
ascending: true,
|
||||
};
|
||||
|
||||
if external_first {
|
||||
local.insert(0, comparator);
|
||||
} else {
|
||||
local.push(comparator);
|
||||
}
|
||||
}
|
||||
|
||||
Ok(results.with_comparators(local).into_sorted())
|
||||
} else {
|
||||
Ok(results.results().iter().collect())
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn sub_query(
|
||||
&self,
|
||||
index: SearchIndex,
|
||||
filters: &[SearchFilter],
|
||||
sort: &[SearchComparator],
|
||||
) -> trc::Result<Vec<u32>> {
|
||||
match self {
|
||||
SearchStore::Store(store) => match store {
|
||||
#[cfg(feature = "postgres")]
|
||||
Store::PostgreSQL(store) => store.query(index, filters, sort).await,
|
||||
#[cfg(feature = "mysql")]
|
||||
Store::MySQL(store) => store.query(index, filters, sort).await,
|
||||
_ => unreachable!(),
|
||||
},
|
||||
SearchStore::ElasticSearch(store) => store.query(index, filters, sort).await,
|
||||
SearchStore::MeiliSearch(store) => store.query(index, filters, sort).await,
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn query_global(&self, query: SearchQuery) -> trc::Result<Vec<u64>> {
|
||||
match self {
|
||||
SearchStore::Store(store) => match store {
|
||||
#[cfg(feature = "postgres")]
|
||||
Store::PostgreSQL(store) => {
|
||||
store
|
||||
.query(query.index, &query.filters, &query.comparators)
|
||||
.await
|
||||
}
|
||||
#[cfg(feature = "mysql")]
|
||||
Store::MySQL(store) => {
|
||||
store
|
||||
.query(query.index, &query.filters, &query.comparators)
|
||||
.await
|
||||
}
|
||||
store => store.query_global(query).await,
|
||||
},
|
||||
SearchStore::ElasticSearch(store) => {
|
||||
store
|
||||
.query(query.index, &query.filters, &query.comparators)
|
||||
.await
|
||||
}
|
||||
SearchStore::MeiliSearch(store) => {
|
||||
store
|
||||
.query(query.index, &query.filters, &query.comparators)
|
||||
.await
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn index(&self, documents: Vec<IndexDocument>) -> trc::Result<()> {
|
||||
match self {
|
||||
SearchStore::Store(store) => match store {
|
||||
#[cfg(feature = "postgres")]
|
||||
Store::PostgreSQL(store) => store.index(documents).await,
|
||||
#[cfg(feature = "mysql")]
|
||||
Store::MySQL(store) => store.index(documents).await,
|
||||
store => store.index(documents).await,
|
||||
},
|
||||
SearchStore::ElasticSearch(store) => store.index(documents).await,
|
||||
SearchStore::MeiliSearch(store) => store.index(documents).await,
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn unindex(&self, query: SearchQuery) -> trc::Result<u64> {
|
||||
match self {
|
||||
SearchStore::Store(store) => match store {
|
||||
#[cfg(feature = "postgres")]
|
||||
Store::PostgreSQL(store) => store.unindex(query).await,
|
||||
#[cfg(feature = "mysql")]
|
||||
Store::MySQL(store) => store.unindex(query).await,
|
||||
store => store.unindex(query).await.map(|_| 0),
|
||||
},
|
||||
SearchStore::ElasticSearch(store) => store.unindex(query).await,
|
||||
SearchStore::MeiliSearch(store) => store.unindex(query).await,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn internal_fts(&self) -> Option<&Store> {
|
||||
match self {
|
||||
SearchStore::Store(store) => match store {
|
||||
#[cfg(feature = "postgres")]
|
||||
Store::PostgreSQL(_) => None,
|
||||
#[cfg(feature = "mysql")]
|
||||
Store::MySQL(_) => None,
|
||||
store => Some(store),
|
||||
},
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn is_mysql(&self) -> bool {
|
||||
match self {
|
||||
#[cfg(feature = "mysql")]
|
||||
SearchStore::Store(Store::MySQL(_)) => true,
|
||||
_ => false,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn is_postgres(&self) -> bool {
|
||||
match self {
|
||||
#[cfg(feature = "postgres")]
|
||||
SearchStore::Store(Store::PostgreSQL(_)) => true,
|
||||
_ => false,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn is_elasticsearch(&self) -> bool {
|
||||
matches!(self, SearchStore::ElasticSearch(_))
|
||||
}
|
||||
|
||||
pub fn is_meilisearch(&self) -> bool {
|
||||
matches!(self, SearchStore::MeiliSearch(_))
|
||||
}
|
||||
|
||||
pub async fn create_indexes(&self) -> trc::Result<()> {
|
||||
match self {
|
||||
SearchStore::Store(store) => match store {
|
||||
#[cfg(feature = "postgres")]
|
||||
Store::PostgreSQL(store) => store.create_search_tables().await,
|
||||
#[cfg(feature = "mysql")]
|
||||
Store::MySQL(store) => store.create_search_tables().await,
|
||||
_ => Ok(()),
|
||||
},
|
||||
SearchStore::ElasticSearch(store) => store.create_indexes().await,
|
||||
SearchStore::MeiliSearch(store) => store.create_indexes().await,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl SearchFilter {
|
||||
pub fn is_external(&self) -> bool {
|
||||
matches!(self, SearchFilter::Operator { .. })
|
||||
}
|
||||
}
|
||||
|
||||
impl SearchComparator {
|
||||
pub fn is_external(&self) -> bool {
|
||||
matches!(self, SearchComparator::Field { .. })
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,373 @@
|
||||
/*
|
||||
* SPDX-FileCopyrightText: 2020 Stalwart Labs LLC <[email protected]>
|
||||
*
|
||||
* SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL
|
||||
*/
|
||||
|
||||
use super::DocumentSet;
|
||||
use crate::{
|
||||
Deserialize, IterateParams, Key, QueryResult, SUBSPACE_COUNTER, SUBSPACE_INDEXES,
|
||||
SUBSPACE_LOGS, Store, U32_LEN, Value, ValueKey,
|
||||
write::{
|
||||
AnyClass, AnyKey, AssignedIds, Batch, BatchBuilder, Operation, ValueClass, ValueOp,
|
||||
key::{DeserializeBigEndian, KeySerializer},
|
||||
},
|
||||
};
|
||||
use compact_str::ToCompactString;
|
||||
use std::time::Instant;
|
||||
use trc::{AddContext, StoreEvent};
|
||||
use types::collection::Collection;
|
||||
|
||||
impl Store {
|
||||
pub async fn get_value<U>(&self, key: impl Key) -> trc::Result<Option<U>>
|
||||
where
|
||||
U: Deserialize + 'static,
|
||||
{
|
||||
match self {
|
||||
#[cfg(feature = "sqlite")]
|
||||
Self::SQLite(store) => store.get_value(key).await,
|
||||
#[cfg(feature = "foundation")]
|
||||
Self::FoundationDb(store) => store.get_value(key).await,
|
||||
#[cfg(feature = "postgres")]
|
||||
Self::PostgreSQL(store) => store.get_value(key).await,
|
||||
#[cfg(feature = "mysql")]
|
||||
Self::MySQL(store) => store.get_value(key).await,
|
||||
#[cfg(feature = "rocks")]
|
||||
Self::RocksDb(store) => store.get_value(key).await,
|
||||
Self::Ephemeral(store) => store.get_value(key).await,
|
||||
Self::None => Err(trc::StoreEvent::NotConfigured.into()),
|
||||
}
|
||||
.caused_by(trc::location!())
|
||||
}
|
||||
|
||||
pub async fn key_exists(&self, key: impl Key) -> trc::Result<bool> {
|
||||
match self {
|
||||
#[cfg(feature = "sqlite")]
|
||||
Self::SQLite(store) => store.key_exists(key).await,
|
||||
#[cfg(feature = "foundation")]
|
||||
Self::FoundationDb(store) => store.key_exists(key).await,
|
||||
#[cfg(feature = "postgres")]
|
||||
Self::PostgreSQL(store) => store.key_exists(key).await,
|
||||
#[cfg(feature = "mysql")]
|
||||
Self::MySQL(store) => store.key_exists(key).await,
|
||||
#[cfg(feature = "rocks")]
|
||||
Self::RocksDb(store) => store.key_exists(key).await,
|
||||
Self::Ephemeral(store) => store.key_exists(key).await,
|
||||
Self::None => Err(trc::StoreEvent::NotConfigured.into()),
|
||||
}
|
||||
.caused_by(trc::location!())
|
||||
}
|
||||
|
||||
pub async fn iterate<T: Key>(
|
||||
&self,
|
||||
params: IterateParams<T>,
|
||||
cb: impl for<'x> FnMut(&'x [u8], &'x [u8]) -> trc::Result<bool> + Sync + Send,
|
||||
) -> trc::Result<()> {
|
||||
let start_time = Instant::now();
|
||||
let result = match self {
|
||||
#[cfg(feature = "sqlite")]
|
||||
Self::SQLite(store) => store.iterate(params, cb).await,
|
||||
#[cfg(feature = "foundation")]
|
||||
Self::FoundationDb(store) => store.iterate(params, cb).await,
|
||||
#[cfg(feature = "postgres")]
|
||||
Self::PostgreSQL(store) => store.iterate(params, cb).await,
|
||||
#[cfg(feature = "mysql")]
|
||||
Self::MySQL(store) => store.iterate(params, cb).await,
|
||||
#[cfg(feature = "rocks")]
|
||||
Self::RocksDb(store) => store.iterate(params, cb).await,
|
||||
Self::Ephemeral(store) => store.iterate(params, cb).await,
|
||||
Self::None => Err(trc::StoreEvent::NotConfigured.into()),
|
||||
}
|
||||
.caused_by(trc::location!());
|
||||
|
||||
trc::event!(
|
||||
Store(StoreEvent::DataIterate),
|
||||
Elapsed = start_time.elapsed(),
|
||||
);
|
||||
|
||||
result
|
||||
}
|
||||
|
||||
pub async fn get_counter(
|
||||
&self,
|
||||
key: impl Into<ValueKey<ValueClass>> + Sync + Send,
|
||||
) -> trc::Result<i64> {
|
||||
match self {
|
||||
#[cfg(feature = "sqlite")]
|
||||
Self::SQLite(store) => store.get_counter(key).await,
|
||||
#[cfg(feature = "foundation")]
|
||||
Self::FoundationDb(store) => store.get_counter(key).await,
|
||||
#[cfg(feature = "postgres")]
|
||||
Self::PostgreSQL(store) => store.get_counter(key).await,
|
||||
#[cfg(feature = "mysql")]
|
||||
Self::MySQL(store) => store.get_counter(key).await,
|
||||
#[cfg(feature = "rocks")]
|
||||
Self::RocksDb(store) => store.get_counter(key).await,
|
||||
Self::Ephemeral(store) => store.get_counter(key).await,
|
||||
Self::None => Err(trc::StoreEvent::NotConfigured.into()),
|
||||
}
|
||||
.caused_by(trc::location!())
|
||||
}
|
||||
|
||||
#[allow(unreachable_patterns)]
|
||||
#[allow(unused_variables)]
|
||||
pub async fn sql_query<T: QueryResult + std::fmt::Debug>(
|
||||
&self,
|
||||
query: &str,
|
||||
params: Vec<Value<'_>>,
|
||||
) -> trc::Result<T> {
|
||||
let result = match self {
|
||||
#[cfg(feature = "sqlite")]
|
||||
Self::SQLite(store) => store.sql_query(query, ¶ms).await,
|
||||
#[cfg(feature = "postgres")]
|
||||
Self::PostgreSQL(store) => store.sql_query(query, ¶ms).await,
|
||||
#[cfg(feature = "mysql")]
|
||||
Self::MySQL(store) => store.sql_query(query, ¶ms).await,
|
||||
_ => Err(trc::StoreEvent::NotSupported.into_err()),
|
||||
};
|
||||
|
||||
trc::event!(
|
||||
Store(trc::StoreEvent::SqlQuery),
|
||||
Details = query.to_compact_string(),
|
||||
Value = params.as_slice(),
|
||||
Result = &result,
|
||||
);
|
||||
|
||||
result.caused_by(trc::location!())
|
||||
}
|
||||
|
||||
pub async fn write(&self, batch: Batch<'_>) -> trc::Result<AssignedIds> {
|
||||
let start_time = Instant::now();
|
||||
let ops = batch.ops.len();
|
||||
|
||||
let result = match self {
|
||||
#[cfg(feature = "sqlite")]
|
||||
Self::SQLite(store) => store.write(batch).await,
|
||||
#[cfg(feature = "foundation")]
|
||||
Self::FoundationDb(store) => store.write(batch).await,
|
||||
#[cfg(feature = "postgres")]
|
||||
Self::PostgreSQL(store) => store.write(batch).await,
|
||||
#[cfg(feature = "mysql")]
|
||||
Self::MySQL(store) => store.write(batch).await,
|
||||
#[cfg(feature = "rocks")]
|
||||
Self::RocksDb(store) => store.write(batch).await,
|
||||
Self::Ephemeral(store) => store.write(batch).await,
|
||||
Self::None => Err(trc::StoreEvent::NotConfigured.into()),
|
||||
};
|
||||
|
||||
trc::event!(
|
||||
Store(StoreEvent::DataWrite),
|
||||
Elapsed = start_time.elapsed(),
|
||||
Total = ops,
|
||||
);
|
||||
|
||||
result
|
||||
}
|
||||
|
||||
pub async fn assign_document_ids(
|
||||
&self,
|
||||
account_id: u32,
|
||||
collection: Collection,
|
||||
num_ids: u64,
|
||||
) -> trc::Result<u32> {
|
||||
// Increment UID next
|
||||
let mut batch = BatchBuilder::new();
|
||||
batch
|
||||
.with_account_id(account_id)
|
||||
.with_collection(collection)
|
||||
.add_and_get(ValueClass::DocumentId, num_ids as i64);
|
||||
self.write(batch.build_all()).await.and_then(|v| {
|
||||
v.last_counter_id().map(|id| {
|
||||
debug_assert!(id >= num_ids as i64, "{} < {}", id, num_ids);
|
||||
id as u32
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
pub async fn purge_store(&self) -> trc::Result<()> {
|
||||
match self {
|
||||
#[cfg(feature = "sqlite")]
|
||||
Self::SQLite(store) => store.purge_store().await,
|
||||
#[cfg(feature = "foundation")]
|
||||
Self::FoundationDb(store) => store.purge_store().await,
|
||||
#[cfg(feature = "postgres")]
|
||||
Self::PostgreSQL(store) => store.purge_store().await,
|
||||
#[cfg(feature = "mysql")]
|
||||
Self::MySQL(store) => store.purge_store().await,
|
||||
#[cfg(feature = "rocks")]
|
||||
Self::RocksDb(store) => store.purge_store().await,
|
||||
Self::Ephemeral(store) => store.purge_store().await,
|
||||
Self::None => Err(trc::StoreEvent::NotConfigured.into()),
|
||||
}
|
||||
.caused_by(trc::location!())
|
||||
}
|
||||
|
||||
pub async fn delete_range(&self, from: impl Key, to: impl Key) -> trc::Result<()> {
|
||||
match self {
|
||||
#[cfg(feature = "sqlite")]
|
||||
Self::SQLite(store) => store.delete_range(from, to).await,
|
||||
#[cfg(feature = "foundation")]
|
||||
Self::FoundationDb(store) => store.delete_range(from, to).await,
|
||||
#[cfg(feature = "postgres")]
|
||||
Self::PostgreSQL(store) => store.delete_range(from, to).await,
|
||||
#[cfg(feature = "mysql")]
|
||||
Self::MySQL(store) => store.delete_range(from, to).await,
|
||||
#[cfg(feature = "rocks")]
|
||||
Self::RocksDb(store) => store.delete_range(from, to).await,
|
||||
Self::Ephemeral(store) => store.delete_range(from, to).await,
|
||||
Self::None => Err(trc::StoreEvent::NotConfigured.into()),
|
||||
}
|
||||
.caused_by(trc::location!())
|
||||
}
|
||||
|
||||
pub async fn delete_documents(
|
||||
&self,
|
||||
subspace: u8,
|
||||
account_id: u32,
|
||||
collection: u8,
|
||||
collection_offset: Option<usize>,
|
||||
document_ids: &impl DocumentSet,
|
||||
) -> trc::Result<()> {
|
||||
// Serialize keys
|
||||
let (from_key, to_key) = if collection_offset.is_some() {
|
||||
(
|
||||
KeySerializer::new(U32_LEN + 2)
|
||||
.write(account_id)
|
||||
.write(collection),
|
||||
KeySerializer::new(U32_LEN + 2)
|
||||
.write(account_id)
|
||||
.write(collection + 1),
|
||||
)
|
||||
} else {
|
||||
(
|
||||
KeySerializer::new(U32_LEN).write(account_id),
|
||||
KeySerializer::new(U32_LEN).write(account_id + 1),
|
||||
)
|
||||
};
|
||||
|
||||
// Find keys to delete
|
||||
let mut delete_keys = Vec::new();
|
||||
self.iterate(
|
||||
IterateParams::new(
|
||||
AnyKey {
|
||||
subspace,
|
||||
key: from_key.finalize(),
|
||||
},
|
||||
AnyKey {
|
||||
subspace,
|
||||
key: to_key.finalize(),
|
||||
},
|
||||
)
|
||||
.no_values(),
|
||||
|key, _| {
|
||||
if collection_offset.is_none_or(|offset| {
|
||||
key.get(key.len() - U32_LEN - offset).copied() == Some(collection)
|
||||
}) {
|
||||
let document_id = key.deserialize_be_u32(key.len() - U32_LEN)?;
|
||||
if document_ids.contains(document_id) {
|
||||
delete_keys.push(key.to_vec());
|
||||
}
|
||||
}
|
||||
|
||||
Ok(true)
|
||||
},
|
||||
)
|
||||
.await
|
||||
.caused_by(trc::location!())?;
|
||||
|
||||
// Remove keys
|
||||
let mut batch = BatchBuilder::new();
|
||||
|
||||
for key in delete_keys {
|
||||
if batch.is_large_batch() {
|
||||
self.write(std::mem::take(&mut batch).build_all())
|
||||
.await
|
||||
.caused_by(trc::location!())?;
|
||||
}
|
||||
batch.any_op(Operation::Value {
|
||||
class: ValueClass::Any(AnyClass { subspace, key }),
|
||||
op: ValueOp::Clear,
|
||||
});
|
||||
}
|
||||
|
||||
if !batch.is_empty() {
|
||||
self.write(batch.build_all())
|
||||
.await
|
||||
.caused_by(trc::location!())?;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn danger_destroy_account(&self, account_id: u32) -> trc::Result<()> {
|
||||
for subspace in [SUBSPACE_LOGS, SUBSPACE_INDEXES, SUBSPACE_COUNTER] {
|
||||
self.delete_range(
|
||||
AnyKey {
|
||||
subspace,
|
||||
key: KeySerializer::new(U32_LEN).write(account_id).finalize(),
|
||||
},
|
||||
AnyKey {
|
||||
subspace,
|
||||
key: KeySerializer::new(U32_LEN).write(account_id + 1).finalize(),
|
||||
},
|
||||
)
|
||||
.await
|
||||
.caused_by(trc::location!())?;
|
||||
}
|
||||
|
||||
self.delete_range(
|
||||
ValueKey {
|
||||
account_id: 0,
|
||||
collection: 0,
|
||||
document_id: 0,
|
||||
class: ValueClass::Acl(account_id),
|
||||
},
|
||||
ValueKey {
|
||||
account_id: 0,
|
||||
collection: 0,
|
||||
document_id: 0,
|
||||
class: ValueClass::Acl(account_id + 1),
|
||||
},
|
||||
)
|
||||
.await
|
||||
.caused_by(trc::location!())?;
|
||||
|
||||
self.delete_range(
|
||||
ValueKey {
|
||||
account_id,
|
||||
collection: 0,
|
||||
document_id: 0,
|
||||
class: ValueClass::Property(0),
|
||||
},
|
||||
ValueKey {
|
||||
account_id: account_id + 1,
|
||||
collection: 0,
|
||||
document_id: 0,
|
||||
class: ValueClass::Property(0),
|
||||
},
|
||||
)
|
||||
.await
|
||||
.caused_by(trc::location!())?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn create_tables(&self) -> trc::Result<()> {
|
||||
match self {
|
||||
#[cfg(feature = "sqlite")]
|
||||
Self::SQLite(store) => store.create_tables(),
|
||||
#[cfg(feature = "postgres")]
|
||||
Self::PostgreSQL(store) => store.create_storage_tables().await,
|
||||
#[cfg(feature = "mysql")]
|
||||
Self::MySQL(store) => store.create_storage_tables().await,
|
||||
_ => Ok(()),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn invalidate_read_snapshot(&self) {
|
||||
#[cfg(feature = "foundation")]
|
||||
if let Self::FoundationDb(store) = self {
|
||||
store.invalidate_read_snapshot();
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user