use std::collections::BTreeMap;
use std::num::{NonZeroU32, NonZeroU64, NonZeroUsize};
use std::os::unix::fs::PermissionsExt;
use std::path::Path;
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, Version,
};
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,
pub written: BTreeMap<String, Version>,
pub patients_only: Version,
directory: tempfile::TempDir,
}
impl Fixture {
pub fn new(documents: &[Value]) -> Result<Self, Box<dyn std::error::Error>> {
Self::with_wards(documents, &[("wrong", json!({"name":"Wrong collection"}))])
}
pub fn with_wards(
documents: &[Value],
wards: &[(&str, 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;
version = storage
.mutate(&Mutation {
request_id: "collection-patients".to_owned(),
branch: "main".to_owned(),
expected: version,
change: Change::Collection(CollectionWrite {
collection: "patients".to_owned(),
schema: None,
}),
})?
.version;
let patients_only = version;
version = storage
.mutate(&Mutation {
request_id: "collection-wards".to_owned(),
branch: "main".to_owned(),
expected: version,
change: Change::Collection(CollectionWrite {
collection: "wards".to_owned(),
schema: None,
}),
})?
.version;
let mut written = BTreeMap::new();
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;
written.insert(format!("patients/p{index}"), version);
}
for (id, document) in wards {
version = storage
.mutate(&Mutation {
request_id: format!("ward-{id}"),
branch: "main".to_owned(),
expected: version,
change: Change::Put(DocumentWrite {
collection: "wards".to_owned(),
document_id: (*id).to_owned(),
document: serde_json::value::to_raw_value(document)?,
}),
})?
.version;
written.insert(format!("wards/{id}"), 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()?,
written,
patients_only,
directory,
})
}
pub fn scan(&mut self, records: NonZeroU64) {
self.settings.scan_records = records;
}
pub fn worker_started(&self) -> bool {
self.directory.path().join("inputs").exists()
}
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)
}
pub fn inputs(&self) -> Result<Vec<Value>, Box<dyn std::error::Error>> {
self.frames()?
.into_iter()
.map(|frame| {
frame
.get("value")
.cloned()
.ok_or_else(|| "worker Input lacks value".into())
})
.collect()
}
pub fn frames(&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 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,
}
}
fn worker(directory: &Path) -> Result<std::path::PathBuf, Box<dyn std::error::Error>> {
let worker = directory.join("worker");
let inputs = directory.join("inputs");
let log = serde_json::to_string(inputs.to_str().ok_or("non-UTF8 fixture 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:])
reading = 'candidate' if start['filter'] in ['.', '.items[]'] else 'alternate'
program = start['filter'] if reading == 'candidate' else start.get('alternate')
if program is None or program not in ['.', '.items[]']:
send(6, dict(code='query_compile', message='fixture requires an alternate echo or expansion'))
sys.exit(2)
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
if frame[0] != 2:
raise ValueError('expected Input frame')
packet = json.loads(frame[1:])
log.write(json.dumps(packet) + '\n')
log.flush()
value = packet['value']
values = value['items'] if program == '.items[]' else [value]
for value in values:
send(4, value)
send(5, dict(outputs=len(values)))
"
);
std::fs::write(&worker, script)?;
std::fs::set_permissions(&worker, std::fs::Permissions::from_mode(0o700))?;
Ok(worker)
}