From 73f09072ab041a75b7266dbac759cb9203f35a72 Mon Sep 17 00:00:00 2001 From: mdecimus <11444311+mdecimus@users.noreply.github.com> Date: Thu, 6 Nov 2025 22:29:03 +0100 Subject: [PATCH] Database schema optimization - part 7 --- Cargo.lock | 345 +++++---- crates/common/src/config/mod.rs | 8 +- crates/common/src/config/network.rs | 2 +- crates/common/src/config/smtp/queue.rs | 8 + crates/common/src/config/telemetry.rs | 3 +- crates/common/src/enterprise/llm.rs | 7 +- crates/common/src/manager/restore.rs | 5 +- crates/common/src/storage/index.rs | 36 +- crates/dav/src/common/mod.rs | 9 +- crates/smtp/src/queue/spool.rs | 60 +- crates/store/Cargo.toml | 2 +- crates/store/src/backend/elastic/main.rs | 319 +++++--- crates/store/src/backend/elastic/search.rs | 376 ++++++++++ .../store/src/backend/foundationdb/write.rs | 147 ++-- crates/store/src/backend/mysql/search.rs | 7 +- crates/store/src/backend/mysql/write.rs | 138 ++-- crates/store/src/backend/postgres/search.rs | 12 +- crates/store/src/backend/postgres/write.rs | 153 ++-- crates/store/src/backend/rocksdb/write.rs | 43 +- crates/store/src/backend/s3/mod.rs | 8 + crates/store/src/backend/sqlite/main.rs | 20 +- crates/store/src/backend/sqlite/write.rs | 96 ++- crates/store/src/dispatch/lookup.rs | 14 +- crates/store/src/lib.rs | 10 +- crates/store/src/search/document.rs | 256 +++++++ crates/store/src/search/fields.rs | 246 ++++++ crates/store/src/search/index.rs | 492 ++++++------ crates/store/src/search/local.rs | 283 +++++-- crates/store/src/search/mod.rs | 708 +----------------- crates/store/src/search/query.rs | 397 ++-------- crates/store/src/search/term.rs | 445 +++++++++++ crates/store/src/write/batch.rs | 39 +- crates/store/src/write/key.rs | 67 +- crates/store/src/write/mod.rs | 258 ++++++- crates/store/src/write/serialize.rs | 51 +- crates/utils/src/cheeky_hash.rs | 52 +- crates/utils/src/config/http.rs | 12 +- 37 files changed, 3209 insertions(+), 1925 deletions(-) create mode 100644 crates/store/src/search/document.rs create mode 100644 crates/store/src/search/fields.rs create mode 100644 crates/store/src/search/term.rs diff --git a/Cargo.lock b/Cargo.lock index c1d3885b..756615e9 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -378,17 +378,18 @@ checksum = "1505bd5d3d116872e7271a6d4e16d81d0c8570876c8de68093a09ac269d8aac0" [[package]] name = "attohttpc" -version = "0.28.5" +version = "0.30.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "07a9b245ba0739fc90935094c29adbaee3f977218b5fb95e822e261cda7f56a3" +checksum = "16e2cdb6d5ed835199484bb92bb8b3edd526effe995c61732580439c1a67e2e9" dependencies = [ + "base64 0.22.1", "http 1.3.1", "log", "rustls 0.23.35", "serde", "serde_json", "url", - "webpki-roots 0.26.11", + "webpki-roots 1.0.4", ] [[package]] @@ -399,28 +400,51 @@ checksum = "c08606f8c3cbf4ce6ec8e28fb0014a2c086708fe954eaa885384a6165172e7e8" [[package]] name = "aws-creds" -version = "0.37.0" +version = "0.39.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7f84143206b9c72b3c5cb65415de60c7539c79cd1559290fddec657939131be0" +checksum = "b13804829a843b3f26e151c97acbb315ee1177a2724690edfcd28f1894146200" dependencies = [ "attohttpc", "home", "log", - "quick-xml 0.32.0", + "quick-xml 0.38.3", "rust-ini", "serde", - "thiserror 1.0.69", + "thiserror 2.0.17", "time", "url", ] [[package]] -name = "aws-region" -version = "0.25.5" +name = "aws-lc-rs" +version = "1.14.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e9aed3f9c7eac9be28662fdb3b0f4d1951e812f7c64fed4f0327ba702f459b3b" +checksum = "879b6c89592deb404ba4dc0ae6b58ffd1795c78991cbb5b8bc441c48a070440d" dependencies = [ - "thiserror 1.0.69", + "aws-lc-sys", + "zeroize", +] + +[[package]] +name = "aws-lc-sys" +version = "0.32.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "107a4e9d9cab9963e04e84bb8dee0e25f2a987f9a8bad5ed054abd439caa8f8c" +dependencies = [ + "bindgen 0.72.1", + "cc", + "cmake", + "dunce", + "fs_extra", +] + +[[package]] +name = "aws-region" +version = "0.28.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5532f65342f789f9c1b7078ea9c9cd9293cd62dcc284fa99adc4a1c9ba43469c" +dependencies = [ + "thiserror 2.0.17", ] [[package]] @@ -619,6 +643,8 @@ dependencies = [ "cexpr", "clang-sys", "itertools 0.13.0", + "log", + "prettyplease", "proc-macro2", "quote", "regex", @@ -1164,7 +1190,7 @@ dependencies = [ "futures", "hashify", "hostname", - "hyper 1.7.0", + "hyper", "idna", "imagesize", "imap_proto", @@ -1176,7 +1202,7 @@ dependencies = [ "mail-builder", "mail-parser", "mail-send", - "md5 0.8.0", + "md5", "nlp", "num_cpus", "opentelemetry", @@ -1647,7 +1673,7 @@ dependencies = [ "groupware", "hashify", "http_proto", - "hyper 1.7.0", + "hyper", "percent-encoding", "rkyv", "store", @@ -1664,7 +1690,7 @@ dependencies = [ "chrono", "compact_str", "hashify", - "hyper 1.7.0", + "hyper", "mail-parser", "quick-xml 0.38.3", "rkyv", @@ -1837,7 +1863,7 @@ dependencies = [ "mail-builder", "mail-parser", "mail-send", - "md5 0.8.0", + "md5", "nlp", "password-hash", "pbkdf2", @@ -1962,6 +1988,12 @@ dependencies = [ "zeroize", ] +[[package]] +name = "dunce" +version = "1.0.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "92773504d58c093f6de2459af4af33faa518c13451eb8f2b5698ed3d36e7c813" + [[package]] name = "dyn-clone" version = "1.0.20" @@ -2457,6 +2489,12 @@ dependencies = [ "uuid", ] +[[package]] +name = "fs_extra" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "42703706b716c37f96a77aea830392ad231f44c9e9a67872fa5548707e11b11c" + [[package]] name = "funty" version = "2.0.0" @@ -3010,7 +3048,7 @@ dependencies = [ "groupware", "http-body-util", "http_proto", - "hyper 1.7.0", + "hyper", "hyper-util", "jmap", "jmap_proto", @@ -3051,17 +3089,6 @@ dependencies = [ "itoa", ] -[[package]] -name = "http-body" -version = "0.4.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7ceab25649e9960c0311ea418d17bee82c0dcec1bd053b5f9a66e265a693bed2" -dependencies = [ - "bytes", - "http 0.2.12", - "pin-project-lite", -] - [[package]] name = "http-body" version = "1.0.1" @@ -3081,7 +3108,7 @@ dependencies = [ "bytes", "futures-core", "http 1.3.1", - "http-body 1.0.1", + "http-body", "pin-project-lite", ] @@ -3113,7 +3140,7 @@ dependencies = [ "compact_str", "form_urlencoded", "http-body-util", - "hyper 1.7.0", + "hyper", "hyper-util", "percent-encoding", "serde", @@ -3145,29 +3172,6 @@ version = "2.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "135b12329e5e3ce057a9f972339ea52bc954fe1e9358ef27f95e89716fbc5424" -[[package]] -name = "hyper" -version = "0.14.32" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "41dfc780fdec9373c01bae43289ea34c972e40ee3c9f6b3c8801a35f35586ce7" -dependencies = [ - "bytes", - "futures-channel", - "futures-core", - "futures-util", - "http 0.2.12", - "http-body 0.4.6", - "httparse", - "httpdate", - "itoa", - "pin-project-lite", - "socket2 0.5.10", - "tokio", - "tower-service", - "tracing", - "want", -] - [[package]] name = "hyper" version = "1.7.0" @@ -3180,7 +3184,7 @@ dependencies = [ "futures-core", "h2 0.4.12", "http 1.3.1", - "http-body 1.0.1", + "http-body", "httparse", "httpdate", "itoa", @@ -3191,20 +3195,6 @@ dependencies = [ "want", ] -[[package]] -name = "hyper-rustls" -version = "0.24.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ec3efd23720e2049821a693cbc7e65ea87c72f1c58ff2f9522ff332b1491e590" -dependencies = [ - "futures-util", - "http 0.2.12", - "hyper 0.14.32", - "rustls 0.21.12", - "tokio", - "tokio-rustls 0.24.1", -] - [[package]] name = "hyper-rustls" version = "0.27.7" @@ -3212,7 +3202,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e3c93eb611681b207e1fe55d5a71ecf91572ec8a6705cdb6857f7d8d5242cf58" dependencies = [ "http 1.3.1", - "hyper 1.7.0", + "hyper", "hyper-util", "rustls 0.23.35", "rustls-pki-types", @@ -3228,7 +3218,7 @@ version = "0.5.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2b90d566bffbce6a75bd8b09a05aa8c2cb1fabb6cb348f8840c9e4c90a0d83b0" dependencies = [ - "hyper 1.7.0", + "hyper", "hyper-util", "pin-project-lite", "tokio", @@ -3247,8 +3237,8 @@ dependencies = [ "futures-core", "futures-util", "http 1.3.1", - "http-body 1.0.1", - "hyper 1.7.0", + "http-body", + "hyper", "ipnet", "libc", "percent-encoding", @@ -3271,7 +3261,7 @@ dependencies = [ "js-sys", "log", "wasm-bindgen", - "windows-core", + "windows-core 0.62.2", ] [[package]] @@ -3419,7 +3409,7 @@ dependencies = [ "indexmap 2.12.0", "mail-parser", "mail-send", - "md5 0.8.0", + "md5", "nlp", "parking_lot", "rand 0.9.2", @@ -3713,7 +3703,7 @@ dependencies = [ "hkdf", "http-body-util", "http_proto", - "hyper 1.7.0", + "hyper", "hyper-util", "jmap-tools", "jmap_proto", @@ -4245,7 +4235,7 @@ checksum = "114a4e27f3cfaf8918783e8fa4149b820c813b1bedc7755e20e12eff4518331e" dependencies = [ "base64 0.22.1", "gethostname", - "md5 0.8.0", + "md5", "rustls 0.23.35", "rustls-pki-types", "smtp-proto", @@ -4268,7 +4258,7 @@ dependencies = [ "jmap_proto", "mail-parser", "mail-send", - "md5 0.8.0", + "md5", "parking_lot", "rkyv", "rustls 0.23.35", @@ -4329,12 +4319,6 @@ dependencies = [ "digest 0.10.7", ] -[[package]] -name = "md5" -version = "0.7.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "490cc448043f947bae3cbee9c203358d62dbee0db12107a74be5c30ccfd09771" - [[package]] name = "md5" version = "0.8.0" @@ -4664,6 +4648,15 @@ dependencies = [ "serde", ] +[[package]] +name = "ntapi" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e8a3895c6391c39d7fe7ebc444a87eb2991b2a0bc718fdabd071eec617fc68e4" +dependencies = [ + "winapi", +] + [[package]] name = "nu-ansi-term" version = "0.50.3" @@ -4782,6 +4775,25 @@ version = "0.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "830b246a0e5f20af87141b25c173cd1b609bd7779a4617d6ec582abaf90870f3" +[[package]] +name = "objc2-core-foundation" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2a180dd8642fa45cdb7dd721cd4c11b1cadd4929ce112ebd8b9f5803cc79d536" +dependencies = [ + "bitflags", +] + +[[package]] +name = "objc2-io-kit" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "33fafba39597d6dc1fb709123dfa8289d39406734be322956a69f0931c73bb15" +dependencies = [ + "libc", + "objc2-core-foundation", +] + [[package]] name = "object" version = "0.37.3" @@ -5666,16 +5678,6 @@ dependencies = [ "serde", ] -[[package]] -name = "quick-xml" -version = "0.32.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1d3a6e5838b60e0e8fa7a43f22ade549a37d61f8bdbe636d0d7816191de969c2" -dependencies = [ - "memchr", - "serde", -] - [[package]] name = "quick-xml" version = "0.38.3" @@ -5683,6 +5685,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "42a232e7487fc2ef313d96dde7948e7a3c05101870d8985e4fd8d26aedd27b89" dependencies = [ "memchr", + "serde", ] [[package]] @@ -6179,10 +6182,10 @@ dependencies = [ "futures-util", "h2 0.4.12", "http 1.3.1", - "http-body 1.0.1", + "http-body", "http-body-util", - "hyper 1.7.0", - "hyper-rustls 0.27.7", + "hyper", + "hyper-rustls", "hyper-util", "js-sys", "log", @@ -6410,9 +6413,9 @@ dependencies = [ [[package]] name = "rust-s3" -version = "0.35.1" +version = "0.37.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c3df3f353b1f4209dcf437d777cda90279c397ab15a0cd6fd06bd32c88591533" +checksum = "94f9b973bd4097f5bb47e5827dcb9fb5dc17e93879e46badc27d2a4e9a4e5588" dependencies = [ "async-trait", "aws-creds", @@ -6420,27 +6423,24 @@ dependencies = [ "base64 0.22.1", "bytes", "cfg-if", - "futures", + "futures-util", "hex", "hmac 0.12.1", - "http 0.2.12", - "hyper 0.14.32", - "hyper-rustls 0.24.2", + "http 1.3.1", "log", "maybe-async", - "md5 0.7.0", + "md5", "percent-encoding", - "quick-xml 0.32.0", - "rustls 0.21.12", - "rustls-native-certs 0.6.3", + "quick-xml 0.38.3", + "reqwest", "serde", "serde_derive", "serde_json", "sha2 0.10.9", - "thiserror 1.0.69", + "sysinfo", + "thiserror 2.0.17", "time", "tokio", - "tokio-rustls 0.24.1", "tokio-stream", "url", ] @@ -6522,6 +6522,7 @@ version = "0.23.35" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "533f54bc6a7d4f647e46ad909549eda97bf5afc1585190ef692b4286b198bd8f" dependencies = [ + "aws-lc-rs", "log", "once_cell", "ring", @@ -6531,18 +6532,6 @@ dependencies = [ "zeroize", ] -[[package]] -name = "rustls-native-certs" -version = "0.6.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a9aace74cb666635c918e9c12bc0d348266037aa8eb599b5cba565709a8dff00" -dependencies = [ - "openssl-probe", - "rustls-pemfile 1.0.4", - "schannel", - "security-framework 2.11.1", -] - [[package]] name = "rustls-native-certs" version = "0.7.3" @@ -6650,6 +6639,7 @@ version = "0.103.8" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2ffdfa2f5286e2247234e03f680868ac2815974dc39e00ea15adc445d0aafe52" dependencies = [ + "aws-lc-rs", "ring", "rustls-pki-types", "untrusted", @@ -7332,14 +7322,14 @@ dependencies = [ "email", "form_urlencoded", "http-body-util", - "hyper 1.7.0", + "hyper", "hyper-util", "lru-cache", "mail-auth", "mail-builder", "mail-parser", "mail-send", - "md5 0.8.0", + "md5", "nlp", "num_cpus", "parking_lot", @@ -7428,7 +7418,7 @@ dependencies = [ "common", "compact_str", "decancer", - "hyper 1.7.0", + "hyper", "idna", "infer 0.19.0", "mail-auth", @@ -7702,6 +7692,20 @@ dependencies = [ "syn 2.0.108", ] +[[package]] +name = "sysinfo" +version = "0.37.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "16607d5caffd1c07ce073528f9ed972d88db15dd44023fa57142963be3feb11f" +dependencies = [ + "libc", + "memchr", + "ntapi", + "objc2-core-foundation", + "objc2-io-kit", + "windows", +] + [[package]] name = "tagptr" version = "0.2.0" @@ -7760,7 +7764,7 @@ dependencies = [ "http 0.14.1", "http-body-util", "http_proto", - "hyper 1.7.0", + "hyper", "hyper-util", "imap", "imap_proto", @@ -8113,9 +8117,9 @@ dependencies = [ "base64 0.22.1", "bytes", "http 1.3.1", - "http-body 1.0.1", + "http-body", "http-body-util", - "hyper 1.7.0", + "hyper", "hyper-timeout", "hyper-util", "percent-encoding", @@ -8192,7 +8196,7 @@ dependencies = [ "bytes", "futures-util", "http 1.3.1", - "http-body 1.0.1", + "http-body", "iri-string", "pin-project-lite", "tower 0.5.2", @@ -8895,6 +8899,41 @@ version = "0.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f" +[[package]] +name = "windows" +version = "0.61.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9babd3a767a4c1aef6900409f85f5d53ce2544ccdfaa86dad48c91782c6d6893" +dependencies = [ + "windows-collections", + "windows-core 0.61.2", + "windows-future", + "windows-link 0.1.3", + "windows-numerics", +] + +[[package]] +name = "windows-collections" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3beeceb5e5cfd9eb1d76b381630e82c4241ccd0d27f1a39ed41b2760b255c5e8" +dependencies = [ + "windows-core 0.61.2", +] + +[[package]] +name = "windows-core" +version = "0.61.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c0fdd3ddb90610c7638aa2b3a3ab2904fb9e5cdbecc643ddb3647212781c4ae3" +dependencies = [ + "windows-implement", + "windows-interface", + "windows-link 0.1.3", + "windows-result 0.3.4", + "windows-strings 0.4.2", +] + [[package]] name = "windows-core" version = "0.62.2" @@ -8904,8 +8943,19 @@ dependencies = [ "windows-implement", "windows-interface", "windows-link 0.2.1", - "windows-result", - "windows-strings", + "windows-result 0.4.1", + "windows-strings 0.5.1", +] + +[[package]] +name = "windows-future" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc6a41e98427b19fe4b73c550f060b59fa592d7d686537eebf9385621bfbad8e" +dependencies = [ + "windows-core 0.61.2", + "windows-link 0.1.3", + "windows-threading", ] [[package]] @@ -8942,6 +8992,25 @@ version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" +[[package]] +name = "windows-numerics" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9150af68066c4c5c07ddc0ce30421554771e528bde427614c61038bc2c92c2b1" +dependencies = [ + "windows-core 0.61.2", + "windows-link 0.1.3", +] + +[[package]] +name = "windows-result" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "56f42bd332cc6c8eac5af113fc0c1fd6a8fd2aa08a0119358686e5160d0586c6" +dependencies = [ + "windows-link 0.1.3", +] + [[package]] name = "windows-result" version = "0.4.1" @@ -8951,6 +9020,15 @@ dependencies = [ "windows-link 0.2.1", ] +[[package]] +name = "windows-strings" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "56e6c93f3a0c3b36176cb1327a4958a0353d5d166c2a35cb268ace15e91d3b57" +dependencies = [ + "windows-link 0.1.3", +] + [[package]] name = "windows-strings" version = "0.5.1" @@ -9077,6 +9155,15 @@ dependencies = [ "windows_x86_64_msvc 0.53.1", ] +[[package]] +name = "windows-threading" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b66463ad2e0ea3bbf808b7f1d371311c80e115c0b71d60efc142cafbcfb057a6" +dependencies = [ + "windows-link 0.1.3", +] + [[package]] name = "windows_aarch64_gnullvm" version = "0.42.2" diff --git a/crates/common/src/config/mod.rs b/crates/common/src/config/mod.rs index fc364eb6..533e8132 100644 --- a/crates/common/src/config/mod.rs +++ b/crates/common/src/config/mod.rs @@ -13,16 +13,12 @@ use crate::{ listener::tls::AcmeProviders, manager::config::ConfigManager, }; use arc_swap::ArcSwap; -use base64::{Engine, engine::general_purpose}; use directory::{Directories, Directory}; use groupware::GroupwareConfig; -use hyper::{ - HeaderMap, - header::{AUTHORIZATION, HeaderName, HeaderValue}, -}; +use hyper::HeaderMap; use ring::signature::{EcdsaKeyPair, RsaKeyPair}; use spamfilter::SpamFilterConfig; -use std::{str::FromStr, sync::Arc}; +use std::sync::Arc; use store::{BlobBackend, BlobStore, InMemoryStore, SearchStore, Store, Stores}; use telemetry::Metrics; use utils::config::{Config, utils::AsKey}; diff --git a/crates/common/src/config/network.rs b/crates/common/src/config/network.rs index 3956648d..34dfc163 100644 --- a/crates/common/src/config/network.rs +++ b/crates/common/src/config/network.rs @@ -8,7 +8,7 @@ use super::*; use crate::expr::{if_block::IfBlock, tokenizer::TokenMap}; use ahash::AHashSet; use std::{hash::Hasher, time::Duration}; -use utils::config::{Config, Rate, utils::ParseValue}; +use utils::config::{Config, Rate, http::parse_http_headers, utils::ParseValue}; use xxhash_rust::xxh3::Xxh3Builder; #[derive(Clone)] diff --git a/crates/common/src/config/smtp/queue.rs b/crates/common/src/config/smtp/queue.rs index d020afa4..78565f95 100644 --- a/crates/common/src/config/smtp/queue.rs +++ b/crates/common/src/config/smtp/queue.rs @@ -933,6 +933,10 @@ impl QueueName { pub fn into_inner(self) -> [u8; 8] { self.0 } + + pub fn as_slice(&self) -> &[u8] { + &self.0 + } } impl ArchivedQueueName { @@ -941,6 +945,10 @@ impl ArchivedQueueName { .unwrap_or_default() .trim_end_matches('\0') } + + pub fn as_slice(&self) -> &[u8] { + self.0.as_ref() + } } impl Default for QueueName { diff --git a/crates/common/src/config/telemetry.rs b/crates/common/src/config/telemetry.rs index f2e0dace..4ea5be00 100644 --- a/crates/common/src/config/telemetry.rs +++ b/crates/common/src/config/telemetry.rs @@ -4,7 +4,6 @@ * SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL */ -use super::parse_http_headers; use ahash::{AHashMap, AHashSet}; use base64::{Engine, engine::general_purpose::STANDARD}; use hyper::{HeaderMap, header::CONTENT_TYPE}; @@ -21,7 +20,7 @@ use opentelemetry_semantic_conventions::resource::SERVICE_VERSION; use std::{collections::HashMap, str::FromStr, sync::Arc, time::Duration}; use store::Stores; use trc::{EventType, Level, TelemetryEvent, ipc::subscriber::Interests}; -use utils::config::{Config, utils::ParseValue}; +use utils::config::{Config, http::parse_http_headers, utils::ParseValue}; #[derive(Debug)] pub struct TelemetrySubscriber { diff --git a/crates/common/src/enterprise/llm.rs b/crates/common/src/enterprise/llm.rs index d7fdf70a..344f3f20 100644 --- a/crates/common/src/enterprise/llm.rs +++ b/crates/common/src/enterprise/llm.rs @@ -8,13 +8,10 @@ * */ -use std::time::Duration; - use hyper::{HeaderMap, header::CONTENT_TYPE}; use serde::{Deserialize, Serialize}; -use utils::config::Config; - -use crate::config::parse_http_headers; +use std::time::Duration; +use utils::config::{Config, http::parse_http_headers}; #[derive(Clone, Debug)] pub struct AiApiConfig { diff --git a/crates/common/src/manager/restore.rs b/crates/common/src/manager/restore.rs index e5ad7c22..cb06cd25 100644 --- a/crates/common/src/manager/restore.rs +++ b/crates/common/src/manager/restore.rs @@ -293,10 +293,7 @@ async fn restore_file(store: Store, blob_store: BlobStore, path: &Path) { } .serialize(0), }), - op: ValueOp::Set { - value, - version_offset: None, - }, + op: ValueOp::Set(value), }); } Family::None => failed("No family specified in file"), diff --git a/crates/common/src/storage/index.rs b/crates/common/src/storage/index.rs index 1d9fd28a..a5cb71ea 100644 --- a/crates/common/src/storage/index.rs +++ b/crates/common/src/storage/index.rs @@ -14,8 +14,8 @@ use std::{borrow::Cow, fmt::Debug}; use store::{ Serialize, SerializeInfallible, write::{ - Archive, Archiver, BatchBuilder, BlobOp, DirectoryClass, IntoOperations, SearchIndex, - TaskQueueClass, ValueClass, now, + Archive, Archiver, BatchBuilder, BlobOp, DirectoryClass, IntoOperations, Params, + SearchIndex, TaskQueueClass, ValueClass, now, }, }; use types::{ @@ -310,7 +310,21 @@ impl IntoOperations } if N::is_versioned() { let (offset, bytes) = Archiver::new(changes).serialize_versioned()?; - batch.set_versioned(Field::ARCHIVE, bytes, offset); + batch.set_fnc( + Field::ARCHIVE, + Params::with_capacity(2).with_bytes(bytes).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) + }, + ); } else { batch.set(Field::ARCHIVE, Archiver::new(changes).serialize()?); } @@ -338,7 +352,21 @@ impl IntoOperations } if N::is_versioned() { let (offset, bytes) = Archiver::new(changes).serialize_versioned()?; - batch.set_versioned(Field::ARCHIVE, bytes, offset); + batch.set_fnc( + Field::ARCHIVE, + Params::with_capacity(2).with_bytes(bytes).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) + }, + ); } else { batch.set(Field::ARCHIVE, Archiver::new(changes).serialize()?); } diff --git a/crates/dav/src/common/mod.rs b/crates/dav/src/common/mod.rs index 99ce741c..385e46cd 100644 --- a/crates/dav/src/common/mod.rs +++ b/crates/dav/src/common/mod.rs @@ -118,11 +118,18 @@ impl ExtractETag for BatchBuilder { match op { Operation::Value { class: ValueClass::Property(p_id), - op: ValueOp::Set { value, .. }, + op: ValueOp::Set(value), } if *p_id == p_value => { return Archive::::extract_hash(value) .map(|hash| format!("\"{}\"", hash)); } + Operation::Value { + class: ValueClass::Property(p_id), + op: ValueOp::SetFnc(set_fnc), + } if *p_id == p_value => { + return Archive::::extract_hash(set_fnc.params().bytes(0)) + .map(|hash| format!("\"{}\"", hash)); + } _ => {} } } diff --git a/crates/smtp/src/queue/spool.rs b/crates/smtp/src/queue/spool.rs index 3e0ec169..1af0e9bb 100644 --- a/crates/smtp/src/queue/spool.rs +++ b/crates/smtp/src/queue/spool.rs @@ -22,8 +22,10 @@ use std::future::Future; use std::net::{IpAddr, Ipv4Addr}; use std::time::SystemTime; use store::write::key::DeserializeBigEndian; +use store::write::serialize::rkyv_deserialize; use store::write::{ - AlignedBytes, Archive, Archiver, BatchBuilder, BlobOp, QueueClass, ValueClass, now, + AlignedBytes, Archive, Archiver, BatchBuilder, BlobOp, MergeResult, Params, QueueClass, + ValueClass, now, }; use store::{Deserialize, IterateParams, Serialize, SerializeInfallible, U64_LEN, ValueKey}; use trc::{AddContext, ServerEvent}; @@ -533,61 +535,75 @@ impl MessageWrapper { ); } + let message_bytes = match Archiver::new(self.message).serialize() { + Ok(data) => data, + Err(err) => { + trc::error!( + err.details("Failed to serialize message.") + .span_id(self.span_id) + .caused_by(trc::location!()) + ); + return false; + } + }; if self.is_multi_queue { - batch.merge( + batch.merge_fnc( ValueClass::Queue(QueueClass::Message(self.queue_id)), - move |bytes| { + Params::with_capacity(3) + .with_u64(self.queue_id) + .with_bytes(self.queue_name.into_inner().to_vec()) + .with_bytes(message_bytes), + |params, _, bytes| { let mut cur_message = as Deserialize>::deserialize( bytes.ok_or_else(|| { trc::StoreEvent::NotFound .into_err() .details("Message no longer exists.") .caused_by(trc::location!()) - .ctx(trc::Key::QueueId, self.queue_id) + .ctx(trc::Key::QueueId, params.u64(0)) })?, ) .and_then(|archive| archive.deserialize::()) .caused_by(trc::location!())?; - if cur_message.blob_hash == self.message.blob_hash - && cur_message.recipients.len() == self.message.recipients.len() + let new_message_ = + as Deserialize>::deserialize(params.bytes(2)) + .caused_by(trc::location!())?; + let new_message = new_message_ + .unarchive::() + .caused_by(trc::location!())?; + + if cur_message.blob_hash.as_slice() == new_message.blob_hash.0.as_slice() + && cur_message.recipients.len() == new_message.recipients.len() { - for (rcpt_idx, rcpt) in self - .message + let queue_name = params.bytes(1); + for (rcpt_idx, rcpt) in new_message .recipients .iter() .enumerate() - .filter(|(_, rcpt)| rcpt.queue == self.queue_name) + .filter(|(_, rcpt)| rcpt.queue.as_slice() == queue_name) { - cur_message.recipients[rcpt_idx] = rcpt.clone(); + cur_message.recipients[rcpt_idx] = + rkyv_deserialize(rcpt).caused_by(trc::location!())?; } Archiver::new(cur_message) .serialize() .caused_by(trc::location!()) + .map(MergeResult::Update) } else { Err(trc::StoreEvent::UnexpectedError .into_err() .details("Message blob hash or recipient count mismatch.") .caused_by(trc::location!()) - .ctx(trc::Key::QueueId, self.queue_id)) + .ctx(trc::Key::QueueId, params.u64(0))) } }, ); } else { batch.set( ValueClass::Queue(QueueClass::Message(self.queue_id)), - match Archiver::new(self.message).serialize() { - Ok(data) => data, - Err(err) => { - trc::error!( - err.details("Failed to serialize message.") - .span_id(self.span_id) - .caused_by(trc::location!()) - ); - return false; - } - }, + message_bytes, ); } diff --git a/crates/store/Cargo.toml b/crates/store/Cargo.toml index 45429fe5..a8180c22 100644 --- a/crates/store/Cargo.toml +++ b/crates/store/Cargo.toml @@ -11,7 +11,7 @@ trc = { path = "../trc" } rocksdb = { version = "0.24", optional = true, features = ["multi-threaded-cf"] } foundationdb = { version = "0.9.2", features = ["embedded-fdb-include", "fdb-7_3"], optional = true } rusqlite = { version = "0.37", features = ["bundled"], optional = true } -rust-s3 = { version = "0.35", default-features = false, features = ["tokio-rustls-tls", "no-verify-ssl"], optional = true } +rust-s3 = { version = "0.37", default-features = false, features = ["tokio-rustls-tls"], optional = true } async-nats = { version = "0.44", default-features = false, features = ["server_2_10", "server_2_11", "ring"], optional = true } azure_core = { version = "0.21.0", optional = true } azure_storage = { version = "0.21.0", default-features = false, features = ["enable_reqwest_rustls", "hmac_rust"], optional = true } diff --git a/crates/store/src/backend/elastic/main.rs b/crates/store/src/backend/elastic/main.rs index 75596dce..49ef2696 100644 --- a/crates/store/src/backend/elastic/main.rs +++ b/crates/store/src/backend/elastic/main.rs @@ -4,37 +4,65 @@ * SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL */ -use crate::backend::elastic::ElasticSearchStore; -use reqwest::Client; -use serde_json::json; -use std::time::Duration; +use crate::{ + backend::elastic::ElasticSearchStore, + search::{ + CalendarSearchField, ContactSearchField, EmailSearchField, FileSearchField, SearchField, + SearchableField, TracingSearchField, + }, + write::SearchIndex, +}; +use reqwest::{Error, Response, Url}; +use serde_json::{Value, json}; use utils::config::{Config, http::build_http_client, utils::AsKey}; -pub(crate) static INDEX_NAMES: &[&str] = &["stalwart_email"]; - impl ElasticSearchStore { pub async fn open(config: &mut Config, prefix: impl AsKey) -> Option { - let client = build_http_client(config, prefix.clone())?; + let client = build_http_client(config, prefix.clone(), "application/json".into())?; + let prefix = prefix.as_key(); let url = config .value_require((&prefix, "url"))? - .trim_end_matches("/"); - Url::parse(url) + .trim_end_matches("/") + .to_string(); + Url::parse(&url) .map_err(|e| config.new_parse_error((&prefix, "url"), format!("Invalid URL: {e}",))) .ok()?; - let es = Self { - client, - url: url.to_string(), - }; + + let es = Self { client, url }; + + let shards = config + .property_or_default((&prefix, "index.shards"), "3") + .unwrap_or(3); + let replicas = config + .property_or_default((&prefix, "index.replicas"), "0") + .unwrap_or(0); + let with_source = config + .property_or_default((&prefix, "index.include-source"), "false") + .unwrap_or(false); if let Err(err) = es - .create_index( - config - .property_or_default((&prefix, "index.shards"), "3") - .unwrap_or(3), - config - .property_or_default((&prefix, "index.replicas"), "0") - .unwrap_or(0), - ) + .create_index::(shards, replicas, with_source) + .await + { + config.new_build_error(prefix.as_str(), err.to_string()); + } + + if let Err(err) = es + .create_index::(shards, replicas, with_source) + .await + { + config.new_build_error(prefix.as_str(), err.to_string()); + } + + if let Err(err) = es + .create_index::(shards, replicas, with_source) + .await + { + config.new_build_error(prefix.as_str(), err.to_string()); + } + + if let Err(err) = es + .create_index::(shards, replicas, with_source) .await { config.new_build_error(prefix.as_str(), err.to_string()); @@ -43,82 +71,54 @@ impl ElasticSearchStore { Some(es) } - async fn create_index(&self, shards: usize, replicas: usize) -> trc::Result<()> { - let exists = self - .index - .indices() - .exists(IndicesExistsParts::Index(&[INDEX_NAMES[0]])) - .send() - .await - .map_err(|err| trc::StoreEvent::ElasticsearchError.reason(err))?; - - if exists.status_code() == StatusCode::NOT_FOUND { - let response = self - .index - .indices() - .create(IndicesCreateParts::Index(INDEX_NAMES[0])) - .body(json!({ - "mappings": { - "properties": { - "document_id": { - "type": "integer" - }, - "account_id": { - "type": "integer" - }, - "header": { - "type": "object", - "properties": { - "name": { - "type": "keyword" - }, - "value": { - "type": "text", - "analyzer": "default_analyzer", - } - } - }, - "body": { - "analyzer": "default_analyzer", - "type": "text" - }, - "attachment": { - "analyzer": "default_analyzer", - "type": "text" - }, - "keyword": { - "type": "keyword" - } - } - }, - "settings": { - "index.number_of_shards": shards, - "index.number_of_replicas": replicas, - "analysis": { - "analyzer": { - "default_analyzer": { - "type": "custom", - "tokenizer": "standard", - "filter": ["lowercase"] - } - } - } - } - })) - .send() - .await; - - assert_success(response).await?; + async fn create_index( + &self, + shards: usize, + replicas: usize, + with_source: bool, + ) -> trc::Result<()> { + let mut mappings = T::primary_keys() + .iter() + .chain(T::all_fields()) + .map(|field| (field.es_field().to_string(), field.es_schema())) + .collect::>(); + if !with_source { + mappings.insert("_source".to_string(), json!({ "enabled": false })); } + let body = json!({ + "mappings": mappings, + "settings": { + "index.number_of_shards": shards, + "index.number_of_replicas": replicas, + "analysis": { + "analyzer": { + "default": { + "type": "custom", + "tokenizer": "standard", + "filter": ["lowercase"] + } + } + } + } + }); + let body = serde_json::to_string(&body).unwrap_or_default(); - Ok(()) + assert_success( + self.client + .put(format!("{}/{}", self.url, T::index().es_index_name())) + .body(body) + .send() + .await, + ) + .await + .map(|_| ()) } } -/*pub(crate) async fn assert_success(response: Result) -> trc::Result { +pub(crate) async fn assert_success(response: Result) -> trc::Result { match response { Ok(response) => { - let status = response.status_code(); + let status = response.status(); if status.is_success() { Ok(response) } else { @@ -130,4 +130,141 @@ impl ElasticSearchStore { Err(err) => Err(trc::StoreEvent::ElasticsearchError.reason(err)), } } -*/ + +impl SearchIndex { + pub fn es_index_name(&self) -> &'static str { + match self { + SearchIndex::Email => "st_email", + SearchIndex::Calendar => "st_calendar", + SearchIndex::Contacts => "st_contact", + SearchIndex::File => "st_file", + SearchIndex::Tracing => "st_tracing", + SearchIndex::InMemory => unreachable!(), + } + } +} + +impl SearchField { + pub fn es_field(&self) -> &'static str { + match self { + SearchField::AccountId => "doc_id", + SearchField::DocumentId => "acc_id", + SearchField::Id => "id", + SearchField::Email(field) => match field { + EmailSearchField::From => "from", + EmailSearchField::To => "to", + EmailSearchField::Cc => "cc", + EmailSearchField::Bcc => "bcc", + EmailSearchField::Subject => "subj", + EmailSearchField::Body => "body", + EmailSearchField::Attachment => "attach", + EmailSearchField::ReceivedAt => "received", + EmailSearchField::SentAt => "sent", + EmailSearchField::Size => "size", + EmailSearchField::HasAttachment => "has_att", + EmailSearchField::Headers => "headers", + }, + SearchField::Calendar(field) => match field { + CalendarSearchField::Title => "title", + CalendarSearchField::Description => "desc", + CalendarSearchField::Location => "loc", + CalendarSearchField::Owner => "owner", + CalendarSearchField::Attendee => "attendee", + CalendarSearchField::Start => "start", + CalendarSearchField::Uid => "uid", + }, + SearchField::Contact(field) => match field { + ContactSearchField::Member => "member", + ContactSearchField::Kind => "kind", + ContactSearchField::Name => "name", + ContactSearchField::Nickname => "nick", + ContactSearchField::Organization => "org", + ContactSearchField::Email => "email", + ContactSearchField::Phone => "phone", + ContactSearchField::OnlineService => "online", + ContactSearchField::Address => "addr", + ContactSearchField::Note => "note", + ContactSearchField::Uid => "uid", + }, + SearchField::File(field) => match field { + FileSearchField::Name => "name", + FileSearchField::Content => "content", + }, + SearchField::Tracing(field) => match field { + TracingSearchField::EventType => "ev_type", + TracingSearchField::QueueId => "queue_id", + TracingSearchField::Keywords => "keywords", + }, + } + } + + pub fn es_schema(&self) -> Value { + match self { + SearchField::AccountId + | SearchField::DocumentId + | SearchField::Email(EmailSearchField::Size) => json!({ + "type": "integer" + }), + SearchField::Id + | SearchField::Email(EmailSearchField::SentAt | EmailSearchField::ReceivedAt) + | SearchField::Calendar(CalendarSearchField::Start) + | SearchField::Tracing(TracingSearchField::QueueId | TracingSearchField::EventType) => { + json!({ + "type": "long" + }) + } + SearchField::Email(EmailSearchField::HasAttachment) => json!({ + "type": "boolean" + }), + SearchField::Calendar(CalendarSearchField::Uid) + | SearchField::Contact(ContactSearchField::Uid) => json!({ + "type": "keyword", + }), + SearchField::Email( + EmailSearchField::From | EmailSearchField::To | EmailSearchField::Subject, + ) => json!({ + "type": "text", + "fields": { + "keyword": { + "type": "keyword" + } + } + }), + SearchField::Email(EmailSearchField::Headers) => { + json!({ + "type": "object", + "enabled": true + }) + } + SearchField::Email( + EmailSearchField::Cc + | EmailSearchField::Bcc + | EmailSearchField::Body + | EmailSearchField::Attachment, + ) + | SearchField::Calendar( + CalendarSearchField::Title + | CalendarSearchField::Description + | CalendarSearchField::Location + | CalendarSearchField::Owner + | CalendarSearchField::Attendee, + ) + | SearchField::Contact( + ContactSearchField::Member + | ContactSearchField::Kind + | ContactSearchField::Name + | ContactSearchField::Nickname + | ContactSearchField::Organization + | ContactSearchField::Email + | ContactSearchField::Phone + | ContactSearchField::OnlineService + | ContactSearchField::Address + | ContactSearchField::Note, + ) + | SearchField::File(FileSearchField::Name | FileSearchField::Content) + | SearchField::Tracing(TracingSearchField::Keywords) => json!({ + "type": "text" + }), + } + } +} diff --git a/crates/store/src/backend/elastic/search.rs b/crates/store/src/backend/elastic/search.rs index 69b01f91..59795f5d 100644 --- a/crates/store/src/backend/elastic/search.rs +++ b/crates/store/src/backend/elastic/search.rs @@ -3,3 +3,379 @@ * * SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL */ + +use crate::{ + backend::elastic::{ElasticSearchStore, main::assert_success}, + search::{ + IndexDocument, SearchComparator, SearchDocumentId, SearchField, SearchFilter, + SearchOperator, SearchQuery, SearchValue, + }, + write::SearchIndex, +}; +use serde::{Deserialize, Deserializer}; +use serde_json::{Map, Value, json}; +use std::fmt::Write; + +#[derive(Debug, Deserialize)] +pub struct SearchResponse { + pub hits: Hits, +} + +#[derive(Debug, Deserialize)] +pub struct Hits { + pub total: Total, + pub hits: Vec, +} + +#[derive(Debug, Deserialize)] +pub struct Total { + pub value: u64, +} + +#[derive(Debug, Deserialize)] +pub struct Hit { + #[serde(rename = "_id", deserialize_with = "deserialize_string_to_u64")] + pub id: u64, +} + +#[derive(Debug, Deserialize)] +pub struct DeleteByQueryResponse { + pub deleted: u64, +} + +impl ElasticSearchStore { + pub async fn index(&self, documents: Vec) -> trc::Result<()> { + let mut request = String::with_capacity(512); + + for document in documents { + let id = if let (Some(SearchValue::Uint(account_id)), Some(SearchValue::Uint(doc_id))) = ( + document.fields.get(&SearchField::AccountId), + document.fields.get(&SearchField::DocumentId), + ) { + *account_id << 32 | *doc_id + } else if let Some(SearchValue::Uint(id)) = document.fields.get(&SearchField::Id) { + *id + } else { + debug_assert!(false, "Document is missing required ID fields"); + continue; + }; + + let _ = writeln!( + &mut request, + "{{\"index\":{{\"_index\":\"{}\",\"_id\":{id}}}}}", + document.index.es_index_name() + ); + json_serialize(&mut request, &document); + request.push('\n'); + } + + assert_success( + self.client + .post(format!("{}/_bulk", self.url)) + .body(request) + .send() + .await, + ) + .await + .map(|_| ()) + } + + pub async fn query( + &self, + index: SearchIndex, + filters: &[SearchFilter], + sort: &[SearchComparator], + ) -> trc::Result> { + let query = Map::from_iter( + [ + Some(("query".to_string(), build_query(filters))), + Some(("size".to_string(), Value::from(10_000))), + Some(("source".to_string(), Value::from(false))), + (!sort.is_empty()).then(|| ("sort".to_string(), build_sort(sort))), + ] + .into_iter() + .flatten(), + ); + let request = serde_json::to_string(&query).unwrap_or_default(); + + let response = assert_success( + self.client + .post(format!("{}/{}/_search", self.url, index.es_index_name())) + .body(request) + .send() + .await, + ) + .await?; + + let text = response + .text() + .await + .map_err(|err| trc::StoreEvent::ElasticsearchError.reason(err))?; + + serde_json::from_str::(&text) + .map(|results| { + results + .hits + .hits + .into_iter() + .map(|hit| R::from_u64(hit.id)) + .collect() + }) + .map_err(|err| { + trc::StoreEvent::ElasticsearchError + .reason(err) + .details(text) + }) + } + + pub async fn unindex(&self, filter: SearchQuery) -> trc::Result { + if filter.filters.is_empty() { + return Err(trc::StoreEvent::ElasticsearchError + .reason("Unindex operation requires at least one filter")); + } + + let query = json!({ + "query": build_query(&filter.filters), + }); + let request = serde_json::to_string(&query).unwrap_or_default(); + + let response = assert_success( + self.client + .post(format!( + "{}/{}/_delete_by_query", + self.url, + filter.index.es_index_name() + )) + .body(request) + .send() + .await, + ) + .await?; + + let response_body = response + .text() + .await + .map_err(|err| trc::StoreEvent::ElasticsearchError.reason(err))?; + + serde_json::from_str::(&response_body) + .map(|delete_response| delete_response.deleted) + .map_err(|err| trc::StoreEvent::ElasticsearchError.reason(err)) + } +} + +fn build_query(filters: &[SearchFilter]) -> Value { + if filters.is_empty() { + return json!({ "match_all": {} }); + } + + let mut stack = Vec::new(); + let mut conditions = Vec::new(); + let mut logical_op = &SearchFilter::And; + + for filter in filters { + match filter { + SearchFilter::Operator { field, op, value } => { + if field.is_text() { + let SearchValue::Text { value, .. } = value else { + debug_assert!(false, "Invalid value type for text field"); + continue; + }; + + if op != &SearchOperator::Equal { + conditions.push(json!({ + "match": { field.es_field(): { + "query": value, + "operator": "and" + } } + })); + } else { + conditions.push(json!({ + "match_phrase": { field.es_field(): value } + })); + } + } else { + let value = match value { + SearchValue::Text { value, .. } => json!(value), + SearchValue::Int(value) => json!(value), + SearchValue::Uint(value) => json!(value), + SearchValue::Boolean(value) => json!(value), + SearchValue::KeyValues(kv) => { + let (key, value) = kv.iter().next().unwrap(); + + let cond = if !value.is_empty() { + if op == &SearchOperator::Equal { + json!({ + "term": { + format!("{}.{}.keyword", field.es_field(), key): value + } + }) + } else { + json!({ + "match": { + format!("{}.{}", field.es_field(), key): value + } + }) + } + } else { + json!({ + "exists": { "field": format!("{}.{}", field.es_field(), key) } + }) + }; + + conditions.push(cond); + continue; + } + }; + + let cond = match op { + SearchOperator::Equal | SearchOperator::Contains => json!({ + "term": { field.es_field(): value } + }), + op => { + let op = match op { + SearchOperator::LowerThan => "lt", + SearchOperator::LowerEqualThan => "lte", + SearchOperator::GreaterThan => "gt", + SearchOperator::GreaterEqualThan => "gte", + _ => unreachable!(), + }; + + json!({ + "range": { field.es_field(): { op: value } } + }) + } + }; + + conditions.push(cond); + } + } + + SearchFilter::And | SearchFilter::Or | SearchFilter::Not => { + stack.push((logical_op, conditions)); + logical_op = filter; + conditions = Vec::new(); + } + SearchFilter::End => { + if let Some((prev_logical_op, mut prev_conditions)) = stack.pop() { + if !conditions.is_empty() { + match logical_op { + SearchFilter::And => { + prev_conditions.push(json!({ "bool": { "must": conditions } })); + } + SearchFilter::Or => { + prev_conditions.push(json!({ "bool": { "should": conditions } })); + } + SearchFilter::Not => { + prev_conditions.push(json!({ "bool": { "must_not": conditions } })); + } + _ => unreachable!(), + } + } + logical_op = prev_logical_op; + conditions = prev_conditions; + } + } + SearchFilter::DocumentSet(_) => { + debug_assert!( + false, + "DocumentSet filters are not supported in this backend" + ); + continue; + } + } + } + + debug_assert!( + !conditions.is_empty(), + "No conditions were built for the query" + ); + + if conditions.len() == 1 { + conditions.pop().unwrap() + } else { + json!({ "bool": { "must": conditions } }) + } +} + +fn build_sort(sort: &[SearchComparator]) -> Value { + Value::Array( + sort.iter() + .filter_map(|comp| match comp { + SearchComparator::Field { field, ascending } => Some(json!({ + field.es_field(): if *ascending { "asc" } else { "desc" } + })), + _ => None, + }) + .collect(), + ) +} + +fn json_serialize(request: &mut String, document: &IndexDocument) { + request.push('{'); + for (idx, (k, v)) in document.fields.iter().enumerate() { + if idx > 0 { + request.push(','); + } + + let _ = write!(request, "{:?}:", k.es_field()); + match v { + SearchValue::Text { value, .. } => { + json_serialize_str(request, value); + } + SearchValue::KeyValues(map) => { + request.push('{'); + for (i, (key, value)) in map.iter().enumerate() { + if i > 0 { + request.push(','); + } + json_serialize_str(request, key); + request.push(':'); + json_serialize_str(request, value); + } + request.push('}'); + } + SearchValue::Int(v) => { + let _ = write!(request, "{}", v); + } + SearchValue::Uint(v) => { + let _ = write!(request, "{}", v); + } + SearchValue::Boolean(v) => { + let _ = write!(request, "{}", v); + } + } + } + request.push('}'); +} + +fn json_serialize_str(request: &mut String, value: &str) { + request.push('"'); + for c in value.chars() { + match c { + '"' => request.push_str("\\\""), + '\\' => request.push_str("\\\\"), + '\n' => request.push_str("\\n"), + '\r' => request.push_str("\\r"), + '\t' => request.push_str("\\t"), + '\u{0008}' => request.push_str("\\b"), // backspace + '\u{000C}' => request.push_str("\\f"), // form feed + _ => { + if !c.is_control() { + request.push(c); + } else { + let _ = write!(request, "\\u{:04x}", c as u32); + } + } + } + } + request.push('"'); +} + +fn deserialize_string_to_u64<'de, D>(deserializer: D) -> Result +where + D: Deserializer<'de>, +{ + <&str>::deserialize(deserializer)? + .parse::() + .map_err(serde::de::Error::custom) +} diff --git a/crates/store/src/backend/foundationdb/write.rs b/crates/store/src/backend/foundationdb/write.rs index 8809ac80..9ca638fc 100644 --- a/crates/store/src/backend/foundationdb/write.rs +++ b/crates/store/src/backend/foundationdb/write.rs @@ -9,12 +9,12 @@ use super::{ read::{ChunkedValue, read_chunked_value}, }; use crate::{ - IndexKey, Key, LogKey, SUBSPACE_COUNTER, SUBSPACE_IN_MEMORY_COUNTER, SUBSPACE_QUOTA, U64_LEN, + IndexKey, Key, LogKey, SUBSPACE_COUNTER, SUBSPACE_IN_MEMORY_COUNTER, SUBSPACE_QUOTA, WITH_SUBSPACE, backend::deserialize_i64_le, write::{ - AssignedIds, Batch, MAX_COMMIT_ATTEMPTS, MAX_COMMIT_TIME, Operation, ValueClass, ValueOp, - key::KeySerializer, + AssignedIds, Batch, MAX_COMMIT_ATTEMPTS, MAX_COMMIT_TIME, MergeResult, Operation, + ValueClass, ValueOp, key::KeySerializer, }, }; use foundationdb::{ @@ -66,7 +66,7 @@ impl FdbStore { } => { account_id = *account_id_; if has_changes { - change_id = result.last_change_id(account_id)?; + change_id = result.set_current_change_id(account_id)?; } } Operation::Collection { @@ -85,41 +85,74 @@ impl FdbStore { let do_chunk = !class.is_counter(collection); match op { - ValueOp::Set { - value, - version_offset, - } => { - if let Some(offset) = version_offset { - value[*offset..*offset + U64_LEN] - .copy_from_slice(&change_id.to_be_bytes()); - } - - if !value.is_empty() && do_chunk { - for (pos, chunk) in value.chunks(MAX_VALUE_SIZE).enumerate() { - match pos.cmp(&1) { - Ordering::Less => {} - Ordering::Equal => { - key.push(0); - } - Ordering::Greater => { - if pos < u8::MAX as usize { - *key.last_mut().unwrap() += 1; - } else { - trx.cancel(); - return Err(trc::StoreEvent::FoundationdbError - .ctx( - trc::Key::Reason, - "Value is too large", - )); - } - } - } - trx.set(&key, chunk); - } - } else { - trx.set(&key, value.as_ref()); + ValueOp::Set(value) => { + if !chunk_value(&trx, &mut key, value) { + trx.cancel(); + return Err(trc::StoreEvent::FoundationdbError + .ctx(trc::Key::Reason, "Value is too large")); } } + ValueOp::SetFnc(set_op) => { + let value = (set_op.fnc)(&set_op.params, &result)?; + if !chunk_value(&trx, &mut key, &value) { + trx.cancel(); + return Err(trc::StoreEvent::FoundationdbError + .ctx(trc::Key::Reason, "Value is too large")); + } + } + ValueOp::MergeFnc(merge_op) => { + let (merge_result, is_chunked) = + match read_chunked_value(&key, &trx, false) + .await + .caused_by(trc::location!())? + { + ChunkedValue::Single(slice) => ( + (merge_op.fnc)( + &merge_op.params, + &result, + Some(slice.as_ref()), + )?, + false, + ), + ChunkedValue::Chunked { bytes, .. } => ( + (merge_op.fnc)( + &merge_op.params, + &result, + Some(bytes.as_ref()), + )?, + true, + ), + ChunkedValue::None => ( + (merge_op.fnc)(&merge_op.params, &result, None)?, + false, + ), + }; + + match merge_result { + MergeResult::Update(value) => { + if !chunk_value(&trx, &mut key, &value) { + trx.cancel(); + return Err(trc::StoreEvent::FoundationdbError + .ctx(trc::Key::Reason, "Value is too large")); + } + } + MergeResult::Delete => { + if is_chunked { + trx.clear_range( + &key, + &KeySerializer::new(key.len() + 1) + .write(key.as_slice()) + .write(u8::MAX) + .finalize(), + ); + } else { + trx.clear(&key); + } + } + MergeResult::Skip => (), + } + } + ValueOp::AtomicAdd(by) => { trx.atomic_op(&key, &by.to_le_bytes()[..], MutationType::Add); } @@ -134,21 +167,6 @@ impl FdbStore { trx.set(&key, &num.to_le_bytes()[..]); result.push_counter_id(num); } - ValueOp::Merge(merge) => { - let value = match read_chunked_value(&key, &trx, false) - .await - .caused_by(trc::location!())? - { - ChunkedValue::Single(slice) => { - (merge.fnc)(Some(slice.as_ref())) - } - ChunkedValue::Chunked { bytes, .. } => { - (merge.fnc)(Some(bytes.as_ref())) - } - ChunkedValue::None => (merge.fnc)(None), - }?; - trx.set(&key, value.as_ref()); - } ValueOp::Clear => { if do_chunk { trx.clear_range( @@ -311,3 +329,28 @@ impl FdbStore { self.commit(trx, false).await.map(|_| ()) } } + +fn chunk_value(trx: &Transaction, key: &mut Vec, value: &[u8]) -> bool { + if !value.is_empty() && value.len() > MAX_VALUE_SIZE { + for (pos, chunk) in value.chunks(MAX_VALUE_SIZE).enumerate() { + match pos.cmp(&1) { + Ordering::Less => {} + Ordering::Equal => { + key.push(0); + } + Ordering::Greater => { + if pos < u8::MAX as usize { + *key.last_mut().unwrap() += 1; + } else { + return false; + } + } + } + trx.set(key, chunk); + } + } else { + trx.set(key, value.as_ref()); + } + + true +} diff --git a/crates/store/src/backend/mysql/search.rs b/crates/store/src/backend/mysql/search.rs index f2046b3a..69586b49 100644 --- a/crates/store/src/backend/mysql/search.rs +++ b/crates/store/src/backend/mysql/search.rs @@ -78,7 +78,7 @@ impl MysqlStore { sort: &[SearchComparator], ) -> trc::Result> { let mut query = format!( - "SELECT {} FROM {} ", + "SELECT {} FROM {}", R::field().column(), index.mysql_table() ); @@ -110,7 +110,10 @@ impl MysqlStore { } fn build_filter(query: &mut String, filters: &[SearchFilter]) -> Vec { - query.push_str("WHERE "); + if filters.is_empty() { + return Vec::new(); + } + query.push_str(" WHERE "); let mut operator_stack = Vec::new(); let mut operator = &SearchFilter::And; let mut is_first = true; diff --git a/crates/store/src/backend/mysql/write.rs b/crates/store/src/backend/mysql/write.rs index e5879aba..59a8ddb0 100644 --- a/crates/store/src/backend/mysql/write.rs +++ b/crates/store/src/backend/mysql/write.rs @@ -6,9 +6,10 @@ use super::{MysqlStore, into_error}; use crate::{ - IndexKey, Key, LogKey, SUBSPACE_COUNTER, SUBSPACE_IN_MEMORY_COUNTER, SUBSPACE_QUOTA, U64_LEN, + IndexKey, Key, LogKey, SUBSPACE_COUNTER, SUBSPACE_IN_MEMORY_COUNTER, SUBSPACE_QUOTA, write::{ - AssignedIds, Batch, MAX_COMMIT_ATTEMPTS, MAX_COMMIT_TIME, Operation, ValueClass, ValueOp, + AssignedIds, Batch, MAX_COMMIT_ATTEMPTS, MAX_COMMIT_TIME, MergeResult, Operation, + ValueClass, ValueOp, }, }; use ahash::AHashMap; @@ -110,7 +111,7 @@ impl MysqlStore { } => { account_id = *account_id_; if has_changes { - change_id = result.last_change_id(account_id)?; + change_id = result.set_current_change_id(account_id)?; } } Operation::Collection { @@ -128,15 +129,7 @@ impl MysqlStore { let table = char::from(class.subspace(collection)); match op { - ValueOp::Set { - value, - version_offset, - } => { - if let Some(offset) = version_offset { - value[*offset..*offset + U64_LEN] - .copy_from_slice(&change_id.to_be_bytes()); - } - + ValueOp::Set(value) => { let exists = asserted_values.get(&key); let s = if let Some(exists) = exists { if *exists { @@ -176,6 +169,94 @@ impl MysqlStore { } } } + ValueOp::SetFnc(set_op) => { + let value = (set_op.fnc)(&set_op.params, &result)?; + let exists = asserted_values.get(&key); + let s = if let Some(exists) = exists { + if *exists { + trx.prep(format!("UPDATE {} SET v = :v WHERE k = :k", table)) + .await? + } else { + trx.prep(format!( + "INSERT INTO {} (k, v) VALUES (:k, :v)", + table + )) + .await? + } + } else { + trx + .prep( + format!("INSERT INTO {} (k, v) VALUES (:k, :v) ON DUPLICATE KEY UPDATE v = VALUES(v)", table), + ) + .await? + }; + + match trx.exec_drop(&s, params! {"k" => key, "v" => &value}).await { + Ok(_) => { + if exists.is_some() && trx.affected_rows() == 0 { + trx.rollback().await?; + return Err(trc::StoreEvent::AssertValueFailed + .into_err() + .caused_by(trc::location!()) + .into()); + } + } + Err(err) => { + trx.rollback().await?; + return Err(err.into()); + } + } + } + ValueOp::MergeFnc(merge_op) => { + let s = trx + .prep(format!("SELECT v FROM {} WHERE k = ? FOR UPDATE", table)) + .await?; + let (exists, merge_result) = trx + .exec_first::, _, _>(&s, (&key,)) + .await? + .map(|bytes| { + (merge_op.fnc)(&merge_op.params, &result, Some(bytes.as_ref())) + .map(|v| (true, v)) + .map_err(CommitError::from) + }) + .unwrap_or_else(|| { + (merge_op.fnc)(&merge_op.params, &result, None) + .map(|v| (false, v)) + .map_err(CommitError::from) + })?; + + let s = if exists { + trx.prep(format!("UPDATE {} SET v = :v WHERE k = :k", table)) + .await? + } else { + trx.prep(format!("INSERT INTO {} (k, v) VALUES (:k, :v)", table)) + .await? + }; + + match merge_result { + MergeResult::Update(value) => { + if let Err(err) = + trx.exec_drop(&s, params! {"k" => key, "v" => &value}).await + { + trx.rollback().await?; + return Err(err.into()); + } + } + MergeResult::Delete if exists => { + // Update asserted value + if let Some(exists) = asserted_values.get_mut(&key) { + *exists = false; + } + + let s = trx + .prep(format!("DELETE FROM {} WHERE k = ?", table)) + .await?; + trx.exec_drop(&s, (key,)).await?; + } + _ => (), + } + } + ValueOp::AtomicAdd(by) => { if *by >= 0 { let s = trx @@ -217,39 +298,6 @@ impl MysqlStore { })?, ); } - ValueOp::Merge(merge) => { - let s = trx - .prep(format!("SELECT v FROM {} WHERE k = ? FOR UPDATE", table)) - .await?; - let (exists, value) = trx - .exec_first::, _, _>(&s, (&key,)) - .await? - .map(|bytes| { - (merge.fnc)(Some(bytes.as_ref())) - .map(|v| (true, v)) - .map_err(CommitError::from) - }) - .unwrap_or_else(|| { - (merge.fnc)(None) - .map(|v| (false, v)) - .map_err(CommitError::from) - })?; - - let s = if exists { - trx.prep(format!("UPDATE {} SET v = :v WHERE k = :k", table)) - .await? - } else { - trx.prep(format!("INSERT INTO {} (k, v) VALUES (:k, :v)", table)) - .await? - }; - - if let Err(err) = - trx.exec_drop(&s, params! {"k" => key, "v" => &value}).await - { - trx.rollback().await?; - return Err(err.into()); - } - } ValueOp::Clear => { // Update asserted value if let Some(exists) = asserted_values.get_mut(&key) { diff --git a/crates/store/src/backend/postgres/search.rs b/crates/store/src/backend/postgres/search.rs index 050d44cf..47ed6381 100644 --- a/crates/store/src/backend/postgres/search.rs +++ b/crates/store/src/backend/postgres/search.rs @@ -116,11 +116,7 @@ impl PostgresStore { filters: &[SearchFilter], sort: &[SearchComparator], ) -> trc::Result> { - let mut query = format!( - "SELECT {} FROM {} ", - R::field().column(), - index.psql_table() - ); + let mut query = format!("SELECT {} FROM {}", R::field().column(), index.psql_table()); let params = self.build_filter(&mut query, filters); if !sort.is_empty() { build_sort(&mut query, sort); @@ -139,6 +135,7 @@ impl PostgresStore { } pub async fn unindex(&self, filter: SearchQuery) -> trc::Result { + debug_assert!(!filter.filters.is_empty()); let mut query = format!("DELETE FROM {} ", filter.index.psql_table()); let params = self.build_filter(&mut query, &filter.filters); let conn = self.conn_pool.get().await.map_err(into_error)?; @@ -154,7 +151,10 @@ impl PostgresStore { query: &mut String, filters: &'x [SearchFilter], ) -> Vec<&'x (dyn ToSql + Sync)> { - query.push_str("WHERE "); + if filters.is_empty() { + return Vec::new(); + } + query.push_str(" WHERE "); let mut operator_stack = Vec::new(); let mut operator = &SearchFilter::And; let mut is_first = true; diff --git a/crates/store/src/backend/postgres/write.rs b/crates/store/src/backend/postgres/write.rs index c19c5c58..dbeaadaa 100644 --- a/crates/store/src/backend/postgres/write.rs +++ b/crates/store/src/backend/postgres/write.rs @@ -6,9 +6,10 @@ use super::{PostgresStore, into_error}; use crate::{ - IndexKey, Key, LogKey, SUBSPACE_COUNTER, SUBSPACE_IN_MEMORY_COUNTER, SUBSPACE_QUOTA, U64_LEN, + IndexKey, Key, LogKey, SUBSPACE_COUNTER, SUBSPACE_IN_MEMORY_COUNTER, SUBSPACE_QUOTA, write::{ - AssignedIds, Batch, MAX_COMMIT_ATTEMPTS, MAX_COMMIT_TIME, Operation, ValueClass, ValueOp, + AssignedIds, Batch, MAX_COMMIT_ATTEMPTS, MAX_COMMIT_TIME, MergeResult, Operation, + ValueClass, ValueOp, }, }; use ahash::AHashMap; @@ -112,7 +113,7 @@ impl PostgresStore { } => { account_id = *account_id_; if has_changes { - change_id = result.last_change_id(account_id)?; + change_id = result.set_current_change_id(account_id)?; } } Operation::Collection { @@ -130,15 +131,7 @@ impl PostgresStore { let table = char::from(class.subspace(collection)); match op { - ValueOp::Set { - value, - version_offset, - } => { - if let Some(offset) = version_offset { - value[*offset..*offset + U64_LEN] - .copy_from_slice(&change_id.to_be_bytes()); - } - + ValueOp::Set(value) => { let s = if let Some(exists) = asserted_values.get(&key) { if *exists { trx.prepare_cached(&format!( @@ -171,6 +164,101 @@ impl PostgresStore { .into()); } } + ValueOp::SetFnc(set_op) => { + let value = (set_op.fnc)(&set_op.params, &result)?; + + let s = if let Some(exists) = asserted_values.get(&key) { + if *exists { + trx.prepare_cached(&format!( + "UPDATE {} SET v = $2 WHERE k = $1", + table + )) + .await? + } else { + trx.prepare_cached(&format!( + "INSERT INTO {} (k, v) VALUES ($1, $2)", + table + )) + .await? + } + } else { + trx.prepare_cached(&format!( + concat!( + "INSERT INTO {} (k, v) VALUES ($1, $2) ", + "ON CONFLICT (k) DO UPDATE SET v = EXCLUDED.v" + ), + table + )) + .await? + }; + + if trx.execute(&s, &[&key, &value]).await? == 0 { + return Err(trc::StoreEvent::AssertValueFailed + .into_err() + .caused_by(trc::location!()) + .into()); + } + } + ValueOp::MergeFnc(merge_op) => { + let s = trx + .prepare_cached(&format!( + "SELECT v FROM {} WHERE k = $1 FOR UPDATE", + table + )) + .await?; + let (exists, merge_result) = trx + .query_opt(&s, &[&key]) + .await? + .map(|row| { + row.try_get::<_, &[u8]>(0) + .map_err(CommitError::from) + .and_then(|v| { + (merge_op.fnc)(&merge_op.params, &result, Some(v)) + .map(|v| (true, v)) + .map_err(CommitError::from) + }) + }) + .unwrap_or_else(|| { + (merge_op.fnc)(&merge_op.params, &result, None) + .map(|v| (false, v)) + .map_err(CommitError::from) + })?; + + match merge_result { + MergeResult::Update(value) => { + let s = if exists { + trx.prepare_cached(&format!( + "UPDATE {} SET v = $2 WHERE k = $1", + table + )) + .await? + } else { + trx.prepare_cached(&format!( + "INSERT INTO {} (k, v) VALUES ($1, $2)", + table + )) + .await? + }; + + trx.execute(&s, &[&key, &value]).await?; + } + MergeResult::Delete if exists => { + let s = trx + .prepare_cached(&format!( + "DELETE FROM {} WHERE k = $1", + table + )) + .await?; + trx.execute(&s, &[&key]).await?; + + // Update asserted value + if let Some(exists) = asserted_values.get_mut(&key) { + *exists = false; + } + } + _ => (), + } + } ValueOp::AtomicAdd(by) => { if *by >= 0 { let s = trx @@ -208,47 +296,6 @@ impl PostgresStore { .and_then(|row| row.try_get::<_, i64>(0))?, ); } - ValueOp::Merge(merge) => { - let s = trx - .prepare_cached(&format!( - "SELECT v FROM {} WHERE k = $1 FOR UPDATE", - table - )) - .await?; - let (exists, value) = trx - .query_opt(&s, &[&key]) - .await? - .map(|row| { - row.try_get::<_, &[u8]>(0) - .map_err(CommitError::from) - .and_then(|v| { - (merge.fnc)(Some(v)) - .map(|v| (true, v)) - .map_err(CommitError::from) - }) - }) - .unwrap_or_else(|| { - (merge.fnc)(None) - .map(|v| (false, v)) - .map_err(CommitError::from) - })?; - - let s = if exists { - trx.prepare_cached(&format!( - "UPDATE {} SET v = $2 WHERE k = $1", - table - )) - .await? - } else { - trx.prepare_cached(&format!( - "INSERT INTO {} (k, v) VALUES ($1, $2)", - table - )) - .await? - }; - - trx.execute(&s, &[&key, &value]).await?; - } ValueOp::Clear => { let s = trx .prepare_cached(&format!("DELETE FROM {} WHERE k = $1", table)) diff --git a/crates/store/src/backend/rocksdb/write.rs b/crates/store/src/backend/rocksdb/write.rs index 1fc2fa85..84418462 100644 --- a/crates/store/src/backend/rocksdb/write.rs +++ b/crates/store/src/backend/rocksdb/write.rs @@ -7,10 +7,11 @@ use super::{CF_INDEXES, CF_LOGS, CfHandle, RocksDbStore, into_error}; use crate::{ Deserialize, IndexKey, Key, LogKey, SUBSPACE_COUNTER, SUBSPACE_IN_MEMORY_COUNTER, - SUBSPACE_QUOTA, U64_LEN, + SUBSPACE_QUOTA, backend::deserialize_i64_le, write::{ - AssignedIds, Batch, MAX_COMMIT_ATTEMPTS, MAX_COMMIT_TIME, Operation, ValueClass, ValueOp, + AssignedIds, Batch, MAX_COMMIT_ATTEMPTS, MAX_COMMIT_TIME, MergeResult, Operation, + ValueClass, ValueOp, }, }; use rand::Rng; @@ -173,7 +174,7 @@ impl RocksDBTransaction<'_, '_> { } => { account_id = *account_id_; if has_changes { - change_id = result.last_change_id(account_id)?; + change_id = result.set_current_change_id(account_id)?; } } Operation::Collection { @@ -191,17 +192,31 @@ impl RocksDBTransaction<'_, '_> { let cf = self.db.subspace_handle(class.subspace(collection)); match op { - ValueOp::Set { - value, - version_offset, - } => { - if let Some(offset) = version_offset { - value[*offset..*offset + U64_LEN] - .copy_from_slice(&change_id.to_be_bytes()); - } + ValueOp::Set(value) => { + txn.put_cf(&cf, &key, value)?; + } + ValueOp::SetFnc(set_op) => { + let value = (set_op.fnc)(&set_op.params, &result)?; txn.put_cf(&cf, &key, value)?; } + ValueOp::MergeFnc(merge_op) => { + let merge_result = (merge_op.fnc)( + &merge_op.params, + &result, + txn.get_pinned_for_update_cf(&cf, &key, true)?.as_deref(), + )?; + + match merge_result { + MergeResult::Update(value) => { + txn.put_cf(&cf, &key, value)?; + } + MergeResult::Delete => { + txn.delete_cf(&cf, &key)?; + } + MergeResult::Skip => (), + } + } ValueOp::AtomicAdd(by) => { txn.merge_cf(&cf, &key, &by.to_le_bytes()[..])?; } @@ -221,12 +236,6 @@ impl RocksDBTransaction<'_, '_> { txn.put_cf(&cf, &key, &num.to_le_bytes()[..])?; result.push_counter_id(num); } - ValueOp::Merge(merge) => { - let value = (merge.fnc)( - txn.get_pinned_for_update_cf(&cf, &key, true)?.as_deref(), - )?; - txn.put_cf(&cf, &key, value)?; - } ValueOp::Clear => { txn.delete_cf(&cf, &key)?; } diff --git a/crates/store/src/backend/s3/mod.rs b/crates/store/src/backend/s3/mod.rs index 9573c07d..eefeb81c 100644 --- a/crates/store/src/backend/s3/mod.rs +++ b/crates/store/src/backend/s3/mod.rs @@ -48,6 +48,9 @@ impl S3Store { let timeout = config .property_or_default::((&prefix, "timeout"), "30s") .unwrap_or_else(|| Duration::from_secs(30)); + let allow_invalid = config + .property_or_default::((&prefix, "tls.allow-invalid"), "false") + .unwrap_or_default(); Some(S3Store { bucket: Bucket::new( @@ -60,6 +63,11 @@ impl S3Store { }) .ok()? .with_path_style() + .set_dangereous_config(allow_invalid, allow_invalid) + .map_err(|err| { + config.new_build_error(prefix.as_str(), format!("Failed to create bucket: {err:?}")) + }) + .ok()? .with_request_timeout(timeout) .map_err(|err| { config.new_build_error(prefix.as_str(), format!("Failed to create bucket: {err:?}")) diff --git a/crates/store/src/backend/sqlite/main.rs b/crates/store/src/backend/sqlite/main.rs index b42e8598..60239283 100644 --- a/crates/store/src/backend/sqlite/main.rs +++ b/crates/store/src/backend/sqlite/main.rs @@ -119,18 +119,16 @@ impl SqliteStore { .map_err(into_error)?; } - for table in [SUBSPACE_INDEXES] { - let table = char::from(table); - conn.execute( - &format!( - "CREATE TABLE IF NOT EXISTS {table} ( + let table = char::from(SUBSPACE_INDEXES); + conn.execute( + &format!( + "CREATE TABLE IF NOT EXISTS {table} ( k BLOB PRIMARY KEY - )" - ), - [], - ) - .map_err(into_error)?; - } + )" + ), + [], + ) + .map_err(into_error)?; for table in [SUBSPACE_COUNTER, SUBSPACE_QUOTA, SUBSPACE_IN_MEMORY_COUNTER] { conn.execute( diff --git a/crates/store/src/backend/sqlite/write.rs b/crates/store/src/backend/sqlite/write.rs index 6f68738c..e96d0a90 100644 --- a/crates/store/src/backend/sqlite/write.rs +++ b/crates/store/src/backend/sqlite/write.rs @@ -6,8 +6,8 @@ use super::{SqliteStore, into_error}; use crate::{ - IndexKey, Key, LogKey, SUBSPACE_COUNTER, SUBSPACE_IN_MEMORY_COUNTER, SUBSPACE_QUOTA, U64_LEN, - write::{AssignedIds, Batch, Operation, ValueClass, ValueOp}, + IndexKey, Key, LogKey, SUBSPACE_COUNTER, SUBSPACE_IN_MEMORY_COUNTER, SUBSPACE_QUOTA, + write::{AssignedIds, Batch, MergeResult, Operation, ValueClass, ValueOp}, }; use rusqlite::{OptionalExtension, TransactionBehavior, params}; use trc::AddContext; @@ -56,7 +56,7 @@ impl SqliteStore { } => { account_id = *account_id_; if has_changes { - change_id = result.last_change_id(account_id)?; + change_id = result.set_current_change_id(account_id)?; } } Operation::Collection { @@ -74,15 +74,7 @@ impl SqliteStore { let table = char::from(class.subspace(collection)); match op { - ValueOp::Set { - value, - version_offset, - } => { - if let Some(offset) = version_offset { - value[*offset..*offset + U64_LEN] - .copy_from_slice(&change_id.to_be_bytes()); - } - + ValueOp::Set(value) => { trx.prepare_cached(&format!( "INSERT OR REPLACE INTO {} (k, v) VALUES (?, ?)", table @@ -93,6 +85,63 @@ impl SqliteStore { .map_err(into_error) .caused_by(trc::location!())?; } + ValueOp::SetFnc(set_op) => { + let value = (set_op.fnc)(&set_op.params, &result)?; + trx.prepare_cached(&format!( + "INSERT OR REPLACE INTO {} (k, v) VALUES (?, ?)", + table + )) + .map_err(into_error) + .caused_by(trc::location!())? + .execute([&key, &value]) + .map_err(into_error) + .caused_by(trc::location!())?; + } + ValueOp::MergeFnc(merge_op) => { + let merge_result = trx + .prepare_cached(&format!("SELECT v FROM {} WHERE k = ?", table)) + .map_err(into_error) + .caused_by(trc::location!())? + .query_row([&key], |row| { + Ok((merge_op.fnc)( + &merge_op.params, + &result, + Some(row.get_ref(0)?.as_bytes()?), + )) + }) + .optional() + .map_err(into_error) + .caused_by(trc::location!())? + .unwrap_or_else(|| { + (merge_op.fnc)(&merge_op.params, &result, None) + })?; + + match merge_result { + MergeResult::Update(value) => { + trx.prepare_cached(&format!( + "INSERT OR REPLACE INTO {} (k, v) VALUES (?, ?)", + table + )) + .map_err(into_error) + .caused_by(trc::location!())? + .execute([&key, &value]) + .map_err(into_error) + .caused_by(trc::location!())?; + } + MergeResult::Delete => { + trx.prepare_cached(&format!( + "DELETE FROM {} WHERE k = ?", + table + )) + .map_err(into_error) + .caused_by(trc::location!())? + .execute([&key]) + .map_err(into_error) + .caused_by(trc::location!())?; + } + MergeResult::Skip => (), + } + } ValueOp::AtomicAdd(by) => { if *by >= 0 { trx.prepare_cached(&format!( @@ -135,29 +184,6 @@ impl SqliteStore { .caused_by(trc::location!())?, ); } - ValueOp::Merge(merge) => { - let value = trx - .prepare_cached(&format!("SELECT v FROM {} WHERE k = ?", table)) - .map_err(into_error) - .caused_by(trc::location!())? - .query_row([&key], |row| { - Ok((merge.fnc)(Some(row.get_ref(0)?.as_bytes()?))) - }) - .optional() - .map_err(into_error) - .caused_by(trc::location!())? - .unwrap_or_else(|| (merge.fnc)(None))?; - - trx.prepare_cached(&format!( - "INSERT OR REPLACE INTO {} (k, v) VALUES (?, ?)", - table - )) - .map_err(into_error) - .caused_by(trc::location!())? - .execute([&key, &value]) - .map_err(into_error) - .caused_by(trc::location!())?; - } ValueOp::Clear => { trx.prepare_cached(&format!("DELETE FROM {} WHERE k = ?", table)) .map_err(into_error) diff --git a/crates/store/src/dispatch/lookup.rs b/crates/store/src/dispatch/lookup.rs index fc4d6c2f..52f086f8 100644 --- a/crates/store/src/dispatch/lookup.rs +++ b/crates/store/src/dispatch/lookup.rs @@ -37,13 +37,12 @@ impl InMemoryStore { let mut batch = BatchBuilder::new(); batch.any_op(Operation::Value { class: ValueClass::InMemory(InMemoryClass::Key(kv.key)), - op: ValueOp::Set { - value: KeySerializer::new(kv.value.len() + U64_LEN) + 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(), - version_offset: None, - }, + ), }); store.write(batch.build_all()).await.map(|_| ()) } @@ -70,13 +69,12 @@ impl InMemoryStore { if let Some(expires) = kv.expires { batch.any_op(Operation::Value { class: ValueClass::InMemory(InMemoryClass::Key(kv.key.clone())), - op: ValueOp::Set { - value: KeySerializer::new(U64_LEN * 2) + op: ValueOp::Set( + KeySerializer::new(U64_LEN * 2) .write(0u64) .write(now() + expires) .finalize(), - version_offset: None, - }, + ), }); } diff --git a/crates/store/src/lib.rs b/crates/store/src/lib.rs index 3955e345..79337682 100644 --- a/crates/store/src/lib.rs +++ b/crates/store/src/lib.rs @@ -103,15 +103,7 @@ pub const SUBSPACE_REPORT_OUT: u8 = b'h'; pub const SUBSPACE_REPORT_IN: u8 = b'r'; pub const SUBSPACE_TELEMETRY_SPAN: u8 = b'o'; pub const SUBSPACE_TELEMETRY_METRIC: u8 = b'x'; - -pub const SUBSPACE_RESERVED_2: u8 = b'z'; -/* -pub const SUBSPACE_BITMAP_ID: u8 = b'b'; -pub const SUBSPACE_BITMAP_TAG: u8 = b'c'; -pub const SUBSPACE_BITMAP_TEXT: u8 = b'v'; -pub const SUBSPACE_FTS_INDEX: u8 = b'g'; -pub const SUBSPACE_TELEMETRY_INDEX: u8 = b'w'; -*/ +pub const SUBSPACE_SEARCH_INDEX: u8 = b'z'; #[derive(Clone)] pub struct IterateParams { diff --git a/crates/store/src/search/document.rs b/crates/store/src/search/document.rs new file mode 100644 index 00000000..ad24121f --- /dev/null +++ b/crates/store/src/search/document.rs @@ -0,0 +1,256 @@ +/* + * SPDX-FileCopyrightText: 2020 Stalwart Labs LLC + * + * SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL + */ + +use crate::search::*; + +impl IndexDocument { + pub fn new(index: SearchIndex) -> Self { + Self { + fields: Default::default(), + index, + } + } + + pub fn with_account_id(mut self, account_id: u32) -> Self { + self.fields + .insert(SearchField::AccountId, SearchValue::Uint(account_id as u64)); + self + } + + pub fn with_document_id(mut self, document_id: u32) -> Self { + self.fields.insert( + SearchField::DocumentId, + SearchValue::Uint(document_id as u64), + ); + self + } + + pub fn with_id(mut self, id: u64) -> Self { + self.fields.insert(SearchField::Id, SearchValue::Uint(id)); + self + } + + pub fn index_text(&mut self, field: impl Into, value: &str, language: Language) { + match self.fields.entry(field.into()) { + Entry::Occupied(mut entry) => { + if let SearchValue::Text { + value: existing_value, + .. + } = entry.get_mut() + { + existing_value.push(' '); + existing_value.push_str(value); + } + } + Entry::Vacant(entry) => { + entry.insert(SearchValue::Text { + value: value.to_string(), + language, + }); + } + } + } + + pub fn index_bool(&mut self, field: impl Into, value: bool) { + self.fields + .insert(field.into(), SearchValue::Boolean(value)); + } + + pub fn index_integer>(&mut self, field: impl Into, value: N) { + self.fields + .insert(field.into(), SearchValue::Int(value.into())); + } + + pub fn index_unsigned>(&mut self, field: impl Into, value: N) { + self.fields + .insert(field.into(), SearchValue::Uint(value.into())); + } + + pub fn insert_key_value( + &mut self, + field: impl Into, + key: impl Into, + value: impl Into, + ) { + let search_field = field.into(); + + match self.fields.entry(search_field) { + Entry::Occupied(mut entry) => { + if let SearchValue::KeyValues(existing_key_values) = entry.get_mut() { + existing_key_values.append(key.into(), value.into()); + } + } + Entry::Vacant(entry) => { + let mut new_key_values = VecMap::new(); + new_key_values.append(key.into(), value.into()); + entry.insert(SearchValue::KeyValues(new_key_values)); + } + } + } + + pub fn is_empty(&self) -> bool { + self.fields.is_empty() + } + + pub fn has_field(&self, field: &SearchField) -> bool { + self.fields.contains_key(field) + } + + pub fn set_unknown_language(&mut self, lang: Language) { + for value in self.fields.values_mut() { + if let SearchValue::Text { language, .. } = value + && language.is_unknown() + { + *language = lang; + } + } + } +} + +impl SearchFilter { + pub fn cond( + field: impl Into, + op: SearchOperator, + value: impl Into, + ) -> Self { + SearchFilter::Operator { + field: field.into(), + op, + value: value.into(), + } + } + + pub fn eq(field: impl Into, value: impl Into) -> Self { + SearchFilter::Operator { + field: field.into(), + op: SearchOperator::Equal, + value: value.into(), + } + } + + pub fn lt(field: impl Into, value: impl Into) -> Self { + SearchFilter::Operator { + field: field.into(), + op: SearchOperator::LowerThan, + value: value.into(), + } + } + + pub fn le(field: impl Into, value: impl Into) -> Self { + SearchFilter::Operator { + field: field.into(), + op: SearchOperator::LowerEqualThan, + value: value.into(), + } + } + + pub fn gt(field: impl Into, value: impl Into) -> Self { + SearchFilter::Operator { + field: field.into(), + op: SearchOperator::GreaterThan, + value: value.into(), + } + } + + pub fn ge(field: impl Into, value: impl Into) -> Self { + SearchFilter::Operator { + field: field.into(), + op: SearchOperator::GreaterEqualThan, + value: value.into(), + } + } + + pub fn has_text_detect( + field: impl Into, + text: impl Into, + default_language: Language, + ) -> Self { + let (text, language) = Language::detect(text.into(), default_language); + Self::has_text(field, text, language) + } + + pub fn has_text( + field: impl Into, + text: impl Into, + language: Language, + ) -> Self { + let text = text.into(); + let (is_exact, text) = if let Some(text) = text + .strip_prefix('"') + .and_then(|t| t.strip_suffix('"')) + .or_else(|| text.strip_prefix('\'').and_then(|t| t.strip_suffix('\''))) + { + (true, text.to_string()) + } else { + (false, text) + }; + + if !matches!(language, Language::None) && is_exact { + SearchFilter::Operator { + field: field.into(), + op: SearchOperator::Equal, + value: SearchValue::Text { + value: text, + language, + }, + } + } else { + SearchFilter::Operator { + field: field.into(), + op: SearchOperator::Contains, + value: SearchValue::Text { + value: text, + language, + }, + } + } + } + + #[inline(always)] + pub fn has_english_text(field: impl Into, text: impl Into) -> Self { + Self::has_text(field, text, Language::English) + } + + #[inline(always)] + pub fn has_unknown_text(field: impl Into, text: impl Into) -> Self { + Self::has_text(field, text, Language::Unknown) + } + + pub fn is_in_set(set: RoaringBitmap) -> Self { + SearchFilter::DocumentSet(set) + } +} + +impl SearchComparator { + pub fn field(field: impl Into, ascending: bool) -> Self { + Self::Field { + field: field.into(), + ascending, + } + } + + pub fn set(set: RoaringBitmap, ascending: bool) -> Self { + Self::DocumentSet { set, ascending } + } + + pub fn sorted_set(set: AHashMap, ascending: bool) -> Self { + Self::SortedSet { set, ascending } + } + + pub fn ascending(field: impl Into) -> Self { + Self::Field { + field: field.into(), + ascending: true, + } + } + + pub fn descending(field: impl Into) -> Self { + Self::Field { + field: field.into(), + ascending: false, + } + } +} diff --git a/crates/store/src/search/fields.rs b/crates/store/src/search/fields.rs new file mode 100644 index 00000000..f238e35d --- /dev/null +++ b/crates/store/src/search/fields.rs @@ -0,0 +1,246 @@ +/* + * SPDX-FileCopyrightText: 2020 Stalwart Labs LLC + * + * SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL + */ + +use crate::search::*; + +impl SearchableField for EmailSearchField { + fn index() -> SearchIndex { + SearchIndex::Email + } + + fn primary_keys() -> &'static [SearchField] { + &[SearchField::AccountId, SearchField::DocumentId] + } + + fn all_fields() -> &'static [SearchField] { + &[ + SearchField::Email(EmailSearchField::From), + SearchField::Email(EmailSearchField::To), + SearchField::Email(EmailSearchField::Cc), + SearchField::Email(EmailSearchField::Bcc), + SearchField::Email(EmailSearchField::Subject), + SearchField::Email(EmailSearchField::Body), + SearchField::Email(EmailSearchField::Attachment), + SearchField::Email(EmailSearchField::ReceivedAt), + SearchField::Email(EmailSearchField::SentAt), + SearchField::Email(EmailSearchField::Size), + SearchField::Email(EmailSearchField::HasAttachment), + SearchField::Email(EmailSearchField::Headers), + ] + } + + fn is_indexed(&self) -> bool { + matches!( + self, + EmailSearchField::From + | EmailSearchField::To + | EmailSearchField::Subject + | EmailSearchField::ReceivedAt + | EmailSearchField::Size + | EmailSearchField::HasAttachment, + ) + } + + fn is_text(&self) -> bool { + matches!( + self, + EmailSearchField::From + | EmailSearchField::To + | EmailSearchField::Cc + | EmailSearchField::Bcc + | EmailSearchField::Subject + | EmailSearchField::Body + | EmailSearchField::Attachment, + ) + } +} + +impl SearchableField for CalendarSearchField { + fn index() -> SearchIndex { + SearchIndex::Calendar + } + + fn primary_keys() -> &'static [SearchField] { + &[SearchField::AccountId, SearchField::DocumentId] + } + + fn all_fields() -> &'static [SearchField] { + &[ + SearchField::Calendar(CalendarSearchField::Title), + SearchField::Calendar(CalendarSearchField::Description), + SearchField::Calendar(CalendarSearchField::Location), + SearchField::Calendar(CalendarSearchField::Owner), + SearchField::Calendar(CalendarSearchField::Attendee), + SearchField::Calendar(CalendarSearchField::Start), + SearchField::Calendar(CalendarSearchField::Uid), + ] + } + + fn is_indexed(&self) -> bool { + matches!(self, CalendarSearchField::Start | CalendarSearchField::Uid) + } + + fn is_text(&self) -> bool { + matches!( + self, + CalendarSearchField::Title + | CalendarSearchField::Description + | CalendarSearchField::Location + | CalendarSearchField::Owner + | CalendarSearchField::Attendee + ) + } +} + +impl SearchableField for ContactSearchField { + fn index() -> SearchIndex { + SearchIndex::Contacts + } + + fn primary_keys() -> &'static [SearchField] { + &[SearchField::AccountId, SearchField::DocumentId] + } + + fn all_fields() -> &'static [SearchField] { + &[ + SearchField::Contact(ContactSearchField::Member), + SearchField::Contact(ContactSearchField::Kind), + SearchField::Contact(ContactSearchField::Name), + SearchField::Contact(ContactSearchField::Nickname), + SearchField::Contact(ContactSearchField::Organization), + SearchField::Contact(ContactSearchField::Email), + SearchField::Contact(ContactSearchField::Phone), + SearchField::Contact(ContactSearchField::OnlineService), + SearchField::Contact(ContactSearchField::Address), + SearchField::Contact(ContactSearchField::Note), + SearchField::Contact(ContactSearchField::Uid), + ] + } + + fn is_indexed(&self) -> bool { + matches!(self, ContactSearchField::Uid | ContactSearchField::Kind) + } + + fn is_text(&self) -> bool { + matches!( + self, + ContactSearchField::Name + | ContactSearchField::Nickname + | ContactSearchField::Organization + | ContactSearchField::Email + | ContactSearchField::Phone + | ContactSearchField::OnlineService + | ContactSearchField::Address + | ContactSearchField::Note + ) + } +} + +impl SearchableField for FileSearchField { + fn index() -> SearchIndex { + SearchIndex::File + } + + fn primary_keys() -> &'static [SearchField] { + &[SearchField::AccountId, SearchField::DocumentId] + } + + fn all_fields() -> &'static [SearchField] { + &[ + SearchField::File(FileSearchField::Name), + SearchField::File(FileSearchField::Content), + ] + } + + fn is_indexed(&self) -> bool { + false + } + + fn is_text(&self) -> bool { + true + } +} + +impl SearchableField for TracingSearchField { + fn index() -> SearchIndex { + SearchIndex::Tracing + } + + fn primary_keys() -> &'static [SearchField] { + &[SearchField::Id] + } + + fn all_fields() -> &'static [SearchField] { + &[ + SearchField::Tracing(TracingSearchField::EventType), + SearchField::Tracing(TracingSearchField::QueueId), + SearchField::Tracing(TracingSearchField::Keywords), + ] + } + + fn is_indexed(&self) -> bool { + matches!( + self, + TracingSearchField::QueueId | TracingSearchField::EventType + ) + } + + fn is_text(&self) -> bool { + matches!(self, TracingSearchField::Keywords) + } +} + +impl SearchField { + pub(crate) fn is_indexed(&self) -> bool { + match self { + SearchField::Email(field) => field.is_indexed(), + SearchField::Calendar(field) => field.is_indexed(), + SearchField::Contact(field) => field.is_indexed(), + SearchField::File(field) => field.is_indexed(), + SearchField::Tracing(field) => field.is_indexed(), + SearchField::AccountId | SearchField::DocumentId | SearchField::Id => false, + } + } + + pub(crate) fn is_text(&self) -> bool { + match self { + SearchField::Email(field) => field.is_text(), + SearchField::Calendar(field) => field.is_text(), + SearchField::Contact(field) => field.is_text(), + SearchField::File(field) => field.is_text(), + SearchField::Tracing(field) => field.is_text(), + SearchField::AccountId | SearchField::DocumentId | SearchField::Id => false, + } + } + + pub(crate) fn is_json(&self) -> bool { + matches!(self, SearchField::Email(EmailSearchField::Headers)) + } +} + +impl SearchIndex { + pub fn all_fields(&self) -> &[SearchField] { + match self { + SearchIndex::Email => EmailSearchField::all_fields(), + SearchIndex::Calendar => CalendarSearchField::all_fields(), + SearchIndex::Contacts => ContactSearchField::all_fields(), + SearchIndex::File => FileSearchField::all_fields(), + SearchIndex::Tracing => TracingSearchField::all_fields(), + SearchIndex::InMemory => unreachable!(), + } + } + + pub fn primary_keys(&self) -> &'static [SearchField] { + match self { + SearchIndex::Email => EmailSearchField::primary_keys(), + SearchIndex::Calendar => CalendarSearchField::primary_keys(), + SearchIndex::Contacts => ContactSearchField::primary_keys(), + SearchIndex::File => FileSearchField::primary_keys(), + SearchIndex::Tracing => TracingSearchField::primary_keys(), + SearchIndex::InMemory => unreachable!(), + } + } +} diff --git a/crates/store/src/search/index.rs b/crates/store/src/search/index.rs index 75398585..5ec19f59 100644 --- a/crates/store/src/search/index.rs +++ b/crates/store/src/search/index.rs @@ -4,249 +4,275 @@ * SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL */ -use std::{borrow::Cow, fmt::Display}; - -use ahash::AHashMap; -use nlp::{ - language::{ - Language, - detect::{LanguageDetector, MIN_LANGUAGE_SCORE}, - stemmer::Stemmer, - }, - tokenizers::word::WordTokenizer, -}; -use trc::AddContext; -use types::collection::Collection; - use crate::{ - IterateParams, SerializeInfallible, Store, U32_LEN, ValueKey, - backend::MAX_TOKEN_LENGTH, - dispatch::DocumentSet, - search::IndexDocument, - write::{BatchBuilder, Operation, ValueClass, ValueOp, key::DeserializeBigEndian}, + Deserialize, IterateParams, Store, U64_LEN, ValueKey, + search::{ + IndexDocument, SearchField, SearchFilter, SearchOperator, SearchQuery, SearchValue, + term::{TermIndex, TermIndexBuilder}, + }, + write::{ + AlignedBytes, Archive, BatchBuilder, SEARCH_INDEX_MAX_FIELD_LEN, SearchIndexClass, + SearchIndexField, SearchIndexId, SearchIndexType, ValueClass, key::DeserializeBigEndian, + }, }; - -pub const TERM_INDEX_VERSION: u8 = 1; +use ahash::AHashMap; +use trc::AddContext; +use utils::cheeky_hash::CheekyHash; impl Store { - pub(crate) async fn index_insert(&self, document: IndexDocument) -> trc::Result<()> { - /*let mut detect = LanguageDetector::new(); - let mut tokens: AHashMap = AHashMap::new(); - let mut parts = Vec::new(); - let mut position = 0; - - for text in document.parts { - match text.typ { - Type::Text(language) => { - let language = if language == Language::Unknown { - detect.detect(&text.text, MIN_LANGUAGE_SCORE) - } else { - language - }; - parts.push((text.field, language, text.text)); - } - Type::Tokenize => { - let field = u8::from(text.field); - for token in WordTokenizer::new(text.text.as_ref(), MAX_TOKEN_LENGTH) { - tokens - .entry(BitmapHash::new(token.word.as_ref())) - .or_default() - .insert(TokenType::word(field), position); - position += 1; - } - position += 10; - } - Type::Keyword => { - let value = text.text.as_ref(); - if !value.is_empty() { - let field = u8::from(text.field); - tokens - .entry(BitmapHash::new(value)) - .or_default() - .insert_keyword(TokenType::word(field)); - } - } - } - } - - let default_language = detect - .most_frequent_language() - .unwrap_or(document.default_language); - - for (field, language, text) in parts.into_iter() { - let language = if language != Language::Unknown { - language - } else { - default_language - }; - let field: u8 = field.into(); - - for token in Stemmer::new(&text, language, MAX_TOKEN_LENGTH) { - tokens - .entry(BitmapHash::new(token.word.as_ref())) - .or_default() - .insert(TokenType::word(field), position); - - if let Some(stemmed_word) = token.stemmed_word { - tokens - .entry(BitmapHash::new(stemmed_word.as_ref())) - .or_default() - .insert_keyword(TokenType::stemmed(field)); - } - - position += 1; - } - - position += 10; - } - - if tokens.is_empty() { - return Ok(()); - } - - // Serialize keys - let mut keys = Vec::with_capacity(tokens.len()); - for (hash, postings) in tokens.into_iter() { - keys.push(Operation::Value { - class: ValueClass::FtsIndex(hash), - op: ValueOp::Set { - value: postings.serialize(), - version_offset: None, - }, - }); - } - - // Commit index - let mut batch = BatchBuilder::new(); - batch - .with_account_id(document.account_id) - .with_collection(document.collection) - .with_document(document.document_id); - - for key in keys.into_iter() { - if batch.is_large_batch() { - self.write(batch.build_all()).await?; - batch = BatchBuilder::new(); - batch - .with_account_id(document.account_id) - .with_collection(document.collection) - .with_document(document.document_id); - } - batch.any_op(key); - } - - if !batch.is_empty() { - self.write(batch.build_all()).await?; - }*/ - - Ok(()) - } - - pub(crate) async fn index_remove( - &self, - account_id: u32, - collection: Collection, - document_ids: &impl DocumentSet, - ) -> trc::Result<()> { - // Find keys to delete - /*let mut delete_keys: AHashMap> = AHashMap::new(); - self.iterate( - IterateParams::new( - ValueKey { - account_id, - collection: collection as u8, - document_id: 0, - class: ValueClass::FtsIndex(BitmapHash { - hash: [0; 8], - len: 1, - }), - }, - ValueKey { - account_id: account_id + 1, - collection: collection as u8, - document_id: 0, - class: ValueClass::FtsIndex(BitmapHash { - hash: [0; 8], - len: 1, - }), - }, - ) - .no_values(), - |key, _| { - let document_id = key.deserialize_be_u32(key.len() - U32_LEN)?; - if document_ids.contains(document_id) { - let mut hash = [0u8; 8]; - let (hash, len) = match key.len() - (U32_LEN * 2) - 1 { - 9 => { - hash[..8].copy_from_slice(&key[U32_LEN..U32_LEN + 8]); - (hash, key[key.len() - U32_LEN - 2]) - } - len @ (1..=7) => { - hash[..len].copy_from_slice(&key[U32_LEN..U32_LEN + len]); - (hash, len as u8) - } - 0 => { - // Temporary fix for empty keywords - (hash, 0) - } - invalid => { - return Err(trc::Error::corrupted_key(key, None, trc::location!()) - .ctx(trc::Key::Reason, "Invalid bitmap key length") - .ctx(trc::Key::Size, invalid)); - } - }; - - delete_keys - .entry(document_id) - .or_default() - .push(ValueClass::FtsIndex(BitmapHash { hash, len })); - } - - Ok(true) - }, - ) - .await - .caused_by(trc::location!())?; - - // Remove keys - let mut batch = BatchBuilder::new(); - batch - .with_account_id(account_id) - .with_collection(collection); - - for (document_id, keys) in delete_keys { - batch.with_document(document_id); - - for key in keys { - if batch.is_large_batch() { - self.write(batch.build_all()) - .await - .caused_by(trc::location!())?; - batch = BatchBuilder::new(); - batch - .with_account_id(account_id) - .with_collection(collection) - .with_document(document_id); - } - batch.any_op(Operation::Value { - class: key, - op: ValueOp::Clear, - }); - } - } - - if !batch.is_empty() { + pub(crate) async fn index(&self, documents: Vec) -> trc::Result<()> { + for document in documents { + let mut batch = BatchBuilder::new(); + let index = document.index; + let term_index_builder = TermIndexBuilder::build(document); + term_index_builder + .index + .write_index(&mut batch, index, term_index_builder.id) + .caused_by(trc::location!())?; self.write(batch.build_all()) .await .caused_by(trc::location!())?; - }*/ - + } Ok(()) } - pub(crate) async fn index_remove_all(&self, _: u32) -> trc::Result<()> { - // No-op - // Term indexes are stored in the same key range as the document + pub(crate) async fn unindex(&self, query: SearchQuery) -> trc::Result<()> { + let index = query.index; + let mut account_documents: AHashMap> = AHashMap::new(); + let mut ids = vec![]; + let mut to_id = None; + let mut last_account_id = None; + + for filter in query.filters { + match filter { + SearchFilter::Operator { field, op, value } => match (field, value) { + (SearchField::AccountId, SearchValue::Uint(id)) + if op == SearchOperator::Equal => + { + last_account_id = Some(id as u32); + account_documents.entry(id as u32).or_default(); + } + (SearchField::DocumentId, SearchValue::Uint(id)) + if op == SearchOperator::Equal && last_account_id.is_some() => + { + account_documents + .get_mut(&last_account_id.unwrap()) + .unwrap() + .push(id as u32); + } + (SearchField::Id, SearchValue::Uint(id)) => match op { + SearchOperator::LowerThan | SearchOperator::LowerEqualThan => { + to_id = Some(id); + } + SearchOperator::Equal => { + ids.push(id); + } + _ => { + return Err(trc::StoreEvent::UnexpectedError + .into_err() + .reason("Unsupported operator for Id field")); + } + }, + _ => { + return Err(trc::StoreEvent::UnexpectedError + .into_err() + .reason("Unsupported filter")); + } + }, + SearchFilter::And | SearchFilter::Or | SearchFilter::End => {} + SearchFilter::Not | SearchFilter::DocumentSet(_) => { + return Err(trc::StoreEvent::UnexpectedError + .into_err() + .reason("Unsupported filter")); + } + } + } + + // Delete by account and document ids + for (account_id, document_ids) in account_documents { + if !document_ids.is_empty() { + for document_id in document_ids { + let Some(archive) = self + .get_value::>(ValueKey::from( + ValueClass::SearchIndex(SearchIndexClass { + index, + typ: SearchIndexType::Document { + id: SearchIndexId::Account { + account_id, + document_id, + }, + }, + }), + )) + .await + .caused_by(trc::location!())? + else { + continue; + }; + let term_index = archive + .unarchive::() + .caused_by(trc::location!())?; + let mut batch = BatchBuilder::new(); + term_index.delete_index( + &mut batch, + index, + SearchIndexId::Account { + account_id, + document_id, + }, + ); + self.write(batch.build_all()) + .await + .caused_by(trc::location!())?; + } + } else { + // Delete all documents for the account + self.delete_range( + ValueKey::from(ValueClass::SearchIndex(SearchIndexClass { + index, + typ: SearchIndexType::Document { + id: SearchIndexId::Account { + account_id, + document_id: 0, + }, + }, + })), + ValueKey::from(ValueClass::SearchIndex(SearchIndexClass { + index, + typ: SearchIndexType::Document { + id: SearchIndexId::Account { + account_id, + document_id: u32::MAX, + }, + }, + })), + ) + .await + .caused_by(trc::location!())?; + + self.delete_range( + ValueKey::from(ValueClass::SearchIndex(SearchIndexClass { + index, + typ: SearchIndexType::Index { + id: SearchIndexId::Account { + account_id, + document_id: 0, + }, + field: SearchIndexField { + field_id: 0, + len: 1, + data: [0; SEARCH_INDEX_MAX_FIELD_LEN], + }, + }, + })), + ValueKey::from(ValueClass::SearchIndex(SearchIndexClass { + index, + typ: SearchIndexType::Index { + id: SearchIndexId::Account { + account_id, + document_id: u32::MAX, + }, + field: SearchIndexField { + field_id: u8::MAX, + len: 1, + data: [u8::MAX; SEARCH_INDEX_MAX_FIELD_LEN], + }, + }, + })), + ) + .await + .caused_by(trc::location!())?; + + self.delete_range( + ValueKey::from(ValueClass::SearchIndex(SearchIndexClass { + index, + typ: SearchIndexType::Term { + account_id: Some(account_id), + hash: CheekyHash::NULL, + }, + })), + ValueKey::from(ValueClass::SearchIndex(SearchIndexClass { + index, + typ: SearchIndexType::Term { + account_id: Some(account_id), + hash: CheekyHash::FULL, + }, + })), + ) + .await + .caused_by(trc::location!())?; + } + } + + // Delete by global ids + for id in ids { + let Some(archive) = self + .get_value::>(ValueKey::from(ValueClass::SearchIndex( + SearchIndexClass { + index, + typ: SearchIndexType::Document { + id: SearchIndexId::Global { id }, + }, + }, + ))) + .await + .caused_by(trc::location!())? + else { + continue; + }; + let term_index = archive + .unarchive::() + .caused_by(trc::location!())?; + let mut batch = BatchBuilder::new(); + term_index.delete_index(&mut batch, index, SearchIndexId::Global { id }); + self.write(batch.build_all()) + .await + .caused_by(trc::location!())?; + } + + // Delete ranges + if let Some(to_id) = to_id { + let mut batches = Vec::new(); + self.iterate( + IterateParams::new( + ValueKey::from(ValueClass::SearchIndex(SearchIndexClass { + index, + typ: SearchIndexType::Document { + id: SearchIndexId::Global { id: 0 }, + }, + })), + ValueKey::from(ValueClass::SearchIndex(SearchIndexClass { + index, + typ: SearchIndexType::Document { + id: SearchIndexId::Global { id: to_id }, + }, + })), + ), + |key, value| { + let archive = as Deserialize>::deserialize(value)?; + let term_index = archive.unarchive::()?; + let mut batch = BatchBuilder::new(); + term_index.delete_index( + &mut batch, + index, + SearchIndexId::Global { + id: key.deserialize_be_u64(key.len() - U64_LEN)?, + }, + ); + batches.push(batch); + + Ok(true) + }, + ) + .await + .caused_by(trc::location!())?; + + for mut batch in batches { + self.write(batch.build_all()) + .await + .caused_by(trc::location!())?; + } + } Ok(()) } diff --git a/crates/store/src/search/local.rs b/crates/store/src/search/local.rs index d18e63ef..f4d4a0d4 100644 --- a/crates/store/src/search/local.rs +++ b/crates/store/src/search/local.rs @@ -4,93 +4,210 @@ * SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL */ -use crate::search::SearchFilter; +use crate::search::*; +use roaring::RoaringBitmap; -/*pub enum FilterGroup { - Fts(Vec), - Store(T), +struct State { + pub op: SearchFilter, + pub bm: Option, } -fn split_local_remote(filter: Vec) -> Vec> { - let mut filter = Vec::with_capacity(self.len()); - let mut iter = self.into_iter(); - let mut logical_op = None; - - while let Some(item) = iter.next() { - if matches!(item.filter_type(), FilterType::Fts) { - let mut store_item = None; - let mut depth = 0; - let mut fts = Vec::with_capacity(5); - - // Add the logical operator if there is one - let in_logical_op = if let Some(op) = logical_op.take() { - fts.push(op); - true - } else { - false - }; - fts.push(item); - - for item in iter.by_ref() { - match item.filter_type() { - FilterType::And | FilterType::Or | FilterType::Not => { - depth += 1; - fts.push(item); - } - FilterType::End if depth > 0 => { - depth -= 1; - fts.push(item); - } - FilterType::Fts => { - fts.push(item); - } - _ => { - store_item = Some(item); - break; - } - } - } - - if in_logical_op { - fts.push(T::from(FilterType::End)); - } - - if depth > 0 { - let mut store = Vec::with_capacity(depth * 2); - while depth > 0 { - let item = fts.pop().unwrap(); - if matches!( - item.filter_type(), - FilterType::And | FilterType::Or | FilterType::Not - ) { - depth -= 1; - } - store.push(FilterGroup::Store(item)); - } - - filter.push(FilterGroup::Fts(fts)); - filter.extend(store); - } else { - filter.push(FilterGroup::Fts(fts)); - } - - if let Some(item) = store_item { - filter.push(FilterGroup::Store(item)); - } - } else { - match item.filter_type() { - FilterType::And | FilterType::Or => { - logical_op = Some(item.clone()); - } - FilterType::Not => { - logical_op = Some(T::from(FilterType::And)); - } - _ => {} - } - filter.push(FilterGroup::Store(item)); +impl SearchQuery { + pub fn new(index: SearchIndex) -> Self { + Self { + index, + filters: Vec::new(), + comparators: Vec::new(), + mask: RoaringBitmap::new(), } } - filter + pub fn with_filters(mut self, filters: Vec) -> Self { + if self.filters.is_empty() { + self.filters = filters; + } else { + self.filters.extend(filters); + } + self + } + + pub fn with_comparators(mut self, comparators: Vec) -> Self { + if self.comparators.is_empty() { + self.comparators = comparators; + } else { + self.comparators.extend(comparators); + } + self + } + + pub fn with_filter(mut self, filter: SearchFilter) -> Self { + self.filters.push(filter); + self + } + + pub fn add_filter(&mut self, filter: SearchFilter) -> &mut Self { + self.filters.push(filter); + self + } + + pub fn with_comparator(mut self, comparator: SearchComparator) -> Self { + self.comparators.push(comparator); + self + } + + pub fn with_mask(mut self, mask: RoaringBitmap) -> Self { + self.mask = mask; + self + } + + pub fn with_account_id(mut self, account_id: u32) -> Self { + self.filters.push(SearchFilter::cond( + SearchField::AccountId, + SearchOperator::Equal, + SearchValue::Uint(account_id as u64), + )); + self + } + + pub fn filter(self) -> QueryResults { + if self.filters.is_empty() { + return QueryResults { + results: self.mask, + comparators: self.comparators, + }; + } + let mut state: State = State { + op: SearchFilter::And, + bm: None, + }; + let mut stack = Vec::new(); + let mut filters = self.filters.into_iter().peekable(); + let mask = self.mask; + + while let Some(filter) = filters.next() { + let mut result = match filter { + SearchFilter::DocumentSet(set) => Some(set), + op @ (SearchFilter::And | SearchFilter::Or | SearchFilter::Not) => { + stack.push(state); + state = State { op, bm: None }; + continue; + } + SearchFilter::End => { + if let Some(prev_state) = stack.pop() { + let bm = state.bm; + state = prev_state; + bm + } else { + break; + } + } + SearchFilter::Operator { .. } => { + continue; + } + }; + + // Apply logical operation + if let Some(dest) = &mut state.bm { + match state.op { + SearchFilter::And => { + if let Some(result) = result { + dest.bitand_assign(result); + } else { + dest.clear(); + } + } + SearchFilter::Or => { + if let Some(result) = result { + dest.bitor_assign(result); + } + } + SearchFilter::Not => { + if let Some(mut result) = result { + result.bitxor_assign(&mask); + dest.bitand_assign(result); + } + } + _ => unreachable!(), + } + } else if let Some(ref mut result_) = result { + if let SearchFilter::Not = state.op { + result_.bitxor_assign(&mask); + } + state.bm = result; + } else if let SearchFilter::Not = state.op { + state.bm = Some(mask.clone()); + } else { + state.bm = Some(RoaringBitmap::new()); + } + + // And short-circuit + if matches!(state.op, SearchFilter::And) && state.bm.as_ref().unwrap().is_empty() { + while let Some(filter) = filters.peek() { + if matches!(filter, SearchFilter::End) { + break; + } else { + filters.next(); + } + } + } + } + + // AND with mask + let mut results = state.bm.unwrap_or_default(); + results.bitand_assign(&mask); + QueryResults { + results, + comparators: self.comparators, + } + } +} + +impl QueryResults { + pub fn results(&self) -> &RoaringBitmap { + &self.results + } + + pub fn update_results(&mut self, results: RoaringBitmap) { + self.results = results; + } + + pub fn into_bitmap(self) -> RoaringBitmap { + self.results + } + + pub fn into_sorted(self) -> Vec { + let comparators = self.comparators; + let mut results = self.results.into_iter().collect::>(); + + if !results.is_empty() && !comparators.is_empty() { + results.sort_by(|a, b| { + for comparator in &comparators { + let (a, b, is_ascending) = match comparator { + SearchComparator::DocumentSet { set, ascending } => { + (set.contains(*a) as u32, set.contains(*b) as u32, *ascending) + } + SearchComparator::SortedSet { set, ascending } => ( + *set.get(a).unwrap_or(&u32::MAX), + *set.get(b).unwrap_or(&u32::MAX), + *ascending, + ), + SearchComparator::Field { .. } => continue, + }; + + let ordering = if is_ascending { + a.cmp(&b).reverse() + } else { + a.cmp(&b) + }; + + if ordering != Ordering::Equal { + return ordering; + } + } + Ordering::Equal + }); + } + + results + } } -*/ diff --git a/crates/store/src/search/mod.rs b/crates/store/src/search/mod.rs index 768ab321..7aeb6acc 100644 --- a/crates/store/src/search/mod.rs +++ b/crates/store/src/search/mod.rs @@ -4,10 +4,14 @@ * SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL */ +pub mod document; +pub mod fields; pub mod index; pub mod local; pub mod query; +pub mod term; +use crate::write::SearchIndex; use ahash::AHashMap; use nlp::language::Language; use roaring::RoaringBitmap; @@ -16,8 +20,6 @@ use std::collections::hash_map::Entry; use std::ops::{BitAndAssign, BitOrAssign, BitXorAssign}; use utils::map::vec_map::VecMap; -use crate::write::SearchIndex; - #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum SearchOperator { LowerThan, @@ -105,7 +107,6 @@ pub enum SearchValue { } pub trait SearchDocumentId: Sized { - fn from_u32(id: u32) -> Self; fn from_u64(id: u64) -> Self; fn field() -> SearchField; } @@ -154,465 +155,11 @@ pub struct IndexDocument { pub(crate) fields: AHashMap, } -impl SearchFilter { - pub fn cond( - field: impl Into, - op: SearchOperator, - value: impl Into, - ) -> Self { - SearchFilter::Operator { - field: field.into(), - op, - value: value.into(), - } - } - - pub fn eq(field: impl Into, value: impl Into) -> Self { - SearchFilter::Operator { - field: field.into(), - op: SearchOperator::Equal, - value: value.into(), - } - } - - pub fn lt(field: impl Into, value: impl Into) -> Self { - SearchFilter::Operator { - field: field.into(), - op: SearchOperator::LowerThan, - value: value.into(), - } - } - - pub fn le(field: impl Into, value: impl Into) -> Self { - SearchFilter::Operator { - field: field.into(), - op: SearchOperator::LowerEqualThan, - value: value.into(), - } - } - - pub fn gt(field: impl Into, value: impl Into) -> Self { - SearchFilter::Operator { - field: field.into(), - op: SearchOperator::GreaterThan, - value: value.into(), - } - } - - pub fn ge(field: impl Into, value: impl Into) -> Self { - SearchFilter::Operator { - field: field.into(), - op: SearchOperator::GreaterEqualThan, - value: value.into(), - } - } - - pub fn has_text_detect( - field: impl Into, - text: impl Into, - default_language: Language, - ) -> Self { - let (text, language) = Language::detect(text.into(), default_language); - Self::has_text(field, text, language) - } - - pub fn has_text( - field: impl Into, - text: impl Into, - language: Language, - ) -> Self { - let text = text.into(); - let (is_exact, text) = if let Some(text) = text - .strip_prefix('"') - .and_then(|t| t.strip_suffix('"')) - .or_else(|| text.strip_prefix('\'').and_then(|t| t.strip_suffix('\''))) - { - (true, text.to_string()) - } else { - (false, text) - }; - - if !matches!(language, Language::None) && is_exact { - SearchFilter::Operator { - field: field.into(), - op: SearchOperator::Equal, - value: SearchValue::Text { - value: text, - language, - }, - } - } else { - SearchFilter::Operator { - field: field.into(), - op: SearchOperator::Contains, - value: SearchValue::Text { - value: text, - language, - }, - } - } - } - - #[inline(always)] - pub fn has_english_text(field: impl Into, text: impl Into) -> Self { - Self::has_text(field, text, Language::English) - } - - #[inline(always)] - pub fn has_unknown_text(field: impl Into, text: impl Into) -> Self { - Self::has_text(field, text, Language::Unknown) - } - - pub fn is_in_set(set: RoaringBitmap) -> Self { - SearchFilter::DocumentSet(set) - } -} - -impl SearchComparator { - pub fn field(field: impl Into, ascending: bool) -> Self { - Self::Field { - field: field.into(), - ascending, - } - } - - pub fn set(set: RoaringBitmap, ascending: bool) -> Self { - Self::DocumentSet { set, ascending } - } - - pub fn sorted_set(set: AHashMap, ascending: bool) -> Self { - Self::SortedSet { set, ascending } - } - - pub fn ascending(field: impl Into) -> Self { - Self::Field { - field: field.into(), - ascending: true, - } - } - - pub fn descending(field: impl Into) -> Self { - Self::Field { - field: field.into(), - ascending: false, - } - } -} - -impl IndexDocument { - pub fn new(index: SearchIndex) -> Self { - Self { - fields: Default::default(), - index, - } - } - - pub fn with_account_id(mut self, account_id: u32) -> Self { - self.fields - .insert(SearchField::AccountId, SearchValue::Uint(account_id as u64)); - self - } - - pub fn with_document_id(mut self, document_id: u32) -> Self { - self.fields.insert( - SearchField::DocumentId, - SearchValue::Uint(document_id as u64), - ); - self - } - - pub fn with_id(mut self, id: u64) -> Self { - self.fields.insert(SearchField::Id, SearchValue::Uint(id)); - self - } - - pub fn index_text(&mut self, field: impl Into, value: &str, language: Language) { - match self.fields.entry(field.into()) { - Entry::Occupied(mut entry) => { - if let SearchValue::Text { - value: existing_value, - .. - } = entry.get_mut() - { - existing_value.push(' '); - existing_value.push_str(value); - } - } - Entry::Vacant(entry) => { - entry.insert(SearchValue::Text { - value: value.to_string(), - language, - }); - } - } - } - - pub fn index_bool(&mut self, field: impl Into, value: bool) { - self.fields - .insert(field.into(), SearchValue::Boolean(value)); - } - - pub fn index_integer>(&mut self, field: impl Into, value: N) { - self.fields - .insert(field.into(), SearchValue::Int(value.into())); - } - - pub fn index_unsigned>(&mut self, field: impl Into, value: N) { - self.fields - .insert(field.into(), SearchValue::Uint(value.into())); - } - - pub fn insert_key_value( - &mut self, - field: impl Into, - key: impl Into, - value: impl Into, - ) { - let search_field = field.into(); - - match self.fields.entry(search_field) { - Entry::Occupied(mut entry) => { - if let SearchValue::KeyValues(existing_key_values) = entry.get_mut() { - existing_key_values.append(key.into(), value.into()); - } - } - Entry::Vacant(entry) => { - let mut new_key_values = VecMap::new(); - new_key_values.append(key.into(), value.into()); - entry.insert(SearchValue::KeyValues(new_key_values)); - } - } - } - - pub fn is_empty(&self) -> bool { - self.fields.is_empty() - } - - pub fn has_field(&self, field: &SearchField) -> bool { - self.fields.contains_key(field) - } - - pub fn set_unknown_language(&mut self, lang: Language) { - for value in self.fields.values_mut() { - if let SearchValue::Text { language, .. } = value - && language.is_unknown() - { - *language = lang; - } - } - } -} - -struct State { - pub op: SearchFilter, - pub bm: Option, -} - -impl SearchQuery { - pub fn new(index: SearchIndex) -> Self { - Self { - index, - filters: Vec::new(), - comparators: Vec::new(), - mask: RoaringBitmap::new(), - } - } - - pub fn with_filters(mut self, filters: Vec) -> Self { - if self.filters.is_empty() { - self.filters = filters; - } else { - self.filters.extend(filters); - } - self - } - - pub fn with_comparators(mut self, comparators: Vec) -> Self { - if self.comparators.is_empty() { - self.comparators = comparators; - } else { - self.comparators.extend(comparators); - } - self - } - - pub fn with_filter(mut self, filter: SearchFilter) -> Self { - self.filters.push(filter); - self - } - - pub fn add_filter(&mut self, filter: SearchFilter) -> &mut Self { - self.filters.push(filter); - self - } - - pub fn with_comparator(mut self, comparator: SearchComparator) -> Self { - self.comparators.push(comparator); - self - } - - pub fn with_mask(mut self, mask: RoaringBitmap) -> Self { - self.mask = mask; - self - } - - pub fn with_account_id(mut self, account_id: u32) -> Self { - self.filters.push(SearchFilter::cond( - SearchField::AccountId, - SearchOperator::Equal, - SearchValue::Uint(account_id as u64), - )); - self - } - - pub fn filter(self) -> QueryResults { - if self.filters.is_empty() { - return QueryResults { - results: self.mask, - comparators: self.comparators, - }; - } - let mut state: State = State { - op: SearchFilter::And, - bm: None, - }; - let mut stack = Vec::new(); - let mut filters = self.filters.into_iter().peekable(); - let mask = self.mask; - - while let Some(filter) = filters.next() { - let mut result = match filter { - SearchFilter::DocumentSet(set) => Some(set), - op @ (SearchFilter::And | SearchFilter::Or | SearchFilter::Not) => { - stack.push(state); - state = State { op, bm: None }; - continue; - } - SearchFilter::End => { - if let Some(prev_state) = stack.pop() { - let bm = state.bm; - state = prev_state; - bm - } else { - break; - } - } - SearchFilter::Operator { .. } => { - continue; - } - }; - - // Apply logical operation - if let Some(dest) = &mut state.bm { - match state.op { - SearchFilter::And => { - if let Some(result) = result { - dest.bitand_assign(result); - } else { - dest.clear(); - } - } - SearchFilter::Or => { - if let Some(result) = result { - dest.bitor_assign(result); - } - } - SearchFilter::Not => { - if let Some(mut result) = result { - result.bitxor_assign(&mask); - dest.bitand_assign(result); - } - } - _ => unreachable!(), - } - } else if let Some(ref mut result_) = result { - if let SearchFilter::Not = state.op { - result_.bitxor_assign(&mask); - } - state.bm = result; - } else if let SearchFilter::Not = state.op { - state.bm = Some(mask.clone()); - } else { - state.bm = Some(RoaringBitmap::new()); - } - - // And short-circuit - if matches!(state.op, SearchFilter::And) && state.bm.as_ref().unwrap().is_empty() { - while let Some(filter) = filters.peek() { - if matches!(filter, SearchFilter::End) { - break; - } else { - filters.next(); - } - } - } - } - - // AND with mask - let mut results = state.bm.unwrap_or_default(); - results.bitand_assign(&mask); - QueryResults { - results, - comparators: self.comparators, - } - } -} - pub struct QueryResults { results: RoaringBitmap, comparators: Vec, } -impl QueryResults { - pub fn results(&self) -> &RoaringBitmap { - &self.results - } - - pub fn update_results(&mut self, results: RoaringBitmap) { - self.results = results; - } - - pub fn into_bitmap(self) -> RoaringBitmap { - self.results - } - - pub fn into_sorted(self) -> Vec { - let comparators = self.comparators; - let mut results = self.results.into_iter().collect::>(); - - if !results.is_empty() && !comparators.is_empty() { - results.sort_by(|a, b| { - for comparator in &comparators { - let (a, b, is_ascending) = match comparator { - SearchComparator::DocumentSet { set, ascending } => { - (set.contains(*a) as u32, set.contains(*b) as u32, *ascending) - } - SearchComparator::SortedSet { set, ascending } => ( - *set.get(a).unwrap_or(&u32::MAX), - *set.get(b).unwrap_or(&u32::MAX), - *ascending, - ), - SearchComparator::Field { .. } => continue, - }; - - let ordering = if is_ascending { - a.cmp(&b).reverse() - } else { - a.cmp(&b) - }; - - if ordering != Ordering::Equal { - return ordering; - } - } - Ordering::Equal - }); - } - - results - } -} - impl From for SearchField { fn from(field: EmailSearchField) -> Self { SearchField::Email(field) @@ -689,10 +236,6 @@ impl From for SearchValue { } impl SearchDocumentId for u32 { - fn from_u32(id: u32) -> Self { - id - } - fn from_u64(id: u64) -> Self { id as u32 } @@ -703,10 +246,6 @@ impl SearchDocumentId for u32 { } impl SearchDocumentId for u64 { - fn from_u32(id: u32) -> Self { - id as u64 - } - fn from_u64(id: u64) -> Self { id } @@ -716,30 +255,6 @@ impl SearchDocumentId for u64 { } } -impl SearchIndex { - pub fn all_fields(&self) -> &[SearchField] { - match self { - SearchIndex::Email => EmailSearchField::all_fields(), - SearchIndex::Calendar => CalendarSearchField::all_fields(), - SearchIndex::Contacts => ContactSearchField::all_fields(), - SearchIndex::File => FileSearchField::all_fields(), - SearchIndex::Tracing => TracingSearchField::all_fields(), - SearchIndex::InMemory => unreachable!(), - } - } - - pub fn primary_keys(&self) -> &'static [SearchField] { - match self { - SearchIndex::Email => EmailSearchField::primary_keys(), - SearchIndex::Calendar => CalendarSearchField::primary_keys(), - SearchIndex::Contacts => ContactSearchField::primary_keys(), - SearchIndex::File => FileSearchField::primary_keys(), - SearchIndex::Tracing => TracingSearchField::primary_keys(), - SearchIndex::InMemory => unreachable!(), - } - } -} - pub trait SearchableField: Sized { fn index() -> SearchIndex; fn primary_keys() -> &'static [SearchField]; @@ -747,218 +262,3 @@ pub trait SearchableField: Sized { fn is_indexed(&self) -> bool; fn is_text(&self) -> bool; } - -impl SearchableField for EmailSearchField { - fn index() -> SearchIndex { - SearchIndex::Email - } - - fn primary_keys() -> &'static [SearchField] { - &[SearchField::AccountId, SearchField::DocumentId] - } - - fn all_fields() -> &'static [SearchField] { - &[ - SearchField::Email(EmailSearchField::From), - SearchField::Email(EmailSearchField::To), - SearchField::Email(EmailSearchField::Cc), - SearchField::Email(EmailSearchField::Bcc), - SearchField::Email(EmailSearchField::Subject), - SearchField::Email(EmailSearchField::Body), - SearchField::Email(EmailSearchField::Attachment), - SearchField::Email(EmailSearchField::ReceivedAt), - SearchField::Email(EmailSearchField::SentAt), - SearchField::Email(EmailSearchField::Size), - SearchField::Email(EmailSearchField::HasAttachment), - SearchField::Email(EmailSearchField::Headers), - ] - } - - fn is_indexed(&self) -> bool { - matches!( - self, - EmailSearchField::From - | EmailSearchField::To - | EmailSearchField::Subject - | EmailSearchField::ReceivedAt - | EmailSearchField::Size - | EmailSearchField::HasAttachment, - ) - } - - fn is_text(&self) -> bool { - matches!( - self, - EmailSearchField::From - | EmailSearchField::To - | EmailSearchField::Cc - | EmailSearchField::Bcc - | EmailSearchField::Subject - | EmailSearchField::Body - | EmailSearchField::Attachment, - ) - } -} - -impl SearchableField for CalendarSearchField { - fn index() -> SearchIndex { - SearchIndex::Calendar - } - - fn primary_keys() -> &'static [SearchField] { - &[SearchField::AccountId, SearchField::DocumentId] - } - - fn all_fields() -> &'static [SearchField] { - &[ - SearchField::Calendar(CalendarSearchField::Title), - SearchField::Calendar(CalendarSearchField::Description), - SearchField::Calendar(CalendarSearchField::Location), - SearchField::Calendar(CalendarSearchField::Owner), - SearchField::Calendar(CalendarSearchField::Attendee), - SearchField::Calendar(CalendarSearchField::Start), - SearchField::Calendar(CalendarSearchField::Uid), - ] - } - - fn is_indexed(&self) -> bool { - matches!(self, CalendarSearchField::Start | CalendarSearchField::Uid) - } - - fn is_text(&self) -> bool { - matches!( - self, - CalendarSearchField::Title - | CalendarSearchField::Description - | CalendarSearchField::Location - | CalendarSearchField::Owner - | CalendarSearchField::Attendee - ) - } -} - -impl SearchableField for ContactSearchField { - fn index() -> SearchIndex { - SearchIndex::Contacts - } - - fn primary_keys() -> &'static [SearchField] { - &[SearchField::AccountId, SearchField::DocumentId] - } - - fn all_fields() -> &'static [SearchField] { - &[ - SearchField::Contact(ContactSearchField::Member), - SearchField::Contact(ContactSearchField::Kind), - SearchField::Contact(ContactSearchField::Name), - SearchField::Contact(ContactSearchField::Nickname), - SearchField::Contact(ContactSearchField::Organization), - SearchField::Contact(ContactSearchField::Email), - SearchField::Contact(ContactSearchField::Phone), - SearchField::Contact(ContactSearchField::OnlineService), - SearchField::Contact(ContactSearchField::Address), - SearchField::Contact(ContactSearchField::Note), - SearchField::Contact(ContactSearchField::Uid), - ] - } - - fn is_indexed(&self) -> bool { - matches!(self, ContactSearchField::Uid | ContactSearchField::Kind) - } - - fn is_text(&self) -> bool { - matches!( - self, - ContactSearchField::Name - | ContactSearchField::Nickname - | ContactSearchField::Organization - | ContactSearchField::Email - | ContactSearchField::Phone - | ContactSearchField::OnlineService - | ContactSearchField::Address - | ContactSearchField::Note - ) - } -} - -impl SearchableField for FileSearchField { - fn index() -> SearchIndex { - SearchIndex::File - } - - fn primary_keys() -> &'static [SearchField] { - &[SearchField::AccountId, SearchField::DocumentId] - } - - fn all_fields() -> &'static [SearchField] { - &[ - SearchField::File(FileSearchField::Name), - SearchField::File(FileSearchField::Content), - ] - } - - fn is_indexed(&self) -> bool { - false - } - - fn is_text(&self) -> bool { - true - } -} - -impl SearchableField for TracingSearchField { - fn index() -> SearchIndex { - SearchIndex::Tracing - } - - fn primary_keys() -> &'static [SearchField] { - &[SearchField::Id] - } - - fn all_fields() -> &'static [SearchField] { - &[ - SearchField::Tracing(TracingSearchField::EventType), - SearchField::Tracing(TracingSearchField::QueueId), - SearchField::Tracing(TracingSearchField::Keywords), - ] - } - - fn is_indexed(&self) -> bool { - matches!( - self, - TracingSearchField::QueueId | TracingSearchField::EventType - ) - } - - fn is_text(&self) -> bool { - matches!(self, TracingSearchField::Keywords) - } -} - -impl SearchField { - pub(crate) fn is_indexed(&self) -> bool { - match self { - SearchField::Email(field) => field.is_indexed(), - SearchField::Calendar(field) => field.is_indexed(), - SearchField::Contact(field) => field.is_indexed(), - SearchField::File(field) => field.is_indexed(), - SearchField::Tracing(field) => field.is_indexed(), - SearchField::AccountId | SearchField::DocumentId | SearchField::Id => false, - } - } - - pub(crate) fn is_text(&self) -> bool { - match self { - SearchField::Email(field) => field.is_text(), - SearchField::Calendar(field) => field.is_text(), - SearchField::Contact(field) => field.is_text(), - SearchField::File(field) => field.is_text(), - SearchField::Tracing(field) => field.is_text(), - SearchField::AccountId | SearchField::DocumentId | SearchField::Id => false, - } - } - - pub(crate) fn is_json(&self) -> bool { - matches!(self, SearchField::Email(EmailSearchField::Headers)) - } -} diff --git a/crates/store/src/search/query.rs b/crates/store/src/search/query.rs index e2d8206f..ea856eb9 100644 --- a/crates/store/src/search/query.rs +++ b/crates/store/src/search/query.rs @@ -5,193 +5,80 @@ */ use crate::{ - IterateParams, Store, U32_LEN, ValueKey, - backend::MAX_TOKEN_LENGTH, - search::{SearchComparator, SearchFilter}, - write::{ValueClass, key::DeserializeBigEndian}, + Store, ValueKey, backend::MAX_TOKEN_LENGTH, search::{SearchFilter, SearchOperator, SearchQuery, SearchValue}, write::{SearchIndexClass, SearchIndexType, ValueClass} }; -use ahash::AHashMap; -use nlp::language::stemmer::Stemmer; +use nlp::language; use roaring::RoaringBitmap; -use std::{ - fmt::Display, - ops::{BitAndAssign, BitOrAssign, BitXorAssign}, -}; use trc::AddContext; -use types::collection::Collection; - -/*struct State { - pub op: FtsTokenized, - pub bm: Option, -} - -enum FtsTokenized { - Exact { - tokens: Vec<(BitmapHash, u8)>, - }, - Contains { - field: u8, - tokens: Vec<(BitmapHash, Option)>, - }, - Keyword { - field: u8, - token: BitmapHash, - }, - And, - Or, - Not, - End, -}*/ +use utils::cheeky_hash::{CheekyHash, CheekyHashMap}; +use std::{collections::hash_map::Entry, ops::{BitAndAssign, BitOrAssign, BitXorAssign}, sync::Arc}; impl Store { - pub(crate) async fn index_query( - &self, - account_id: u32, - collection: Collection, - filters: Vec, - comparators: Vec, - ) -> trc::Result> { - todo!() - // Tokenize text - /*let mut tokenized_filters = Vec::with_capacity(filters.len()); - let mut token_count = AHashMap::new(); - for filter in filters { - let filter = match filter { - FtsFilter::Exact { - field, - text, - language, - } => { - let mut tokens = Vec::new(); - let field = TokenType::word(field.into()); - - for token in language.tokenize_text(text.as_ref(), MAX_TOKEN_LENGTH) { - let hash = BitmapHash::new(token.word.as_ref()); - token_count.entry(hash).and_modify(|c| *c += 1).or_insert(1); - tokens.push((hash, field)); - } - FtsTokenized::Exact { tokens } - } - FtsFilter::Contains { - field, - text, - language, - } => { - let mut tokens = Vec::new(); - for token in Stemmer::new(text.as_ref(), language, MAX_TOKEN_LENGTH) { - let hash = BitmapHash::new(token.word.as_ref()); - let stemmed_hash = token.stemmed_word.as_deref().map(BitmapHash::new); - - token_count.entry(hash).and_modify(|c| *c += 1).or_insert(1); - if let Some(stemmed_hash) = stemmed_hash { - token_count - .entry(stemmed_hash) - .and_modify(|c| *c += 1) - .or_insert(1); - } - - tokens.push((hash, stemmed_hash)); - } - FtsTokenized::Contains { - field: field.into(), - tokens, - } - } - FtsFilter::Keyword { field, text } => { - let hash = BitmapHash::new(text); - token_count.entry(hash).and_modify(|c| *c += 1).or_insert(1); - - FtsTokenized::Keyword { - field: field.into(), - token: hash, - } - } - FtsFilter::And => FtsTokenized::And, - FtsFilter::Or => FtsTokenized::Or, - FtsFilter::Not => FtsTokenized::Not, - FtsFilter::End => FtsTokenized::End, - }; - - tokenized_filters.push(filter); + pub(crate) async fn query_account(&self, query: SearchQuery) -> trc::Result> { + struct State { + pub op: SearchFilter, + pub bm: Option, } - - let mut not_mask = RoaringBitmap::new(); - let mut not_fetch = false; - - let mut state: State = FtsTokenized::And.into(); + let mut state: State = State { + op: SearchFilter::And, + bm: None, + }; let mut stack = Vec::new(); - let mut token_cache = AHashMap::with_capacity(token_count.len()); - let mut filters = tokenized_filters.into_iter().peekable(); - let collection = u8::from(collection); + let mask = query.mask; + let mut filters = query.filters.into_iter().peekable(); + let mut token_cache : CheekyHashMap> = CheekyHashMap::default(); + let account_id = None; while let Some(filter) = filters.next() { let mut result = match filter { - FtsTokenized::Exact { tokens } => { - self.get_postings( - account_id, - collection, - &tokens, - &token_count, - &mut token_cache, - true, - ) - .await? - } - FtsTokenized::Contains { field, tokens } => { - let mut result = RoaringBitmap::new(); + SearchFilter::Operator { field, op, value } => { + if field.is_text() { + let (value, language) = match value { + SearchValue::Text { value, language } => (value, language), + _ => return Err(trc::Error::InvalidInput("Expected text value for text field".into())), + }; - for (token, stemmed_token) in tokens { - match self - .get_postings( - account_id, - collection, - &[ - (token, TokenType::word(field)), - (stemmed_token.unwrap_or(token), TokenType::stemmed(field)), - ], - &token_count, - &mut token_cache, - false, - ) - .await? - { - Some(b) if !b.is_empty() => { - if !result.is_empty() { - result &= b; - if result.is_empty() { - break; - } - } else { - result = b; + if op == &SearchOperator::Equal { + for token in language.tokenize_text(&value, MAX_TOKEN_LENGTH) { + let hash = CheekyHash::new(token.word.as_bytes()); + match token_cache.entry(hash) { + Entry::Occupied(entry) => { + entry.get().clone() + }, + Entry::Vacant(entry) => { + let value = self.get_value::(ValueKey::from(ValueClass::SearchIndex(SearchIndexClass { + index: query.index, + typ: SearchIndexType::Term { account_id, hash }, + }))).await.caused_by(trc::location!())?.map(Arc::new); + entry.insert(value.clone()); + value + + }, } + + } else { + todo!() } - _ => break, + + } else { + todo!() + } + + + } else { + todo!() + } - if !result.is_empty() { - Some(result) - } else { - None - } } - FtsTokenized::Keyword { field, token } => { - self.get_postings( - account_id, - collection, - &[(token, TokenType::word(field))], - &token_count, - &mut token_cache, - false, - ) - .await? - } - op @ (FtsTokenized::And | FtsTokenized::Or | FtsTokenized::Not) => { + SearchFilter::DocumentSet(bitmap) => Some(Arc::new(bitmap)), + op @ (SearchFilter::And | SearchFilter::Or | SearchFilter::Not) => { stack.push(state); - state = op.into(); + state = State { op, bm: None }; continue; } - FtsTokenized::End => { + SearchFilter::End => { if let Some(prev_state) = stack.pop() { let bm = state.bm; state = prev_state; @@ -202,53 +89,44 @@ impl Store { } }; - // Only fetch not mask if we need it - if matches!(state.op, FtsTokenized::Not) && !not_fetch { - not_mask = self - .get_bitmap(BitmapKey::document_ids(account_id, collection)) - .await? - .unwrap_or_else(RoaringBitmap::new); - not_fetch = true; - } - // Apply logical operation if let Some(dest) = &mut state.bm { match state.op { - FtsTokenized::And => { + SearchFilter::And => { if let Some(result) = result { dest.bitand_assign(result); } else { dest.clear(); } } - FtsTokenized::Or => { + SearchFilter::Or => { if let Some(result) = result { dest.bitor_assign(result); } } - FtsTokenized::Not => { + SearchFilter::Not => { if let Some(mut result) = result { - result.bitxor_assign(¬_mask); + result.bitxor_assign(&mask); dest.bitand_assign(result); } } _ => unreachable!(), } } else if let Some(ref mut result_) = result { - if let FtsTokenized::Not = state.op { - result_.bitxor_assign(¬_mask); + if let SearchFilter::Not = state.op { + result_.bitxor_assign(&mask); } state.bm = result; - } else if let FtsTokenized::Not = state.op { - state.bm = Some(not_mask.clone()); + } else if let SearchFilter::Not = state.op { + state.bm = Some(mask.clone()); } else { state.bm = Some(RoaringBitmap::new()); } // And short circuit - if matches!(state.op, FtsTokenized::And) && state.bm.as_ref().unwrap().is_empty() { + if matches!(state.op, SearchFilter::And) && state.bm.as_ref().unwrap().is_empty() { while let Some(filter) = filters.peek() { - if matches!(filter, FtsTokenized::End) { + if matches!(filter, SearchFilter::End) { break; } else { filters.next(); @@ -257,151 +135,6 @@ impl Store { } } - Ok(state.bm.unwrap_or_default())*/ + todo!() } - - /*async fn get_postings( - &self, - account_id: u32, - collection: u8, - tokens: &[(BitmapHash, u8)], - token_count: &AHashMap, - token_cache: &mut AHashMap>>>, - is_intersect: bool, - ) -> trc::Result> { - let mut result_bm = RoaringBitmap::new(); - let mut position_candidates = AHashMap::new(); - let num_tokens = tokens.len(); - - for (pos, (token, field)) in tokens.iter().enumerate() { - let needs_caching = token_count[token] > 1; - let is_first = pos == 0; - let mut bm = RoaringBitmap::new(); - - if needs_caching { - // Try to fetch from cache - if let Some(postings) = token_cache.get(token) { - for (document_id, postings) in postings { - if postings.has_field(*field) { - if is_intersect { - if is_first { - if num_tokens > 1 { - position_candidates - .insert(*document_id, postings.positions()); - } - bm.insert(*document_id); - } else if position_candidates.get(document_id).is_some_and( - |positions| postings.matches_positions(positions, pos as u32), - ) { - bm.insert(*document_id); - } - } else { - result_bm.insert(*document_id); - } - } - } - - if is_intersect { - if is_first { - result_bm = bm; - } else { - result_bm &= bm; - } - if result_bm.is_empty() { - return Ok(None); - } - } - - continue; - } - - // Insert empty cache entry - token_cache.insert(*token, AHashMap::new()); - } - - // Fetch from store - let key_len = ValueClass::FtsIndex(*token).serialized_size(); - self.iterate( - IterateParams::new( - ValueKey { - account_id, - collection, - document_id: 0, - class: ValueClass::FtsIndex(*token), - }, - ValueKey { - account_id, - collection, - document_id: u32::MAX, - class: ValueClass::FtsIndex(*token), - }, - ), - |key, value| { - if key.len() != key_len { - return Ok(true); - } - - // Make sure this document contain the field - let document_id = key.deserialize_be_u32(key.len() - U32_LEN)?; - let postings = SerializedPostings::new(value); - if postings.has_field(*field) { - if is_intersect { - if is_first { - if num_tokens > 1 { - position_candidates.insert(document_id, postings.positions()); - } - bm.insert(document_id); - } else if position_candidates.get(&document_id).is_some_and( - |positions| postings.matches_positions(positions, pos as u32), - ) { - bm.insert(document_id); - } - } else { - result_bm.insert(document_id); - } - } - - // Cache the postings if needed - if needs_caching { - token_cache - .entry(*token) - .or_default() - .insert(document_id, SerializedPostings::new(value.to_vec())); - } - - Ok(true) - }, - ) - .await - .caused_by(trc::location!())?; - - if is_intersect { - if is_first { - result_bm = bm; - } else { - result_bm &= bm; - } - if result_bm.is_empty() { - return Ok(None); - } - } - } - - Ok(if !result_bm.is_empty() { - Some(result_bm) - } else { - None - }) - } - - */ } - -/*impl From for State { - fn from(value: FtsTokenized) -> Self { - Self { - op: value, - bm: None, - } - } -}*/ diff --git a/crates/store/src/search/term.rs b/crates/store/src/search/term.rs new file mode 100644 index 00000000..83fba133 --- /dev/null +++ b/crates/store/src/search/term.rs @@ -0,0 +1,445 @@ +/* + * SPDX-FileCopyrightText: 2020 Stalwart Labs LLC + * + * SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL + */ + +use crate::{ + Deserialize, Serialize, U64_LEN, + backend::MAX_TOKEN_LENGTH, + search::*, + write::{ + Archiver, BatchBuilder, MergeResult, Params, SEARCH_INDEX_MAX_FIELD_LEN, SearchIndexClass, + SearchIndexField, SearchIndexId, SearchIndexType, ValueClass, + }, +}; +use nlp::{language::stemmer::Stemmer, tokenizers::word::WordTokenizer}; +use roaring::RoaringTreemap; +use utils::cheeky_hash::{CheekyBTreeMap, CheekyHash}; + +#[derive(Debug, PartialEq, Eq, rkyv::Serialize, rkyv::Deserialize, rkyv::Archive)] +pub(crate) struct TermIndex { + terms: Vec, + fields: Vec, +} + +#[derive(Debug, PartialEq, Eq, rkyv::Serialize, rkyv::Deserialize, rkyv::Archive)] +pub(crate) struct Term { + hash: CheekyHash, + fields: u32, +} + +pub(crate) struct TermIndexBuilder { + pub(crate) index: TermIndex, + pub(crate) id: SearchIndexId, +} + +impl TermIndexBuilder { + pub fn build(document: IndexDocument) -> Self { + let mut terms: CheekyBTreeMap = CheekyBTreeMap::new(); + let mut fields: Vec = Vec::new(); + let mut account_id = None; + let mut document_id = None; + let mut id = None; + + for (field, value) in document.fields { + match field { + SearchField::Id => { + if let SearchValue::Uint(v) = value { + id = Some(v); + } + continue; + } + SearchField::AccountId => { + if let SearchValue::Uint(v) = value { + account_id = Some(v); + } + continue; + } + SearchField::DocumentId => { + if let SearchValue::Uint(v) = value { + document_id = Some(v); + } + continue; + } + _ => {} + } + + let field_id = 1 << (field.u8_id() as u32); + + let field = match value { + SearchValue::Text { value, language } => { + if field.is_text() { + if !matches!(language, Language::Unknown | Language::None) { + for token in Stemmer::new(&value, language, MAX_TOKEN_LENGTH) { + *terms + .entry(CheekyHash::new(token.word.as_bytes())) + .or_default() |= field_id; + + if let Some(stemmed_word) = token.stemmed_word { + *terms + .entry(CheekyHash::new( + format!("{}*", stemmed_word).as_bytes(), + )) + .or_default() |= field_id; + } + } + } else { + for token in WordTokenizer::new(value.as_str(), MAX_TOKEN_LENGTH) { + *terms + .entry(CheekyHash::new(token.word.as_bytes())) + .or_default() |= field_id; + } + } + } + + if field.is_indexed() { + let bytes = value.as_bytes(); + let len = bytes.len().min(SEARCH_INDEX_MAX_FIELD_LEN); + let mut data = [0u8; SEARCH_INDEX_MAX_FIELD_LEN]; + + data[..len].copy_from_slice(&bytes[..len]); + + SearchIndexField { + field_id: field.u8_id(), + len: len as u8, + data, + } + } else { + continue; + } + } + SearchValue::KeyValues(map) => { + for (key, value) in map { + *terms.entry(CheekyHash::new(key.as_bytes())).or_default() |= field_id; + for token in value.split_ascii_whitespace() { + *terms + .entry(CheekyHash::new(format!("{key} {token}").as_bytes())) + .or_default() |= field_id; + } + } + + continue; + } + SearchValue::Int(v) => { + let mut data = [0u8; SEARCH_INDEX_MAX_FIELD_LEN]; + data[..U64_LEN].copy_from_slice(&v.to_be_bytes()); + + SearchIndexField { + field_id: field.u8_id(), + len: U64_LEN as u8, + data, + } + } + SearchValue::Uint(v) => { + let mut data = [0u8; SEARCH_INDEX_MAX_FIELD_LEN]; + data[..U64_LEN].copy_from_slice(&v.to_be_bytes()); + + SearchIndexField { + field_id: field.u8_id(), + len: U64_LEN as u8, + data, + } + } + SearchValue::Boolean(v) if v => SearchIndexField { + field_id: field.u8_id(), + len: 1, + data: [1u8; SEARCH_INDEX_MAX_FIELD_LEN], + }, + _ => continue, + }; + + fields.push(field); + } + + TermIndexBuilder { + index: TermIndex { + terms: terms + .into_iter() + .map(|(k, v)| Term { hash: k, fields: v }) + .collect(), + fields, + }, + id: match (account_id, document_id, id) { + (Some(account_id), Some(document_id), None) => SearchIndexId::Account { + account_id: account_id as u32, + document_id: document_id as u32, + }, + (None, None, Some(id)) => SearchIndexId::Global { id }, + _ => { + debug_assert!( + false, + "Invalid combination of AccountId, DocumentId and Id fields" + ); + SearchIndexId::Global { id: 0 } + } + }, + } + } +} + +impl TermIndex { + pub fn write_index( + self, + batch: &mut BatchBuilder, + index: SearchIndex, + id: SearchIndexId, + ) -> trc::Result<()> { + let archive = Archiver::new(self); + batch.set( + ValueClass::SearchIndex(SearchIndexClass { + index, + typ: SearchIndexType::Document { id }, + }), + archive.serialize()?, + ); + + match id { + SearchIndexId::Account { + account_id, + document_id, + } => { + for term in archive.inner.terms { + batch.merge_fnc( + ValueClass::SearchIndex(SearchIndexClass { + index, + typ: SearchIndexType::Term { + account_id: Some(account_id), + hash: term.hash, + }, + }), + Params::with_capacity(1).with_u64(document_id as u64), + |params, _, bytes| { + let document_id = params.u64(0) as u32; + + if let Some(bytes) = bytes { + let mut bitmap = RoaringBitmap::deserialize(bytes)?; + if bitmap.insert(document_id) { + Ok(MergeResult::Update(bitmap.serialize()?)) + } else { + Ok(MergeResult::Skip) + } + } else { + Ok(MergeResult::Update( + RoaringBitmap::from_iter([document_id]).serialize()?, + )) + } + }, + ); + } + } + SearchIndexId::Global { id } => { + for term in archive.inner.terms { + batch.merge_fnc( + ValueClass::SearchIndex(SearchIndexClass { + index, + typ: SearchIndexType::Term { + account_id: None, + hash: term.hash, + }, + }), + Params::with_capacity(1).with_u64(id), + |params, _, bytes| { + let id = params.u64(0); + + if let Some(bytes) = bytes { + let mut bitmap = RoaringTreemap::deserialize(bytes)?; + if bitmap.insert(id) { + Ok(MergeResult::Update(bitmap.serialize()?)) + } else { + Ok(MergeResult::Skip) + } + } else { + Ok(MergeResult::Update( + RoaringTreemap::from_iter([id]).serialize()?, + )) + } + }, + ); + } + } + } + + for field in archive.inner.fields { + batch.set( + ValueClass::SearchIndex(SearchIndexClass { + index, + typ: SearchIndexType::Index { id, field }, + }), + vec![], + ); + } + + Ok(()) + } +} + +impl ArchivedTermIndex { + pub fn has_term(&self, hash: &CheekyHash, field: &SearchField) -> bool { + let hash = hash.as_raw_bytes(); + self.terms + .binary_search_by(|term| term.hash.as_raw_bytes().cmp(hash)) + .is_ok_and(|idx| { + (self.terms[idx].fields.to_native() & (1 << (field.u8_id() as u32))) != 0 + }) + } + + pub fn delete_index(&self, batch: &mut BatchBuilder, index: SearchIndex, id: SearchIndexId) { + batch.clear(ValueClass::SearchIndex(SearchIndexClass { + index, + typ: SearchIndexType::Document { id }, + })); + + match id { + SearchIndexId::Account { + account_id, + document_id, + } => { + for term in self.terms.iter() { + batch.merge_fnc( + ValueClass::SearchIndex(SearchIndexClass { + index, + typ: SearchIndexType::Term { + account_id: Some(account_id), + hash: term.hash.to_native(), + }, + }), + Params::with_capacity(1).with_u64(document_id as u64), + |params, _, bytes| { + let document_id = params.u64(0) as u32; + + if let Some(bytes) = bytes { + let mut bitmap = RoaringBitmap::deserialize(bytes)?; + if bitmap.remove(document_id) { + if !bitmap.is_empty() { + Ok(MergeResult::Update(bitmap.serialize()?)) + } else { + Ok(MergeResult::Delete) + } + } else { + Ok(MergeResult::Skip) + } + } else { + Ok(MergeResult::Skip) + } + }, + ); + } + } + SearchIndexId::Global { id } => { + for term in self.terms.iter() { + batch.merge_fnc( + ValueClass::SearchIndex(SearchIndexClass { + index, + typ: SearchIndexType::Term { + account_id: None, + hash: term.hash.to_native(), + }, + }), + Params::with_capacity(1).with_u64(id), + |params, _, bytes| { + let id = params.u64(0); + + if let Some(bytes) = bytes { + let mut bitmap = RoaringTreemap::deserialize(bytes)?; + if bitmap.remove(id) { + if !bitmap.is_empty() { + Ok(MergeResult::Update(bitmap.serialize()?)) + } else { + Ok(MergeResult::Delete) + } + } else { + Ok(MergeResult::Skip) + } + } else { + Ok(MergeResult::Skip) + } + }, + ); + } + } + } + + for field in self.fields.iter() { + batch.clear(ValueClass::SearchIndex(SearchIndexClass { + index, + typ: SearchIndexType::Index { + id, + field: SearchIndexField { + field_id: field.field_id, + len: field.len, + data: field.data, + }, + }, + })); + } + } +} + +impl SearchIndex { + pub(crate) fn as_u8(&self) -> u8 { + match self { + SearchIndex::Email => 0, + SearchIndex::Calendar => 1, + SearchIndex::Contacts => 2, + SearchIndex::File => 3, + SearchIndex::Tracing => 4, + SearchIndex::InMemory => unreachable!(), + } + } +} + +impl SearchField { + pub(crate) fn u8_id(&self) -> u8 { + match self { + SearchField::AccountId => 0, + SearchField::DocumentId => 1, + SearchField::Id => 2, + SearchField::Email(field) => match field { + EmailSearchField::From => 3, + EmailSearchField::To => 4, + EmailSearchField::Cc => 5, + EmailSearchField::Bcc => 6, + EmailSearchField::Subject => 7, + EmailSearchField::Body => 8, + EmailSearchField::Attachment => 9, + EmailSearchField::ReceivedAt => 10, + EmailSearchField::SentAt => 11, + EmailSearchField::Size => 12, + EmailSearchField::HasAttachment => 13, + EmailSearchField::Headers => 14, + }, + SearchField::Calendar(field) => match field { + CalendarSearchField::Title => 3, + CalendarSearchField::Description => 4, + CalendarSearchField::Location => 5, + CalendarSearchField::Owner => 6, + CalendarSearchField::Attendee => 7, + CalendarSearchField::Start => 8, + CalendarSearchField::Uid => 9, + }, + SearchField::Contact(field) => match field { + ContactSearchField::Member => 3, + ContactSearchField::Kind => 4, + ContactSearchField::Name => 5, + ContactSearchField::Nickname => 6, + ContactSearchField::Organization => 7, + ContactSearchField::Email => 8, + ContactSearchField::Phone => 9, + ContactSearchField::OnlineService => 10, + ContactSearchField::Address => 11, + ContactSearchField::Note => 12, + ContactSearchField::Uid => 13, + }, + SearchField::File(field) => match field { + FileSearchField::Name => 3, + FileSearchField::Content => 4, + }, + SearchField::Tracing(field) => match field { + TracingSearchField::EventType => 3, + TracingSearchField::QueueId => 4, + TracingSearchField::Keywords => 5, + }, + } + } +} diff --git a/crates/store/src/write/batch.rs b/crates/store/src/write/batch.rs index d631537b..a61c8ac1 100644 --- a/crates/store/src/write/batch.rs +++ b/crates/store/src/write/batch.rs @@ -10,7 +10,7 @@ use super::{ }; use crate::{ SerializeInfallible, U32_LEN, - write::{LogCollection, MergeFn}, + write::{LogCollection, MergeFnc, MergeOperation, Params, SetFnc, SetOperation}, }; use types::{ collection::{Collection, SyncCollection, VanishedCollection}, @@ -143,47 +143,35 @@ impl BatchBuilder { self.batch_size += class.serialized_size() + value.len(); self.ops.push(Operation::Value { class, - op: ValueOp::Set { - value, - version_offset: None, - }, + op: ValueOp::Set(value), }); self.batch_ops += 1; self } - pub fn merge( + pub fn set_fnc( &mut self, class: impl Into, - value: impl Fn(Option<&[u8]>) -> trc::Result> + Sync + Send + 'static, + params: Params, + fnc: SetFnc, ) -> &mut Self { self.ops.push(Operation::Value { class: class.into(), - op: ValueOp::Merge(MergeFn { - fnc: Box::new(value), - fnc_id: rand::random::(), - }), + op: ValueOp::SetFnc(SetOperation { fnc, params }), }); self } - pub fn set_versioned( + pub fn merge_fnc( &mut self, class: impl Into, - value: impl Into>, - version_offset: usize, + params: Params, + fnc: MergeFnc, ) -> &mut Self { - let class = class.into(); - let value = value.into(); - self.batch_size += class.serialized_size() + value.len(); self.ops.push(Operation::Value { - class, - op: ValueOp::Set { - value, - version_offset: Some(version_offset), - }, + class: class.into(), + op: ValueOp::MergeFnc(MergeOperation { fnc, params }), }); - self.batch_ops += 1; self } @@ -202,10 +190,7 @@ impl BatchBuilder { self.batch_size += (U32_LEN * 3) + op.len(); self.ops.push(Operation::Value { class: ValueClass::Acl(grant_account_id), - op: ValueOp::Set { - value: op, - version_offset: None, - }, + op: ValueOp::Set(op), }); self.batch_ops += 1; self diff --git a/crates/store/src/write/key.rs b/crates/store/src/write/key.rs index 199f85d6..4414ff8d 100644 --- a/crates/store/src/write/key.rs +++ b/crates/store/src/write/key.rs @@ -13,9 +13,10 @@ use crate::{ SUBSPACE_BLOB_RESERVE, SUBSPACE_COUNTER, SUBSPACE_DIRECTORY, SUBSPACE_IN_MEMORY_COUNTER, SUBSPACE_IN_MEMORY_VALUE, SUBSPACE_INDEXES, SUBSPACE_LOGS, SUBSPACE_PROPERTY, SUBSPACE_QUEUE_EVENT, SUBSPACE_QUEUE_MESSAGE, SUBSPACE_QUOTA, SUBSPACE_REPORT_IN, - SUBSPACE_REPORT_OUT, SUBSPACE_SETTINGS, SUBSPACE_TASK_QUEUE, SUBSPACE_TELEMETRY_METRIC, - SUBSPACE_TELEMETRY_SPAN, U16_LEN, U32_LEN, U64_LEN, ValueKey, WITH_SUBSPACE, - write::{IndexPropertyClass, SearchIndex}, + SUBSPACE_REPORT_OUT, SUBSPACE_SEARCH_INDEX, SUBSPACE_SETTINGS, SUBSPACE_TASK_QUEUE, + SUBSPACE_TELEMETRY_METRIC, SUBSPACE_TELEMETRY_SPAN, U16_LEN, U32_LEN, U64_LEN, ValueKey, + WITH_SUBSPACE, + write::{IndexPropertyClass, SearchIndex, SearchIndexId, SearchIndexType}, }; use std::convert::TryInto; use types::{blob_hash::BLOB_HASH_LEN, collection::SyncCollection}; @@ -440,6 +441,51 @@ impl ValueClass { .write(*notify_account_id) .write(u8::from(SyncCollection::ShareNotification)) .write(*notification_id), + ValueClass::SearchIndex(index) => match &index.typ { + SearchIndexType::Term { account_id, hash } => { + let class = index.index.as_u8(); + if let Some(account_id) = account_id { + serializer + .write(class) + .write(*account_id) + .write(hash.as_bytes()) + } else { + serializer.write(class).write(hash.as_bytes()) + } + } + SearchIndexType::Index { id, field } => { + let class = index.index.as_u8() | 1 << 6; + match id { + SearchIndexId::Account { + account_id, + document_id, + } => serializer + .write(class) + .write(*account_id) + .write(field.field_id) + .write(&field.data[..field.len as usize]) + .write(*document_id), + SearchIndexId::Global { id } => serializer + .write(class) + .write(field.field_id) + .write(&field.data[..field.len as usize]) + .write(*id), + } + } + SearchIndexType::Document { id } => { + let class = index.index.as_u8() | 2 << 6; + match id { + SearchIndexId::Account { + account_id, + document_id, + } => serializer + .write(class) + .write(*account_id) + .write(*document_id), + SearchIndexId::Global { id } => serializer.write(class).write(*id), + } + } + }, ValueClass::Any(any) => serializer.write(any.key.as_slice()), } .finalize() @@ -543,6 +589,20 @@ impl ValueClass { ValueClass::DocumentId => U32_LEN + 1, ValueClass::ChangeId => U32_LEN, ValueClass::ShareNotification { .. } => U32_LEN + U64_LEN + 1, + ValueClass::SearchIndex(v) => match &v.typ { + SearchIndexType::Term { account_id, hash } => { + if account_id.is_some() { + 1 + U32_LEN + hash.len() + } else { + 1 + hash.len() + } + } + SearchIndexType::Index { field, .. } => 1 + field.len as usize + U64_LEN, + SearchIndexType::Document { id } => match id { + SearchIndexId::Account { .. } => 1 + U32_LEN * 2, + SearchIndexId::Global { .. } => 1 + U64_LEN, + }, + }, ValueClass::Any(v) => v.key.len(), } } @@ -590,6 +650,7 @@ impl ValueClass { }, ValueClass::DocumentId | ValueClass::ChangeId => SUBSPACE_COUNTER, ValueClass::ShareNotification { .. } => SUBSPACE_LOGS, + ValueClass::SearchIndex(_) => SUBSPACE_SEARCH_INDEX, ValueClass::Any(any) => any.subspace, } } diff --git a/crates/store/src/write/mod.rs b/crates/store/src/write/mod.rs index 2f1a09cb..d0ca25a9 100644 --- a/crates/store/src/write/mod.rs +++ b/crates/store/src/write/mod.rs @@ -11,6 +11,7 @@ use nlp::tokenizers::word::WordTokenizer; use rkyv::util::AlignedVec; use std::{ collections::HashSet, + hash::Hash, time::{Duration, SystemTime}, }; use types::{ @@ -73,6 +74,7 @@ where #[derive(Debug, Default)] pub struct AssignedIds { pub ids: Vec, + current_change_id: Option, } #[derive(Debug)] @@ -173,6 +175,7 @@ pub enum ValueClass { Queue(QueueClass), Report(ReportClass), Telemetry(TelemetryClass), + SearchIndex(SearchIndexClass), Any(AnyClass), ShareNotification { notification_id: u64, @@ -188,6 +191,42 @@ pub enum IndexPropertyClass { Integer { property: u8, value: u64 }, } +#[derive(Debug, PartialEq, Clone, Eq, Hash)] +pub struct SearchIndexClass { + pub index: SearchIndex, + pub typ: SearchIndexType, +} + +#[derive(Debug, PartialEq, Clone, Eq, Hash)] +pub enum SearchIndexType { + Term { + account_id: Option, + hash: CheekyHash, + }, + Index { + id: SearchIndexId, + field: SearchIndexField, + }, + Document { + id: SearchIndexId, + }, +} + +pub(crate) const SEARCH_INDEX_MAX_FIELD_LEN: usize = 16; + +#[derive(Debug, PartialEq, Eq, Clone, Hash, rkyv::Serialize, rkyv::Deserialize, rkyv::Archive)] +pub struct SearchIndexField { + pub(crate) field_id: u8, + pub(crate) len: u8, + pub(crate) data: [u8; SEARCH_INDEX_MAX_FIELD_LEN], +} + +#[derive(Debug, PartialEq, Clone, Copy, Eq, Hash)] +pub enum SearchIndexId { + Account { account_id: u32, document_id: u32 }, + Global { id: u64 }, +} + #[derive(Debug, PartialEq, Clone, Eq, Hash)] pub enum TaskQueueClass { UpdateIndex { @@ -295,21 +334,47 @@ pub struct ReportEvent { #[derive(Debug, PartialEq, Eq, Hash, Default)] pub enum ValueOp { - Set { - value: Vec, - version_offset: Option, - }, + Set(Vec), + SetFnc(SetOperation), + MergeFnc(MergeOperation), AtomicAdd(i64), AddAndGet(i64), - Merge(MergeFn), #[default] Clear, } -#[allow(clippy::type_complexity)] -pub struct MergeFn { - pub fnc: Box) -> trc::Result> + Send + Sync>, - pub fnc_id: u64, +pub enum MergeResult { + Update(Vec), + Skip, + Delete, +} + +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub enum Param { + I64(i64), + U64(u64), + String(String), + Bytes(Vec), + Bool(bool), +} + +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +#[repr(transparent)] +pub struct Params(Vec); + +pub type SetFnc = fn(&Params, &AssignedIds) -> trc::Result>; +pub type MergeFnc = fn(&Params, &AssignedIds, Option<&[u8]>) -> trc::Result; + +#[derive(Debug, Clone)] +pub struct MergeOperation { + pub(crate) fnc: MergeFnc, + pub(crate) params: Params, +} + +#[derive(Debug, Clone)] +pub struct SetOperation { + pub(crate) fnc: SetFnc, + pub(crate) params: Params, } #[derive(Debug, PartialEq, Clone, Eq, Hash)] @@ -391,6 +456,20 @@ impl AssignedIds { }) } + pub fn current_change_id(&self) -> trc::Result { + self.current_change_id.ok_or_else(|| { + trc::StoreEvent::UnexpectedError + .caused_by(trc::location!()) + .ctx(trc::Key::Reason, "No current change id is set") + }) + } + + pub(crate) fn set_current_change_id(&mut self, account_id: u32) -> trc::Result { + let change_id = self.last_change_id(account_id)?; + self.current_change_id = Some(change_id); + Ok(change_id) + } + pub fn last_counter_id(&self) -> trc::Result { self.ids .iter() @@ -449,28 +528,6 @@ impl From for u8 { } } -impl std::fmt::Debug for MergeFn { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("MergeFn") - .field("fnc_id", &self.fnc_id) - .finish() - } -} - -impl PartialEq for MergeFn { - fn eq(&self, other: &Self) -> bool { - self.fnc_id == other.fnc_id - } -} - -impl Eq for MergeFn {} - -impl std::hash::Hash for MergeFn { - fn hash(&self, state: &mut H) { - self.fnc_id.hash(state); - } -} - impl From for ValueClass { fn from(value: ContactField) -> Self { ValueClass::Property(value.into()) @@ -524,3 +581,142 @@ impl From for ValueClass { ValueClass::Property(value.into()) } } + +impl PartialEq for MergeOperation { + fn eq(&self, other: &Self) -> bool { + self.params == other.params + } +} + +impl Eq for MergeOperation {} + +impl PartialEq for SetOperation { + fn eq(&self, other: &Self) -> bool { + self.params == other.params + } +} + +impl Eq for SetOperation {} + +impl Hash for MergeOperation { + fn hash(&self, state: &mut H) { + self.params.hash(state); + } +} + +impl Hash for SetOperation { + fn hash(&self, state: &mut H) { + self.params.hash(state); + } +} + +impl SetOperation { + pub fn params(&self) -> &Params { + &self.params + } +} + +impl MergeOperation { + pub fn params(&self) -> &Params { + &self.params + } +} + +impl Params { + pub fn with_capacity(capacity: usize) -> Self { + Self(Vec::with_capacity(capacity)) + } + + pub fn new() -> Self { + Self(Vec::new()) + } + + pub fn with_i64(mut self, value: i64) -> Self { + self.0.push(Param::I64(value)); + self + } + + pub fn with_u64(mut self, value: u64) -> Self { + self.0.push(Param::U64(value)); + self + } + + pub fn with_string(mut self, value: String) -> Self { + self.0.push(Param::String(value)); + self + } + + pub fn with_str(mut self, value: &str) -> Self { + self.0.push(Param::String(value.to_string())); + self + } + + pub fn with_bytes(mut self, value: Vec) -> Self { + self.0.push(Param::Bytes(value)); + self + } + + pub fn with_bool(mut self, value: bool) -> Self { + self.0.push(Param::Bool(value)); + self + } + + pub fn i64(&self, idx: usize) -> i64 { + match &self.0[idx] { + Param::I64(v) => *v, + _ => panic!("Param at index {} is not an i64", idx), + } + } + + pub fn u64(&self, idx: usize) -> u64 { + match &self.0[idx] { + Param::U64(v) => *v, + _ => panic!("Param at index {} is not a u64", idx), + } + } + + pub fn string(&self, idx: usize) -> &str { + match &self.0[idx] { + Param::String(v) => v.as_str(), + _ => panic!("Param at index {} is not a String", idx), + } + } + + pub fn bytes(&self, idx: usize) -> &[u8] { + match &self.0[idx] { + Param::Bytes(v) => v.as_slice(), + _ => panic!("Param at index {} is not Bytes", idx), + } + } + + pub fn bool(&self, idx: usize) -> bool { + match &self.0[idx] { + Param::Bool(v) => *v, + _ => panic!("Param at index {} is not a bool", idx), + } + } + + pub fn len(&self) -> usize { + self.0.len() + } + + pub fn is_empty(&self) -> bool { + self.0.is_empty() + } + + pub fn as_slice(&self) -> &[Param] { + &self.0 + } +} + +impl Default for Params { + fn default() -> Self { + Self::new() + } +} + +impl AsRef<[Param]> for Params { + fn as_ref(&self) -> &[Param] { + &self.0 + } +} diff --git a/crates/store/src/write/serialize.rs b/crates/store/src/write/serialize.rs index 66b73495..9e142139 100644 --- a/crates/store/src/write/serialize.rs +++ b/crates/store/src/write/serialize.rs @@ -8,6 +8,7 @@ use super::{ARCHIVE_ALIGNMENT, AlignedBytes, Archive, ArchiveVersion, Archiver}; use crate::{Deserialize, Serialize, SerializeInfallible, U32_LEN, U64_LEN, Value}; use compact_str::format_compact; use rkyv::util::AlignedVec; +use roaring::{RoaringBitmap, RoaringTreemap}; const MAGIC_MARKER: u8 = 1 << 7; const VERSIONED: u8 = 1 << 6; @@ -419,10 +420,10 @@ where } } - pub fn serialize_versioned(self) -> trc::Result<(usize, Vec)> { + pub fn serialize_versioned(self) -> trc::Result<(u64, Vec)> { self.with_version() .serialize() - .map(|bytes| (bytes.len() - U64_LEN - 1, bytes)) + .map(|bytes| ((bytes.len() - U64_LEN - 1) as u64, bytes)) } } @@ -582,3 +583,49 @@ impl Default for Archive { } } } + +impl Serialize for RoaringBitmap { + fn serialize(&self) -> trc::Result> { + let mut bytes = Vec::with_capacity(self.serialized_size()); + self.serialize_into(&mut bytes) + .map_err(|err| { + trc::StoreEvent::UnexpectedError + .caused_by(trc::location!()) + .reason(err) + }) + .map(|_| bytes) + } +} + +impl Deserialize for RoaringBitmap { + fn deserialize(bytes: &[u8]) -> trc::Result { + RoaringBitmap::deserialize_from(bytes).map_err(|err| { + trc::StoreEvent::DeserializeError + .caused_by(trc::location!()) + .reason(err) + }) + } +} + +impl Serialize for RoaringTreemap { + fn serialize(&self) -> trc::Result> { + let mut bytes = Vec::with_capacity(self.serialized_size()); + self.serialize_into(&mut bytes) + .map_err(|err| { + trc::StoreEvent::UnexpectedError + .caused_by(trc::location!()) + .reason(err) + }) + .map(|_| bytes) + } +} + +impl Deserialize for RoaringTreemap { + fn deserialize(bytes: &[u8]) -> trc::Result { + RoaringTreemap::deserialize_from(bytes).map_err(|err| { + trc::StoreEvent::DeserializeError + .caused_by(trc::location!()) + .reason(err) + }) + } +} diff --git a/crates/utils/src/cheeky_hash.rs b/crates/utils/src/cheeky_hash.rs index 53c54da8..f2a67447 100644 --- a/crates/utils/src/cheeky_hash.rs +++ b/crates/utils/src/cheeky_hash.rs @@ -6,25 +6,42 @@ use nohash_hasher::IsEnabled; use std::{ - collections::{HashMap, HashSet}, + collections::{BTreeMap, HashMap, HashSet}, hash::Hash, }; -#[derive(Debug, Copy, Clone, PartialEq, Eq)] +// A hash that can cheekily store small inputs directly without hashing them. +#[derive( + Debug, + Copy, + Clone, + PartialEq, + Eq, + PartialOrd, + Ord, + rkyv::Serialize, + rkyv::Deserialize, + rkyv::Archive, +)] #[repr(transparent)] pub struct CheekyHash([u8; HASH_SIZE]); + const HASH_SIZE: usize = std::mem::size_of::() * 2; const HASH_PAYLOAD: usize = HASH_SIZE - 1; pub type CheekyHashSet = HashSet>; pub type CheekyHashMap = HashMap>; +pub type CheekyBTreeMap = BTreeMap; impl CheekyHash { + pub const NULL: CheekyHash = CheekyHash([0u8; HASH_SIZE]); + pub const FULL: CheekyHash = CheekyHash([u8::MAX; HASH_SIZE]); + pub fn new(bytes: impl AsRef<[u8]>) -> Self { let mut hash = [0u8; HASH_SIZE]; let bytes = bytes.as_ref(); - if bytes.len() < HASH_PAYLOAD { + if bytes.len() <= HASH_PAYLOAD { hash[0] = bytes.len() as u8; hash[1..1 + bytes.len()].copy_from_slice(bytes); } else { @@ -58,6 +75,15 @@ impl CheekyHash { pub fn as_bytes(&self) -> &[u8] { &self.0[..self.len()] } + + #[inline(always)] + pub fn as_raw_bytes(&self) -> &[u8; HASH_SIZE] { + &self.0 + } + + pub fn into_inner(self) -> [u8; HASH_SIZE] { + self.0 + } } impl AsRef<[u8]> for CheekyHash { @@ -69,7 +95,7 @@ impl AsRef<[u8]> for CheekyHash { impl Hash for CheekyHash { fn hash(&self, state: &mut H) { let len = self.0[0] as usize; - if len < HASH_PAYLOAD { + if len <= HASH_PAYLOAD { state.write_u64(xxhash_rust::xxh3::xxh3_64(&self.0[1..1 + len])); } else { state.write_u64(u64::from_be_bytes( @@ -83,6 +109,24 @@ impl Hash for CheekyHash { impl IsEnabled for CheekyHash {} +impl ArchivedCheekyHash { + #[inline(always)] + pub fn as_raw_bytes(&self) -> &[u8; HASH_SIZE] { + &self.0 + } + + #[inline(always)] + pub fn as_bytes(&self) -> &[u8] { + let len = self.0[0] as usize; + &self.0[..1 + len.min(HASH_PAYLOAD)] + } + + #[inline(always)] + pub fn to_native(&self) -> CheekyHash { + CheekyHash(self.0) + } +} + #[cfg(test)] mod tests { use super::*; diff --git a/crates/utils/src/config/http.rs b/crates/utils/src/config/http.rs index fa036dd0..4475ddc7 100644 --- a/crates/utils/src/config/http.rs +++ b/crates/utils/src/config/http.rs @@ -8,14 +8,22 @@ use crate::config::{Config, utils::AsKey}; use base64::{Engine, engine::general_purpose}; use reqwest::{ Client, - header::{AUTHORIZATION, HeaderMap, HeaderName, HeaderValue, USER_AGENT}, + header::{AUTHORIZATION, CONTENT_TYPE, HeaderMap, HeaderName, HeaderValue, USER_AGENT}, }; use std::{str::FromStr, time::Duration}; -pub fn build_http_client(config: &mut Config, prefix: impl AsKey) -> Option { +pub fn build_http_client( + config: &mut Config, + prefix: impl AsKey, + content_type: Option<&str>, +) -> Option { let mut headers = parse_http_headers(config, prefix.clone()); headers.insert(USER_AGENT, "Stalwart/1.0.0".parse().unwrap()); + if let Some(content_type) = content_type { + headers.insert(CONTENT_TYPE, HeaderValue::from_str(content_type).unwrap()); + } + let prefix = prefix.as_key(); match Client::builder() .connect_timeout(