entity_sled/
lib.rs

1use entity::{
2    Database, DatabaseError, DatabaseResult, EdgeDeletionPolicy, Ent, Filter, Id, IdAllocator,
3    Predicate, Primitive, Query, Value, EPHEMERAL_ID,
4};
5use std::collections::HashSet;
6
7type EntIdSet = HashSet<Id>;
8
9/// Represents a sled database that performs synchronous insertion,
10/// retrieval, and removal. Sled maintains disk-backed data, so the `serde`
11/// feature has no purpose with this database.
12///
13/// Sled itself is thread-safe, maintaining an internal `Arc` for each tree;
14/// therefore, this database can be cloned to increment those counters.
15#[derive(Clone)]
16pub struct SledDatabase(sled::Db);
17
18fn id_to_ivec(id: Id) -> sled::IVec {
19    id.to_be_bytes().as_ref().into()
20}
21
22fn ivec_to_id(ivec: sled::IVec) -> Option<Id> {
23    use std::convert::TryInto;
24    let (bytes, _) = ivec.as_ref().split_at(std::mem::size_of::<Id>());
25    bytes.try_into().map(Id::from_be_bytes).ok()
26}
27
28const ENTS_OF_TYPE: &str = "ents_of_type";
29const ID_ALLOCATOR: &str = "id_allocator";
30
31impl SledDatabase {
32    /// Creates a new instance of the database wrapping a `sled::Db`
33    pub fn new(db: sled::Db) -> Self {
34        Self(db)
35    }
36
37    /// Returns ids of all ents stored in the database
38    pub fn ids(&self) -> EntIdSet {
39        self.0
40            .iter()
41            .keys()
42            .filter_map(Result::ok)
43            .filter_map(ivec_to_id)
44            .collect()
45    }
46
47    /// Returns true if database contains the provided id
48    pub fn has_id(&self, id: Id) -> bool {
49        self.0.contains_key(id_to_ivec(id)).ok().unwrap_or_default()
50    }
51
52    /// Returns ids of all ents for the given type
53    pub fn ids_for_type(&self, r#type: &str) -> EntIdSet {
54        fn inner(this: &SledDatabase, r#type: &str) -> DatabaseResult<EntIdSet> {
55            match this
56                .0
57                .open_tree(ENTS_OF_TYPE)
58                .map_err(|e| DatabaseError::Connection {
59                    source: Box::from(e),
60                })?
61                .get(r#type)
62                .map_err(|e| DatabaseError::Connection {
63                    source: Box::from(e),
64                })? {
65                Some(ivec) => match bincode::deserialize::<EntIdSet>(&ivec) {
66                    Ok(x) => Ok(x),
67                    Err(x) => Err(DatabaseError::Connection {
68                        source: Box::from(x),
69                    }),
70                },
71                None => Ok(HashSet::new()),
72            }
73        }
74
75        inner(self, r#type).ok().unwrap_or_default()
76    }
77
78    /// Returns sled tree for id allocator
79    fn id_allocator_tree(&self) -> DatabaseResult<sled::Tree> {
80        self.0
81            .open_tree(ID_ALLOCATOR)
82            .map_err(|e| DatabaseError::Connection {
83                source: Box::from(e),
84            })
85    }
86
87    /// Provides a mutable reference to the id allocator, returning an optional
88    /// id in the case that we want to return the next id from the allocator.
89    ///
90    /// Any changes made to the allocator are persisted back to disk.
91    fn with_id_allocator<F: Fn(&mut IdAllocator) -> Option<Id>>(
92        &self,
93        f: F,
94    ) -> DatabaseResult<Option<Id>> {
95        self.id_allocator_tree()?
96            .transaction(move |tx_db| {
97                let mut id_alloc = match tx_db.get([0])? {
98                    Some(ivec) => match bincode::deserialize::<IdAllocator>(&ivec) {
99                        Ok(x) => x,
100                        Err(x) => {
101                            sled::transaction::abort(x)?;
102                            return Ok(None);
103                        }
104                    },
105                    None => IdAllocator::new(),
106                };
107
108                let maybe_id = f(&mut id_alloc);
109
110                let id_alloc_bytes = match bincode::serialize(&id_alloc) {
111                    Ok(x) => x,
112                    Err(x) => {
113                        sled::transaction::abort(x)?;
114                        return Ok(maybe_id);
115                    }
116                };
117
118                tx_db.insert(&[0], id_alloc_bytes)?;
119                Ok(maybe_id)
120            })
121            .map_err(|e| DatabaseError::Connection {
122                source: Box::from(e),
123            })
124    }
125
126    /// Returns sled tree for ent types
127    fn ent_type_tree(&self) -> DatabaseResult<sled::Tree> {
128        self.0
129            .open_tree(ENTS_OF_TYPE)
130            .map_err(|e| DatabaseError::Connection {
131                source: Box::from(e),
132            })
133    }
134
135    /// Provides a mutable reference to the id set associated with an ent type.
136    /// Any changes made to the set are persisted back to disk.
137    fn with_ent_type_set<F: Fn(&mut EntIdSet)>(&self, r#type: &str, f: F) -> DatabaseResult<()> {
138        self.ent_type_tree()?
139            .transaction(move |tx_db| {
140                let mut set = match tx_db.get(r#type)? {
141                    Some(ivec) => match bincode::deserialize::<EntIdSet>(&ivec) {
142                        Ok(x) => x,
143                        Err(x) => {
144                            sled::transaction::abort(x)?;
145                            return Ok(());
146                        }
147                    },
148                    None => HashSet::new(),
149                };
150
151                f(&mut set);
152
153                let set_bytes = match bincode::serialize(&set) {
154                    Ok(x) => x,
155                    Err(x) => {
156                        sled::transaction::abort(x)?;
157                        return Ok(());
158                    }
159                };
160
161                tx_db.insert(r#type, set_bytes)?;
162                Ok(())
163            })
164            .map_err(|e| DatabaseError::Connection {
165                source: Box::from(e),
166            })
167    }
168}
169
170impl Database for SledDatabase {
171    fn get_all(&self, ids: Vec<Id>) -> DatabaseResult<Vec<Box<dyn Ent>>> {
172        ids.into_iter()
173            .filter_map(|id| self.get(id).transpose())
174            .collect()
175    }
176
177    fn find_all(&self, query: Query) -> DatabaseResult<Vec<Box<dyn Ent>>> {
178        let mut pipeline: Option<EntIdSet> = None;
179
180        for filter in query {
181            let mut_pipeline = pipeline.get_or_insert_with(|| prefill_ids(self, &filter));
182
183            // If our filter is the special IntoEdge case, we don't want to
184            // actually filter out ids but rather transform them into the ids
185            // of their edge
186            match filter {
187                Filter::IntoEdge(name) => {
188                    pipeline = Some(
189                        mut_pipeline
190                            .iter()
191                            .flat_map(|id| {
192                                self.get(*id)
193                                    .map(|maybe_ent| {
194                                        maybe_ent
195                                            .and_then(|ent| {
196                                                ent.edge(&name).map(|edge| edge.to_ids())
197                                            })
198                                            .unwrap_or_default()
199                                    })
200                                    .unwrap_or_default()
201                            })
202                            .collect(),
203                    )
204                }
205                // Otherwise, the filter is a traditional case where we will
206                // strip out ids by the filter
207                f => {
208                    mut_pipeline.retain(|id| filter_id(self, id, &f));
209                }
210            }
211        }
212
213        pipeline
214            .unwrap_or_default()
215            .into_iter()
216            .filter_map(|id| self.get(id).transpose())
217            .collect()
218    }
219
220    fn get(&self, id: Id) -> DatabaseResult<Option<Box<dyn Ent>>> {
221        let maybe_ivec = self
222            .0
223            .get(id_to_ivec(id))
224            .map_err(|e| DatabaseError::Connection {
225                source: Box::from(e),
226            })?;
227
228        let result: Result<Option<Box<dyn Ent>>, DatabaseError> = maybe_ivec
229            .map(|ivec| bincode::deserialize(ivec.as_ref()))
230            .transpose()
231            .map_err(|e| DatabaseError::CorruptedEnt {
232                id,
233                source: Box::from(e),
234            });
235
236        // If we found an ent without a database connection, attempt to fill
237        // it in with the global database if it exists
238        match result {
239            Ok(Some(mut ent)) => {
240                if !ent.is_connected() {
241                    ent.connect(entity::global::db());
242                }
243                Ok(Some(ent))
244            }
245            x => x,
246        }
247    }
248
249    fn remove(&self, id: Id) -> DatabaseResult<bool> {
250        if let Some(ent) = self
251            .0
252            .remove(id_to_ivec(id))
253            .map_err(|e| DatabaseError::Connection {
254                source: Box::from(e),
255            })?
256            .map(|ivec| bincode::deserialize::<Box<dyn Ent>>(ivec.as_ref()))
257            .transpose()
258            .map_err(|e| DatabaseError::CorruptedEnt {
259                id,
260                source: Box::from(e),
261            })?
262        {
263            for edge in ent.edges() {
264                match edge.deletion_policy() {
265                    // If shallow deletion, we only want to remove the connections
266                    // back to this ent from the corresponding ents
267                    EdgeDeletionPolicy::ShallowDelete => {
268                        for edge_id in edge.to_ids() {
269                            self.0
270                                .transaction(|tx_db| {
271                                    let maybe_ivec = tx_db.get(id_to_ivec(id))?;
272                                    let result = maybe_ivec
273                                        .map(|ivec| {
274                                            bincode::deserialize::<Box<dyn Ent>>(ivec.as_ref())
275                                        })
276                                        .transpose()
277                                        .map_err(|e| DatabaseError::CorruptedEnt {
278                                            id,
279                                            source: Box::from(e),
280                                        });
281                                    match result {
282                                        Ok(Some(mut ent)) => {
283                                            for mut edge in ent.edges() {
284                                                let _ = edge.value_mut().remove_ids(Some(edge_id));
285                                                let name = edge.name().to_string();
286                                                let _ = ent.update_edge(&name, edge.into_value());
287                                            }
288                                            match bincode::serialize(&ent) {
289                                                Ok(bytes) => tx_db.insert(id_to_ivec(id), bytes)?,
290                                                Err(x) => sled::transaction::abort(
291                                                    DatabaseError::CorruptedEnt {
292                                                        id: ent.id(),
293                                                        source: x,
294                                                    },
295                                                )?,
296                                            };
297                                        }
298                                        Ok(None) => {}
299                                        Err(x) => {
300                                            sled::transaction::abort(x)?;
301                                        }
302                                    };
303                                    Ok(())
304                                })
305                                .map_err(|e| DatabaseError::Connection {
306                                    source: Box::from(e),
307                                })?;
308                        }
309                    }
310                    // If deep deletion, we want to remove the ents connected
311                    // by the edge
312                    EdgeDeletionPolicy::DeepDelete => {
313                        for id in edge.to_ids() {
314                            let _ = self.remove(id);
315                        }
316                    }
317                    // If deletion policy is nothing, then do nothing
318                    EdgeDeletionPolicy::Nothing => {}
319                }
320            }
321
322            // Remove the id from our type mapping if it is there
323            self.with_ent_type_set(ent.r#type(), |set| {
324                set.remove(&id);
325            })?;
326
327            // Add the id to the freed ids available in the allocator
328            self.with_id_allocator(|alloc| {
329                alloc.extend(vec![id]);
330                None
331            })?;
332
333            Ok(true)
334        } else {
335            Ok(false)
336        }
337    }
338
339    fn insert(&self, mut ent: Box<dyn Ent>) -> DatabaseResult<Id> {
340        // Get the id of the ent, swapping out the ephemeral id
341        let id = ent.id();
342        let id = self
343            .with_id_allocator(move |alloc| {
344                if id == EPHEMERAL_ID {
345                    alloc.next()
346                } else {
347                    alloc.mark_external_id(id);
348                    Some(id)
349                }
350            })?
351            .ok_or(DatabaseError::EntCapacityReached)?;
352
353        // Update the ent's id to match what is actually to be used
354        ent.set_id(id);
355
356        // Clear any cache before saving the ent
357        ent.clear_cache();
358
359        // Update the ent's last_updated to be the current time
360        ent.mark_updated().map_err(|e| DatabaseError::Other {
361            source: Box::from(e),
362        })?;
363
364        // Add our ent's id to the set of ids associated with the ent's type
365        self.with_ent_type_set(ent.r#type(), |set| {
366            set.insert(id);
367        })?;
368
369        // Add our ent to the primary database
370        let ent_bytes = bincode::serialize(&ent).map_err(|e| DatabaseError::CorruptedEnt {
371            id,
372            source: Box::from(e),
373        })?;
374        self.0
375            .insert(id_to_ivec(id), ent_bytes)
376            .map_err(|e| DatabaseError::Connection {
377                source: Box::from(e),
378            })?;
379
380        Ok(id)
381    }
382}
383
384/// Called once when first beginning to filter to determine which ent ids
385/// to start with based on the leading filter
386///
387/// 1. If lead filter by id equality, will only include those ids that match
388///    the predicate
389/// 2. If lead filter by type equality, will only include those ids that equal
390///    the type (or many types if wrapped in Or)
391/// 3. Any other variation of id/type filter or other kind of filter will
392///    result in the more expensive pulling of all ids
393fn prefill_ids(db: &SledDatabase, filter: &Filter) -> EntIdSet {
394    fn from_id_predicate(db: &SledDatabase, p: &Predicate, mut ids: EntIdSet) -> Option<EntIdSet> {
395        match p {
396            Predicate::Equals(Value::Primitive(Primitive::Number(id))) => Some({
397                ids.insert(id.to_usize());
398                ids
399            }),
400            Predicate::Or(list) => list.iter().fold(Some(ids), |ids, p| match ids {
401                Some(ids) => from_id_predicate(db, p, ids),
402                None => None,
403            }),
404            _ => None,
405        }
406    }
407
408    fn from_type_predicate(
409        db: &SledDatabase,
410        p: &Predicate,
411        mut ids: EntIdSet,
412    ) -> Option<EntIdSet> {
413        match p {
414            Predicate::Equals(Value::Text(t)) => Some({
415                ids.extend(db.ids_for_type(t));
416                ids
417            }),
418            Predicate::Or(list) => list.iter().fold(Some(ids), |ids, p| match ids {
419                Some(ids) => from_type_predicate(db, p, ids),
420                None => None,
421            }),
422            _ => None,
423        }
424    }
425
426    match filter {
427        // If leading with id, support Equals and Or(Equals(...), ...) for
428        // specific ids; otherwise, too hard to figure out so we pull in all ids
429        Filter::Id(p) => {
430            from_id_predicate(db, p.as_untyped(), EntIdSet::new()).unwrap_or_else(|| db.ids())
431        }
432
433        // If leading with type, support Equals and Or(Equals(...), ...) for
434        // specific ids; otherwise, too hard to figure out so we pull in all ids
435        Filter::Type(p) => {
436            from_type_predicate(db, p.as_untyped(), EntIdSet::new()).unwrap_or_else(|| db.ids())
437        }
438
439        // Otherwise, currently no cached/indexed way to look up (yet)
440        // TODO: Support database field indexing so equality of a field can
441        //       be used for faster id lookup; do the same for timestamp fields
442        _ => db.ids(),
443    }
444}
445
446fn filter_id(db: &SledDatabase, id: &Id, filter: &Filter) -> bool {
447    match filter {
448        Filter::Id(p) => p.check(*id),
449        Filter::Type(p) => with_ent(db, id, |ent| p.check(ent.r#type().to_string())),
450        Filter::Created(p) => with_ent(db, id, |ent| p.check(ent.created())),
451        Filter::LastUpdated(p) => with_ent(db, id, |ent| p.check(ent.last_updated())),
452        Filter::Field(name, p) => with_ent(db, id, |ent| match ent.field(name) {
453            Some(value) => p.check(&value),
454            None => false,
455        }),
456        Filter::Edge(name, f) => with_ent(db, id, |ent| match ent.edge(name) {
457            Some(edge) => edge.to_ids().iter().any(|id| filter_id(db, id, f)),
458            None => false,
459        }),
460
461        // NOTE: Logically, this should be impossible to reach since we only
462        //       call this when we know that the filter is not a transformation
463        Filter::IntoEdge(_) => unreachable!("Bug: Transformation in filter"),
464    }
465}
466
467fn with_ent<F: Fn(Box<dyn Ent>) -> bool>(db: &SledDatabase, id: &Id, f: F) -> bool {
468    db.get(*id)
469        .map(|maybe_ent| maybe_ent.map(f).unwrap_or_default())
470        .unwrap_or_default()
471}
472
473#[cfg(test)]
474mod tests {
475    use super::*;
476    use entity::{Predicate as P, TypedPredicate as TP, *};
477    use std::collections::HashMap;
478
479    fn new_db() -> SledDatabase {
480        let config = sled::Config::new().temporary(true);
481        let db = config.open().expect("Failed to create database");
482        SledDatabase::new(db)
483    }
484
485    /// Creates a new database with some test entries used throughout
486    ///
487    /// IDs: 1-3 ~ are type1 with no fields or edges
488    /// IDs: 4-6 ~ are type2 with value fields and no edges
489    /// IDs: 7-9 ~ are type3 with collection fields and no edges
490    /// IDs: 10-12 ~ are type4 with edges to 1-9 and no fields
491    fn new_test_database() -> SledDatabase {
492        let db = new_db();
493
494        // 1-3 have no fields or edges
495        let _ = db
496            .insert(Box::from(UntypedEnt::from_collections(1, vec![], vec![])))
497            .unwrap();
498        let _ = db
499            .insert(Box::from(UntypedEnt::from_collections(2, vec![], vec![])))
500            .unwrap();
501        let _ = db
502            .insert(Box::from(UntypedEnt::from_collections(3, vec![], vec![])))
503            .unwrap();
504
505        // 4-6 have value fields only
506        let _ = db
507            .insert(Box::from(UntypedEnt::from_collections(
508                4,
509                vec![Field::new("a", 1), Field::new("b", 2)],
510                vec![],
511            )))
512            .unwrap();
513        let _ = db
514            .insert(Box::from(UntypedEnt::from_collections(
515                5,
516                vec![Field::new("a", 3), Field::new("b", 4)],
517                vec![],
518            )))
519            .unwrap();
520        let _ = db
521            .insert(Box::from(UntypedEnt::from_collections(
522                6,
523                vec![Field::new("a", 5), Field::new("b", 6)],
524                vec![],
525            )))
526            .unwrap();
527
528        // 7-9 have collection fields only
529        let _ = db
530            .insert(Box::from(UntypedEnt::from_collections(
531                7,
532                vec![Field::new(
533                    "f",
534                    Value::from(
535                        vec![(String::from("a"), 3), (String::from("b"), 5)]
536                            .into_iter()
537                            .collect::<HashMap<String, u8>>(),
538                    ),
539                )],
540                vec![],
541            )))
542            .unwrap();
543        let _ = db
544            .insert(Box::from(UntypedEnt::from_collections(
545                8,
546                vec![Field::new("f", vec![1, 2])],
547                vec![],
548            )))
549            .unwrap();
550        let _ = db
551            .insert(Box::from(UntypedEnt::from_collections(
552                9,
553                vec![Field::new(
554                    "f",
555                    Value::from(
556                        vec![
557                            (String::from("a"), Value::from(vec![1, 2])),
558                            (String::from("b"), Value::from(vec![3, 4])),
559                        ]
560                        .into_iter()
561                        .collect::<HashMap<String, Value>>(),
562                    ),
563                )],
564                vec![],
565            )))
566            .unwrap();
567
568        // 10-12 have edges only
569        let _ = db
570            .insert(Box::from(UntypedEnt::from_collections(
571                10,
572                vec![],
573                vec![
574                    Edge::new("a", 1),
575                    Edge::new("b", vec![3, 4, 5]),
576                    Edge::new("c", None),
577                ],
578            )))
579            .unwrap();
580        let _ = db
581            .insert(Box::from(UntypedEnt::from_collections(
582                11,
583                vec![],
584                vec![Edge::new("a", 2), Edge::new("b", vec![1, 2, 3, 4, 5, 6])],
585            )))
586            .unwrap();
587        let _ = db
588            .insert(Box::from(UntypedEnt::from_collections(
589                12,
590                vec![],
591                vec![
592                    Edge::new("a", 3),
593                    Edge::new("b", vec![]),
594                    Edge::new("c", Some(8)),
595                ],
596            )))
597            .unwrap();
598
599        db
600    }
601
602    fn query_and_assert<Q: Into<Query>>(db: &SledDatabase, query: Q, expected: &[Id]) {
603        let query = query.into();
604        let results = db
605            .find_all(query.clone())
606            .expect("Failed to retrieve ents")
607            .iter()
608            .map(|ent| ent.id())
609            .collect::<HashSet<Id>>();
610        assert_eq!(
611            results,
612            expected.into_iter().copied().collect(),
613            "{:?}\nExpected: {:?}, Actual: {:?}",
614            query,
615            expected,
616            results
617        );
618    }
619
620    #[test]
621    fn insert_should_replace_ephemeral_id_with_allocator_id() {
622        let db = new_db();
623
624        let ent = UntypedEnt::empty_with_id(EPHEMERAL_ID);
625        let id = db.insert(Box::from(ent)).expect("Failed to insert ent");
626        assert_ne!(id, EPHEMERAL_ID);
627
628        let ent = db.get(id).expect("Failed to get ent").expect("Ent missing");
629        assert_eq!(ent.id(), id);
630    }
631
632    #[test]
633    fn insert_should_update_the_last_updated_time_with_the_current_time() {
634        let db = new_db();
635
636        let ent = UntypedEnt::empty_with_id(EPHEMERAL_ID);
637        let last_updated = ent.last_updated();
638        std::thread::sleep(std::time::Duration::from_millis(10));
639
640        let id = db.insert(Box::from(ent)).expect("Failed to insert ent");
641        let ent = db.get(id).expect("Failed to get ent").expect("Ent missing");
642        assert!(ent.last_updated() > last_updated);
643    }
644
645    #[test]
646    fn insert_should_add_a_new_ent_using_its_id() {
647        let db = new_db();
648
649        let ent = UntypedEnt::empty_with_id(999);
650        let id = db.insert(Box::from(ent)).expect("Failed to insert ent");
651        assert_eq!(id, 999);
652
653        let ent = db
654            .get(999)
655            .expect("Failed to get ent")
656            .expect("Ent missing");
657        assert_eq!(ent.id(), 999);
658        assert_eq!(db.with_id_allocator(Iterator::next).unwrap().unwrap(), 1000);
659    }
660
661    #[test]
662    fn insert_should_overwrite_an_existing_ent_with_the_same_id() {
663        let db = new_db();
664
665        let ent = UntypedEnt::from_collections(999, vec![Field::new("field1", 3)], vec![]);
666        let _ = db.insert(Box::from(ent)).expect("Failed to insert ent");
667
668        let ent = db
669            .get(999)
670            .expect("Failed to get ent")
671            .expect("Ent missing");
672        assert_eq!(ent.field("field1").expect("Field missing"), Value::from(3));
673    }
674
675    #[test]
676    fn insert_should_reset_all_computed_field_caches_to_none() {
677        let db = new_db();
678
679        // Verify that a computed field is reset to None
680        let ent = UntypedEnt::from_collections(
681            999,
682            vec![Field::new_with_attributes(
683                "field1",
684                Some(3),
685                vec![FieldAttribute::Computed],
686            )],
687            vec![],
688        );
689        let _ = db.insert(Box::from(ent)).expect("Failed to insert ent");
690
691        let ent = db
692            .get(999)
693            .expect("Failed to get ent")
694            .expect("Ent missing");
695        assert_eq!(
696            ent.field("field1").expect("Field missing"),
697            Value::Optional(None)
698        );
699    }
700
701    #[test]
702    fn get_should_return_an_ent_by_id() {
703        let db = new_db();
704
705        let result = db.get(999).expect("Failed to get ent");
706        assert!(result.is_none(), "Unexpectedly acquired ent");
707
708        let _ = db
709            .insert(Box::from(UntypedEnt::empty_with_id(999)))
710            .unwrap();
711
712        let result = db.get(999).expect("Failed to get ent");
713        assert!(result.is_some(), "Unexpectedly missing ent");
714        assert!(
715            !result.unwrap().is_connected(),
716            "Ent unexpectedly connected to database"
717        );
718
719        // Verify that if a global database is available, it will populate
720        let db = DatabaseRc::new(Box::new(db));
721        entity::global::with_db_from_rc(DatabaseRc::clone(&db), || {
722            let result = db.get(999).expect("Failed to get ent");
723            assert!(result.is_some(), "Unexpectedly missing ent");
724            assert!(
725                result.unwrap().is_connected(),
726                "Ent unexpectedly not connected to database"
727            );
728        });
729    }
730
731    #[test]
732    fn remove_should_remove_an_ent_by_id() {
733        let db = new_db();
734
735        let _ = db.remove(999).expect("Failed to remove ent");
736
737        let _ = db
738            .insert(Box::from(UntypedEnt::empty_with_id(999)))
739            .unwrap();
740        assert!(db.get(999).unwrap().is_some(), "Failed to set up ent");
741
742        let _ = db.remove(999).expect("Failed to remove ent");
743        assert!(db.get(999).unwrap().is_none(), "Did not remove ent");
744
745        // Id allocator should indicate that id has been freed
746        assert_eq!(
747            db.with_id_allocator(|alloc| alloc.freed().first().copied())
748                .unwrap(),
749            Some(999),
750        );
751    }
752
753    #[test]
754    fn get_all_should_return_all_ents_with_associated_ids() {
755        let db = new_db();
756
757        let _ = db.insert(Box::from(UntypedEnt::empty_with_id(1))).unwrap();
758        let _ = db.insert(Box::from(UntypedEnt::empty_with_id(2))).unwrap();
759        let _ = db.insert(Box::from(UntypedEnt::empty_with_id(3))).unwrap();
760
761        let results = db
762            .get_all(vec![1, 2, 3])
763            .expect("Failed to retrieve ents")
764            .iter()
765            .map(|ent| ent.id())
766            .collect::<HashSet<Id>>();
767        assert_eq!(results, [1, 2, 3].iter().copied().collect());
768
769        let results = db
770            .get_all(vec![1, 3])
771            .expect("Failed to retrieve ents")
772            .iter()
773            .map(|ent| ent.id())
774            .collect::<HashSet<Id>>();
775        assert_eq!(results, [1, 3].iter().copied().collect());
776
777        let results = db
778            .get_all(vec![2, 3, 4, 5, 6, 7, 8])
779            .expect("Failed to retrieve ents")
780            .iter()
781            .map(|ent| ent.id())
782            .collect::<HashSet<Id>>();
783        assert_eq!(results, [2, 3].iter().copied().collect());
784    }
785
786    #[test]
787    fn find_all_should_return_no_ents_by_default() {
788        let db = new_test_database();
789
790        let q = Query::default();
791        query_and_assert(&db, q, &[]);
792    }
793
794    #[test]
795    fn find_all_should_support_filtering_by_id() {
796        let db = new_test_database();
797
798        // If ent with id exists, we expect it to be available
799        let q = Query::default().where_id(TP::equals(1));
800        query_and_assert(&db, q, &[1]);
801
802        // If ent with either id exists, we expect it to be available
803        let q = Query::default().where_id(TP::equals(1) | TP::equals(2));
804        query_and_assert(&db, q, &[1, 2]);
805
806        // If ent with id does not exist, we expect empty
807        let q = Query::default().where_id(TP::equals(999));
808        query_and_assert(&db, q, &[]);
809
810        // If already in a pipeline, should only filter the existing ids
811        let q = Query::default()
812            .where_id(TP::equals(1) | TP::equals(2))
813            .where_id(TP::equals(1) | TP::equals(3));
814        query_and_assert(&db, q, &[1]);
815    }
816
817    #[test]
818    fn find_all_should_support_filtering_by_type() {
819        let db = new_test_database();
820        let _ = db.insert(Box::from(TestEnt::new(20))).unwrap();
821        let _ = db.insert(Box::from(TestEnt::new(21))).unwrap();
822        let _ = db.insert(Box::from(TestEnt::new(22))).unwrap();
823
824        // If ent with type exists, we expect it to be available
825        let ts = <UntypedEnt as EntType>::type_str();
826        let q = Query::default().where_type(TP::equals(ts.to_string()));
827        query_and_assert(&db, q, &[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]);
828
829        // If ent with either type exists, we expect it to be available
830        let q = Query::default().where_type(TP::or(vec![
831            TP::equals(<UntypedEnt as EntType>::type_str().to_string()),
832            TP::equals(<TestEnt as EntType>::type_str().to_string()),
833        ]));
834        query_and_assert(&db, q, &[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 20, 21, 22]);
835
836        // If ent with type does not exist, we expect empty
837        let q = Query::default().where_type(TP::equals(String::from("unknown")));
838        query_and_assert(&db, q, &[]);
839
840        // If already in a pipeline, should only filter the existing ids
841        let q = Query::default()
842            .where_id(TP::equals(1) | TP::equals(2) | TP::equals(4))
843            .where_type(TP::equals(ts.to_string()));
844        query_and_assert(&db, q, &[1, 2, 4]);
845    }
846
847    #[test]
848    fn find_all_should_support_filtering_by_created_timestamp() {
849        let db = new_test_database();
850
851        // Re-create all ents with enough time split between them for us to
852        // properly test creation time
853        for i in 1..=12 {
854            let ent = UntypedEnt::empty_with_id(i);
855            db.insert(Box::from(ent))
856                .expect(&format!("Failed to replace ent {}", i));
857            std::thread::sleep(std::time::Duration::from_millis(1));
858        }
859
860        // Get all ents created after our third ent
861        let time = db.get(3).unwrap().expect("Missing ent 3").created();
862        let q = Query::default().where_created(TP::greater_than(time));
863        query_and_assert(&db, q, &[4, 5, 6, 7, 8, 9, 10, 11, 12]);
864
865        // If already in a pipeline, should only filter the existing ids
866        let time = db.get(3).unwrap().expect("Missing ent 3").created();
867        let q = Query::default()
868            .where_id(TP::less_than(8))
869            .where_created(TP::greater_than(time));
870        query_and_assert(&db, q, &[4, 5, 6, 7]);
871    }
872
873    #[test]
874    fn find_all_should_support_filtering_by_last_updated_timestamp() {
875        let db = new_test_database();
876
877        // Update all ents with enough time split between them for us to
878        // properly test last updated time
879        for i in (1..=12).rev() {
880            let mut ent = db
881                .get_typed::<UntypedEnt>(i)
882                .unwrap()
883                .expect(&format!("Missing ent {}", i));
884            ent.mark_updated().unwrap();
885            db.insert(Box::from(ent))
886                .expect(&format!("Failed to update ent {}", i));
887            std::thread::sleep(std::time::Duration::from_millis(1));
888        }
889
890        // Get all ents updated after our third ent
891        let time = db.get(3).unwrap().expect("Missing ent 3").last_updated();
892        let q = Query::default().where_last_updated(TP::greater_than(time));
893        query_and_assert(&db, q, &[1, 2]);
894
895        // If already in a pipeline, should only filter the existing ids
896        let time = db.get(3).unwrap().expect("Missing ent 3").created();
897        let q = Query::default()
898            .where_id(TP::equals(2))
899            .where_last_updated(TP::greater_than(time));
900        query_and_assert(&db, q, &[2]);
901    }
902
903    #[test]
904    fn find_all_should_support_filtering_by_field() {
905        let db = new_test_database();
906
907        // If ent's field passes condition, it will be included in return
908        let q = Query::default().where_field("a", P::equals(3));
909        query_and_assert(&db, q, &[5]);
910
911        // If already have ents in pipeline, they will be filtered by "field"
912        let q = Query::default()
913            .where_id(TP::equals(4) | TP::equals(6))
914            .where_field("a", P::greater_than(1));
915        query_and_assert(&db, q, &[6]);
916    }
917
918    #[test]
919    fn find_all_should_support_filtering_by_edge() {
920        let db = new_test_database();
921
922        // If ent's edge passes condition, it will be included in return
923        let q = Query::default().where_edge("a", Filter::Id(TP::equals(3)));
924        query_and_assert(&db, q, &[12]);
925
926        // If already have ents in pipeline, they will be filtered by "edge"
927        let q = Query::default()
928            .where_id(TP::equals(10) | TP::equals(12))
929            .where_edge("a", Filter::Id(TP::always()));
930        query_and_assert(&db, q, &[10, 12]);
931    }
932
933    #[test]
934    fn find_all_should_support_transforming_into_edge() {
935        let db = new_test_database();
936
937        // Will take the ids of each ent with the given edge and use
938        // them going forward; in this example, ents #10 and #11 have
939        // overlapping ids for edge b
940        let q = Query::default().where_into_edge("b");
941        query_and_assert(&db, q, &[1, 2, 3, 4, 5, 6]);
942
943        // If already have ents in pipeline, their edge's ids will
944        // be used specifically; in this example, ent #12 has no ents
945        // for edge b
946        let q = Query::default()
947            .where_id(TP::equals(10) | TP::equals(12))
948            .where_into_edge("b");
949        query_and_assert(&db, q, &[3, 4, 5]);
950    }
951
952    #[derive(Clone, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
953    struct TestEnt(Id);
954
955    impl TestEnt {
956        pub fn new(id: Id) -> Self {
957            Self(id)
958        }
959    }
960
961    impl EntType for TestEnt {
962        fn type_data() -> EntTypeData {
963            EntTypeData::Concrete {
964                ty: concat!(module_path!(), "::TestEnt"),
965            }
966        }
967    }
968
969    #[typetag::serde]
970    impl Ent for TestEnt {
971        fn id(&self) -> Id {
972            self.0
973        }
974
975        fn set_id(&mut self, id: Id) {
976            self.0 = id;
977        }
978
979        fn r#type(&self) -> &str {
980            Self::type_str()
981        }
982
983        fn created(&self) -> u64 {
984            0
985        }
986
987        fn last_updated(&self) -> u64 {
988            0
989        }
990
991        fn mark_updated(&mut self) -> Result<(), EntMutationError> {
992            Ok(())
993        }
994
995        fn field_definitions(&self) -> Vec<FieldDefinition> {
996            Vec::new()
997        }
998
999        fn field_names(&self) -> Vec<String> {
1000            Vec::new()
1001        }
1002
1003        fn field(&self, _name: &str) -> Option<Value> {
1004            None
1005        }
1006
1007        fn update_field(&mut self, name: &str, _value: Value) -> Result<Value, EntMutationError> {
1008            Err(EntMutationError::NoField {
1009                name: name.to_string(),
1010            })
1011        }
1012
1013        fn edge_definitions(&self) -> Vec<EdgeDefinition> {
1014            Vec::new()
1015        }
1016
1017        fn edge_names(&self) -> Vec<String> {
1018            Vec::new()
1019        }
1020
1021        fn edge(&self, _name: &str) -> Option<EdgeValue> {
1022            None
1023        }
1024
1025        fn update_edge(
1026            &mut self,
1027            name: &str,
1028            _value: EdgeValue,
1029        ) -> Result<EdgeValue, EntMutationError> {
1030            Err(EntMutationError::NoEdge {
1031                name: name.to_string(),
1032            })
1033        }
1034
1035        fn connect(&mut self, _database: WeakDatabaseRc) {}
1036
1037        fn disconnect(&mut self) {}
1038
1039        fn is_connected(&self) -> bool {
1040            false
1041        }
1042
1043        fn load_edge(&self, _name: &str) -> DatabaseResult<Vec<Box<dyn Ent>>> {
1044            Err(DatabaseError::Disconnected)
1045        }
1046
1047        fn clear_cache(&mut self) {}
1048
1049        fn refresh(&mut self) -> DatabaseResult<()> {
1050            Err(DatabaseError::Disconnected)
1051        }
1052
1053        fn commit(&mut self) -> DatabaseResult<()> {
1054            Err(DatabaseError::Disconnected)
1055        }
1056
1057        fn remove(&self) -> DatabaseResult<bool> {
1058            Err(DatabaseError::Disconnected)
1059        }
1060    }
1061}