#![forbid(unsafe_code)]
use std::collections::BTreeMap;
use std::error::Error;
use std::num::NonZeroU32;
use haematite::db::DatabaseConfig;
use haematite::store::NodeCacheBudget;
use haematite_cli::serve::protocol::{
BranchPage, Cardinality, Change, CollectionWrite, DocumentImport, DocumentPage, DocumentRead,
Edge, Fork, ImportedDocument, Mutation, Node, Operation, Relation, ResponseFailure, Selection,
StoreRead, VectorDelete, VectorWrite, Version,
};
use haematite_cli::serve::service::Service;
use haematite_cli::serve::storage::{Failure, Receipt, Storage};
use serde_json::{Value, json, value::RawValue};
type TestResult = Result<(), Box<dyn Error>>;
struct Fixture {
service: Service,
directory: tempfile::TempDir,
versions: BTreeMap<String, Version>,
request_count: u64,
}
impl Fixture {
fn new(max_page: u32, schema: Option<Value>) -> Result<Self, Box<dyn Error>> {
let directory = tempfile::tempdir()?;
let storage = Storage::create(
DatabaseConfig {
data_dir: directory.path().join("database"),
shard_count: 2,
distributed: None,
executor_threads: Some(2),
node_cache_budget: Some(NodeCacheBudget::Unlimited),
retention: haematite::Retention::KeepAll,
},
NodeCacheBudget::Unlimited,
long_wait_interval(),
)?;
let version = storage.capture(&selection("main", None))?.version;
let mut fixture = Self {
service: Service::new(storage, NonZeroU32::new(max_page).ok_or("zero page limit")?)?,
directory,
versions: BTreeMap::from([("main".to_owned(), version)]),
request_count: 0,
};
fixture.apply(
"main",
Change::Collection(CollectionWrite {
collection: "records".to_owned(),
schema,
}),
)?;
Ok(fixture)
}
fn version(&self, branch: &str) -> Result<Version, Box<dyn Error>> {
self.versions
.get(branch)
.copied()
.ok_or_else(|| "fixture branch missing".into())
}
fn request(&mut self, branch: &str, change: Change) -> Result<Mutation, Box<dyn Error>> {
self.request_count += 1;
Ok(Mutation {
request_id: format!("request-{}", self.request_count),
branch: branch.to_owned(),
expected: self.version(branch)?,
change,
})
}
fn apply_request(&mut self, request: Mutation) -> Result<Receipt, Box<dyn Error>> {
let receipt: Receipt =
serde_json::from_value(self.service.dispatch(Operation::Mutate(request))?)?;
if self.versions.get(&receipt.branch).is_none_or(|current| {
current.generation != receipt.version.generation
|| current.sequence < receipt.version.sequence
}) {
self.versions
.insert(receipt.branch.clone(), receipt.version);
}
Ok(receipt)
}
fn apply(&mut self, branch: &str, change: Change) -> Result<Receipt, Box<dyn Error>> {
let request = self.request(branch, change)?;
self.apply_request(request)
}
fn get(&self, branch: &str, id: &str) -> Result<Value, Box<dyn Error>> {
Ok(self.service.dispatch(Operation::Get(DocumentRead {
selection: selection(branch, None),
collection: "records".to_owned(),
document_id: id.to_owned(),
}))?)
}
fn stats(&self, branch: &str, version: Option<Version>) -> Result<Value, Box<dyn Error>> {
Ok(self.service.dispatch(Operation::StoreStats(StoreRead {
selection: selection(branch, version),
}))?)
}
fn fields(&self, branch: &str, version: Option<Version>) -> Result<Value, Box<dyn Error>> {
Ok(self.service.dispatch(Operation::VectorFields(StoreRead {
selection: selection(branch, version),
}))?)
}
fn refuse(
&mut self,
change: Change,
code: &str,
path: &str,
) -> Result<ResponseFailure, Box<dyn Error>> {
let before = self.version("main")?;
let documents = self.stats("main", None)?["documents"].clone();
let request = self.request("main", change)?;
let failure = self
.service
.dispatch(Operation::Mutate(request))
.err()
.ok_or("invalid import accepted")?;
assert_eq!(failure.code, code);
assert_eq!(failure.path.as_deref(), Some(path));
let branches = self.service.dispatch(Operation::Branches(BranchPage {
after: None,
limit: NonZeroU32::MIN,
}))?;
assert_eq!(
serde_json::from_value::<Version>(branches["entries"][0]["version"].clone())?,
before
);
assert_eq!(self.stats("main", None)?["documents"], documents);
Ok(failure)
}
}
fn selection(branch: &str, version: Option<Version>) -> Selection {
Selection {
branch: branch.to_owned(),
version,
}
}
fn row(id: &str, raw: &str) -> Result<ImportedDocument, serde_json::Error> {
Ok(ImportedDocument {
document_id: id.to_owned(),
document: RawValue::from_string(raw.to_owned())?,
})
}
fn import(documents: Vec<ImportedDocument>) -> Change {
Change::PutMany(DocumentImport {
collection: "records".to_owned(),
documents,
})
}
fn required_name() -> Value {
json!({"type":"object", "required":["name"], "properties":{"name":{"type":"string"}}})
}
fn node(id: &str) -> Node {
Node {
collection: "records".to_owned(),
document_id: id.to_owned(),
}
}
fn vector(id: &str, field: &str, values: Vec<f32>) -> Change {
Change::VectorPut(VectorWrite {
collection: "records".to_owned(),
document_id: id.to_owned(),
field: field.to_owned(),
vector: values,
})
}
#[test]
fn import_is_one_version_with_one_revision_and_replays_after_a_later_write() -> TestResult {
let mut fixture = Fixture::new(8, None)?;
let before = fixture.version("main")?;
let raw = r#"{ "name": "kept", "number":18446744073709551616001 }"#;
let request = fixture.request(
"main",
import(vec![row("a", raw)?, row("b", "{}")?, row("c", "null")?]),
)?;
let replay = serde_json::to_vec(&request)?;
let receipt = fixture.apply_request(request)?;
assert_eq!(receipt.version.generation, before.generation);
assert_eq!(receipt.version.sequence, before.sequence + 1);
let listed = fixture.service.dispatch(Operation::List(DocumentPage {
selection: selection("main", Some(receipt.version)),
collection: "records".to_owned(),
prefix: None,
content: None,
after: None,
limit: NonZeroU32::new(8).ok_or("zero page")?,
}))?;
let documents = listed["entries"].as_array().ok_or("entries absent")?;
assert_eq!(documents.len(), 3);
assert!(listed["next_after"].is_null());
assert_eq!(documents[0]["raw_json"], raw);
for document in documents {
assert_eq!(
serde_json::from_value::<Version>(document["revision"]["written_at"].clone())?,
receipt.version
);
assert_eq!(document["revision"]["written_in"], "main");
}
let later = fixture
.apply("main", import(vec![row("a", r#"{"name":"later"}"#)?]))?
.version;
let replayed = fixture.apply_request(serde_json::from_slice(&replay)?)?;
assert_eq!(replayed.version, receipt.version);
assert_eq!(replayed.fingerprint, receipt.fingerprint);
assert_eq!(
fixture.get("main", "a")?["document"]["raw_json"],
r#"{"name":"later"}"#
);
let current = fixture.stats("main", None)?;
assert_eq!(
serde_json::from_value::<Version>(current["selection"]["version"].clone())?,
later
);
assert_eq!(current["documents"], 3);
Ok(())
}
#[test]
fn every_import_refusal_preserves_the_whole_branch() -> TestResult {
let mut fixture = Fixture::new(8, Some(required_name()))?;
fixture.refuse(import(Vec::new()), "import_empty", "/documents")?;
let duplicate = fixture.refuse(
import(vec![
row("a", r#"{"name":"first"}"#)?,
row("a", r#"{"name":"duplicate"}"#)?,
]),
"import_duplicate",
"/documents/1/document_id",
)?;
assert_eq!(duplicate.message, "repeated document_id a");
for id in ["", "bad\0id"] {
fixture.refuse(
import(vec![
row("a", r#"{"name":"valid"}"#)?,
row(id, r#"{"name":"valid"}"#)?,
]),
"invalid_identifier",
"/documents/1/document_id",
)?;
}
fixture.refuse(
Change::PutMany(DocumentImport {
collection: "missing".to_owned(),
documents: vec![row("a", "{}")?],
}),
"collection_missing",
"/collection",
)?;
fixture.refuse(
Change::PutMany(DocumentImport {
collection: String::new(),
documents: vec![row("a", "{}")?],
}),
"invalid_identifier",
"/collection",
)?;
assert!(fixture.get("main", "a")?["document"].is_null());
assert_eq!(fixture.stats("main", None)?["documents"], 0);
Ok(())
}
#[test]
fn schema_failure_names_the_row_and_leaves_earlier_overwrites_and_inserts_absent() -> TestResult {
let mut fixture = Fixture::new(8, Some(required_name()))?;
fixture.apply("main", import(vec![row("a", r#"{"name":"original"}"#)?]))?;
fixture.refuse(
import(vec![
row("a", r#"{"name":"overwrite"}"#)?,
row("b", r#"{"name":"insert"}"#)?,
row("c", "{}")?,
]),
"schema_refused",
"/documents/2/document/name",
)?;
assert_eq!(
fixture.get("main", "a")?["document"]["raw_json"],
r#"{"name":"original"}"#
);
assert!(fixture.get("main", "b")?["document"].is_null());
assert!(fixture.get("main", "c")?["document"].is_null());
Ok(())
}
#[test]
fn import_bound_is_sixty_four_pages_and_an_oversize_list_changes_nothing() -> TestResult {
let mut fixture = Fixture::new(1, None)?;
let before = fixture.version("main")?;
let documents = (0..64)
.map(|index| row(&format!("document-{index:02}"), "{}"))
.collect::<Result<Vec<_>, _>>()?;
let receipt = fixture.apply("main", import(documents))?;
assert_eq!(receipt.version.sequence, before.sequence + 1);
assert_eq!(fixture.stats("main", None)?["documents"], 64);
let documents = (0..65)
.map(|index| row(&format!("refused-{index:02}"), "{}"))
.collect::<Result<Vec<_>, _>>()?;
let request = fixture.request("main", import(documents))?;
let failure = fixture
.service
.dispatch(Operation::Mutate(request))
.err()
.ok_or("oversize import accepted")?;
assert_eq!(failure.code, "import_limit");
assert_eq!(failure.path.as_deref(), Some("/documents/64"));
assert!(failure.message.contains("64"));
assert_eq!(fixture.stats("main", None)?["documents"], 64);
assert!(fixture.get("main", "refused-00")?["document"].is_null());
let head: Version =
serde_json::from_value(fixture.stats("main", None)?["selection"]["version"].clone())?;
assert_eq!(head, receipt.version);
Ok(())
}
fn counted_source() -> Result<(Fixture, Version), Box<dyn Error>> {
let mut fixture = Fixture::new(8, None)?;
fixture.apply(
"main",
Change::Collection(CollectionWrite {
collection: "empty".to_owned(),
schema: None,
}),
)?;
fixture.apply(
"main",
import(vec![row("a", "{}")?, row("b", "{}")?, row("c", "{}")?]),
)?;
fixture.apply(
"main",
Change::RelationPut(Relation {
relation: "links".to_owned(),
from_collection: "records".to_owned(),
to_collection: "records".to_owned(),
cardinality: Cardinality::ManyToMany,
}),
)?;
fixture.apply(
"main",
Change::EdgePut(Edge {
relation: "links".to_owned(),
from: node("a"),
to: node("b"),
properties: json!({}),
}),
)?;
for (id, field, values) in [
("a", "embedding", vec![1.0, 2.0]),
("b", "embedding", vec![3.0, 4.0]),
("a", "quality", vec![1.0]),
] {
fixture.apply("main", vector(id, field, values))?;
}
let source = fixture.version("main")?;
Ok((fixture, source))
}
fn counted_fork(fixture: &mut Fixture, source: Version) -> TestResult {
fixture.apply(
"main",
Change::Fork(Fork {
name: "experiment".to_owned(),
from: source,
}),
)?;
fixture.apply("experiment", import(vec![row("d", "{}")?, row("e", "{}")?]))?;
fixture.apply(
"experiment",
Change::EdgePut(Edge {
relation: "links".to_owned(),
from: node("c"),
to: node("d"),
properties: json!({}),
}),
)?;
fixture.apply("experiment", vector("d", "embedding", vec![5.0, 6.0]))?;
fixture.apply("experiment", vector("d", "local", vec![1.0, 2.0, 3.0]))?;
fixture.apply(
"experiment",
Change::VectorDelete(VectorDelete {
collection: "records".to_owned(),
document_id: "a".to_owned(),
field: "quality".to_owned(),
}),
)?;
Ok(())
}
#[test]
fn counts_and_fields_are_exact_at_captured_main_and_fork_roots() -> TestResult {
let (mut fixture, source) = counted_source()?;
counted_fork(&mut fixture, source)?;
for (branch, documents, edges, vectors) in [("main", 3, 1, 3), ("experiment", 5, 2, 4)] {
let stats = fixture.stats(branch, None)?;
assert_eq!(stats["documents"], documents);
assert_eq!(stats["edges"], edges);
assert_eq!(stats["vectors"], vectors);
assert_eq!(
stats["collections"],
json!([{"name":"empty", "documents":0}, {"name":"records", "documents":documents}])
);
for key in ["live_bytes", "history_bytes", "reclaimable_bytes"] {
assert!(stats[key].is_null());
}
}
assert_eq!(
fixture.fields("main", None)?["fields"],
json!([
{"collection":"records", "field":"embedding", "dimensions":2, "count":2},
{"collection":"records", "field":"quality", "dimensions":1, "count":1},
])
);
assert_eq!(
fixture.fields("experiment", None)?["fields"],
json!([
{"collection":"records", "field":"embedding", "dimensions":2, "count":3},
{"collection":"records", "field":"local", "dimensions":3, "count":1},
{"collection":"records", "field":"quality", "dimensions":1, "count":0},
])
);
fixture.apply("main", import(vec![row("later", "{}")?]))?;
assert_eq!(fixture.stats("main", None)?["documents"], 4);
assert_eq!(fixture.stats("main", Some(source))?["documents"], 3);
assert_eq!(
fixture.fields("main", Some(source))?["fields"],
fixture.fields("main", None)?["fields"]
);
assert_eq!(fixture.stats("experiment", None)?["documents"], 5);
Ok(())
}
#[test]
fn import_replay_survives_restart_and_a_smaller_current_page_bound() -> TestResult {
let mut fixture = Fixture::new(2, None)?;
let documents = (0..65)
.map(|index| row(&format!("document-{index:02}"), "{}"))
.collect::<Result<Vec<_>, _>>()?;
let request = fixture.request("main", import(documents))?;
let saved = serde_json::to_vec(&request)?;
let receipt = fixture.apply_request(request)?;
let Fixture {
service,
directory,
versions,
request_count,
} = fixture;
drop(service);
let storage = Storage::open(
&directory.path().join("database"),
NodeCacheBudget::Unlimited,
long_wait_interval(),
)?;
let mut fixture = Fixture {
service: Service::new(storage, NonZeroU32::MIN)?,
directory,
versions,
request_count,
};
let replayed = fixture.apply_request(serde_json::from_slice(&saved)?)?;
assert_eq!(replayed.version, receipt.version);
assert_eq!(replayed.fingerprint, receipt.fingerprint);
assert_eq!(fixture.stats("main", None)?["documents"], 65);
Ok(())
}
#[test]
fn direct_import_requires_an_explicit_bound_and_studio_shapes_deserialize() -> TestResult {
let fixture = Fixture::new(8, None)?;
let Fixture {
service,
directory,
versions: _,
request_count: _,
} = fixture;
drop(service);
let mut storage = Storage::open(
&directory.path().join("database"),
NodeCacheBudget::Unlimited,
long_wait_interval(),
)?;
let before = storage.capture(&selection("main", None))?.version;
let request = Mutation {
request_id: "direct".to_owned(),
branch: "main".to_owned(),
expected: before,
change: import(vec![row("a", "{}")?]),
};
let refused: Failure = storage
.mutate(&request)
.err()
.ok_or("unbounded direct import accepted")?;
assert_eq!(refused.code, "import_limit_required");
assert_eq!(storage.capture(&selection("main", None))?.version, before);
storage.mutate_bounded(&request, NonZeroU32::MIN)?;
for method in ["store.stats", "vectors.fields"] {
let operation: Operation = serde_json::from_value(
json!({"method":method, "params":{"selection":{"branch":"main"}}}),
)?;
assert!(matches!(
operation,
Operation::StoreStats(_) | Operation::VectorFields(_)
));
}
let change: Change = serde_json::from_str(
r#"{"kind":"documents.put_many","value":{"collection":"records","documents":[{"document_id":"a","document":{ "number":1.00 }}]}}"#,
)?;
let Change::PutMany(imported) = change else {
return Err("wrong import variant".into());
};
assert_eq!(imported.documents[0].document.get(), r#"{ "number":1.00 }"#);
Ok(())
}
const fn long_wait_interval() -> std::time::Duration {
const LONG_WAIT_INTERVAL: std::time::Duration = std::time::Duration::from_secs(1);
LONG_WAIT_INTERVAL
}