1use std::{collections::BTreeSet, path::Path, time::Instant};
4
5use hyphae_query::{
6 ExecutionLimits, Query, QueryError, QueryResult, Record, execute, validate_query,
7};
8use hyphae_retrieval::{
9 RetrievalError, RetrievalLimits, RetrievalOutcome, RetrievalRequest, VectorRecord, retrieve,
10};
11use hyphae_storage::{
12 AppendOutcome, BackupError, BackupInfo, CompactionOutcome, MAX_SCAN_PAGE_ENTRIES, Mutation,
13 RestoreInfo, SnapshotInfo, StorageEngine, StorageError, StorageRecoveryReport, restore_backup,
14 verify_backup,
15};
16use thiserror::Error;
17use uuid::Uuid;
18
19use crate::{
20 DocumentError, ProofError, ResultProof, ResultProofArtifact, decode_document, encode_document,
21};
22
23#[derive(Debug, Error)]
25pub enum EngineError {
26 #[error(transparent)]
28 Storage(#[from] StorageError),
29
30 #[error(transparent)]
32 Backup(#[from] BackupError),
33
34 #[error(transparent)]
36 Document(#[from] DocumentError),
37
38 #[error(transparent)]
40 Query(#[from] QueryError),
41
42 #[error(transparent)]
44 Retrieval(#[from] RetrievalError),
45
46 #[error(transparent)]
48 Proof(#[from] ProofError),
49
50 #[error("atomic document batch contains a duplicate key")]
52 DuplicateDocumentKey,
53}
54
55#[derive(Debug)]
57pub struct OpenedEngine {
58 pub engine: HyphaeEngine,
60 pub recovery: StorageRecoveryReport,
62}
63
64#[derive(Debug)]
66pub struct HyphaeEngine {
67 storage: StorageEngine,
68}
69
70impl HyphaeEngine {
71 pub fn open(path: impl AsRef<Path>) -> Result<OpenedEngine, EngineError> {
78 let opened = StorageEngine::open(path)?;
79 Ok(OpenedEngine {
80 engine: Self {
81 storage: opened.storage,
82 },
83 recovery: opened.recovery,
84 })
85 }
86
87 pub fn data_path(&self) -> &Path {
89 self.storage.data_path()
90 }
91
92 pub fn put_record(
98 &mut self,
99 transaction_id: Uuid,
100 record: &Record,
101 ) -> Result<AppendOutcome, EngineError> {
102 self.put_records(transaction_id, std::slice::from_ref(record))
103 }
104
105 pub fn put_records(
115 &mut self,
116 transaction_id: Uuid,
117 records: &[Record],
118 ) -> Result<AppendOutcome, EngineError> {
119 let mut keys = BTreeSet::new();
120 let mut mutations = Vec::with_capacity(records.len());
121 for record in records {
122 if !keys.insert(record.key.as_slice()) {
123 return Err(EngineError::DuplicateDocumentKey);
124 }
125 mutations.push(Mutation::put(
126 record.key.clone(),
127 encode_document(&record.value)?,
128 ));
129 }
130 Ok(self.storage.write(transaction_id, &mutations)?)
131 }
132
133 pub fn delete_record(
139 &mut self,
140 transaction_id: Uuid,
141 key: &[u8],
142 ) -> Result<AppendOutcome, EngineError> {
143 self.delete_records(transaction_id, &[key])
144 }
145
146 pub fn delete_records(
156 &mut self,
157 transaction_id: Uuid,
158 keys: &[&[u8]],
159 ) -> Result<AppendOutcome, EngineError> {
160 let mut unique = BTreeSet::new();
161 let mut mutations = Vec::with_capacity(keys.len());
162 for key in keys {
163 if !unique.insert(*key) {
164 return Err(EngineError::DuplicateDocumentKey);
165 }
166 mutations.push(Mutation::delete(*key));
167 }
168 Ok(self.storage.write(transaction_id, &mutations)?)
169 }
170
171 pub fn get_record(&self, key: &[u8]) -> Result<Option<Record>, EngineError> {
177 self.storage
178 .get(key)?
179 .map(|encoded| {
180 Ok(Record {
181 key: key.to_vec(),
182 value: decode_document(&encoded)?,
183 })
184 })
185 .transpose()
186 }
187
188 pub fn get_record_with_proof(&self, key: &[u8]) -> Result<ResultProofArtifact, EngineError> {
195 let result = self.get_record(key)?;
196 let snapshot = self.snapshot()?;
197 let proof = ResultProof::for_get(&snapshot, key.to_vec(), result)?;
198 Ok(ResultProofArtifact { proof, snapshot })
199 }
200
201 pub fn query(
211 &self,
212 query: &Query,
213 limits: &ExecutionLimits,
214 ) -> Result<QueryResult, EngineError> {
215 validate_query(query, limits)?;
216 let started = Instant::now();
217 let mut records = Vec::new();
218 let mut after = None;
219 loop {
220 if started.elapsed() >= limits.timeout {
221 return Err(QueryError::TimedOut.into());
222 }
223 let loaded = u64::try_from(records.len()).unwrap_or(u64::MAX);
224 let remaining = limits.max_scanned_records.saturating_sub(loaded);
225 let remaining_entries = match usize::try_from(remaining) {
226 Ok(value) => value,
227 Err(_) => usize::MAX,
228 };
229 let page_limit = remaining_entries
230 .saturating_add(1)
231 .min(MAX_SCAN_PAGE_ENTRIES);
232 let page = self.storage.scan_page(after.as_deref(), page_limit)?;
233 for entry in page.entries {
234 if u64::try_from(records.len()).unwrap_or(u64::MAX) >= limits.max_scanned_records {
235 return Err(QueryError::ScannedBudgetExceeded {
236 maximum: limits.max_scanned_records,
237 }
238 .into());
239 }
240 records.push(Record {
241 key: entry.key,
242 value: decode_document(&entry.value)?,
243 });
244 }
245 let Some(next_after) = page.next_after else {
246 break;
247 };
248 after = Some(next_after);
249 }
250
251 let elapsed = started.elapsed();
252 let Some(timeout) = limits.timeout.checked_sub(elapsed) else {
253 return Err(QueryError::TimedOut.into());
254 };
255 if timeout.is_zero() {
256 return Err(QueryError::TimedOut.into());
257 }
258 let execution_limits = ExecutionLimits {
259 timeout,
260 ..limits.clone()
261 };
262 Ok(execute(&[records.as_slice()], query, &execution_limits)?)
263 }
264
265 pub fn query_with_proof(
273 &self,
274 query: &Query,
275 limits: &ExecutionLimits,
276 ) -> Result<ResultProofArtifact, EngineError> {
277 let result = self.query(query, limits)?;
278 let snapshot = self.snapshot()?;
279 let proof = ResultProof::for_query(&snapshot, query.clone(), result)?;
280 Ok(ResultProofArtifact { proof, snapshot })
281 }
282
283 pub fn retrieve_vectors(
290 shards: &[&[VectorRecord]],
291 request: &RetrievalRequest,
292 limits: &RetrievalLimits,
293 ) -> Result<RetrievalOutcome, EngineError> {
294 Ok(retrieve(shards, request, limits)?)
295 }
296
297 pub fn snapshot(&self) -> Result<SnapshotInfo, EngineError> {
303 Ok(self.storage.snapshot()?)
304 }
305
306 pub fn compact(&mut self) -> Result<CompactionOutcome, EngineError> {
312 Ok(self.storage.compact()?)
313 }
314
315 pub fn backup(&self, destination: impl AsRef<Path>) -> Result<BackupInfo, EngineError> {
321 Ok(self.storage.backup(destination)?)
322 }
323
324 pub fn verify_backup(path: impl AsRef<Path>) -> Result<BackupInfo, EngineError> {
331 Ok(verify_backup(path)?)
332 }
333
334 pub fn restore_backup(
341 backup: impl AsRef<Path>,
342 destination: impl AsRef<Path>,
343 ) -> Result<RestoreInfo, EngineError> {
344 Ok(restore_backup(backup, destination)?)
345 }
346}
347
348#[cfg(test)]
349mod tests {
350 use std::{collections::BTreeMap, fs, path::PathBuf};
351
352 use hyphae_query::{
353 AggregationPlan, CompareOperator, FieldPath, Filter, Metric, MetricValue, NamedMetric,
354 NullPlacement, SortDirection, SortField, Value,
355 };
356 use uuid::Uuid;
357
358 use super::{EngineError, ExecutionLimits, HyphaeEngine, Query, Record};
359
360 struct TestDirectory {
361 path: PathBuf,
362 }
363
364 impl TestDirectory {
365 fn new(name: &str) -> std::io::Result<Self> {
366 let path = std::env::temp_dir().join(format!(
367 "hyphae-engine-{name}-{}-{}",
368 std::process::id(),
369 Uuid::now_v7()
370 ));
371 fs::create_dir_all(&path)?;
372 Ok(Self { path })
373 }
374
375 fn path(&self) -> &std::path::Path {
376 &self.path
377 }
378 }
379
380 impl Drop for TestDirectory {
381 fn drop(&mut self) {
382 let _ignored = fs::remove_dir_all(&self.path);
383 }
384 }
385
386 fn value(score: i64, group: &str) -> Value {
387 Value::Object(BTreeMap::from([
388 ("group".to_owned(), Value::String(group.to_owned())),
389 ("score".to_owned(), Value::Integer(score)),
390 ]))
391 }
392
393 #[test]
394 fn durable_documents_query_identically_after_compaction_and_reopen()
395 -> Result<(), Box<dyn std::error::Error>> {
396 let temporary = TestDirectory::new("engine-query-reopen")?;
397 let root = temporary.path().join("data");
398 let mut opened = HyphaeEngine::open(&root)?;
399 opened.engine.put_records(
400 Uuid::now_v7(),
401 &[
402 Record::new(b"a", value(10, "x")),
403 Record::new(b"b", value(8, "x")),
404 Record::new(b"c", value(7, "y")),
405 Record::new(b"d", value(2, "y")),
406 ],
407 )?;
408 let request = Query {
409 filter: Filter::Compare {
410 path: FieldPath::field("score"),
411 operator: CompareOperator::GreaterOrEqual,
412 value: Value::Integer(7),
413 },
414 sort: vec![SortField {
415 path: FieldPath::field("score"),
416 direction: SortDirection::Descending,
417 nulls: NullPlacement::Last,
418 }],
419 cursor: None,
420 limit: 2,
421 aggregation: Some(AggregationPlan {
422 group_by: Vec::new(),
423 metrics: vec![NamedMetric {
424 name: "count".to_owned(),
425 metric: Metric::Count,
426 }],
427 }),
428 };
429 let before = opened.engine.query(&request, &ExecutionLimits::default())?;
430 assert_eq!(before.rows.len(), 2);
431 assert_eq!(
432 before
433 .aggregation
434 .as_ref()
435 .map(|aggregation| { aggregation.groups[0].metrics[0].value.clone() }),
436 Some(MetricValue::Count(3))
437 );
438 opened.engine.compact()?;
439 drop(opened);
440
441 let reopened = HyphaeEngine::open(&root)?;
442 let after = reopened
443 .engine
444 .query(&request, &ExecutionLimits::default())?;
445 assert_eq!(before, after);
446 assert_eq!(
447 reopened.engine.get_record(b"a")?.map(|record| record.value),
448 Some(value(10, "x"))
449 );
450 Ok(())
451 }
452
453 #[test]
454 fn facade_enforces_scan_budget_before_building_a_partial_page()
455 -> Result<(), Box<dyn std::error::Error>> {
456 let temporary = TestDirectory::new("engine-query-budget")?;
457 let mut opened = HyphaeEngine::open(temporary.path().join("data"))?;
458 opened.engine.put_records(
459 Uuid::now_v7(),
460 &[
461 Record::new(b"a", Value::Null),
462 Record::new(b"b", Value::Null),
463 ],
464 )?;
465 let limits = ExecutionLimits {
466 max_scanned_records: 1,
467 ..ExecutionLimits::default()
468 };
469 let result = opened.engine.query(
470 &Query {
471 filter: Filter::MatchAll,
472 sort: Vec::new(),
473 cursor: None,
474 limit: 1,
475 aggregation: None,
476 },
477 &limits,
478 );
479 assert!(matches!(
480 result,
481 Err(EngineError::Query(
482 hyphae_query::QueryError::ScannedBudgetExceeded { maximum: 1 }
483 ))
484 ));
485 Ok(())
486 }
487}