use std::collections::BTreeMap;
use std::num::{NonZeroU32, NonZeroU64, NonZeroUsize};
use std::os::unix::fs::PermissionsExt;
use std::time::Duration;
use haematite::db::{DatabaseConfig, ReadOnlyDatabase};
use haematite::store::NodeCacheBudget;
use serde_json::Value;
use serde_json::value::RawValue;
use crate::serve::access::{self, TestSource};
use crate::serve::protocol::{
Change, CollectionWrite, DocumentWrite, Edge, Mutation, Node, QueryPage, QueryRequest,
Selection,
};
use crate::serve::query::admission::Slots;
use crate::serve::query::run::{self, Context};
use crate::serve::query::{Rows, Seal, Settings};
use crate::serve::storage::{Failure, Storage, document_key, reads, revision_key};
pub type TestResult = Result<(), Box<dyn std::error::Error>>;
pub struct Fixture {
storage: Storage,
reader: ReadOnlyDatabase,
settings: Settings,
slots: Slots,
seal: Seal,
directory: tempfile::TempDir,
}
impl Fixture {
pub fn new() -> Result<Self, Box<dyn std::error::Error>> {
let directory = tempfile::tempdir()?;
let worker = worker(directory.path())?;
let mut storage = Storage::create(
DatabaseConfig {
data_dir: directory.path().join("database"),
shard_count: 4,
distributed: None,
executor_threads: Some(1),
node_cache_budget: Some(NodeCacheBudget::Unlimited),
retention: haematite::Retention::KeepAll,
},
NodeCacheBudget::Unlimited,
Duration::from_secs(1),
)?;
for collection in ["patients", "wards", "nurses", "empty", "other"] {
apply(
&mut storage,
collection,
Change::Collection(CollectionWrite {
collection: collection.to_owned(),
schema: None,
}),
)?;
}
for (collection, id, text) in [
("patients", "p0", r#"{ "name": "Ada", "precise": 1.2300 }"#),
("patients", "p1", r#"{"name":"Bea"}"#),
("wards", "w0", r#"{ "ward" : "A" }"#),
("wards", "w1", r#"{"ward":"B"}"#),
("nurses", "n0", r#"{"name":"Nia"}"#),
("nurses", "n1", r#"{"name":"Noor"}"#),
("nurses", "n2", r#"{"name":"Nell"}"#),
] {
apply(
&mut storage,
&format!("{collection}-{id}"),
Change::Put(DocumentWrite {
collection: collection.to_owned(),
document_id: id.to_owned(),
document: RawValue::from_string(text.to_owned())?,
}),
)?;
}
for (index, (from, to, relation)) in [
("p1", "w0", "admitted_to"),
("p0", "w1", "admitted_to"),
("p0", "w0", "admitted_to"),
("p1", "w1", "visits"),
]
.iter()
.enumerate()
{
apply(
&mut storage,
&format!("edge-{index}"),
Change::EdgePut(Edge {
from: node("patients", from),
to: node("wards", to),
relation: (*relation).to_owned(),
properties: serde_json::json!({}),
}),
)?;
}
for (index, (from, to)) in [("n1", "w0"), ("n0", "w0"), ("n2", "w1")]
.iter()
.enumerate()
{
apply(
&mut storage,
&format!("staffing-{index}"),
Change::EdgePut(Edge {
from: node("nurses", from),
to: node("wards", to),
relation: "staffed_by".to_owned(),
properties: serde_json::json!({}),
}),
)?;
}
let reader = storage.reader()?;
Ok(Self {
storage,
reader,
settings: Settings {
time_ms: NonZeroU64::MIN.saturating_add(59_999),
memory_bytes: NonZeroU64::MIN.saturating_add((1 << 26) - 1),
outputs_per_record: NonZeroU64::MIN.saturating_add(31),
scan_records: NonZeroU64::MIN.saturating_add(99),
max_running: NonZeroUsize::MIN,
worker,
},
slots: Slots::new(NonZeroUsize::MIN),
seal: Seal::new()?,
directory,
})
}
pub fn scan(&mut self, records: u64) -> TestResult {
self.settings.scan_records = NonZeroU64::new(records).ok_or("zero scan")?;
Ok(())
}
pub fn run_raw(
&self,
request: &QueryRequest,
denied: &[(&str, &str)],
) -> Result<Box<RawValue>, Failure> {
let access = access::request(&TestSource::denying(denied))?;
let catalogue = self.storage.catalogue();
run::run(
&Context {
access: &access,
reader: &self.reader,
catalogue: &catalogue,
slots: &self.slots,
settings: Some(&self.settings),
max_page: NonZeroU32::MIN.saturating_add(99),
seal: &self.seal,
},
request,
None,
)
}
pub fn run(&self, request: &QueryRequest, denied: &[(&str, &str)]) -> Result<Value, Failure> {
serde_json::from_str(self.run_raw(request, denied)?.get()).map_err(Failure::storage)
}
pub fn bytes(&self, collection: &str, id: &str) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
let at = self.storage.capture(&selection())?;
Ok(
reads::get_bytes(&self.reader, &at, &document_key(collection, id))?
.ok_or("stored record missing")?,
)
}
pub fn value(&self, collection: &str, id: &str) -> Result<Value, Box<dyn std::error::Error>> {
Ok(serde_json::from_slice(&self.bytes(collection, id)?)?)
}
pub fn revision(
&self,
collection: &str,
id: &str,
) -> Result<Value, Box<dyn std::error::Error>> {
let at = self.storage.capture(&selection())?;
let bytes = reads::get_bytes(&self.reader, &at, &revision_key(collection, id))?
.ok_or("stored revision missing")?;
let revision: crate::serve::storage::DocumentRevision = serde_json::from_slice(&bytes)?;
Ok(serde_json::to_value(revision.written_at)?)
}
pub fn inputs(&self) -> Result<Vec<Value>, Box<dyn std::error::Error>> {
std::fs::read_to_string(self.directory.path().join("inputs"))?
.lines()
.map(|line| serde_json::from_str(line).map_err(Into::into))
.collect()
}
}
fn node(collection: &str, id: &str) -> Node {
Node {
collection: collection.to_owned(),
document_id: id.to_owned(),
}
}
fn apply(storage: &mut Storage, id: &str, change: Change) -> Result<(), Failure> {
let expected = storage.capture(&selection())?.version;
storage.mutate(&Mutation {
request_id: id.to_owned(),
branch: "main".to_owned(),
expected,
change,
})?;
Ok(())
}
fn selection() -> Selection {
Selection {
branch: "main".to_owned(),
version: None,
}
}
pub fn request(text: &str, rows: NonZeroU64) -> QueryRequest {
QueryRequest {
selection: selection(),
query: text.to_owned(),
parameters: BTreeMap::new(),
page: QueryPage {
rows,
bytes: NonZeroU64::MIN.saturating_add((1 << 20) - 1),
},
continuation: None,
rows: Rows::Values,
}
}
pub fn values(answer: &Value) -> Result<Vec<Value>, Box<dyn std::error::Error>> {
answer
.get("rows")
.and_then(Value::as_array)
.ok_or("rows missing")?
.iter()
.map(|row| {
row.get("value")
.cloned()
.ok_or_else(|| "value missing".into())
})
.collect()
}
fn worker(directory: &std::path::Path) -> Result<std::path::PathBuf, Box<dyn std::error::Error>> {
let worker = directory.join("worker");
let log = serde_json::to_string(
directory
.join("inputs")
.to_str()
.ok_or("non-UTF8 worker path")?,
)?;
let worker_version = env!("CARGO_PKG_VERSION");
let script = format!(
r"#!/usr/bin/python3
import json, struct, sys
log = open({log}, 'a')
def receive():
header = sys.stdin.buffer.read(4)
if not header:
return None
size = struct.unpack('>I', header)[0]
data = sys.stdin.buffer.read(size)
if len(data) != size:
raise ValueError('incomplete frame')
return data
def send(tag, value):
data = bytes([tag]) + json.dumps(value, separators=(',', ':')).encode()
sys.stdout.buffer.write(struct.pack('>I', len(data)) + data)
sys.stdout.buffer.flush()
start = json.loads(receive()[1:])
program = start.get('alternate')
if program not in ['.', '.c.name', '$score']:
raise ValueError('unserved match fixture program')
send(3, dict(protocol=start['protocol'], worker='{worker_version}', core_limit=0,
descriptors_closed=0, environment=0, engine='fixture',
jaq_core='0', jaq_std='0', jaq_json='0', reading='alternate'))
while True:
frame = receive()
if frame is None:
break
if frame[0] != 2:
raise ValueError('expected Input')
packet = json.loads(frame[1:])
log.write(json.dumps(packet) + '\n')
log.flush()
value = packet['value']
if program == '.c.name':
value = value['c']['name']
elif program == '$score':
value = packet['bindings']['score']
send(4, value)
send(5, dict(outputs=1))
"
);
std::fs::write(&worker, script)?;
std::fs::set_permissions(&worker, std::fs::Permissions::from_mode(0o700))?;
Ok(worker)
}