inbuxa/inbuxa-server
1
0
Fork 0
Code Issues Pull Requests 2 Actions 4 Packages Projects Releases 2 Wiki Activity

Import upstream v0.16.22, stripped

Browse Source 
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:
jcoffey-dev
2026-09-18 10:21:56 -07:00
commit 7dae9b29fd
1650 changed files with 485521 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/src/store/blob.rs
+533
View File
@@ -0,0 +1,533 @@
/*
* SPDX-FileCopyrightText: 2020 Stalwart Labs LLC <[email protected]>
*
* SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL
*/
use crate::utils::{cleanup::store_destroy, server::TestServerBuilder};
use ahash::AHashMap;
use email::message::metadata::MessageMetadata;
use registry::{
schema::{enums::CompressionAlgo, structs::Jmap},
types::duration::Duration,
};
use services::task_manager::destroy_account::destroy_account_blobs;
use store::{
BlobStore, Serialize, SerializeInfallible,
write::{Archiver, BatchBuilder, BlobLink, BlobOp, ValueClass, now},
};
use types::{blob::BlobClass, blob_hash::BlobHash, collection::Collection, field::EmailField};
#[tokio::test]
pub async fn blob_tests() {
let test = TestServerBuilder::new("blob_tests")
.await
.with_object(Jmap {
upload_quota: 1024,
upload_ttl: Duration::from_millis(1000),
..Default::default()
})
.await
.build()
.await;
let store = test.server.core.storage.data.clone();
let blob_store = test.server.core.storage.blob.clone();
println!(
"Testing blob store {} with data store {}...",
std::env::var("BLOB_STORE").unwrap_or_else(|_| "default".to_string()),
std::env::var("STORE").unwrap()
);
// Test blob quota
assert!(test.server.blob_has_quota(0, 1024).await.unwrap().allowed);
assert!(!test.server.blob_has_quota(0, 1024).await.unwrap().allowed);
tokio::time::sleep(std::time::Duration::from_secs(1)).await;
assert!(test.server.blob_has_quota(0, 1024).await.unwrap().allowed);
// Test and reset store
test_store(blob_store.clone()).await;
store_destroy(&store).await;
// Blob hash exists
let hash = BlobHash::generate(b"abc".as_slice());
assert!(!store.blob_exists(&hash).await.unwrap());
// Reserve blob
let until = now() + 1;
store
.write(
BatchBuilder::new()
.with_account_id(0)
.set(
BlobOp::Link {
to: BlobLink::Temporary { until },
hash: hash.clone(),
},
1024u32.serialize(),
)
.build_all(),
)
.await
.unwrap();
// Uncommitted blob, should not exist
assert!(!store.blob_exists(&hash).await.unwrap());
// Write blob to store
blob_store
.put_blob(hash.as_ref(), b"abc", CompressionAlgo::Lz4)
.await
.unwrap();
// Commit blob
store
.write(
BatchBuilder::new()
.set(BlobOp::Commit { hash: hash.clone() }, Vec::new())
.build_all(),
)
.await
.unwrap();
// Blob hash should now exist
assert!(store.blob_exists(&hash).await.unwrap());
assert!(
blob_store
.get_blob(hash.as_ref(), 0..usize::MAX)
.await
.unwrap()
.is_some()
);
// AccountId 0 should be able to read blob
assert!(
store
.blob_has_access(
&hash,
BlobClass::Reserved {
account_id: 0,
expires: until
}
)
.await
.unwrap()
);
// AccountId 1 should not be able to read blob
assert!(
!store
.blob_has_access(
&hash,
BlobClass::Reserved {
account_id: 1,
expires: until
}
)
.await
.unwrap()
);
// Purge expired blobs
tokio::time::sleep(std::time::Duration::from_secs(1)).await;
store
.purge_blobs_all_shards(blob_store.clone())
.await
.unwrap();
// Blob hash should no longer exist
assert!(!store.blob_exists(&hash).await.unwrap());
// AccountId 0 should not be able to read blob
assert!(
!store
.blob_has_access(
&hash,
BlobClass::Reserved {
account_id: 0,
expires: until
}
)
.await
.unwrap()
);
// Blob should no longer be in store
assert!(
blob_store
.get_blob(hash.as_ref(), 0..usize::MAX)
.await
.unwrap()
.is_none()
);
// Upload one linked blob to accountId 1, two linked blobs to accountId 0, and three unlinked (reserved) blobs to accountId 2
let expiry_times = AHashMap::from_iter([
(b"abc", now() - 10),
(b"efg", now() + 10),
(b"hij", now() + 10),
]);
for (document_id, (blob, _)) in [
(b"123", vec![]),
(b"456", vec![]),
(b"789", vec![]),
(b"abc", 5000u32.serialize()),
(b"efg", 1000u32.serialize()),
(b"hij", 2000u32.serialize()),
]
.into_iter()
.enumerate()
{
let hash = BlobHash::generate(blob.as_slice());
let mut batch = BatchBuilder::new();
batch
.with_account_id(if document_id > 0 { 0 } else { 1 })
.with_collection(Collection::Email)
.with_document(document_id as u32);
if let Some(until) = expiry_times.get(blob) {
batch.set(
BlobOp::Link {
hash: hash.clone(),
to: BlobLink::Temporary { until: *until },
},
vec![],
);
} else {
batch
.set(
BlobOp::Link {
hash: hash.clone(),
to: BlobLink::Document,
},
vec![],
)
.set(
ValueClass::Property(EmailField::Metadata.into()),
Archiver::new(MessageMetadata {
contents: Default::default(),
rcvd_attach: Default::default(),
blob_hash: hash.clone(),
blob_body_offset: Default::default(),
preview: Default::default(),
raw_headers: Default::default(),
})
.serialize()
.unwrap(),
);
};
batch.set(BlobOp::Commit { hash: hash.clone() }, vec![]);
store.write(batch.build_all()).await.unwrap();
blob_store
.put_blob(hash.as_ref(), blob.as_slice(), CompressionAlgo::Lz4)
.await
.unwrap();
}
// Purge expired blobs and make sure nothing else is deleted
store
.purge_blobs_all_shards(blob_store.clone())
.await
.unwrap();
for (pos, (blob, blob_class)) in [
(
b"abc",
BlobClass::Reserved {
account_id: 0,
expires: expiry_times[&b"abc"],
},
),
(
b"123",
BlobClass::Linked {
account_id: 1,
collection: 0,
document_id: 0,
},
),
(
b"456",
BlobClass::Linked {
account_id: 0,
collection: 0,
document_id: 1,
},
),
(
b"789",
BlobClass::Linked {
account_id: 0,
collection: 0,
document_id: 2,
},
),
(
b"efg",
BlobClass::Reserved {
account_id: 0,
expires: expiry_times[&b"efg"],
},
),
(
b"hij",
BlobClass::Reserved {
account_id: 0,
expires: expiry_times[&b"hij"],
},
),
]
.into_iter()
.enumerate()
{
let hash = BlobHash::generate(blob.as_slice());
let ct = pos == 0;
assert!(store.blob_has_access(&hash, blob_class).await.unwrap() ^ ct);
assert!(store.blob_exists(&hash).await.unwrap() ^ ct);
assert!(
blob_store
.get_blob(hash.as_ref(), 0..usize::MAX)
.await
.unwrap()
.is_some()
^ ct
);
}
// AccountId 0 should not have access to accountId 1's blobs
assert!(
!store
.blob_has_access(
BlobHash::generate(b"123".as_slice()),
BlobClass::Linked {
account_id: 0,
collection: 0,
document_id: 0,
}
)
.await
.unwrap()
);
// Unlink blob
store
.write(
BatchBuilder::new()
.with_account_id(0)
.with_collection(Collection::Email)
.with_document(2)
.clear(BlobOp::Link {
hash: BlobHash::generate(b"789".as_slice()),
to: BlobLink::Document,
})
.build_all(),
)
.await
.unwrap();
// Purge and make sure blob is deleted
store
.purge_blobs_all_shards(blob_store.clone())
.await
.unwrap();
for (pos, (blob, blob_class)) in [
(
b"789",
BlobClass::Linked {
account_id: 0,
collection: 0,
document_id: 2,
},
),
(
b"123",
BlobClass::Linked {
account_id: 1,
collection: 0,
document_id: 0,
},
),
(
b"456",
BlobClass::Linked {
account_id: 0,
collection: 0,
document_id: 1,
},
),
(
b"efg",
BlobClass::Reserved {
account_id: 0,
expires: expiry_times[&b"efg"],
},
),
(
b"hij",
BlobClass::Reserved {
account_id: 0,
expires: expiry_times[&b"hij"],
},
),
]
.into_iter()
.enumerate()
{
let ct = pos == 0;
let hash = BlobHash::generate(blob.as_slice());
assert!(store.blob_has_access(&hash, blob_class).await.unwrap() ^ ct);
assert!(store.blob_exists(&hash).await.unwrap() ^ ct);
assert!(
blob_store
.get_blob(hash.as_ref(), 0..usize::MAX)
.await
.unwrap()
.is_some()
^ ct
);
}
// Unlink all blobs from accountId 1 and purge
destroy_account_blobs(&test.server, 1).await.unwrap();
store
.purge_blobs_all_shards(blob_store.clone())
.await
.unwrap();
// Make sure only accountId 0's blobs are left
for (pos, (blob, blob_class)) in [
(
b"123",
BlobClass::Linked {
account_id: 1,
collection: 0,
document_id: 0,
},
),
(
b"456",
BlobClass::Linked {
account_id: 0,
collection: 0,
document_id: 1,
},
),
(
b"efg",
BlobClass::Reserved {
account_id: 0,
expires: expiry_times[&b"efg"],
},
),
(
b"hij",
BlobClass::Reserved {
account_id: 0,
expires: expiry_times[&b"hij"],
},
),
]
.into_iter()
.enumerate()
{
let ct = pos == 0;
let hash = BlobHash::generate(blob.as_slice());
assert!(store.blob_has_access(&hash, blob_class).await.unwrap() ^ ct);
assert!(store.blob_exists(&hash).await.unwrap() ^ ct);
assert!(
blob_store
.get_blob(hash.as_ref(), 0..usize::MAX)
.await
.unwrap()
.is_some()
^ ct
);
}
test.temp_dir.delete();
}
async fn test_store(store: BlobStore) {
// Test small blob
const DATA: &[u8] = b"Lorem ipsum dolor sit amet, consectetur adipiscing elit. Fusce erat nisl, dignissim a porttitor id, varius nec arcu. Sed mauris.";
let hash = BlobHash::generate(DATA);
store
.put_blob(hash.as_slice(), DATA, CompressionAlgo::Lz4)
.await
.unwrap();
assert_eq!(
String::from_utf8(
store
.get_blob(hash.as_slice(), 0..usize::MAX)
.await
.unwrap()
.unwrap()
)
.unwrap(),
std::str::from_utf8(DATA).unwrap()
);
assert_eq!(
String::from_utf8(
store
.get_blob(hash.as_slice(), 11..57)
.await
.unwrap()
.unwrap()
)
.unwrap(),
std::str::from_utf8(&DATA[11..57]).unwrap()
);
assert!(store.delete_blob(hash.as_slice()).await.unwrap());
assert!(
store
.get_blob(hash.as_slice(), 0..usize::MAX)
.await
.unwrap()
.is_none()
);
// Test large blob
let mut data = Vec::with_capacity(50 * 1024 * 1024);
while data.len() < 50 * 1024 * 1024 {
data.extend_from_slice(DATA);
let marker = format!(" [{}] ", data.len());
data.extend_from_slice(marker.as_bytes());
}
let hash = BlobHash::generate(&data);
store
.put_blob(hash.as_slice(), &data, CompressionAlgo::Lz4)
.await
.unwrap();
assert_eq!(
String::from_utf8(
store
.get_blob(hash.as_slice(), 0..usize::MAX)
.await
.unwrap()
.unwrap()
)
.unwrap(),
std::str::from_utf8(&data).unwrap()
);
assert_eq!(
String::from_utf8(
store
.get_blob(hash.as_slice(), 3000111..4000999)
.await
.unwrap()
.unwrap()
)
.unwrap(),
std::str::from_utf8(&data[3000111..4000999]).unwrap()
);
assert!(store.delete_blob(hash.as_slice()).await.unwrap());
assert!(
store
.get_blob(hash.as_slice(), 0..usize::MAX)
.await
.unwrap()
.is_none()
);
}
tests/src/store/import_export.rs
+328
View File
@@ -0,0 +1,328 @@
/*
* SPDX-FileCopyrightText: 2020 Stalwart Labs LLC <[email protected]>
*
* SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL
*/
use crate::utils::{
cleanup::{store_assert_is_empty, store_destroy},
server::TestServer,
temp_dir::TempDir,
};
use ::registry::schema::enums::CompressionAlgo;
use ahash::AHashSet;
use common::{DATABASE_SCHEMA_VERSION, manager::backup::BackupParams};
use store::{
rand,
write::{
AnyClass, AnyKey, BatchBuilder, BlobLink, BlobOp, Operation, QueueClass, QueueEvent,
RegistryClass, ValueClass, key::KeySerializer,
},
*,
};
use types::{
blob_hash::BlobHash,
collection::{Collection, SyncCollection},
field::{Field, MailboxField},
};
pub async fn test(test: &TestServer) {
// Make sure the store is empty
store_assert_is_empty(test.server.store(), test.server.blob_store().clone(), true).await;
let db = test.server.store().clone();
// Create blobs
println!("Creating blobs...");
let mut batch = BatchBuilder::new();
batch.set(
ValueClass::Any(AnyClass {
subspace: SUBSPACE_PROPERTY,
key: vec![0u8],
}),
DATABASE_SCHEMA_VERSION.serialize(),
);
let mut blob_hashes = Vec::new();
for blob_size in [16, 128, 1024, 2056, 102400] {
let data = random_bytes(blob_size);
let hash = BlobHash::generate(data.as_slice());
blob_hashes.push(hash.clone());
test.server
.blob_store()
.put_blob(hash.as_ref(), &data, CompressionAlgo::Lz4)
.await
.unwrap();
batch.set(ValueClass::Blob(BlobOp::Commit { hash }), vec![]);
}
db.write(batch.build_all()).await.unwrap();
// Create account data
println!("Creating account data...");
for account_id in 0u32..10u32 {
let mut batch = BatchBuilder::new();
batch.with_account_id(account_id);
// Create properties of different sizes
for collection in [
Collection::Email,
Collection::Mailbox,
Collection::Thread,
Collection::Identity,
] {
batch.with_collection(collection);
for document_id in [0, 10, 20, 30, 40] {
batch.with_document(document_id);
if collection == Collection::Mailbox {
batch
.set(
ValueClass::Property(Field::ARCHIVE.into()),
random_bytes(10),
)
.add(
ValueClass::Property(MailboxField::UidCounter.into()),
rand::random(),
);
}
for (idx, value_size) in [16, 128, 1024, 2056, 102400].into_iter().enumerate() {
batch.set(ValueClass::Property(idx as u8), random_bytes(value_size));
}
for grant_account_id in 0u32..10u32 {
if account_id != grant_account_id {
batch.set(
ValueClass::Acl(grant_account_id),
vec![account_id as u8, grant_account_id as u8, document_id as u8],
);
}
}
for hash in &blob_hashes {
batch.set(
ValueClass::Blob(BlobOp::Link {
hash: hash.clone(),
to: BlobLink::Document,
}),
vec![],
);
}
batch.log_item_insert(SyncCollection::from(collection), None);
for field in 0..5 {
batch.any_op(Operation::Index {
field,
key: random_bytes(field as usize + 2),
set: true,
});
}
}
}
db.write(batch.build_all()).await.unwrap();
}
// Create queue, config and lookup data
println!("Creating queue, config and lookup data...");
let mut batch = BatchBuilder::new();
for idx in [1, 2, 3, 4, 5] {
batch.set(
ValueClass::Queue(QueueClass::Message(rand::random())),
random_bytes(idx),
);
batch.set(
ValueClass::Queue(QueueClass::MessageEvent(QueueEvent {
due: rand::random(),
queue_id: rand::random(),
queue_name: rand::random(),
})),
random_bytes(idx),
);
batch.set(
ValueClass::Registry(RegistryClass::Item {
object_id: 0,
item_id: 1,
}),
random_bytes(idx + 10),
);
batch.set(
ValueClass::Registry(RegistryClass::IndexId {
object_id: idx as u16,
item_id: idx as u64 * 100,
}),
vec![],
);
for index_id in 0u16..3 {
batch.set(
ValueClass::Registry(RegistryClass::Index {
object_id: idx as u16,
index_id,
key: random_bytes(idx + index_id as usize),
item_id: idx as u64 * 100,
}),
vec![],
);
}
}
db.write(batch.build_all()).await.unwrap();
// Create directory data
println!("Creating directory data...");
let mut batch = BatchBuilder::new();
batch
.with_account_id(u32::MAX)
.with_collection(Collection::Principal);
for account_id in [1, 2, 3, 4, 5] {
batch
.with_document(account_id)
.add(ValueClass::Quota, account_id as i64 * 1000);
}
db.write(batch.build_all()).await.unwrap();
// Obtain store hash
println!("Calculating store hash...");
let snapshot = Snapshot::new(&db).await;
assert!(!snapshot.keys.is_empty(), "Store hash counts are empty",);
// Export store
println!("Exporting store...");
let temp_dir = TempDir::new("art_vandelay_tests", true);
test.server
.core
.backup(BackupParams::new(temp_dir.path.clone()))
.await;
// Destroy store
println!("Destroying store...");
store_destroy(&db).await;
store_assert_is_empty(&db, db.clone().into(), true).await;
// Import store over a node id lease
println!("Importing store...");
let mut batch = BatchBuilder::new();
batch.set(
ValueClass::NodeId(0),
KeySerializer::new(U64_LEN + 9)
.write(0u64)
.write("localhost")
.finalize(),
);
db.write(batch.build_all()).await.unwrap();
test.server.core.restore(temp_dir.path.clone()).await;
let mut batch = BatchBuilder::new();
batch.clear(ValueClass::NodeId(0));
db.write(batch.build_all()).await.unwrap();
// Verify hash
print!("Verifying store hash...");
snapshot.assert_is_eq(&Snapshot::new(&db).await);
println!(" GREAT SUCCESS!");
// Destroy store
store_destroy(&db).await;
store_assert_is_empty(&db, db.clone().into(), true).await;
temp_dir.delete();
}
#[derive(Debug, PartialEq, Eq)]
struct Snapshot {
keys: AHashSet<KeyValue>,
}
#[derive(Debug, PartialEq, Eq, Hash)]
struct KeyValue {
subspace: u8,
key: Vec<u8>,
value: Vec<u8>,
}
impl Snapshot {
async fn new(db: &Store) -> Self {
let is_sql = db.is_sql();
let mut keys = AHashSet::new();
for (subspace, with_values) in [
(SUBSPACE_ACL, true),
(SUBSPACE_TASK_QUEUE, true),
(SUBSPACE_INDEXES, false),
(SUBSPACE_DELETED_ITEMS, true),
(SUBSPACE_SPAM_SAMPLES, true),
(SUBSPACE_BLOB_LINK, true),
(SUBSPACE_BLOBS, true),
(SUBSPACE_LOGS, true),
(SUBSPACE_COUNTER, !is_sql),
(SUBSPACE_IN_MEMORY_COUNTER, !is_sql),
(SUBSPACE_IN_MEMORY_VALUE, true),
(SUBSPACE_PROPERTY, true),
(SUBSPACE_REGISTRY, true),
(SUBSPACE_REGISTRY_IDX, !is_sql),
(SUBSPACE_REGISTRY_PK, true),
(SUBSPACE_QUEUE_MESSAGE, true),
(SUBSPACE_QUEUE_EVENT, true),
(SUBSPACE_QUOTA, !is_sql),
(SUBSPACE_REPORT_OUT, true),
(SUBSPACE_REPORT_IN, true),
] {
let from_key = AnyKey {
subspace,
key: vec![0u8],
};
let to_key = AnyKey {
subspace,
key: vec![u8::MAX; 10],
};
db.iterate(
IterateParams::new(from_key, to_key).set_values(with_values),
|key, value| {
keys.insert(KeyValue {
subspace,
key: key.to_vec(),
value: value.to_vec(),
});
Ok(true)
},
)
.await
.unwrap();
}
Snapshot { keys }
}
fn assert_is_eq(&self, other: &Self) {
let mut is_err = false;
for key in &self.keys {
if !other.keys.contains(key) {
println!(
"Subspace {}, Key {:?} not found in restored snapshot",
char::from(key.subspace),
key.key,
);
is_err = true;
}
}
for key in &other.keys {
if !self.keys.contains(key) {
println!(
"Subspace {}, Key {:?} not found in original snapshot",
char::from(key.subspace),
key.key,
);
is_err = true;
}
}
if is_err {
panic!("Snapshot mismatch");
}
}
}
fn random_bytes(len: usize) -> Vec<u8> {
(0..len).map(|_| rand::random::<u8>()).collect()
}
tests/src/store/lookup.rs
+300
View File
@@ -0,0 +1,300 @@
/*
* SPDX-FileCopyrightText: 2020 Stalwart Labs LLC <[email protected]>
*
* SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL
*/
use crate::utils::{
cleanup::{store_assert_is_empty, store_destroy},
server::TestServerBuilder,
};
use registry::schema::structs::Rate;
use registry::types::duration::Duration;
use store::{InMemoryStore, dispatch::lookup::KeyValue};
#[tokio::test]
pub async fn lookup_tests() {
let test = TestServerBuilder::new("lookup_tests").await.build().await;
let store = test.server.in_memory_store().clone();
let rate = Rate {
count: 1,
period: Duration::from_millis(1000),
};
println!(
"Testing in-memory store {}...",
std::env::var("MEMORY_STORE").unwrap_or_else(|_| "default".to_string())
);
if let InMemoryStore::Store(store) = &store {
store_destroy(store).await;
} else {
// Reset redis counter
store
.key_set(KeyValue::new("abc", "0".as_bytes().to_vec()))
.await
.unwrap();
}
// Test key
let key = "xyz".as_bytes().to_vec();
store
.key_set(KeyValue::new(key.clone(), "world".to_string().into_bytes()))
.await
.unwrap();
store.purge_in_memory_store().await.unwrap();
assert_eq!(
store.key_get::<String>(key.clone()).await.unwrap(),
Some("world".to_string())
);
// Test value expiry
store
.key_set(KeyValue::new(key.clone(), "hello".to_string().into_bytes()).expires(1))
.await
.unwrap();
assert_eq!(
store.key_get::<String>(key.clone()).await.unwrap(),
Some("hello".to_string())
);
tokio::time::sleep(tokio::time::Duration::from_secs(2)).await;
assert_eq!(None, store.key_get::<String>(key.clone()).await.unwrap());
store.purge_in_memory_store().await.unwrap();
if let InMemoryStore::Store(store) = &store {
store_assert_is_empty(store, store.clone().into(), false).await;
}
// Test counter
let key = "abc".as_bytes().to_vec();
store
.counter_incr(KeyValue::new(key.clone(), 1), true)
.await
.unwrap();
assert_eq!(1, store.counter_get(key.clone()).await.unwrap());
store
.counter_incr(KeyValue::new(key.clone(), 2), true)
.await
.unwrap();
assert_eq!(3, store.counter_get(key.clone()).await.unwrap());
store
.counter_incr(KeyValue::new(key.clone(), -3), false)
.await
.unwrap();
assert_eq!(0, store.counter_get(key.clone()).await.unwrap());
// Test counter expiry
let key = "fgh".as_bytes().to_vec();
store
.counter_incr(KeyValue::new(key.clone(), 1).expires(1), false)
.await
.unwrap();
assert_eq!(1, store.counter_get(key.clone()).await.unwrap());
tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
store.purge_in_memory_store().await.unwrap();
assert_eq!(0, store.counter_get(key.clone()).await.unwrap());
// Test rate limiter
assert!(
store
.is_rate_allowed(0, "rate".as_bytes(), &rate, false)
.await
.unwrap()
.is_none()
);
assert!(
store
.is_rate_allowed(0, "rate".as_bytes(), &rate, false)
.await
.unwrap()
.is_some()
);
tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
assert!(
store
.is_rate_allowed(0, "rate".as_bytes(), &rate, false)
.await
.unwrap()
.is_none()
);
tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
store.purge_in_memory_store().await.unwrap();
if let InMemoryStore::Store(store) = &store {
store_assert_is_empty(store, store.clone().into(), false).await;
}
// Test locking
for iteration in [1, 2] {
let mut tasks = Vec::new();
for _ in 0..100 {
let store = store.clone();
tasks.push(tokio::spawn(async move {
store.try_lock(0, "lock".as_bytes(), 1).await.unwrap()
}));
}
// Only one should return true
let mut count = 0;
for task in tasks {
if task.await.unwrap() {
count += 1;
}
}
assert_eq!(1, count, "Iteration {}", iteration);
// Wait 2 seconds for the lock to expire
tokio::time::sleep(tokio::time::Duration::from_secs(2)).await;
}
store.purge_in_memory_store().await.unwrap();
// Failed lock attempts must not extend the lock expiry
assert!(store.try_lock(0, "lock".as_bytes(), 2).await.unwrap());
let mut acquired = false;
for _ in 0..12 {
tokio::time::sleep(tokio::time::Duration::from_millis(500)).await;
if store.try_lock(0, "lock".as_bytes(), 2).await.unwrap() {
acquired = true;
break;
}
}
assert!(acquired, "Abandoned lock was never released");
store.remove_lock(0, "lock".as_bytes()).await.unwrap();
store.purge_in_memory_store().await.unwrap();
if let InMemoryStore::Store(store) = &store {
store_assert_is_empty(store, store.clone().into(), false).await;
}
// Test prefix delete
store
.key_set(KeyValue::with_prefix(
1,
[0],
"hello".to_string().into_bytes(),
))
.await
.unwrap();
for v in 0u32..2020u32 {
store
.key_set(KeyValue::with_prefix(
0,
pack_u32(0, v),
"world".to_string().into_bytes(),
))
.await
.unwrap();
store
.counter_incr(
KeyValue::with_prefix(0, pack_u32(1, v), 123).expires(3600),
false,
)
.await
.unwrap();
}
// Make sure the keys are there
assert_eq!(
Some("hello"),
store
.key_get::<String>(KeyValue::<()>::build_key(1, [0]))
.await
.unwrap()
.as_deref()
);
for v in [0, 1000, 1001, 2000, 2001] {
assert_eq!(
Some("world"),
store
.key_get::<String>(KeyValue::<()>::build_key(0, pack_u32(0, v)))
.await
.unwrap()
.as_deref()
);
}
for v in [0, 1000, 1001, 2000, 2001] {
assert_ne!(
0,
store
.counter_get(KeyValue::<()>::build_key(0, pack_u32(1, v)))
.await
.unwrap()
);
}
// Delete [0, 0, 0, 0, 1] prefix and make sure only the keys with that prefix are gone
store
.key_delete_prefix(&KeyValue::<()>::build_key(0, 1u32.to_be_bytes()))
.await
.unwrap();
assert_eq!(
Some("hello"),
store
.key_get::<String>(KeyValue::<()>::build_key(1, [0]))
.await
.unwrap()
.as_deref()
);
for v in [0, 1000, 1001, 2000, 2001] {
assert_eq!(
Some("world"),
store
.key_get::<String>(KeyValue::<()>::build_key(0, pack_u32(0, v)))
.await
.unwrap()
.as_deref()
);
}
for v in [0, 1000, 1001, 2000, 2001] {
assert_eq!(
0,
store
.counter_get(KeyValue::<()>::build_key(0, pack_u32(1, v)))
.await
.unwrap()
);
}
// Delete [0, 0, 0, 0, 0] prefix and make sure only the keys with that prefix are gone
store
.key_delete_prefix(&KeyValue::<()>::build_key(0, 0u32.to_be_bytes()))
.await
.unwrap();
assert_eq!(
Some("hello"),
store
.key_get::<String>(KeyValue::<()>::build_key(1, [0]))
.await
.unwrap()
.as_deref()
);
for v in [0, 1000, 1001, 2000, 2001] {
assert_eq!(
None,
store
.key_get::<String>(KeyValue::<()>::build_key(0, pack_u32(0, v)))
.await
.unwrap()
.as_deref()
);
}
// Delete [1, ...] prefix and make sure it's all gone
store.key_delete_prefix(&[1u8]).await.unwrap();
assert_eq!(
None,
store
.key_get::<String>(KeyValue::<()>::build_key(1, [0]))
.await
.unwrap()
.as_deref()
);
if let InMemoryStore::Store(store) = &store {
store_assert_is_empty(store, store.clone().into(), false).await;
}
}
fn pack_u32(a: u32, b: u32) -> Vec<u8> {
(((a as u64) << 32) | b as u64).to_be_bytes().to_vec()
}
tests/src/store/mod.rs
+68
View File
@@ -0,0 +1,68 @@
/*
* SPDX-FileCopyrightText: 2020 Stalwart Labs LLC <[email protected]>
*
* SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL
*/
pub mod blob;
pub mod import_export;
pub mod lookup;
pub mod ops;
pub mod query;
pub mod registry;
#[cfg(any(feature = "postgres", feature = "mysql"))]
pub mod sql_timeout;
use crate::utils::server::TestServerBuilder;
use std::io::Read;
#[tokio::test(flavor = "multi_thread")]
pub async fn store_tests() {
let test = TestServerBuilder::new("store_tests").await.build().await;
println!("Testing store {}...", std::env::var("STORE").unwrap());
test.destroy_store().await;
registry::test(&test).await;
import_export::test(&test).await;
ops::test(&test).await;
#[cfg(any(feature = "postgres", feature = "mysql"))]
sql_timeout::test(&test).await;
if test.is_reset() {
test.temp_dir.delete();
}
}
#[tokio::test(flavor = "multi_thread")]
pub async fn search_tests() {
let test = TestServerBuilder::new("search_store_tests")
.await
.build()
.await;
println!(
"Testing search store {}...",
std::env::var("SEARCH_STORE").unwrap_or("default".to_string())
);
query::test(&test).await;
if test.is_reset() {
test.temp_dir.delete();
}
}
pub fn deflate_test_resource(name: &str) -> Vec<u8> {
let mut csv_path = std::path::PathBuf::from(env!("CARGO_MANIFEST_DIR"));
csv_path.push("resources");
csv_path.push(name);
let mut decoder = flate2::bufread::GzDecoder::new(std::io::BufReader::new(
std::fs::File::open(csv_path).unwrap(),
));
let mut result = Vec::new();
decoder.read_to_end(&mut result).unwrap();
result
}
tests/src/store/ops.rs
+801
View File
@@ -0,0 +1,801 @@
/*
* SPDX-FileCopyrightText: 2020 Stalwart Labs LLC <[email protected]>
*
* SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL
*/
use crate::utils::{cleanup::store_assert_is_empty, server::TestServer};
use ahash::AHashSet;
use std::collections::HashSet;
use store::{
ValueKey,
rand::{self, RngExt},
write::{AlignedBytes, Archive, Archiver, BatchBuilder, MergeResult, Params, ValueClass},
};
use types::collection::Collection;
use types::collection::SyncCollection;
// FDB max value
const MAX_VALUE_SIZE: usize = 100000;
#[cfg(feature = "foundationdb")]
fn value_gen(chunks: impl IntoIterator<Item = (u8, usize)>) -> Vec<u8> {
let mut value = Vec::new();
for (byte, size) in chunks {
value.extend(std::iter::repeat_n(byte, size));
}
value
}
pub async fn test(test: &TestServer) {
let db = test.server.store().clone();
#[cfg(feature = "foundationdb")]
if matches!(db, store::Store::FoundationDb(_)) {
use store::write::RegistryClass;
println!("Running FoundationDB chunked iterator test...");
let kvs = [
(1, value_gen([(b'a', 1)])),
(2, value_gen([(b'b', MAX_VALUE_SIZE), (b'0', 1)])),
(
3,
value_gen([
(b'c', MAX_VALUE_SIZE),
(b'1', MAX_VALUE_SIZE),
(b'2', MAX_VALUE_SIZE),
]),
),
(
4,
value_gen([(b'd', MAX_VALUE_SIZE), (b'3', MAX_VALUE_SIZE)]),
),
(5, value_gen([(b'e', 1)])),
];
let mut batch = BatchBuilder::new();
batch
.with_account_id(0)
.with_collection(Collection::Email)
.with_document(0);
for (key, value) in &kvs {
batch.set(
ValueClass::Registry(RegistryClass::Item {
object_id: *key,
item_id: 0,
}),
value.clone(),
);
}
db.write(batch.build_all()).await.unwrap();
// Iterate over all keys
let mut results = Vec::new();
db.iterate(
store::IterateParams::new(
ValueKey {
account_id: 0,
collection: 0,
document_id: 0,
class: ValueClass::Registry(RegistryClass::Item {
object_id: 0,
item_id: 0,
}),
},
ValueKey {
account_id: 0,
collection: 0,
document_id: 0,
class: ValueClass::Registry(RegistryClass::Item {
object_id: u16::MAX,
item_id: u64::MAX,
}),
},
),
|key, value| {
results.push((String::from_utf8(key.to_vec()).unwrap(), value.to_vec()));
Ok(true)
},
)
.await
.unwrap();
assert_eq!(results.len(), kvs.len());
db.delete_range(
ValueKey {
account_id: 0,
collection: 0,
document_id: 0,
class: ValueClass::Registry(RegistryClass::Item {
object_id: 0,
item_id: 0,
}),
},
ValueKey {
account_id: 0,
collection: 0,
document_id: 0,
class: ValueClass::Registry(RegistryClass::Item {
object_id: u16::MAX,
item_id: u64::MAX,
}),
},
)
.await
.unwrap();
// Read-your-writes through the cached read version: overwrite a key in a tight loop
println!("Running FoundationDB read-your-writes test...");
for n in 0u64..200 {
db.write(
BatchBuilder::new()
.with_account_id(0)
.with_collection(Collection::Email)
.with_document(0)
.set(
ValueClass::Registry(RegistryClass::Item {
object_id: 100,
item_id: 0,
}),
n.to_be_bytes().to_vec(),
)
.build_all(),
)
.await
.unwrap();
let got = db
.get_value::<u64>(ValueKey {
account_id: 0,
collection: 0,
document_id: 0,
class: ValueClass::Registry(RegistryClass::Item {
object_id: 100,
item_id: 0,
}),
})
.await
.unwrap()
.unwrap();
assert_eq!(got, n, "stale read: wrote {n} but read back {got}");
}
db.write(
BatchBuilder::new()
.with_account_id(0)
.with_collection(Collection::Email)
.with_document(0)
.clear(ValueClass::Registry(RegistryClass::Item {
object_id: 100,
item_id: 0,
}))
.build_all(),
)
.await
.unwrap();
// Read-version cache monotonicity under concurrency: while a writer increments a counter
println!("Running FoundationDB read-version monotonicity test...");
let n_increments = 500u64;
let writer = {
let db = db.clone();
tokio::spawn(async move {
for _ in 0..n_increments {
db.write(
BatchBuilder::new()
.with_account_id(0)
.with_collection(Collection::Email)
.with_document(5000)
.add_and_get(ValueClass::Quota, 1)
.build_all(),
)
.await
.unwrap();
}
})
};
let mut readers = Vec::new();
for _ in 0..16 {
let db = db.clone();
readers.push(tokio::spawn(async move {
let deadline = std::time::Instant::now() + std::time::Duration::from_millis(1500);
let mut last = 0i64;
while std::time::Instant::now() < deadline {
let current = db
.get_counter(ValueKey {
account_id: 0,
collection: 0,
document_id: 5000,
class: ValueClass::Quota,
})
.await
.unwrap();
assert!(
current >= last,
"read version regressed: counter went from {last} to {current}"
);
last = current;
}
}));
}
writer.await.unwrap();
for reader in readers {
reader.await.unwrap();
}
assert_eq!(
db.get_counter(ValueKey {
account_id: 0,
collection: 0,
document_id: 5000,
class: ValueClass::Quota,
})
.await
.unwrap(),
n_increments as i64,
"counter did not reach the expected total"
);
db.write(
BatchBuilder::new()
.with_account_id(0)
.with_collection(Collection::Email)
.with_document(5000)
.clear(ValueClass::Quota)
.build_all(),
)
.await
.unwrap();
// Overwriting a chunked value with a shorter one must not leave orphaned chunks behind
println!("Running FoundationDB orphaned chunk test...");
const ORPHAN_FIELD: u8 = 200;
let orphan_range = |document_id: u32| ValueKey {
account_id: 0,
collection: 0,
document_id,
class: ValueClass::Property(ORPHAN_FIELD),
};
let marker = value_gen([(b'z', 16)]);
db.write(
BatchBuilder::new()
.with_account_id(0)
.with_collection(Collection::Email)
.with_document(1)
.set(ValueClass::Property(ORPHAN_FIELD), marker.clone())
.build_all(),
)
.await
.unwrap();
for (byte, size) in [
(b'a', MAX_VALUE_SIZE * 3),
(b'b', MAX_VALUE_SIZE * 2),
(b'c', MAX_VALUE_SIZE / 2),
(b'd', MAX_VALUE_SIZE * 3 / 2),
(b'e', MAX_VALUE_SIZE),
(b'f', 1),
(b'g', MAX_VALUE_SIZE * 4),
(b'h', 0),
] {
let value = value_gen([(byte, size)]);
db.write(
BatchBuilder::new()
.with_account_id(0)
.with_collection(Collection::Email)
.with_document(0)
.set(ValueClass::Property(ORPHAN_FIELD), value.clone())
.build_all(),
)
.await
.unwrap();
let mut results = Vec::new();
db.iterate(
store::IterateParams::new(orphan_range(0), orphan_range(u32::MAX)),
|key, value| {
results.push((key.to_vec(), value.to_vec()));
Ok(true)
},
)
.await
.unwrap();
assert_eq!(
results.len(),
2,
"orphaned chunks surfaced as extra rows after writing {size} bytes"
);
assert_eq!(
results[0].1.len(),
value.len(),
"stale chunk spliced onto the value after writing {size} bytes"
);
assert_eq!(results[0].1, value, "value mismatch for {size} bytes");
assert_eq!(
results[1].1, marker,
"neighbouring document corrupted after writing {size} bytes"
);
assert_eq!(
results[0].0.len(),
results[1].0.len(),
"chunk key returned as a document key after writing {size} bytes"
);
}
// A short key sharing a subspace with longer structured keys must never clear them,
// as the database schema version does in the property subspace
println!("Running FoundationDB short key test...");
for document_id in [0u32, 1, 0xFFFF, 0x10000] {
db.write(
BatchBuilder::new()
.with_account_id(document_id)
.with_collection(Collection::Email)
.with_document(document_id)
.set(ValueClass::Property(ORPHAN_FIELD), marker.clone())
.build_all(),
)
.await
.unwrap();
}
db.write(
BatchBuilder::new()
.set(
ValueClass::Any(store::write::AnyClass {
subspace: store::SUBSPACE_PROPERTY,
key: vec![0u8],
}),
vec![1u8],
)
.build_all(),
)
.await
.unwrap();
for document_id in [0u32, 1, 0xFFFF, 0x10000] {
let key = ValueKey {
account_id: document_id,
collection: 0,
document_id,
class: ValueClass::Property(ORPHAN_FIELD),
};
let mut found = Vec::new();
db.iterate(
store::IterateParams::new(key.clone(), key.clone()),
|_, value| {
found.push(value.to_vec());
Ok(true)
},
)
.await
.unwrap();
assert_eq!(
found,
vec![marker.clone()],
"property key for account {document_id} was cleared by a shorter key"
);
db.write(
BatchBuilder::new()
.with_account_id(document_id)
.with_collection(Collection::Email)
.with_document(document_id)
.clear(ValueClass::Property(ORPHAN_FIELD))
.build_all(),
)
.await
.unwrap();
}
db.write(
BatchBuilder::new()
.clear(ValueClass::Any(store::write::AnyClass {
subspace: store::SUBSPACE_PROPERTY,
key: vec![0u8],
}))
.build_all(),
)
.await
.unwrap();
db.delete_range(orphan_range(0), orphan_range(u32::MAX))
.await
.unwrap();
if std::env::var("SLOW_FDB_TRX").is_ok() {
println!("Running FoundationDB slow transaction tests...");
// Create 900000 keys
let mut batch = BatchBuilder::new();
batch
.with_account_id(0)
.with_collection(Collection::Email)
.with_document(0);
for n in 0..900000 {
batch.set(
ValueClass::Registry(RegistryClass::Item {
object_id: 0,
item_id: n,
}),
format!("value{n:10}").into_bytes(),
);
if n % 10000 == 0 {
db.write(batch.build_all()).await.unwrap();
batch = BatchBuilder::new();
batch
.with_account_id(0)
.with_collection(Collection::Email)
.with_document(0);
}
}
db.write(batch.build_all()).await.unwrap();
println!("Created 900.000 keys...");
// Iterate over all keys
let mut n = 0;
db.iterate(
store::IterateParams::new(
ValueKey {
account_id: 0,
collection: 0,
document_id: 0,
class: ValueClass::Registry(RegistryClass::Item {
object_id: 0,
item_id: 0,
}),
},
ValueKey {
account_id: 0,
collection: 0,
document_id: 0,
class: ValueClass::Registry(RegistryClass::Item {
object_id: 0,
item_id: u64::MAX,
}),
},
),
|key, value| {
assert_eq!(std::str::from_utf8(key).unwrap(), format!("key{n:10}"));
assert_eq!(std::str::from_utf8(value).unwrap(), format!("value{n:10}"));
n += 1;
if n % 10000 == 0 {
println!("Iterated over {n} keys");
std::thread::sleep(std::time::Duration::from_millis(1000));
}
Ok(true)
},
)
.await
.unwrap();
// Delete 100 keys
let mut batch = BatchBuilder::new();
batch
.with_account_id(0)
.with_collection(Collection::Email)
.with_document(0);
for n in 0..900000 {
batch.clear(ValueClass::Registry(RegistryClass::Item {
object_id: 0,
item_id: n,
}));
if n % 10000 == 0 {
db.write(batch.build_all()).await.unwrap();
batch = BatchBuilder::new();
batch
.with_account_id(0)
.with_collection(Collection::Email)
.with_document(0);
}
}
db.write(batch.build_all()).await.unwrap();
}
}
// Merge values 1000 times concurrently
let mut handles = Vec::new();
println!("Merge values 1000 times concurrently...");
for _ in 0..1000 {
handles.push({
let db = db.clone();
tokio::spawn(async move {
for _ in 0..5 {
let mut builder = BatchBuilder::new();
builder
.with_account_id(0)
.with_collection(Collection::Email)
.with_document(0)
.merge_fnc(
ValueClass::Property(3),
Params::with_capacity(0),
|_, _, bytes| {
if let Some(bytes) = bytes {
Ok(MergeResult::Update(
(u64::from_be_bytes(bytes.try_into().unwrap()) + 1)
.to_be_bytes()
.to_vec(),
))
} else {
Ok(MergeResult::Update(0u64.to_be_bytes().to_vec()))
}
},
);
match db.write(builder.build_all()).await {
Ok(_) => {
break;
}
Err(e) if e.is_assertion_failure() => {
// Retry on assertion failures
continue;
}
Err(e) => {
panic!("Merge failed: {:?}", e);
}
}
}
})
});
}
for handle in handles {
handle.await.unwrap();
}
assert_eq!(
999,
db.get_value::<u64>(ValueKey {
account_id: 0,
collection: 0,
document_id: 0,
class: ValueClass::Property(3),
})
.await
.unwrap()
.unwrap()
);
// Increment a counter 1000 times concurrently
let mut handles = Vec::new();
let mut assigned_ids = HashSet::new();
println!("Incrementing counter 1000 times concurrently...");
for _ in 0..1000 {
handles.push({
let db = db.clone();
tokio::spawn(async move {
let mut builder = BatchBuilder::new();
builder
.with_account_id(0)
.with_collection(Collection::Email)
.with_document(0)
.add_and_get(ValueClass::Quota, 1);
db.write(builder.build_all())
.await
.unwrap()
.last_counter_id()
.unwrap()
})
});
}
for handle in handles {
let assigned_id = handle.await.unwrap();
assert!(
assigned_ids.insert(assigned_id),
"counter assigned {assigned_id} twice or more times."
);
}
assert_eq!(assigned_ids.len(), 1000);
assert_eq!(
db.get_counter(ValueKey {
account_id: 0,
collection: 0,
document_id: 0,
class: ValueClass::Quota,
})
.await
.unwrap(),
1000
);
// Concurrent changelog
let mut handles = Vec::new();
let mut assigned_ids = AHashSet::new();
print!("Incrementing changeId 1000 times concurrently...");
let time = std::time::Instant::now();
for document_id in 0..1000 {
handles.push({
let db = db.clone();
tokio::spawn(async move {
let mut builder = BatchBuilder::new();
let value = if document_id != 0 {
(0..rand::rng().random_range(1..=100))
.map(|_| rand::rng().random_range(0..=255))
.collect::<Vec<u8>>()
} else {
vec![0u8; 100000]
};
let (offset, archived_value) = Archiver::new(value).serialize_versioned().unwrap();
builder
.with_account_id(0)
.with_collection(Collection::Email)
.with_document(document_id)
.set_fnc(
ValueClass::Property(5),
Params::with_capacity(2)
.with_bytes(archived_value)
.with_u64(offset),
|params, ids| {
let change_id = ids.current_change_id()?;
let archive = params.bytes(0);
let offset = params.u64(1);
let mut bytes = Vec::with_capacity(archive.len());
bytes.extend_from_slice(&archive[..offset as usize]);
bytes.extend_from_slice(&change_id.to_be_bytes()[..]);
bytes.push(archive.last().copied().unwrap()); // Marker
Ok(bytes)
},
)
.log_container_insert(SyncCollection::Email);
db.write(builder.build_all())
.await
.unwrap()
.last_change_id(0)
.unwrap()
})
});
}
for handle in handles {
let assigned_id = handle.await.unwrap();
assert!(
assigned_ids.insert(assigned_id),
"counter assigned {assigned_id} twice or more times: {:?}.",
assigned_ids
);
}
assert_eq!(assigned_ids.len(), 1000);
println!(" done in {:?}ms", time.elapsed().as_millis());
let mut change_ids = AHashSet::new();
for document_id in 0..1000 {
let archive = db
.get_value::<Archive<AlignedBytes>>(ValueKey {
account_id: 0,
collection: 0,
document_id,
class: ValueClass::Property(5),
})
.await
.unwrap()
.unwrap();
change_ids.insert(archive.version.change_id().unwrap());
archive.unarchive_untrusted::<Vec<u8>>().unwrap();
}
assert_eq!(change_ids, assigned_ids);
println!("Running chunking tests...");
for (test_num, value) in [
vec![b'A'; 0],
vec![b'A'; 1],
vec![b'A'; 100],
vec![b'A'; MAX_VALUE_SIZE],
vec![b'B'; MAX_VALUE_SIZE + 1],
vec![b'C'; MAX_VALUE_SIZE]
.into_iter()
.chain(vec![b'D'; MAX_VALUE_SIZE])
.chain(vec![b'E'; MAX_VALUE_SIZE])
.collect::<Vec<_>>(),
vec![b'F'; MAX_VALUE_SIZE]
.into_iter()
.chain(vec![b'G'; MAX_VALUE_SIZE])
.chain(vec![b'H'; MAX_VALUE_SIZE + 1])
.collect::<Vec<_>>(),
]
.into_iter()
.enumerate()
{
// Write value
let test_len = value.len();
db.write(
BatchBuilder::new()
.with_account_id(0)
.with_collection(Collection::Email)
.with_document(0)
.set(ValueClass::Property(1), value.as_slice())
.set(ValueClass::Property(0), "check1".as_bytes())
.set(ValueClass::Property(2), "check2".as_bytes())
.build_all(),
)
.await
.unwrap();
// Fetch value
assert_eq!(
String::from_utf8(value).unwrap(),
db.get_value::<String>(ValueKey {
account_id: 0,
collection: 0,
document_id: 0,
class: ValueClass::Property(1),
})
.await
.unwrap()
.unwrap_or_else(|| panic!("no value for test {test_num} with value length {test_len}")),
"failed for test {test_num} with value length {test_len}"
);
// Delete value
db.write(
BatchBuilder::new()
.with_account_id(0)
.with_collection(Collection::Email)
.with_document(0)
.clear(ValueClass::Property(1))
.build_all(),
)
.await
.unwrap();
// Make sure value is deleted
assert_eq!(
None,
db.get_value::<String>(ValueKey {
account_id: 0,
collection: 0,
document_id: 0,
class: ValueClass::Property(1),
})
.await
.unwrap()
);
// Make sure other values are still there
for (class, value) in [
(ValueClass::Property(0), "check1"),
(ValueClass::Property(2), "check2"),
] {
assert_eq!(
Some(value.to_string()),
db.get_value::<String>(ValueKey {
account_id: 0,
collection: 0,
document_id: 0,
class,
})
.await
.unwrap()
);
}
// Delete everything
let mut batch = BatchBuilder::new();
batch
.with_account_id(0)
.with_collection(Collection::Email)
.with_account_id(0)
.with_document(0)
.clear(ValueClass::Property(0))
.clear(ValueClass::Property(2))
.clear(ValueClass::Property(3))
.clear(ValueClass::Quota)
.clear(ValueClass::ChangeId);
for document_id in 0..1000 {
batch
.with_document(document_id)
.clear(ValueClass::Property(5));
}
db.write(batch.build_all()).await.unwrap();
// Make sure everything is deleted
store_assert_is_empty(&db, db.clone().into(), false).await;
}
}
tests/src/store/query.rs
+811
View File
@@ -0,0 +1,811 @@
/*
* SPDX-FileCopyrightText: 2020 Stalwart Labs LLC <[email protected]>
*
* SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL
*/
use crate::{store::deflate_test_resource, utils::server::TestServer};
use ahash::AHashSet;
use nlp::language::Language;
use std::{
io::Write,
sync::{Arc, Mutex},
time::Instant,
};
use store::{
SearchStore,
ahash::AHashMap,
rand::{self, RngExt, distr::Alphanumeric},
roaring::RoaringBitmap,
search::{
EmailSearchField, IndexDocument, SearchComparator, SearchField, SearchFilter,
SearchOperator, SearchQuery, SearchValue, TracingSearchField,
},
write::SearchIndex,
};
use utils::map::vec_map::VecMap;
pub const FIELDS: [&str; 20] = [
"id",
"accession_number",
"artist",
"artistRole",
"artistId",
"title",
"dateText",
"medium",
"creditLine",
"year",
"acquisitionYear",
"dimensions",
"width",
"height",
"depth",
"units",
"inscription",
"thumbnailCopyright",
"thumbnailUrl",
"url",
];
/*
"title", // Subject
"year". // ReceivedAt
"width", // Size
"height", // SentAt
"artist" // Headers
"artistRole" // Cc
"medium", // From
"creditLine" // Body
"acquisitionYear" // Bcc
"accession_number" // To
*/
const FIELD_MAPPINGS: [EmailSearchField; 20] = [
EmailSearchField::HasAttachment, // "id",
EmailSearchField::To, // "accession_number",
EmailSearchField::Headers, // "artist",
EmailSearchField::Cc, // "artistRole",
EmailSearchField::HasAttachment, // "artistId",
EmailSearchField::Subject, // "title",
EmailSearchField::HasAttachment, // "dateText",
EmailSearchField::From, // "medium",
EmailSearchField::Body, // "creditLine",
EmailSearchField::ReceivedAt, // "year",
EmailSearchField::Bcc, // "acquisitionYear",
EmailSearchField::HasAttachment, // "dimensions",
EmailSearchField::Size, // "width",
EmailSearchField::SentAt, // "height",
EmailSearchField::HasAttachment, // "depth",
EmailSearchField::HasAttachment, // "units",
EmailSearchField::HasAttachment, // "inscription",
EmailSearchField::HasAttachment, // "thumbnailCopyright",
EmailSearchField::HasAttachment, // "thumbnailUrl",
EmailSearchField::HasAttachment, // "url",
];
const ALL_IDS: &[&str] = &[
"p11293", "p79426", "p79427", "p79428", "p79429", "p79430", "d05503", "d00399", "d05352",
"p01764", "t05843", "n02478", "n02479", "n03568", "n03658", "n04327", "n04328", "n04721",
"n04739", "n05095", "n05096", "n05145", "n05157", "n05158", "n05159", "n05298", "n05303",
"n06070", "t01181", "t03571", "t05805", "t05806", "t12147", "t12154", "t12155", "ar00039",
"t12600", "p80203", "t13209", "t13560", "t13561", "t13655", "t13811", "p13352", "p13351",
"p13350", "p13349", "p13348", "p13347", "p13346", "p13345", "p13344", "p13342", "p13341",
"p13340", "p13339", "p13338", "p13337", "p13336", "p13335", "p13334", "p13333", "p13332",
"p13331", "p13330", "p13329", "p13328", "p13327", "p13326", "p13325", "p13324", "p13323",
"t13786", "p13322", "p13321", "p13320", "p13319", "p13318", "p13317", "p13316", "p13315",
"p13314", "t13588", "t13587", "t13586", "t13585", "t13584", "t13540", "t13444", "ar01154",
"ar01153", "t03681", "t12601", "ar00166", "t12625", "t12915", "p04182", "t06483", "ar00703",
"t07671", "ar00021", "t05557", "t07918", "p06298", "p05465", "p06640", "t12855", "t01355",
"t12800", "t12557", "t02078", "ar00052", "ar00627", "t00352", "t07275", "t12318", "t04931",
"t13683", "t13686", "t13687", "t13688", "t13689", "t13690", "t13691", "t13769", "t13773",
"t07151", "t13684", "t07523", "t12369", "t12567", "ar00627", "ar00052", "t00352", "t07275",
"t12318", "t04931", "t13683", "t13686", "t13687", "t13688", "t13689", "t13690", "t13691",
"t07766", "t07918", "t12993", "ar00044", "t13326", "t07614", "t12414",
];
#[allow(clippy::mutex_atomic)]
pub async fn test(test: &TestServer) {
let store = test.server.search_store().clone();
println!("Running Store query tests...");
let pool = rayon::ThreadPoolBuilder::new()
.num_threads(8)
.build()
.unwrap();
let now = Instant::now();
let documents = Arc::new(Mutex::new(Vec::new()));
let mut mask = RoaringBitmap::new();
let mut fields = AHashMap::new();
// Global ids test
println!("Running global id filtering tests...");
test_global(store.clone()).await;
// Large document insert test
println!("Running large document insert tests...");
let mut large_text = String::with_capacity(20 * 1024 * 1024);
while large_text.len() < 20 * 1024 * 1024 {
let word = rand::rng()
.sample_iter(&Alphanumeric)
.take(rand::rng().random_range(3..10))
.map(char::from)
.collect::<String>();
large_text.push_str(&word);
large_text.push(' ');
}
let mut document = IndexDocument::new(SearchIndex::Email)
.with_account_id(1)
.with_document_id(1);
for field in [
EmailSearchField::From,
EmailSearchField::To,
EmailSearchField::Cc,
EmailSearchField::Bcc,
EmailSearchField::Subject,
] {
document.index_text(field, &large_text[..10 * 1024], Language::English);
}
for field in [EmailSearchField::Body, EmailSearchField::Attachment] {
document.index_text(field, &large_text, Language::English);
}
for field in [
EmailSearchField::ReceivedAt,
EmailSearchField::SentAt,
EmailSearchField::Size,
] {
document.index_unsigned(field, rand::rng().random_range(100u64..1_000_000u64));
}
store.index(vec![document]).await.unwrap();
// Refresh
if let SearchStore::ElasticSearch(store) = &store {
store.refresh_index(SearchIndex::Email).await.unwrap();
}
println!("Running account filtering tests...");
let filter_ids = std::env::var("QUICK_TEST").is_ok().then(|| {
let mut ids = AHashSet::new();
for &id in ALL_IDS {
ids.insert(id.to_string());
let id = id.as_bytes();
if id.last().unwrap() > &b'0' {
let mut alt_id = id.to_vec();
*alt_id.last_mut().unwrap() -= 1;
ids.insert(String::from_utf8(alt_id).unwrap());
}
if id.last().unwrap() < &b'9' {
let mut alt_id = id.to_vec();
*alt_id.last_mut().unwrap() += 1;
ids.insert(String::from_utf8(alt_id).unwrap());
}
}
ids
});
pool.scope_fifo(|s| {
for (document_id, record) in csv::ReaderBuilder::new()
.has_headers(true)
.from_reader(&deflate_test_resource("artwork_data.csv.gz")[..])
.records()
.enumerate()
{
let record = record.unwrap();
let documents = documents.clone();
if let Some(filter_ids) = &filter_ids {
let id = record.get(1).unwrap().to_lowercase();
if !filter_ids.contains(&id) {
continue;
}
}
s.spawn_fifo(move |_| {
let mut document = IndexDocument::new(SearchIndex::Email)
.with_account_id(0)
.with_document_id(document_id as u32);
for (pos, field) in record.iter().enumerate() {
match FIELD_MAPPINGS[pos] {
EmailSearchField::From
| EmailSearchField::To
| EmailSearchField::Cc
| EmailSearchField::Bcc => {
document.index_text(
FIELD_MAPPINGS[pos].clone(),
&field.to_lowercase(),
Language::None,
);
}
EmailSearchField::Subject
| EmailSearchField::Body
| EmailSearchField::Attachment => {
document.index_text(
FIELD_MAPPINGS[pos].clone(),
&field
.replace(|ch: char| !ch.is_alphanumeric(), " ")
.to_lowercase(),
Language::English,
);
}
EmailSearchField::Headers => {
document.insert_key_value(
EmailSearchField::Headers,
"artist",
field.to_lowercase(),
);
}
EmailSearchField::ReceivedAt
| EmailSearchField::SentAt
| EmailSearchField::Size => {
document.index_unsigned(
FIELD_MAPPINGS[pos].clone(),
field.parse::<u64>().unwrap_or(0),
);
}
_ => {
continue;
}
};
}
documents.lock().unwrap().push(document);
});
}
});
println!(
"Parsed {} entries in {} ms.",
documents.lock().unwrap().len(),
now.elapsed().as_millis()
);
let now = Instant::now();
let batches = documents.lock().unwrap().drain(..).collect::<Vec<_>>();
print!("Inserting... ",);
let mut chunks = Vec::new();
let mut chunk = Vec::new();
for document in batches {
let mut document_id = None;
let mut to_field = None;
for (key, value) in document.fields() {
if key == &SearchField::DocumentId {
if let SearchValue::Uint(id) = value {
document_id = Some(*id as u32);
}
} else if key == &SearchField::Email(EmailSearchField::To)
&& let SearchValue::Text { value, .. } = value
{
to_field = Some(value.to_string());
}
}
let document_id = document_id.unwrap();
let to_field = to_field.unwrap();
mask.insert(document_id);
fields.insert(document_id, to_field);
chunk.push(document);
if chunk.len() == 10 {
chunks.push(chunk);
chunk = Vec::new();
}
}
if !chunk.is_empty() {
chunks.push(chunk);
}
if test.is_reset() {
let mut tasks = Vec::new();
for chunk in chunks {
let chunk_instance = Instant::now();
tasks.push({
let db = store.clone();
tokio::spawn(async move { db.index(chunk).await })
});
if tasks.len() == 100 {
for handle in tasks {
handle.await.unwrap().unwrap();
}
print!(" [{} ms]", chunk_instance.elapsed().as_millis());
std::io::stdout().flush().unwrap();
tasks = Vec::new();
}
}
if !tasks.is_empty() {
for handle in tasks {
handle.await.unwrap().unwrap();
}
}
// Refresh
if let SearchStore::ElasticSearch(store) = &store {
store.refresh_index(SearchIndex::Email).await.unwrap();
}
println!("\nInsert took {} ms.", now.elapsed().as_millis());
}
if store.internal_fts().is_none() {
let ids = store
.query_account(
SearchQuery::new(SearchIndex::Email)
.with_filters(vec![SearchFilter::eq(SearchField::AccountId, 0u32)])
.with_comparator(SearchComparator::ascending(EmailSearchField::ReceivedAt))
.with_mask(mask.clone()),
)
.await
.unwrap()
.into_iter()
.collect::<RoaringBitmap>();
assert_eq!(ids, mask);
let ids = store
.query_account(
SearchQuery::new(SearchIndex::Email)
.with_filters(vec![
SearchFilter::eq(SearchField::AccountId, 0u32),
SearchFilter::ge(SearchField::DocumentId, 0u32),
])
.with_mask(mask.clone()),
)
.await
.unwrap()
.into_iter()
.collect::<RoaringBitmap>();
assert_eq!(ids, mask);
}
println!("Running account filter tests...");
let now = Instant::now();
test_filter(store.clone(), &fields, &mask).await;
println!("Filtering took {} ms.", now.elapsed().as_millis());
println!("Running account sort tests...");
let now = Instant::now();
test_sort(store.clone(), &fields, &mask).await;
println!("Sorting took {} ms.", now.elapsed().as_millis());
println!("Running unindex tests...");
let now = Instant::now();
test_unindex(store.clone(), &fields).await;
println!("Unindexing took {} ms.", now.elapsed().as_millis());
}
async fn test_filter(store: SearchStore, fields: &AHashMap<u32, String>, mask: &RoaringBitmap) {
let can_stem = !store.is_mysql();
let can_negate_text = !store.is_meilisearch();
let tests = [
(
vec![
SearchFilter::eq(SearchField::AccountId, 0u32),
SearchFilter::has_english_text(EmailSearchField::Subject, "water"),
SearchFilter::eq(EmailSearchField::ReceivedAt, 1979u32),
],
vec!["p11293"],
),
(
vec![
SearchFilter::eq(SearchField::AccountId, 0u32),
SearchFilter::has_keyword(EmailSearchField::From, "gelatin"),
SearchFilter::gt(EmailSearchField::ReceivedAt, 2000u32),
SearchFilter::lt(EmailSearchField::Size, 180u32),
SearchFilter::gt(EmailSearchField::Size, 0u32),
],
vec!["p79426", "p79427", "p79428", "p79429", "p79430"],
),
(
vec![
SearchFilter::eq(SearchField::AccountId, 0u32),
SearchFilter::has_english_text(EmailSearchField::Subject, "'rustic bridge'"),
],
vec!["d05503"],
),
(
vec![
SearchFilter::eq(SearchField::AccountId, 0u32),
SearchFilter::has_english_text(EmailSearchField::Subject, "'rustic'"),
SearchFilter::has_english_text(
EmailSearchField::Subject,
if can_stem { "study" } else { "studies" },
),
],
vec!["d00399", "d05352"],
),
(
vec![
SearchFilter::eq(SearchField::AccountId, 0u32),
SearchFilter::cond(
EmailSearchField::Headers,
SearchOperator::Contains,
SearchValue::KeyValues(VecMap::from_iter([(
"artist".to_string(),
"kunst, mauro".to_string(),
)])),
),
SearchFilter::has_keyword(EmailSearchField::Cc, "artist"),
SearchFilter::Or,
SearchFilter::eq(EmailSearchField::ReceivedAt, 1969u32),
SearchFilter::eq(EmailSearchField::ReceivedAt, 1971u32),
SearchFilter::End,
],
vec!["p01764", "t05843"],
),
(
vec![
SearchFilter::eq(SearchField::AccountId, 0u32),
SearchFilter::Not,
SearchFilter::has_keyword(EmailSearchField::From, "oil"),
SearchFilter::End,
SearchFilter::has_english_text(
EmailSearchField::Body,
if can_stem { "bequeath" } else { "bequeathed" },
),
SearchFilter::Or,
SearchFilter::And,
SearchFilter::ge(EmailSearchField::ReceivedAt, 1900u32),
SearchFilter::lt(EmailSearchField::ReceivedAt, 1910u32),
SearchFilter::End,
SearchFilter::And,
SearchFilter::ge(EmailSearchField::ReceivedAt, 2000u32),
SearchFilter::lt(EmailSearchField::ReceivedAt, 2010u32),
SearchFilter::End,
SearchFilter::End,
],
vec![
"n02478", "n02479", "n03568", "n03658", "n04327", "n04328", "n04721", "n04739",
"n05095", "n05096", "n05145", "n05157", "n05158", "n05159", "n05298", "n05303",
"n06070", "t01181", "t03571", "t05805", "t05806", "t12147", "t12154", "t12155",
],
),
(
vec![
SearchFilter::And,
SearchFilter::eq(SearchField::AccountId, 0u32),
SearchFilter::cond(
EmailSearchField::Headers,
SearchOperator::Contains,
SearchValue::KeyValues(VecMap::from_iter([(
"artist".to_string(),
"warhol".to_string(),
)])),
),
SearchFilter::Not,
SearchFilter::has_english_text(EmailSearchField::Subject, "'campbell'"),
SearchFilter::End,
SearchFilter::Not,
SearchFilter::Or,
SearchFilter::gt(EmailSearchField::ReceivedAt, 1980u32),
SearchFilter::And,
SearchFilter::gt(EmailSearchField::Size, 500u32),
SearchFilter::gt(EmailSearchField::SentAt, 500u32),
SearchFilter::End,
SearchFilter::End,
SearchFilter::End,
SearchFilter::eq(EmailSearchField::Bcc, "2008".to_string()),
SearchFilter::End,
],
vec!["ar00039", "t12600"],
),
(
if can_stem {
vec![
SearchFilter::eq(SearchField::AccountId, 0u32),
SearchFilter::has_english_text(EmailSearchField::Subject, "study"),
SearchFilter::has_keyword(EmailSearchField::From, "paper"),
SearchFilter::has_english_text(EmailSearchField::Body, "'purchased'"),
SearchFilter::Not,
SearchFilter::Or,
SearchFilter::has_english_text(EmailSearchField::Subject, "'anatomical'"),
SearchFilter::has_english_text(EmailSearchField::Subject, "'discarded'"),
SearchFilter::has_english_text(EmailSearchField::Subject, "'untitled'"),
SearchFilter::has_english_text(EmailSearchField::Subject, "'girl'"),
SearchFilter::End,
SearchFilter::End,
SearchFilter::gt(EmailSearchField::ReceivedAt, 1900u32),
SearchFilter::gt(EmailSearchField::Bcc, "2008".to_string()),
]
} else {
vec![
SearchFilter::eq(SearchField::AccountId, 0u32),
SearchFilter::Or,
SearchFilter::has_english_text(EmailSearchField::Subject, "study"),
SearchFilter::has_english_text(EmailSearchField::Subject, "studies"),
SearchFilter::End,
SearchFilter::has_keyword(EmailSearchField::From, "paper"),
SearchFilter::has_english_text(EmailSearchField::Body, "'purchased'"),
SearchFilter::Not,
SearchFilter::Or,
SearchFilter::has_english_text(EmailSearchField::Subject, "'anatomical'"),
SearchFilter::has_english_text(EmailSearchField::Subject, "'discarded'"),
SearchFilter::has_english_text(EmailSearchField::Subject, "'untitled'"),
SearchFilter::has_english_text(EmailSearchField::Subject, "'girl'"),
SearchFilter::End,
SearchFilter::End,
SearchFilter::gt(EmailSearchField::ReceivedAt, 1900u32),
SearchFilter::gt(EmailSearchField::Bcc, "2008".to_string()),
]
},
vec!["p80203", "t13209", "t13560", "t13561"],
),
];
for (filters, expected_results) in tests {
if !can_negate_text && has_negated_text(&filters) {
continue;
}
//println!("Running test: {:?}", filter);
let ids = store
.query_account(
SearchQuery::new(SearchIndex::Email)
.with_filters(filters)
.with_comparator(SearchComparator::ascending(EmailSearchField::To))
.with_mask(mask.clone()),
)
.await
.unwrap();
let mut results = Vec::new();
for document_id in ids {
results.push(fields.get(&document_id).unwrap());
}
assert_eq!(results, expected_results);
}
}
fn has_negated_text(filters: &[SearchFilter]) -> bool {
let mut stack = Vec::new();
let mut negated = 0;
for filter in filters {
match filter {
SearchFilter::Not => {
stack.push(true);
negated += 1;
}
SearchFilter::And | SearchFilter::Or => {
stack.push(false);
}
SearchFilter::End => {
if stack.pop().unwrap_or(false) {
negated -= 1;
}
}
SearchFilter::Operator {
op: SearchOperator::Equal | SearchOperator::Contains,
value: SearchValue::Text { .. },
..
} if negated > 0 => {
return true;
}
_ => (),
}
}
false
}
async fn test_sort(store: SearchStore, fields: &AHashMap<u32, String>, mask: &RoaringBitmap) {
let is_reversed = store.is_postgres();
let tests = [
(
vec![
SearchFilter::eq(SearchField::AccountId, 0u32),
SearchFilter::gt(EmailSearchField::ReceivedAt, 0u32),
SearchFilter::gt(EmailSearchField::Bcc, "0000".to_string()),
SearchFilter::gt(EmailSearchField::Size, 0u32),
],
vec![
SearchComparator::descending(EmailSearchField::ReceivedAt),
SearchComparator::ascending(EmailSearchField::Bcc),
SearchComparator::ascending(EmailSearchField::Size),
SearchComparator::descending(EmailSearchField::To),
],
vec![
"t13655", "t13811", "p13352", "p13351", "p13350", "p13349", "p13348", "p13347",
"p13346", "p13345", "p13344", "p13342", "p13341", "p13340", "p13339", "p13338",
"p13337", "p13336", "p13335", "p13334", "p13333", "p13332", "p13331", "p13330",
"p13329", "p13328", "p13327", "p13326", "p13325", "p13324", "p13323", "t13786",
"p13322", "p13321", "p13320", "p13319", "p13318", "p13317", "p13316", "p13315",
"p13314", "t13588", "t13587", "t13586", "t13585", "t13584", "t13540", "t13444",
"ar01154", "ar01153",
],
),
(
vec![
SearchFilter::eq(SearchField::AccountId, 0u32),
SearchFilter::gt(EmailSearchField::Size, 0u32),
SearchFilter::gt(EmailSearchField::SentAt, 0u32),
],
vec![
SearchComparator::descending(EmailSearchField::Size),
SearchComparator::ascending(EmailSearchField::SentAt),
],
vec![
"t03681", "t12601", "ar00166", "t12625", "t12915", "p04182", "t06483", "ar00703",
"t07671", "ar00021", "t05557", "t07918", "p06298", "p05465", "p06640", "t12855",
"t01355", "t12800", "t12557", "t02078",
],
),
(
vec![SearchFilter::eq(SearchField::AccountId, 0u32)],
vec![
SearchComparator::descending(EmailSearchField::From),
SearchComparator::descending(EmailSearchField::Cc),
SearchComparator::ascending(EmailSearchField::To),
],
if is_reversed {
vec![
"ar00052", "ar00627", "t00352", "t07275", "t12318", "t04931", "t13683",
"t13686", "t13687", "t13688", "t13689", "t13690", "t13691", "t13769", "t13773",
"t07151", "t13684", "t07523", "t12369", "t12567",
]
} else {
vec![
"ar00627", "ar00052", "t00352", "t07275", "t12318", "t04931", "t13683",
"t13686", "t13687", "t13688", "t13689", "t13690", "t13691", "t07766", "t07918",
"t12993", "ar00044", "t13326", "t07614", "t12414",
]
},
),
];
for (filters, comparators, expected_results) in tests {
//println!("Running test: {:?}", sort);
let ids = store
.query_account(
SearchQuery::new(SearchIndex::Email)
.with_filters(filters)
.with_comparators(comparators)
.with_mask(mask.clone()),
)
.await
.unwrap();
let mut results = Vec::new();
for document_id in ids.into_iter().take(expected_results.len()) {
results.push(fields.get(&document_id).unwrap());
}
assert_eq!(results, expected_results);
}
}
async fn test_unindex(store: SearchStore, fields: &AHashMap<u32, String>) {
let ids = store
.query_account(
SearchQuery::new(SearchIndex::Email)
.with_mask(RoaringBitmap::from_iter(fields.keys().copied()))
.with_filters(vec![
SearchFilter::has_keyword(EmailSearchField::From, "gelatin"),
SearchFilter::gt(EmailSearchField::ReceivedAt, 2000u32),
SearchFilter::lt(EmailSearchField::Size, 180u32),
SearchFilter::gt(EmailSearchField::Size, 0u32),
])
.with_account_id(0),
)
.await
.unwrap();
assert!(!ids.is_empty());
let expected_count = ids.len().saturating_sub(10);
let mut query = SearchQuery::new(SearchIndex::Email)
.with_account_id(0)
.with_filter(SearchFilter::Or);
for id in ids.into_iter().take(10) {
query = query.with_filter(SearchFilter::eq(SearchField::DocumentId, id));
}
query = query.with_filter(SearchFilter::End);
store.unindex(query).await.unwrap();
// Refresh
if let SearchStore::ElasticSearch(store) = &store {
store.refresh_index(SearchIndex::Email).await.unwrap();
}
assert_eq!(
store
.query_account(
SearchQuery::new(SearchIndex::Email)
.with_filters(vec![
SearchFilter::has_keyword(EmailSearchField::From, "gelatin"),
SearchFilter::gt(EmailSearchField::ReceivedAt, 2000u32),
SearchFilter::lt(EmailSearchField::Size, 180u32),
SearchFilter::gt(EmailSearchField::Size, 0u32),
])
.with_account_id(0)
.with_mask(RoaringBitmap::from_iter(fields.keys().copied())),
)
.await
.unwrap()
.len(),
expected_count
);
}
async fn test_global(store: SearchStore) {
// Insert global ids
for (id, queue_id, etyp, keywords) in [
(0, 1000u64, 1u64, "init start"),
(1, 1000u64, 2u64, "init complete"),
(2, 1001u64, 1u64, "process start"),
(3, 1001u64, 2u64, "process complete"),
(4, 1002u64, 1u64, "cleanup start"),
(5, 1002u64, 2u64, "cleanup complete"),
] {
let mut document = IndexDocument::new(SearchIndex::Tracing).with_id(id);
document.index_unsigned(TracingSearchField::QueueId, queue_id);
document.index_unsigned(TracingSearchField::EventType, etyp);
document.index_text(TracingSearchField::Keywords, keywords, Language::None);
store.index(vec![document]).await.unwrap();
}
// Refresh
if let SearchStore::ElasticSearch(store) = &store {
store.refresh_index(SearchIndex::Tracing).await.unwrap();
}
// Query all
assert_eq!(
store
.query_global(
SearchQuery::new(SearchIndex::Tracing)
.with_filter(SearchFilter::ge(SearchField::Id, 0u64))
)
.await
.unwrap()
.into_iter()
.collect::<AHashSet<_>>(),
AHashSet::from_iter([0, 1, 2, 3, 4, 5])
);
// Query with filter
assert_eq!(
store
.query_global(
SearchQuery::new(SearchIndex::Tracing)
.with_filter(SearchFilter::gt(SearchField::Id, 1u64))
.with_filter(SearchFilter::lt(SearchField::Id, 5u64))
.with_filter(SearchFilter::has_keyword(
TracingSearchField::Keywords,
"start",
)),
)
.await
.unwrap()
.into_iter()
.collect::<AHashSet<_>>(),
AHashSet::from_iter([2, 4])
);
// Delete by filter
store
.unindex(
SearchQuery::new(SearchIndex::Tracing)
.with_filter(SearchFilter::lt(SearchField::Id, 3u64)),
)
.await
.unwrap();
// Refresh
if let SearchStore::ElasticSearch(store) = &store {
store.refresh_index(SearchIndex::Tracing).await.unwrap();
}
assert_eq!(
store
.query_global(
SearchQuery::new(SearchIndex::Tracing)
.with_filter(SearchFilter::ge(SearchField::Id, 0u64))
)
.await
.unwrap()
.into_iter()
.collect::<AHashSet<_>>(),
AHashSet::from_iter([3, 4, 5])
);
}
tests/src/store/registry.rs
+804
View File
@@ -0,0 +1,804 @@
/*
* SPDX-FileCopyrightText: 2020 Stalwart Labs LLC <[email protected]>
*
* SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL
*/
use crate::utils::{registry::UnwrapRegistryId, server::TestServer};
use jmap_tools::JsonPointer;
use registry::{
jmap::{IntoValue, JmapValue, JsonPointerPatch, MaybeUnpatched, RegistryJsonPatch},
pickle::{Pickle, PickledStream},
schema::{
enums::{AccountType, Locale, Permission, StorageQuota},
prelude::{Object, ObjectType, Property},
structs::{
Account, CertificateManagement, Credential, CredentialPermissions,
CredentialPermissionsList, CustomRoles, DkimManagement, DnsManagement, Domain,
EmailAlias, EncryptionAtRest, EncryptionSettings, GroupAccount, MailingList,
PasswordCredential, Permissions, PermissionsList, PublicKey, SecondaryCredential,
SieveUserScript, UserAccount, UserRoles,
},
},
types::{
EnumImpl, ObjectImpl, datetime::UTCDateTime, id::ObjectId, ipmask::IpAddrOrMask,
list::List, map::Map,
},
};
use std::str::FromStr;
use store::{
registry::{
RegistryQuery,
write::{RegistryWrite, RegistryWriteResult},
},
write::now,
};
use types::id::Id;
use utils::map::vec_map::VecMap;
pub async fn test(test: &TestServer) {
let r = test.server.registry();
println!("Registry tests...");
test_patch_regressions();
// Pickle-unpickle test
let mut account = Account::User(UserAccount {
aliases: List::from_iter([
EmailAlias {
description: "Test Alias 1".to_string().into(),
domain_id: 1000u64.into(),
enabled: true,
name: "alias1".into(),
},
EmailAlias {
description: "Test Alias 2".to_string().into(),
domain_id: 1001u64.into(),
enabled: true,
name: "alias2".into(),
},
]),
created_at: UTCDateTime::now(),
credentials: List::from_iter([
Credential::Password(PasswordCredential {
allowed_ips: Map::new(vec![IpAddrOrMask::from_str("192.168.1.1").unwrap()]),
credential_id: 3u64.into(),
expires_at: None,
otp_auth: "otpauth://totp/test?secret=SECRET".to_string().into(),
secret: "secret".into(),
}),
Credential::AppPassword(SecondaryCredential {
allowed_ips: Map::new(vec![IpAddrOrMask::from_str("192.168.1.0/24").unwrap()]),
created_at: UTCDateTime::now(),
credential_id: 4u64.into(),
description: "App Password".into(),
expires_at: Some(UTCDateTime::from_timestamp((now() + 1000) as i64)),
permissions: CredentialPermissions::Disable(CredentialPermissionsList {
permissions: Map::new(vec![
Permission::Authenticate,
Permission::ActionClassifySpam,
]),
}),
secret: "app_password_secret".into(),
}),
]),
description: "This is a test Account".to_string().into(),
domain_id: 1004u64.into(),
encryption_at_rest: EncryptionAtRest::Aes128(EncryptionSettings {
allow_spam_training: true,
encrypt_on_append: false,
public_key: 0u64.into(),
}),
external_id: "8f7c1e2a-4b3d-4f1a-9c2e-7d5b6a8f0e11".to_string().into(),
locale: Locale::EnUS,
member_group_ids: Map::new(vec![2000u64.into(), 2001u64.into()]),
member_tenant_id: None,
name: "user".into(),
permissions: Permissions::Merge(PermissionsList {
disabled_permissions: Map::new(vec![Permission::Impersonate]),
enabled_permissions: Map::new(vec![Permission::JmapBlobGet]),
}),
quotas: VecMap::from_iter([
(StorageQuota::MaxDiskQuota, 1024u64),
(StorageQuota::MaxApiKeys, 3u64),
]),
roles: UserRoles::Custom(CustomRoles {
role_ids: Map::new(vec![5000u64.into()]),
}),
time_zone: None,
});
let account_pickle = account.to_pickled_vec();
assert_eq!(
account,
Account::unpickle(&mut PickledStream::new(&account_pickle).unwrap()).unwrap()
);
// Pickle compression test
let script = SieveUserScript {
contents: "A".repeat(100_000),
description: "B".repeat(100_000).into(),
is_active: true,
name: "C".repeat(100_000),
};
let script_pickle = script.to_pickled_vec();
assert!(
script_pickle.len() < 8_192,
"Pickle was not compressed: {} bytes",
script_pickle.len()
);
assert_eq!(
script,
SieveUserScript::unpickle(&mut PickledStream::new(&script_pickle).unwrap()).unwrap()
);
// Create a domain and a group
let domain_id = r
.write(RegistryWrite::insert(
&Domain {
name: "test.org".into(),
certificate_management: CertificateManagement::Manual,
dns_management: DnsManagement::Manual,
dkim_management: DkimManagement::Manual,
is_enabled: true,
..Default::default()
}
.into(),
))
.await
.unwrap()
.unwrap_id(trc::location!());
let domain_id_2 = r
.write(RegistryWrite::insert(
&Domain {
name: "test.net".into(),
certificate_management: CertificateManagement::Manual,
dns_management: DnsManagement::Manual,
dkim_management: DkimManagement::Manual,
is_enabled: true,
..Default::default()
}
.into(),
))
.await
.unwrap()
.unwrap_id(trc::location!());
let group_id = r
.write(RegistryWrite::insert(
&Account::Group(GroupAccount {
name: "group".into(),
domain_id,
..Default::default()
})
.into(),
))
.await
.unwrap()
.unwrap_id(trc::location!());
// Inserting an account linking non-existing ids should fail
test.assert_registry_insert_error(
account.clone(),
RegistryWriteResult::InvalidForeignKey {
object_id: ObjectId::new(ObjectType::Account, Id::new(2000)),
},
trc::location!(),
)
.await;
account.assert_patch(
&format!("memberGroupIds/{}", Id::new(2000)),
false,
trc::location!(),
);
account.assert_patch(
&format!("memberGroupIds/{}", Id::new(2001)),
false,
trc::location!(),
);
account.assert_patch(
&format!("memberGroupIds/{}", group_id),
true,
trc::location!(),
);
test.assert_registry_insert_error(
account.clone(),
RegistryWriteResult::InvalidForeignKey {
object_id: ObjectId::new(ObjectType::Domain, Id::new(1000)),
},
trc::location!(),
)
.await;
account.assert_patch("aliases/0/domainId", domain_id, trc::location!());
account.assert_patch("aliases/1/domainId", domain_id, trc::location!());
test.assert_registry_insert_error(
account.clone(),
RegistryWriteResult::InvalidForeignKey {
object_id: ObjectId::new(ObjectType::Domain, Id::new(1004)),
},
trc::location!(),
)
.await;
account.assert_patch("domainId", domain_id, trc::location!());
test.assert_registry_insert_error(
account.clone(),
RegistryWriteResult::InvalidForeignKey {
object_id: ObjectId::new(ObjectType::PublicKey, Id::new(0)),
},
trc::location!(),
)
.await;
account.assert_patch(
"encryptionAtRest",
EncryptionAtRest::Disabled.into_value(),
trc::location!(),
);
test.assert_registry_insert_error(
account.clone(),
RegistryWriteResult::InvalidForeignKey {
object_id: ObjectId::new(ObjectType::Role, Id::new(5000)),
},
trc::location!(),
)
.await;
account.assert_patch("roles", UserRoles::User.into_value(), trc::location!());
let account_id = r
.write(RegistryWrite::insert(&account.into()))
.await
.unwrap()
.unwrap_id(trc::location!());
// Deleting linked objects should fail
test.assert_registry_delete_error(
ObjectType::Domain,
domain_id,
RegistryWriteResult::CannotDeleteLinked {
object_id: ObjectId::new(ObjectType::Domain, domain_id),
linked_objects: vec![
ObjectId::new(ObjectType::Account, group_id),
ObjectId::new(ObjectType::Account, account_id),
],
},
trc::location!(),
)
.await;
// Primary key violations should not be allowed
test.assert_registry_insert_error(
Domain {
name: "test.org".into(),
is_enabled: true,
certificate_management: CertificateManagement::Manual,
dns_management: DnsManagement::Manual,
dkim_management: DkimManagement::Manual,
..Default::default()
},
RegistryWriteResult::PrimaryKeyConflict {
property: Property::Name,
existing_id: ObjectId::new(ObjectType::Domain, domain_id),
},
trc::location!(),
)
.await;
test.assert_registry_insert_error(
Account::Group(GroupAccount {
name: "group".into(),
domain_id,
..Default::default()
}),
RegistryWriteResult::PrimaryKeyConflict {
property: Property::Email,
existing_id: ObjectId::new(ObjectType::Account, group_id),
},
trc::location!(),
)
.await;
test.assert_registry_insert_error(
MailingList {
name: "user".into(),
domain_id,
recipients: Map::new(vec!["[email protected]".into()]),
..Default::default()
},
RegistryWriteResult::PrimaryKeyConflict {
property: Property::Email,
existing_id: ObjectId::new(ObjectType::Account, account_id),
},
trc::location!(),
)
.await;
test.assert_registry_insert_error(
MailingList {
name: "mailing-list".into(),
domain_id,
aliases: List::from_iter([EmailAlias {
description: "Test Alias 1".to_string().into(),
domain_id,
enabled: true,
name: "alias1".into(),
}]),
recipients: Map::new(vec!["[email protected]".into()]),
..Default::default()
},
RegistryWriteResult::PrimaryKeyConflict {
property: Property::Email,
existing_id: ObjectId::new(ObjectType::Account, account_id),
},
trc::location!(),
)
.await;
// Create a public key and link it to the account
let pk_id = r
.write(RegistryWrite::insert(
&PublicKey {
account_id,
key: "secret".into(),
description: "Test Key".into(),
..Default::default()
}
.into(),
))
.await
.unwrap()
.unwrap_id(trc::location!());
let old_account = r
.get(ObjectId::new(ObjectType::Account, account_id))
.await
.unwrap()
.unwrap();
let mut account = old_account.clone();
assert_obj_patch(
&mut account,
"encryptionAtRest",
EncryptionAtRest::Aes128(EncryptionSettings {
allow_spam_training: true,
encrypt_on_append: false,
public_key: pk_id,
})
.into_value(),
trc::location!(),
);
r.write(RegistryWrite::update(account_id, &account, &old_account))
.await
.unwrap()
.unwrap_id(trc::location!());
// Search tests
assert_eq!(
r.query::<Vec<Id>>(RegistryQuery::new(ObjectType::Domain))
.await
.unwrap(),
vec![domain_id, domain_id_2]
);
assert_eq!(
r.query::<Vec<Id>>(RegistryQuery::new(ObjectType::Domain).equal_pk(
Property::Name,
"test.org".to_string(),
true,
))
.await
.unwrap(),
vec![domain_id]
);
assert_eq!(
r.query::<Vec<Id>>(RegistryQuery::new(ObjectType::Account))
.await
.unwrap(),
vec![group_id, account_id]
);
assert_eq!(
r.query::<Vec<Id>>(
RegistryQuery::new(ObjectType::Account)
.equal(Property::Type, AccountType::User.to_id())
.text(Property::Text, "this is a test")
.equal(Property::Name, "user")
)
.await
.unwrap(),
vec![account_id]
);
// Sort test
assert_eq!(
r.sort_by_index(ObjectType::Account, Property::Type, None, true)
.await
.unwrap(),
vec![account_id, group_id]
);
assert_eq!(
r.sort_by_index(
ObjectType::Account,
Property::Type,
Some(vec![group_id, account_id]),
true
)
.await
.unwrap(),
vec![account_id, group_id]
);
assert_eq!(
r.sort_by_index(ObjectType::Account, Property::Name, None, true)
.await
.unwrap(),
vec![group_id, account_id]
);
assert_eq!(
r.sort_by_pk(ObjectType::Domain, Property::Name, None, true)
.await
.unwrap(),
vec![domain_id_2, domain_id]
);
assert_eq!(
r.sort_by_pk(
ObjectType::Domain,
Property::Name,
Some(vec![domain_id, domain_id_2]),
true
)
.await
.unwrap(),
vec![domain_id_2, domain_id]
);
// Delete everything
let old_account = r
.get(ObjectId::new(ObjectType::Account, account_id))
.await
.unwrap()
.unwrap();
let mut account = old_account.clone();
assert_obj_patch(
&mut account,
"encryptionAtRest",
EncryptionAtRest::Disabled.into_value(),
trc::location!(),
);
r.write(RegistryWrite::update(account_id, &account, &old_account))
.await
.unwrap()
.unwrap_id(trc::location!());
r.write(RegistryWrite::delete(ObjectId::new(
ObjectType::PublicKey,
pk_id,
)))
.await
.unwrap()
.unwrap_id(trc::location!());
r.write(RegistryWrite::delete(ObjectId::new(
ObjectType::Account,
account_id,
)))
.await
.unwrap()
.unwrap_id(trc::location!());
r.write(RegistryWrite::delete(ObjectId::new(
ObjectType::Account,
group_id,
)))
.await
.unwrap()
.unwrap_id(trc::location!());
r.write(RegistryWrite::delete(ObjectId::new(
ObjectType::Domain,
domain_id,
)))
.await
.unwrap()
.unwrap_id(trc::location!());
r.write(RegistryWrite::delete(ObjectId::new(
ObjectType::Domain,
domain_id_2,
)))
.await
.unwrap()
.unwrap_id(trc::location!());
test.assert_is_empty().await;
}
impl TestServer {
pub async fn assert_registry_insert_error(
&self,
obj: impl Into<Object>,
result: RegistryWriteResult,
location: &str,
) {
let obj = obj.into();
assert_eq!(
self.server
.registry()
.write(RegistryWrite::insert(&obj))
.await
.unwrap(),
result,
"{}",
location
);
}
pub async fn assert_registry_delete_error(
&self,
object_type: ObjectType,
id: Id,
result: RegistryWriteResult,
location: &str,
) {
assert_eq!(
self.server
.registry()
.write(RegistryWrite::delete(ObjectId::new(object_type, id)))
.await
.unwrap(),
result,
"{}",
location
);
}
}
fn test_patch_regressions() {
fn fresh_account() -> Account {
Account::User(UserAccount {
credentials: List::from_iter([
Credential::Password(PasswordCredential {
allowed_ips: Map::new(vec![
IpAddrOrMask::from_str("192.168.1.1").unwrap(),
IpAddrOrMask::from_str("192.168.1.2").unwrap(),
]),
credential_id: 3u64.into(),
expires_at: None,
otp_auth: None,
secret: "secret".into(),
}),
Credential::Password(PasswordCredential {
allowed_ips: Map::new(vec![IpAddrOrMask::from_str("10.0.0.1").unwrap()]),
credential_id: 4u64.into(),
expires_at: None,
otp_auth: None,
secret: "another".into(),
}),
]),
domain_id: 1u64.into(),
name: "patch-target".into(),
..Default::default()
})
}
fn user(account: &Account) -> &UserAccount {
match account {
Account::User(u) => u,
_ => panic!("expected user account"),
}
}
fn user_mut(account: &mut Account) -> &mut UserAccount {
match account {
Account::User(u) => u,
_ => panic!("expected user account"),
}
}
fn password_at(account: &Account, idx: u32) -> &PasswordCredential {
let cred = user(account)
.credentials
.0
.get(&idx)
.expect("credential at index");
match cred {
Credential::Password(p) => p,
_ => panic!("expected password credential at idx {idx}"),
}
}
// Leaf-null patch into a List<T> entry removes only the leaf not the whole entry.
let mut account = fresh_account();
account.assert_patch(
"credentials/0/allowedIps/192.168.1.1",
JmapValue::Null,
trc::location!(),
);
{
let cred = password_at(&account, 0);
assert_eq!(cred.allowed_ips.len(), 1, "one ip should remain");
assert!(
cred.allowed_ips
.contains(&IpAddrOrMask::from_str("192.168.1.2").unwrap()),
"remaining ip survived"
);
assert!(
!cred
.allowed_ips
.contains(&IpAddrOrMask::from_str("192.168.1.1").unwrap()),
"targeted ip removed"
);
}
// The sibling credential is untouched.
{
let cred = password_at(&account, 1);
assert_eq!(cred.allowed_ips.len(), 1);
assert!(
cred.allowed_ips
.contains(&IpAddrOrMask::from_str("10.0.0.1").unwrap())
);
}
// Removing every leaf still leaves the entry in place with an empty map.
let mut account = fresh_account();
account.assert_patch(
"credentials/0/allowedIps/192.168.1.1",
JmapValue::Null,
trc::location!(),
);
account.assert_patch(
"credentials/0/allowedIps/192.168.1.2",
JmapValue::Null,
trc::location!(),
);
{
assert_eq!(
user(&account).credentials.len(),
2,
"credential entry retained"
);
let cred = password_at(&account, 0);
assert!(cred.allowed_ips.is_empty(), "leaf map drained");
}
// Direct removal of a list entry with no remaining segments still works.
let mut account = fresh_account();
account.assert_patch("credentials/0", JmapValue::Null, trc::location!());
{
assert_eq!(user(&account).credentials.len(), 1, "credential 0 removed");
let cred = password_at(&account, 1);
assert_eq!(cred.allowed_ips.len(), 1);
}
// Leaf-null patch into a scalar property of a list entry clears only that property.
let mut account = fresh_account();
{
let cred = user_mut(&mut account)
.credentials
.inner_mut()
.get_mut(&0)
.expect("credential at 0");
if let Credential::Password(p) = cred {
p.expires_at = Some(UTCDateTime::from_timestamp(now() as i64));
}
}
account.assert_patch("credentials/0/expiresAt", JmapValue::Null, trc::location!());
{
let cred = password_at(&account, 0);
assert!(cred.expires_at.is_none(), "expiresAt cleared");
assert_eq!(cred.allowed_ips.len(), 2, "siblings untouched");
}
// Map<T> set-style patches
fn account_with_groups() -> Account {
let mut account = match fresh_account() {
Account::User(u) => u,
_ => unreachable!(),
};
account.member_group_ids = Map::new(vec![Id::new(2000), Id::new(2001)]);
Account::User(account)
}
let mut account = account_with_groups();
account.assert_patch(
&format!("memberGroupIds/{}", Id::new(2000)),
JmapValue::Null,
trc::location!(),
);
assert_eq!(
user(&account).member_group_ids.len(),
1,
"one member removed"
);
assert!(
user(&account).member_group_ids.contains(&Id::new(2001)),
"sibling preserved"
);
let mut account = account_with_groups();
let extra_path = format!("memberGroupIds/{}/extra", Id::new(2000));
let ptr = JsonPointer::parse(&extra_path);
let outcome = account.patch(JsonPointerPatch::new(&ptr), JmapValue::Null);
assert!(outcome.is_err(), "extra segments must error on remove");
assert_eq!(
user(&account).member_group_ids.len(),
2,
"membership unchanged after rejected patch"
);
let mut account = account_with_groups();
let ptr = JsonPointer::parse(&format!("memberGroupIds/{}/extra", Id::new(2002)));
let outcome = account.patch(JsonPointerPatch::new(&ptr), JmapValue::Bool(true));
assert!(outcome.is_err(), "extra segments must error on add");
assert_eq!(
user(&account).member_group_ids.len(),
2,
"membership unchanged after rejected add"
);
// Direct adds and removes still work.
let mut account = account_with_groups();
account.assert_patch(
&format!("memberGroupIds/{}", Id::new(2002)),
true,
trc::location!(),
);
assert_eq!(user(&account).member_group_ids.len(), 3, "member added");
assert!(user(&account).member_group_ids.contains(&Id::new(2002)));
}
trait AssertPatch {
fn assert_patch(&mut self, patch: &str, value: impl Into<JmapValue<'static>>, location: &str);
}
impl<T: RegistryJsonPatch> AssertPatch for T {
fn assert_patch(&mut self, patch: &str, value: impl Into<JmapValue<'static>>, location: &str) {
let ptr = JsonPointer::parse(patch);
let patch = JsonPointerPatch::new(&ptr);
let value = value.into();
match self.patch(patch, value) {
Ok(maybe_unpatched) => {
match maybe_unpatched {
MaybeUnpatched::Patched => {
// Patch succeeded
}
MaybeUnpatched::Unpatched { property, value } => {
panic!(
"Expected patch to succeed but it was unpatched at {}: property: {}, value: {:?}",
location, property, value
);
}
MaybeUnpatched::UnpatchedMany { properties } => {
panic!(
"Expected patch to succeed but it was unpatched at {}: properties: {:?}",
location, properties
);
}
}
}
Err(err) => panic!("Patch failed at {}: {:?}", location, err),
}
}
}
fn assert_obj_patch(
obj: &mut Object,
patch: &str,
value: impl Into<JmapValue<'static>>,
location: &str,
) {
let ptr = JsonPointer::parse(patch);
let patch = JsonPointerPatch::new(&ptr);
let value = value.into();
match obj.patch(patch, value) {
Ok(maybe_unpatched) => {
match maybe_unpatched {
MaybeUnpatched::Patched => {
// Patch succeeded
}
MaybeUnpatched::Unpatched { property, value } => {
panic!(
"Expected patch to succeed but it was unpatched at {}: property: {}, value: {:?}",
location, property, value
);
}
MaybeUnpatched::UnpatchedMany { properties } => {
panic!(
"Expected patch to succeed but it was unpatched at {}: properties: {:?}",
location, properties
);
}
}
}
Err(err) => panic!("Patch failed at {}: {:?}", location, err),
}
}
tests/src/store/sql_timeout.rs
+385
View File
@@ -0,0 +1,385 @@
/*
* SPDX-FileCopyrightText: 2020 Stalwart Labs LLC <[email protected]>
*
* SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL
*/
use crate::utils::{server::TestServer, storage::build_data_store};
use futures::FutureExt;
use registry::schema::structs::DataStore;
use std::{panic::AssertUnwindSafe, time::Duration};
use store::{
IterateParams, Key, Rows, SUBSPACE_COUNTER, SUBSPACE_PROPERTY, Store, U32_LEN, Value, ValueKey,
write::ValueClass,
};
use types::collection::Collection;
const ACCOUNT_ID: u32 = 90210;
const NUM_KEYS: u32 = 300000;
const VALUE_SIZE: usize = 64;
const PROPERTY: u8 = 128;
const STATEMENT_TIMEOUT: Duration = Duration::from_millis(10);
#[derive(Clone, Copy)]
enum Backend {
#[cfg(feature = "postgres")]
Postgres,
#[cfg(feature = "mysql")]
MariaDb,
}
pub async fn test(test: &TestServer) {
let Some(backend) = Backend::detect(test.server.store()) else {
return;
};
println!("Running SQL statement timeout tests...");
let admin = Store::build(backend.data_store(false).await)
.await
.expect("Failed to connect to the statement timeout store");
admin.create_tables().await.unwrap();
admin.delete_range(key(0), key(u32::MAX)).await.unwrap();
for query in [backend.populate_query(), backend.populate_counters_query()] {
admin
.sql_query::<usize>(&query, vec![])
.await
.expect("Failed to populate the statement timeout store");
}
backend.set_timeout(&admin, true).await;
let slow = Store::build(backend.data_store(true).await)
.await
.expect("Failed to open the statement timeout connection pool");
let result = AssertUnwindSafe(scenarios(&slow, backend))
.catch_unwind()
.await;
backend.set_timeout(&admin, false).await;
let mut remaining = 0;
let mut counters = String::new();
if result.is_ok() {
admin
.iterate(
IterateParams::new(key(0), key(u32::MAX)).no_values(),
|_, _| {
remaining += 1;
Ok(true)
},
)
.await
.expect("Failed to count remaining keys");
counters = admin
.sql_query::<Rows>(&backend.count_counters_query(), range_params())
.await
.expect("Failed to count remaining counters")
.rows
.into_iter()
.next()
.and_then(|row| row.values.into_iter().next())
.map(|value| value.to_str().into_owned())
.unwrap_or_default();
}
admin.delete_range(key(0), key(u32::MAX)).await.unwrap();
admin
.sql_query::<usize>(&backend.purge_query(), vec![])
.await
.unwrap();
if let Err(err) = result {
std::panic::resume_unwind(err);
}
assert_eq!(remaining, 0, "Keys left behind by delete_range");
assert_eq!(counters, "0", "Counters left behind by purge_store");
}
async fn scenarios(db: &Store, backend: Backend) {
match db
.sql_query::<Rows>(&backend.scan_query(), range_params())
.await
{
Ok(rows) => panic!(
"Unable to reproduce a statement timeout: scanning {} keys completed within a \
{STATEMENT_TIMEOUT:?} statement timeout",
rows.rows.len()
),
Err(err) => backend.assert_timeout(&err),
}
let mut ids = Vec::with_capacity(NUM_KEYS as usize);
db.iterate(IterateParams::new(key(0), key(u32::MAX)), |key, value| {
let document_id = document_id(key);
assert_eq!(value.len(), VALUE_SIZE, "document {document_id}");
assert_eq!(
value.get(..U32_LEN),
Some(document_id.to_be_bytes().as_slice()),
"document {document_id}"
);
ids.push(document_id);
Ok(true)
})
.await
.expect("Failed to iterate in ascending order");
assert_ids(&ids, 0..NUM_KEYS);
ids.clear();
db.iterate(
IterateParams::new(key(0), key(u32::MAX)).descending(),
|key, _| {
ids.push(document_id(key));
Ok(true)
},
)
.await
.expect("Failed to iterate in descending order");
assert_ids(&ids, (0..NUM_KEYS).rev());
ids.clear();
db.iterate(
IterateParams::new(key(0), key(u32::MAX)).no_values(),
|key, value| {
assert!(value.is_empty());
ids.push(document_id(key));
Ok(true)
},
)
.await
.expect("Failed to iterate over keys");
assert_ids(&ids, 0..NUM_KEYS);
ids.clear();
db.iterate(IterateParams::new(key(0), key(u32::MAX)), |key, _| {
ids.push(document_id(key));
Ok(ids.len() < 10)
})
.await
.expect("Failed to stop iterating");
assert_ids(&ids, 0..10);
match db
.sql_query::<usize>(&backend.delete_query(), range_params())
.await
{
Ok(deleted) => panic!(
"Unable to reproduce a statement timeout: deleting {deleted} keys completed within \
a {STATEMENT_TIMEOUT:?} statement timeout"
),
Err(err) => backend.assert_timeout(&err),
}
db.delete_range(key(0), key(u32::MAX))
.await
.expect("Failed to delete range");
match db.sql_query::<usize>(&backend.purge_query(), vec![]).await {
Ok(deleted) => panic!(
"Unable to reproduce a statement timeout: purging {deleted} counters completed \
within a {STATEMENT_TIMEOUT:?} statement timeout"
),
Err(err) => backend.assert_timeout(&err),
}
db.purge_store().await.expect("Failed to purge store");
}
fn assert_ids(ids: &[u32], expected: impl ExactSizeIterator<Item = u32>) {
assert_eq!(ids.len(), expected.len(), "Unexpected number of keys");
for (position, (id, expected)) in ids.iter().zip(expected).enumerate() {
assert_eq!(*id, expected, "Unexpected key at position {position}");
}
}
fn key(document_id: u32) -> ValueKey<ValueClass> {
ValueKey {
account_id: ACCOUNT_ID,
collection: Collection::Email.into(),
document_id,
class: ValueClass::Property(PROPERTY),
}
}
fn document_id(key: &[u8]) -> u32 {
u32::from_be_bytes(key[key.len() - U32_LEN..].try_into().unwrap())
}
fn key_prefix() -> String {
let key = key(0).serialize(0);
key[..key.len() - U32_LEN]
.iter()
.map(|byte| format!("{byte:02x}"))
.collect()
}
fn range_params() -> Vec<Value<'static>> {
vec![
Value::Blob(key(0).serialize(0).into()),
Value::Blob(key(u32::MAX).serialize(0).into()),
]
}
impl Backend {
fn detect(db: &Store) -> Option<Self> {
match db {
#[cfg(feature = "postgres")]
Store::PostgreSQL(_) => Some(Backend::Postgres),
#[cfg(feature = "mysql")]
Store::MySQL(_) => Some(Backend::MariaDb),
_ => None,
}
}
async fn data_store(&self, slow: bool) -> DataStore {
let mut config = build_data_store(
match self {
#[cfg(feature = "postgres")]
Backend::Postgres => "PostgreSql",
#[cfg(feature = "mysql")]
Backend::MariaDb => "MariaDb",
},
"",
)
.await;
if slow {
match &mut config {
DataStore::PostgreSql(config) => {
config.options = Some(format!(
"-c statement_timeout={}ms",
STATEMENT_TIMEOUT.as_millis()
));
}
DataStore::MySql(_) => (),
_ => unreachable!(),
}
}
config
}
async fn set_timeout(&self, _admin: &Store, _enable: bool) {
match self {
#[cfg(feature = "postgres")]
Backend::Postgres => (),
#[cfg(feature = "mysql")]
Backend::MariaDb => {
let timeout = if _enable {
STATEMENT_TIMEOUT.as_secs_f64()
} else {
0.0
};
_admin
.sql_query::<usize>(
&format!("SET GLOBAL max_statement_time = {timeout}"),
vec![],
)
.await
.unwrap();
}
}
}
fn populate_query(&self) -> String {
let table = char::from(SUBSPACE_PROPERTY);
let prefix = key_prefix();
let padding = VALUE_SIZE - U32_LEN;
let last = NUM_KEYS - 1;
match self {
#[cfg(feature = "postgres")]
Backend::Postgres => format!(
"INSERT INTO {table} (k, v) SELECT \
decode('{prefix}' || lpad(to_hex(id), 8, '0'), 'hex'), \
decode(lpad(to_hex(id), 8, '0') || repeat('76', {padding}), 'hex') \
FROM generate_series(0, {last}) id"
),
#[cfg(feature = "mysql")]
Backend::MariaDb => format!(
"INSERT INTO {table} (k, v) SELECT \
UNHEX(CONCAT('{prefix}', LPAD(HEX(seq), 8, '0'))), \
UNHEX(CONCAT(LPAD(HEX(seq), 8, '0'), REPEAT('76', {padding}))) \
FROM seq_0_to_{last}"
),
}
}
fn populate_counters_query(&self) -> String {
let table = char::from(SUBSPACE_COUNTER);
let prefix = key_prefix();
let last = NUM_KEYS - 1;
match self {
#[cfg(feature = "postgres")]
Backend::Postgres => format!(
"INSERT INTO {table} (k, v) SELECT \
decode('{prefix}' || lpad(to_hex(id), 8, '0'), 'hex'), 0 \
FROM generate_series(0, {last}) id"
),
#[cfg(feature = "mysql")]
Backend::MariaDb => format!(
"INSERT INTO {table} (k, v) SELECT \
UNHEX(CONCAT('{prefix}', LPAD(HEX(seq), 8, '0'))), 0 FROM seq_0_to_{last}"
),
}
}
fn purge_query(&self) -> String {
format!("DELETE FROM {} WHERE v = 0", char::from(SUBSPACE_COUNTER))
}
fn count_counters_query(&self) -> String {
let table = char::from(SUBSPACE_COUNTER);
match self {
#[cfg(feature = "postgres")]
Backend::Postgres => format!("SELECT COUNT(*) FROM {table} WHERE k >= $1 AND k <= $2"),
#[cfg(feature = "mysql")]
Backend::MariaDb => format!("SELECT COUNT(*) FROM {table} WHERE k >= ? AND k <= ?"),
}
}
fn scan_query(&self) -> String {
let table = char::from(SUBSPACE_PROPERTY);
match self {
#[cfg(feature = "postgres")]
Backend::Postgres => {
format!("SELECT k, v FROM {table} WHERE k >= $1 AND k <= $2 ORDER BY k ASC")
}
#[cfg(feature = "mysql")]
Backend::MariaDb => {
format!("SELECT k, v FROM {table} WHERE k >= ? AND k <= ? ORDER BY k ASC")
}
}
}
fn delete_query(&self) -> String {
let table = char::from(SUBSPACE_PROPERTY);
match self {
#[cfg(feature = "postgres")]
Backend::Postgres => format!("DELETE FROM {table} WHERE k >= $1 AND k <= $2"),
#[cfg(feature = "mysql")]
Backend::MariaDb => format!("DELETE FROM {table} WHERE k >= ? AND k <= ?"),
}
}
fn assert_timeout(&self, err: &trc::Error) {
let (event_type, marker) = match self {
#[cfg(feature = "postgres")]
Backend::Postgres => (trc::StoreEvent::PostgresqlError, "57014"),
#[cfg(feature = "mysql")]
Backend::MariaDb => (
trc::StoreEvent::MysqlError,
"Query execution was interrupted",
),
};
let details = format!("{err:?}");
assert!(
err.matches(trc::EventType::Store(event_type)) && details.contains(marker),
"Expected a statement timeout, got: {details}"
);
}
}