use std::collections::BTreeMap;
use std::num::{NonZeroU32, NonZeroU64, NonZeroUsize};
use std::os::unix::fs::PermissionsExt;
use std::path::PathBuf;
use std::time::Duration;
use haematite::db::{DatabaseConfig, ReadOnlyDatabase};
use haematite::store::NodeCacheBudget;
use serde_json::value::RawValue;
use serde_json::{Value, json};
use crate::serve::access::{self, TestSource};
use crate::serve::protocol::{
Change, CollectionWrite, DocumentWrite, Mutation, QueryPage, QueryRequest, Selection,
VectorWrite,
};
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};
pub type TestResult = Result<(), Box<dyn std::error::Error>>;
pub struct Fixture {
storage: Storage,
reader: ReadOnlyDatabase,
settings: Settings,
slots: Slots,
pub seal: Seal,
directory: tempfile::TempDir,
}
impl Fixture {
pub fn new(
schema: Option<Value>,
documents: &[Value],
) -> 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),
)?;
let mut version = storage.capture(&selection())?.version;
let change = Change::Collection(CollectionWrite {
collection: "patients".to_owned(),
schema,
});
version = storage
.mutate(&Mutation {
request_id: "collection".to_owned(),
branch: "main".to_owned(),
expected: version,
change,
})?
.version;
for (index, document) in documents.iter().enumerate() {
let id = format!("p{index}");
version = storage
.mutate(&Mutation {
request_id: id.clone(),
branch: "main".to_owned(),
expected: version,
change: Change::Put(DocumentWrite {
collection: "patients".to_owned(),
document_id: id,
document: serde_json::value::to_raw_value(document)?,
}),
})?
.version;
}
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 vectors(&mut self, field: &str, vectors: &[(usize, [f32; 2])]) -> TestResult {
for (index, vector) in vectors {
let expected = self.storage.capture(&selection())?.version;
self.storage.mutate(&Mutation {
request_id: format!("vector-{field}-{index}"),
branch: "main".to_owned(),
expected,
change: Change::VectorPut(VectorWrite {
collection: "patients".to_owned(),
document_id: format!("p{index}"),
field: field.to_owned(),
vector: vector.to_vec(),
}),
})?;
}
Ok(())
}
pub const fn scan(&mut self, records: NonZeroU64) {
self.settings.scan_records = records;
}
pub const fn memory(&mut self, bytes: NonZeroU64) {
self.settings.memory_bytes = bytes;
}
pub fn no_worker(&mut self) {
self.settings.worker = self.directory.path().join("absent");
}
pub fn run(&self, request: &QueryRequest, denied: &[(&str, &str)]) -> Result<Value, Failure> {
let access = access::request(&TestSource::denying(denied))?;
let catalogue = self.storage.catalogue();
let context = 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,
};
serde_json::from_str(run::run(&context, request, None)?.get()).map_err(Failure::storage)
}
}
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::<String, Box<RawValue>>::new(),
page: QueryPage {
rows,
bytes: NonZeroU64::MIN.saturating_add((1 << 20) - 1),
},
continuation: None,
rows: Rows::Values,
}
}
pub fn patients() -> Vec<Value> {
vec![
json!({"name":"Ada", "ward":"4B", "admitted":2}),
json!({"name":"Bea", "ward":"3A", "admitted":4}),
json!({"name":"Cy", "ward":"4B", "admitted":3}),
json!({"name":"Dee", "ward":"4B", "admitted":1}),
]
}
pub fn values(answer: &Value) -> Result<Vec<Value>, Box<dyn std::error::Error>> {
Ok(answer["rows"]
.as_array()
.ok_or("missing rows")?
.iter()
.map(|row| row["value"].clone())
.collect())
}
fn worker(directory: &std::path::Path) -> Result<PathBuf, Box<dyn std::error::Error>> {
let worker = directory.join("worker");
std::fs::write(
&worker,
r#"#!/usr/bin/python3
import json, struct, sys
def receive():
header = sys.stdin.buffer.read(4)
if not header:
return None
size = struct.unpack('>I', header)[0]
return sys.stdin.buffer.read(size)
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:])
reading = 'alternate' if start.get('alternate') is not None else 'candidate'
program = start['alternate'] if reading == 'alternate' else start['filter']
if program not in ['.', '.name', '.names[]', '$score', 'select(.ward == "4B")']:
raise ValueError('unserved fixture filter: ' + 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=reading))
while True:
frame = receive()
if frame is None:
break
incoming = json.loads(frame[1:])
value = incoming['value']
if program == '$score':
values = [incoming['bindings']['score']]
elif program == 'select(.ward == "4B")':
values = [value] if value.get('ward') == '4B' else []
elif program == '.names[]':
values = value['names']
else:
values = [value['name'] if program == '.name' else value]
for value in values:
send(4, value)
send(5, dict(outputs=len(values)))
"#
.replace("@WORKER_VERSION@", env!("CARGO_PKG_VERSION")),
)?;
std::fs::set_permissions(&worker, std::fs::Permissions::from_mode(0o700))?;
Ok(worker)
}