+2

ufos/src/error.rs

Reviewed

··· 34 34 FjallLsmError(#[from] fjall::LsmError), 35 35 #[error("Bytes encoding error")] 36 36 EncodingError(#[from] EncodingError), 37 37 + #[error("If you ever see this, there's a bug in the code. The error was stolen")] 38 38 + Stolen, 37 39 }

+94 -12

ufos/src/storage_fjall.rs

Reviewed

··· 426 426 ) -> StorageResult<Vec<UFOsRecord>> { 427 427 if collections.is_empty() { 428 428 return Ok(vec![]); 429 429 - } else if collections.len() > 1 { 430 430 - todo!() 431 429 } 432 432 - 433 433 - let collection = collections[0]; 434 434 - 435 435 - let mut out = Vec::new(); 436 436 - for rec in RecordIterator::new(&self.feeds, self.records.clone(), collection, limit)? { 437 437 - if let Some(r) = rec? { 438 438 - out.push(r) 439 439 - } else { 440 440 - break; 430 430 + let mut record_iterators = Vec::new(); 431 431 + for collection in collections { 432 432 + let iter = RecordIterator::new(&self.feeds, self.records.clone(), collection, limit)?; 433 433 + record_iterators.push(iter.peekable()); 434 434 + } 435 435 + let mut merged = Vec::new(); 436 436 + loop { 437 437 + let mut latest: Option<(Cursor, usize)> = None; // ugh 438 438 + for (i, iter) in record_iterators.iter_mut().enumerate() { 439 439 + let Some(it) = iter.peek_mut() else { 440 440 + continue; 441 441 + }; 442 442 + let it = match it { 443 443 + Ok(v) => v, 444 444 + Err(e) => Err(std::mem::replace(e, StorageError::Stolen))?, 445 445 + }; 446 446 + let Some(rec) = it else { 447 447 + break; 448 448 + }; 449 449 + if let Some((cursor, _)) = latest { 450 450 + if rec.cursor > cursor { 451 451 + latest = Some((rec.cursor, i)) 452 452 + } 453 453 + } else { 454 454 + latest = Some((rec.cursor, i)); 455 455 + } 441 456 } 457 457 + let Some((_, idx)) = latest else { 458 458 + break; 459 459 + }; 460 460 + // yeah yeah whateverrrrrrrrrrrrrrrr 461 461 + merged.push(record_iterators[idx].next().unwrap().unwrap().unwrap()); 442 462 } 443 443 - Ok(out) 463 463 + Ok(merged) 444 464 } 445 465 } 446 466 ··· 1648 1668 let records = 1649 1669 read.get_records_by_collections(&[&Nsid::new("d.e.f".to_string()).unwrap()], 2)?; 1650 1670 assert_eq!(records.len(), 0); 1671 1671 + 1672 1672 + Ok(()) 1673 1673 + } 1674 1674 + 1675 1675 + #[test] 1676 1676 + fn test_get_multi_collection() -> anyhow::Result<()> { 1677 1677 + let (read, mut write) = fjall_db(); 1678 1678 + 1679 1679 + let mut batch = TestBatch::default(); 1680 1680 + batch.create( 1681 1681 + "did:plc:inze6wrmsm7pjl7yta3oig77", 1682 1682 + "a.a.a", 1683 1683 + "aaa", 1684 1684 + r#""earliest""#, 1685 1685 + Some("rev-a"), 1686 1686 + None, 1687 1687 + 100, 1688 1688 + ); 1689 1689 + batch.create( 1690 1690 + "did:plc:inze6wrmsm7pjl7yta3oig77", 1691 1691 + "a.a.b", 1692 1692 + "aab", 1693 1693 + r#""in between""#, 1694 1694 + Some("rev-ab"), 1695 1695 + None, 1696 1696 + 101, 1697 1697 + ); 1698 1698 + batch.create( 1699 1699 + "did:plc:inze6wrmsm7pjl7yta3oig77", 1700 1700 + "a.a.a", 1701 1701 + "aaa-2", 1702 1702 + r#""last""#, 1703 1703 + Some("rev-a-2"), 1704 1704 + None, 1705 1705 + 102, 1706 1706 + ); 1707 1707 + write.insert_batch(batch.batch)?; 1708 1708 + 1709 1709 + let records = read.get_records_by_collections( 1710 1710 + &[ 1711 1711 + &Nsid::new("a.a.a".to_string()).unwrap(), 1712 1712 + &Nsid::new("a.a.b".to_string()).unwrap(), 1713 1713 + &Nsid::new("a.a.c".to_string()).unwrap(), 1714 1714 + ], 1715 1715 + 100, 1716 1716 + )?; 1717 1717 + assert_eq!(records.len(), 3); 1718 1718 + assert_eq!(records[0].record.get(), r#""last""#); 1719 1719 + assert_eq!( 1720 1720 + records[0].collection, 1721 1721 + Nsid::new("a.a.a".to_string()).unwrap() 1722 1722 + ); 1723 1723 + assert_eq!(records[1].record.get(), r#""in between""#); 1724 1724 + assert_eq!( 1725 1725 + records[1].collection, 1726 1726 + Nsid::new("a.a.b".to_string()).unwrap() 1727 1727 + ); 1728 1728 + assert_eq!(records[2].record.get(), r#""earliest""#); 1729 1729 + assert_eq!( 1730 1730 + records[2].collection, 1731 1731 + Nsid::new("a.a.a".to_string()).unwrap() 1732 1732 + ); 1651 1733 1652 1734 Ok(()) 1653 1735 }