use std::error::Error;
use super::{
Captured, Storage, attribution, branch_name, collection_key, document_key, reads, revision_key,
};
use crate::serve::contracts::{
MutationRequest, ResolvedMutation,
attribution::{Changed, Row},
envelopes::{Attributed, Collection, Document, Edge, Vector, attribution_key},
};
use crate::serve::protocol::{Selection, Version};
use haematite::branch::{CommitDurability, CommitRequest, current_timestamp, open};
use haematite::db::DatabaseConfig;
use haematite::store::NodeCacheBudget;
use serde_json::{Value, json};
type TestResult = Result<(), Box<dyn Error>>;
type RawRow = (Vec<u8>, Vec<u8>);
fn setup() -> Result<(tempfile::TempDir, Storage), 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,
fixture_long_wait_interval(),
)?;
Ok((directory, storage))
}
fn capture(storage: &Storage) -> Result<Captured, Box<dyn Error>> {
Ok(storage.capture(&Selection {
branch: "main".to_owned(),
version: None,
})?)
}
const fn next(origin: Version) -> Version {
Version {
generation: origin.generation,
sequence: origin.sequence + 1,
}
}
fn request(
id: &str,
caller: &Value,
change: &Value,
origin: Version,
) -> Result<ResolvedMutation, Box<dyn Error>> {
Ok(MutationRequest::decode(&serde_json::to_vec(&json!({
"request_id":id, "branch":"main", "caller":caller,
"contract":{"kind":"dependencies", "origin":origin, "reads":[]}, "change":change
}))?)?
.resolve("operator-not-the-author")?)
}
fn person(subject: &str) -> Value {
json!({"issuer":"lys", "subject":subject})
}
fn install(storage: &mut Storage, rows: Vec<RawRow>) -> TestResult {
let name = branch_name("main");
let branch = open(&name, &storage.refs, &storage.registry)?;
for (key, value) in rows {
branch.put(storage.writer.shard_for(&key), key, value)?;
}
storage.writer.commit(
&branch,
&storage.registry,
CommitRequest {
durability: CommitDurability::Durable {
refs: &mut storage.refs,
},
extra_parents: &[],
timestamp: current_timestamp(),
},
)?;
storage.publish_record(&name)?;
Ok(())
}
fn row_bytes(request: &ResolvedMutation, written: Version) -> Result<RawRow, Box<dyn Error>> {
let row = Row::for_request(request, written)?;
Ok((row.key()?, row.encode()?))
}
fn stored(storage: &Storage, at: &Captured, key: &[u8]) -> Result<Vec<u8>, Box<dyn Error>> {
Ok(reads::get_bytes(&storage.reader()?, at, key)?.ok_or("missing fixture row")?)
}
fn verify_pointer(
storage: &Storage,
at: &Captured,
envelope: &impl Attributed,
id: &str,
subject: &str,
) -> TestResult {
let reader = storage.reader()?;
let row = attribution::get(&reader, at, envelope.written_at(), envelope.request_id())?
.ok_or("pointer did not find row")?;
assert_eq!(row.key()?, attribution_key(envelope)?);
assert_eq!(row.request_id(), id);
assert_eq!(row.caller().subject(), subject);
Ok(())
}
#[test]
fn each_stored_envelope_finds_its_own_request_and_not_another_caller() -> TestResult {
let (_directory, mut storage) = setup()?;
let origin = capture(&storage)?.version;
let written = next(origin);
let doc = request(
"document",
&person("doctor-a"),
&json!({"kind":"document.put", "value":{"collection":"charts", "document_id":"child", "document":{"allergy":"penicillin"}}}),
origin,
)?;
let collection = request(
"collection",
&person("doctor-b"),
&json!({"kind":"collection.put", "value":{"collection":"charts", "schema":null}}),
origin,
)?;
let edge = request(
"edge",
&person("doctor-c"),
&json!({"kind":"edge.put", "value":{"relation":"assigned", "from":{"collection":"charts", "document_id":"child"}, "to":{"collection":"doctors", "document_id":"doctor-c"}, "properties":{}}}),
origin,
)?;
let vector = request(
"vector",
&person("doctor-d"),
&json!({"kind":"vector.put", "value":{"collection":"charts", "document_id":"child", "field":"embedding", "vector":[1.0]}}),
origin,
)?;
let d = Document::for_request(&doc, written, "main".to_owned(), origin)?;
let c = Collection::for_request(&collection, written)?;
let e = Edge::for_request(&edge, written)?;
let v = Vector::for_request(&vector, written)?;
let body = b"{ \"allergy\" : \"penicillin\" }".to_vec();
let edge_key = b"e\0charts\0child\0assigned\0doctors\0doctor-c".to_vec();
let incoming = b"i\0doctors\0doctor-c\0assigned\0charts\0child".to_vec();
let vector_key = b"v\0charts\0embedding\0child".to_vec();
let mut rows = vec![
(document_key("charts", "child"), body.clone()),
(revision_key("charts", "child"), d.encode()?),
(collection_key("charts"), c.encode()?),
(edge_key.clone(), e.encode()?),
(incoming.clone(), e.encode()?),
(vector_key.clone(), v.encode()?),
(b"w\0charts\0embedding".to_vec(), b"1".to_vec()),
];
for request in [&doc, &collection, &edge, &vector] {
rows.push(row_bytes(request, written)?);
}
install(&mut storage, rows)?;
let at = capture(&storage)?;
assert_eq!(at.version, written);
assert_eq!(
stored(&storage, &at, &document_key("charts", "child"))?,
body
);
let d = Document::decode(&stored(&storage, &at, &revision_key("charts", "child"))?)?;
let c = Collection::decode(&stored(&storage, &at, &collection_key("charts"))?)?;
let e = Edge::decode(&stored(&storage, &at, &edge_key)?)?;
let mirror = Edge::decode(&stored(&storage, &at, &incoming)?)?;
let v = Vector::decode(&stored(&storage, &at, &vector_key)?)?;
verify_pointer(&storage, &at, &d, "document", "doctor-a")?;
verify_pointer(&storage, &at, &c, "collection", "doctor-b")?;
verify_pointer(&storage, &at, &e, "edge", "doctor-c")?;
verify_pointer(&storage, &at, &v, "vector", "doctor-d")?;
assert_eq!(attribution_key(&e)?, attribution_key(&mirror)?);
assert_eq!(
attribution::version_rows(&storage.reader()?, &at, written)?.len(),
4
);
Ok(())
}
#[test]
fn three_imported_records_point_at_one_row_after_reopen() -> TestResult {
let (directory, mut storage) = setup()?;
let origin = capture(&storage)?.version;
let written = next(origin);
let request = request(
"import",
&person("doctor-a"),
&json!({"kind":"documents.put_many", "value":{"collection":"charts", "documents":[
{"document_id":"a", "document":{"allergy":"a"}}, {"document_id":"b", "document":{"allergy":"b"}}, {"document_id":"c", "document":{"allergy":"c"}}
]}}),
origin,
)?;
let mut rows = vec![row_bytes(&request, written)?];
for id in ["a", "b", "c"] {
rows.push((
document_key("charts", id),
serde_json::to_vec(&json!({"allergy":id}))?,
));
rows.push((
revision_key("charts", id),
Document::for_request(&request, written, "main".to_owned(), origin)?.encode()?,
));
}
install(&mut storage, rows)?;
drop(storage);
let storage = Storage::open(
&directory.path().join("database"),
NodeCacheBudget::Unlimited,
fixture_long_wait_interval(),
)?;
let at = capture(&storage)?;
let rows = attribution::version_rows(&storage.reader()?, &at, written)?;
assert_eq!(rows.len(), 1);
assert!(matches!(
rows[0].changed(),
Changed::Import { count: 3, .. }
));
for id in ["a", "b", "c"] {
let envelope = Document::decode(&stored(&storage, &at, &revision_key("charts", id))?)?;
verify_pointer(&storage, &at, &envelope, "import", "doctor-a")?;
assert_eq!(attribution_key(&envelope)?, rows[0].key()?);
assert_eq!(
stored(&storage, &at, &document_key("charts", id))?,
serde_json::to_vec(&json!({"allergy":id}))?
);
}
Ok(())
}
#[test]
fn vector_replacement_points_at_the_jobs_row_without_rewriting_document_authorship() -> TestResult {
let (_directory, mut storage) = setup()?;
let origin = capture(&storage)?.version;
let first = next(origin);
let doc_request = request(
"document",
&person("doctor-a"),
&json!({"kind":"document.put", "value":{"collection":"charts", "document_id":"child", "document":{}}}),
origin,
)?;
let vector_request = request(
"supplied-vector",
&person("doctor-a"),
&json!({"kind":"vector.put", "value":{"collection":"charts", "document_id":"child", "field":"embedding", "vector":[1.0]}}),
origin,
)?;
let doc = Document::for_request(&doc_request, first, "main".to_owned(), origin)?;
let vector_key = b"v\0charts\0embedding\0child";
install(
&mut storage,
vec![
row_bytes(&doc_request, first)?,
row_bytes(&vector_request, first)?,
(document_key("charts", "child"), b"{}".to_vec()),
(revision_key("charts", "child"), doc.encode()?),
(b"w\0charts\0embedding".to_vec(), b"1".to_vec()),
(
vector_key.to_vec(),
Vector::for_request(&vector_request, first)?.encode()?,
),
],
)?;
let before = capture(&storage)?;
let second = next(before.version);
let job = request(
"embedding-job",
&json!({"issuer":"local", "subject":"embedder", "responsible":"embedding-owner"}),
&json!({"kind":"vector.put", "value":{"collection":"charts", "document_id":"child", "field":"embedding", "vector":[2.0]}}),
before.version,
)?;
install(
&mut storage,
vec![
row_bytes(&job, second)?,
(
vector_key.to_vec(),
Vector::for_request(&job, second)?.encode()?,
),
],
)?;
let after = capture(&storage)?;
let old = Vector::decode(&stored(&storage, &before, vector_key)?)?;
let new = Vector::decode(&stored(&storage, &after, vector_key)?)?;
verify_pointer(&storage, &before, &old, "supplied-vector", "doctor-a")?;
verify_pointer(&storage, &after, &new, "embedding-job", "embedder")?;
assert_ne!(attribution_key(&old)?, attribution_key(&new)?);
let row = attribution::get(
&storage.reader()?,
&after,
new.written_at(),
new.request_id(),
)?
.ok_or("job row missing")?;
assert_eq!(row.caller().issuer(), "local");
assert_eq!(row.caller().responsible(), Some("embedding-owner"));
let unchanged = Document::decode(&stored(&storage, &after, &revision_key("charts", "child"))?)?;
verify_pointer(&storage, &after, &unchanged, "document", "doctor-a")?;
assert_eq!(attribution_key(&unchanged)?, attribution_key(&doc)?);
Ok(())
}
const fn fixture_long_wait_interval() -> std::time::Duration {
const LONG_WAIT_INTERVAL: std::time::Duration = std::time::Duration::from_secs(1);
LONG_WAIT_INTERVAL
}