Skip to main content

mentat/
store.rs

1// Copyright 2016-2018 Mozilla
2//
3// Licensed under the Apache License, Version 2.0 (the "License"); you may not use
4// this file except in compliance with the License. You may obtain a copy of the
5// License at http://www.apache.org/licenses/LICENSE-2.0
6// Unless required by applicable law or agreed to in writing, software distributed
7// under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
8// CONDITIONS OF ANY KIND, either express or implied. See the License for the
9// specific language governing permissions and limitations under the License.
10
11#![allow(dead_code)]
12
13use std::collections::BTreeMap;
14
15use std::sync::Arc;
16
17use core_traits::{Entid, StructuredMap, TypedValue};
18
19use mentat_core::{HasSchema, Keyword, TxReport, ValueRc};
20use mentat_db::TxObserver;
21
22use mentat_transaction::{
23    CacheAction, CacheDirection, InProgress, InProgressRead, Pullable, Queryable,
24};
25
26use crate::conn::Conn;
27
28use public_traits::errors::Result;
29
30use mentat_transaction::query::{PreparedResult, QueryExplanation, QueryInputs, QueryOutput};
31
32/// A convenience wrapper around a single SQLite connection and a Conn. This is suitable
33/// for applications that don't require complex connection management.
34/// The `:db/fn` attribute: the mino source of a user transaction function.
35const DB_FN_ATTRIBUTE: &str =
36    "[{:db/ident :db/fn :db/valueType :db.type/string :db/cardinality :db.cardinality/one}]";
37
38/// The live `:db/fn` source on the entity whose `:db/ident` is `ident`, if any.
39#[cfg(feature = "mino")]
40pub(crate) fn db_fn_source(
41    store: &Store,
42    ident: &str,
43) -> std::result::Result<Option<String>, String> {
44    let db_fn = Keyword::namespaced("db", "fn");
45    let schema = store.conn.current_schema();
46    if schema.get_entid(&db_fn).is_none() {
47        return Ok(None);
48    }
49    let Ok(edn::Value::Keyword(kw)) = edn::parse::value(ident).map(|v| v.without_spans()) else {
50        return Ok(None);
51    };
52    let Some(e) = schema.get_entid(&kw) else {
53        return Ok(None); // not an ident at all
54    };
55    match store
56        .lookup_value_for_attribute(e.0, &db_fn)
57        .map_err(|err| err.to_string())?
58    {
59        Some(TypedValue::String(s)) => Ok(Some(s.as_ref().clone())),
60        _ => Ok(None),
61    }
62}
63
64pub struct Store {
65    conn: Conn,
66    sqlite: rusqlite::Connection,
67    advisor: std::cell::RefCell<crate::auto_index::Advisor>,
68}
69
70impl Store {
71    /// Open a store at the supplied path, ensuring that it includes the bootstrap schema.
72    pub fn open(path: &str) -> Result<Store> {
73        let mut connection = crate::new_connection(path)?;
74        let conn = Conn::connect(&mut connection)?;
75        Ok(Store {
76            conn,
77            sqlite: connection,
78            advisor: Default::default(),
79        })
80    }
81
82    pub fn transact(&mut self, transaction: &str) -> Result<TxReport> {
83        let transaction = self.expand_tx_fns(transaction)?;
84        let mut ip = self.begin_transaction()?;
85        let report = ip.transact(transaction.as_str())?;
86        ip.commit()?;
87        Ok(report)
88    }
89
90    /// Prepare `transaction` for the transactor: replace every invocation of a
91    /// user-defined transaction function (`[:my/fn arg ...]` on an ident that
92    /// carries `:db/fn`) with the tx-data it returns, and define the `:db/fn`
93    /// attribute itself the first time a transaction installs a function.
94    /// [`Store::transact`] and [`Store::with_speculative`] call this; a host
95    /// that drives [`InProgress::transact`] directly calls it first.
96    ///
97    /// Functions run in the sandboxed mino interpreter (`mino` feature), see
98    /// the database as of the start of the transaction, and nest at most
99    /// `MENTAT_TXFN_MAX_DEPTH` (default 16) deep. Without the `mino` feature
100    /// only the `:db/fn` attribute bootstrap applies.
101    pub fn expand_tx_fns(&mut self, transaction: &str) -> Result<String> {
102        let has_db_fn_attr = self
103            .conn
104            .current_schema()
105            .get_entid(&Keyword::namespaced("db", "fn"))
106            .is_some();
107        if !has_db_fn_attr && transaction.contains(":db/fn") {
108            // The embedded core schema predates transaction functions; define
109            // `:db/fn` alongside the first function installed, rather than
110            // changing every store's bootstrap.
111            self.transact_raw(DB_FN_ATTRIBUTE)?;
112        }
113        // Only a store that has (or is about to have) `:db/fn` can hold a
114        // transaction function, so other stores skip the interpreter entirely.
115        #[cfg(feature = "mino")]
116        if has_db_fn_attr || transaction.contains(":db/fn") {
117            // ponytail: swaps in a throwaway in-memory Store (~ms) to move this
118            // one into the script backend; only paid by transactions that could
119            // invoke a function. Borrowing it instead needs a non-'static
120            // ScriptBackend.
121            let store = std::mem::replace(self, Store::open("")?);
122            let (store, expanded) = crate::script::expand_tx_fns(store, transaction);
123            *self = store;
124            return expanded.map_err(public_traits::errors::MentatError::TransactionFunction);
125        }
126        Ok(transaction.to_string())
127    }
128
129    /// Transact without transaction-function expansion.
130    fn transact_raw(&mut self, transaction: &str) -> Result<TxReport> {
131        let mut ip = self.begin_transaction()?;
132        let report = ip.transact(transaction)?;
133        ip.commit()?;
134        Ok(report)
135    }
136
137    /// Speculatively transact `transaction` and return its `TxReport` WITHOUT
138    /// committing (Datomic `d/with`). Delegates to [`Conn::transact_speculative`],
139    /// which opens a real SQLite transaction and rolls it back on drop, so the
140    /// store is left unchanged. This is the thin `&mut Store` wrapper the
141    /// scripting layer needs, since `transact_speculative` lives on `Conn` and
142    /// wants both the `Conn` and the `rusqlite::Connection` (both private here).
143    pub fn with_speculative(&mut self, transaction: &str) -> Result<TxReport> {
144        let transaction = self.expand_tx_fns(transaction)?;
145        self.conn
146            .transact_speculative(&mut self.sqlite, transaction.as_str())
147    }
148}
149
150#[cfg(feature = "sqlcipher")]
151impl Store {
152    /// Variant of `open` that allows a key (for encryption/decryption) to be
153    /// supplied. Fails unless linked against sqlcipher (or something else that
154    /// supports the Sqlite Encryption Extension).
155    pub fn open_with_key(path: &str, encryption_key: &str) -> Result<Store> {
156        let mut connection = crate::new_connection_with_key(path, encryption_key)?;
157        let conn = Conn::connect(&mut connection)?;
158        Ok(Store {
159            conn,
160            sqlite: connection,
161            advisor: Default::default(),
162        })
163    }
164
165    /// Change the key for a database that was opened using `open_with_key` (using `PRAGMA
166    /// rekey`). Fails unless linked against sqlcipher (or something else that supports the Sqlite
167    /// Encryption Extension).
168    pub fn change_encryption_key(&mut self, new_encryption_key: &str) -> Result<()> {
169        crate::change_encryption_key(&self.sqlite, new_encryption_key)?;
170        Ok(())
171    }
172}
173
174impl Store {
175    /// Intended for use from tests.
176    pub fn sqlite_mut(&mut self) -> &mut rusqlite::Connection {
177        &mut self.sqlite
178    }
179
180    /// A shared borrow of the underlying SQLite connection, for read-only
181    /// helpers (e.g. `mentat_db::debug::datoms`) that only need `&Connection`.
182    pub fn sqlite_ref(&self) -> &rusqlite::Connection {
183        &self.sqlite
184    }
185
186    #[cfg(test)]
187    pub fn is_registered_as_observer(&self, key: &str) -> bool {
188        self.conn
189            .tx_observer_service
190            .lock()
191            .unwrap()
192            .is_registered(key)
193    }
194}
195
196impl Store {
197    pub fn dismantle(self) -> (rusqlite::Connection, Conn) {
198        (self.sqlite, self.conn)
199    }
200
201    pub fn conn(&self) -> &Conn {
202        &self.conn
203    }
204
205    pub fn begin_read(&mut self) -> Result<InProgressRead<'_, '_>> {
206        self.conn.begin_read(&mut self.sqlite)
207    }
208
209    pub fn begin_transaction(&mut self) -> Result<InProgress<'_, '_>> {
210        self.conn.begin_transaction(&mut self.sqlite)
211    }
212
213    pub fn cache(&mut self, attr: &Keyword, direction: CacheDirection) -> Result<()> {
214        let schema = &self.conn.current_schema();
215        self.conn.cache(
216            &mut self.sqlite,
217            schema,
218            attr,
219            direction,
220            CacheAction::Register,
221        )
222    }
223
224    pub fn register_observer(&mut self, key: String, observer: Arc<TxObserver>) {
225        self.conn.register_observer(key, observer);
226    }
227
228    pub fn unregister_observer(&mut self, key: &str) {
229        self.conn.unregister_observer(key);
230    }
231
232    pub fn last_tx_id(&self) -> Entid {
233        self.conn.last_tx_id()
234    }
235
236    /// How mentat manages secondary indexes for this `Store` (see [`AutoIndex`]).
237    /// The default comes from `MENTAT_AUTO_INDEX` (`off`/`schema`/`adaptive`),
238    /// else `Schema`.
239    ///
240    /// [`AutoIndex`]: crate::AutoIndex
241    pub fn set_auto_index(&mut self, mode: crate::AutoIndex) {
242        self.advisor.get_mut().mode = mode;
243    }
244
245    /// Adaptive thresholds: create an index once an attribute is value-filtered
246    /// by `min_uses` queries within one tuning period; drop a mentat-created
247    /// index unused for a whole period and for at least `idle`.
248    pub fn set_index_tuning(&mut self, min_uses: u32, idle: std::time::Duration) {
249        let a = self.advisor.get_mut();
250        a.min_uses = min_uses.max(1);
251        a.idle = idle;
252    }
253
254    /// Apply (or, with `dry_run`, only report) index changes per the current
255    /// [`AutoIndex`](crate::AutoIndex) mode. Only ever drops indexes mentat
256    /// created (listed in `mentat_managed_indexes`).
257    pub fn tune_indexes(&mut self, dry_run: bool) -> Result<Vec<crate::IndexAction>> {
258        self.tune(dry_run)
259    }
260
261    fn tune(&self, dry_run: bool) -> Result<Vec<crate::IndexAction>> {
262        let schema = self.conn.current_schema();
263        self.advisor
264            .borrow_mut()
265            .tune(&self.sqlite, &schema, dry_run)
266    }
267
268    /// Run a Datalog query against a historical basis reconstructed as of
269    /// transaction `tx` (inclusive) -- the state that was current at `tx`.
270    pub fn q_once_as_of<T>(&self, query: &str, inputs: T, tx: Entid) -> Result<QueryOutput>
271    where
272        T: Into<Option<QueryInputs>>,
273    {
274        self.conn.q_once_as_of(&self.sqlite, query, inputs, tx)
275    }
276
277    /// `EXPLAIN QUERY PLAN` for a query, optionally against a historical basis.
278    pub fn q_explain_temporal<T>(
279        &self,
280        query: &str,
281        inputs: T,
282        temporal: Option<crate::TemporalBound>,
283    ) -> Result<QueryExplanation>
284    where
285        T: Into<Option<QueryInputs>>,
286    {
287        self.conn
288            .q_explain_temporal(&self.sqlite, query, inputs, temporal)
289    }
290
291    /// Pull `pattern` (EDN, e.g. `[:person/name {:person/friend [:person/name]}]`)
292    /// for `entity` on the current basis.
293    pub fn pull(&self, entity: Entid, pattern: &str) -> Result<StructuredMap> {
294        self.conn.pull_on(&self.sqlite, entity, pattern, None)
295    }
296
297    /// Pull `pattern` for `entity` as of transaction `tx` (inclusive): the
298    /// values asserted at or before `tx` and not retracted at or before `tx`.
299    pub fn pull_as_of(&self, entity: Entid, pattern: &str, tx: Entid) -> Result<StructuredMap> {
300        let basis = Some(crate::TemporalBound::AsOf(tx));
301        self.conn.pull_on(&self.sqlite, entity, pattern, basis)
302    }
303
304    /// Pull `pattern` for `entity` seeing only the current datoms transacted
305    /// after transaction `tx` (Datomic's `since`).
306    pub fn pull_since(&self, entity: Entid, pattern: &str, tx: Entid) -> Result<StructuredMap> {
307        let basis = Some(crate::TemporalBound::Since(tx));
308        self.conn.pull_on(&self.sqlite, entity, pattern, basis)
309    }
310
311    /// Pull `pattern` for `entity` on `basis` (`None` = current; as-of, since,
312    /// or both: `TemporalBound::AsOfSince`).
313    pub fn pull_temporal(
314        &self,
315        entity: Entid,
316        pattern: &str,
317        basis: Option<crate::TemporalBound>,
318    ) -> Result<StructuredMap> {
319        self.conn.pull_on(&self.sqlite, entity, pattern, basis)
320    }
321
322    /// Run a Datalog query on `basis` (`None` = current, via `q_once`; as-of,
323    /// since, or both).
324    pub fn q_once_temporal<T>(
325        &self,
326        query: &str,
327        inputs: T,
328        basis: Option<crate::TemporalBound>,
329    ) -> Result<QueryOutput>
330    where
331        T: Into<Option<QueryInputs>>,
332    {
333        match basis {
334            Some(b) => self.conn.q_once_temporal(&self.sqlite, query, inputs, b),
335            None => self.q_once(query, inputs),
336        }
337    }
338
339    /// Run a Datalog query on Datomic's `since tx`: a pattern sees the current
340    /// datoms transacted after `tx`; a history pattern `[?e ?a ?v ?tx ?added]`
341    /// sees every datom after `tx`, retractions too.
342    pub fn q_once_since<T>(&self, query: &str, inputs: T, tx: Entid) -> Result<QueryOutput>
343    where
344        T: Into<Option<QueryInputs>>,
345    {
346        self.conn.q_once_since(&self.sqlite, query, inputs, tx)
347    }
348}
349
350impl Queryable for Store {
351    fn q_once<T>(&self, query: &str, inputs: T) -> Result<QueryOutput>
352    where
353        T: Into<Option<QueryInputs>>,
354    {
355        if self.advisor.borrow().mode != crate::AutoIndex::Adaptive {
356            return self.conn.q_once(&self.sqlite, query, inputs);
357        }
358        let (out, attrs) = self
359            .conn
360            .q_once_value_filtered(&self.sqlite, query, inputs)?;
361        if self.advisor.borrow_mut().record(&attrs) {
362            // Best effort: e.g. another connection may hold the write lock.
363            if let Err(e) = self.tune(false) {
364                log::warn!("automatic index tuning failed: {e}");
365            }
366        }
367        Ok(out)
368    }
369
370    fn q_prepare<T>(&self, query: &str, inputs: T) -> PreparedResult<'_>
371    where
372        T: Into<Option<QueryInputs>>,
373    {
374        self.conn.q_prepare(&self.sqlite, query, inputs)
375    }
376
377    fn q_explain<T>(&self, query: &str, inputs: T) -> Result<QueryExplanation>
378    where
379        T: Into<Option<QueryInputs>>,
380    {
381        self.conn.q_explain(&self.sqlite, query, inputs)
382    }
383
384    fn lookup_values_for_attribute<E>(
385        &self,
386        entity: E,
387        attribute: &edn::Keyword,
388    ) -> Result<Vec<TypedValue>>
389    where
390        E: Into<Entid>,
391    {
392        self.conn
393            .lookup_values_for_attribute(&self.sqlite, entity.into(), attribute)
394    }
395
396    fn lookup_value_for_attribute<E>(
397        &self,
398        entity: E,
399        attribute: &edn::Keyword,
400    ) -> Result<Option<TypedValue>>
401    where
402        E: Into<Entid>,
403    {
404        self.conn
405            .lookup_value_for_attribute(&self.sqlite, entity.into(), attribute)
406    }
407}
408
409impl Pullable for Store {
410    fn pull_attributes_for_entities<E, A>(
411        &self,
412        entities: E,
413        attributes: A,
414    ) -> Result<BTreeMap<Entid, ValueRc<StructuredMap>>>
415    where
416        E: IntoIterator<Item = Entid>,
417        A: IntoIterator<Item = Entid>,
418    {
419        self.conn
420            .pull_attributes_for_entities(&self.sqlite, entities, attributes)
421    }
422
423    fn pull_attributes_for_entity<A>(&self, entity: Entid, attributes: A) -> Result<StructuredMap>
424    where
425        A: IntoIterator<Item = Entid>,
426    {
427        self.conn
428            .pull_attributes_for_entity(&self.sqlite, entity, attributes)
429    }
430}
431
432#[cfg(test)]
433mod tests {
434    use super::*;
435
436    use uuid::Uuid;
437
438    use std::collections::BTreeSet;
439    use std::path::{Path, PathBuf};
440    use std::sync::mpsc;
441    use std::sync::Mutex;
442    use std::time::Duration;
443
444    use mentat_db::cache::SQLiteAttributeCache;
445
446    use core_traits::{TypedValue, ValueType};
447
448    use mentat_core::{CachedAttributes, HasSchema};
449
450    use mentat_transaction::entity_builder::BuildTerms;
451
452    use mentat_transaction::query::PreparedQuery;
453
454    use mentat_query_algebrizer::QueryInputs;
455
456    use crate::vocabulary::{AttributeBuilder, Definition, VersionedStore};
457
458    use core_traits::attribute::Unique;
459
460    fn fixture_path(rest: &str) -> PathBuf {
461        let fixtures = Path::new("../fixtures/");
462        fixtures.join(Path::new(rest))
463    }
464
465    #[test]
466    fn test_prepared_query_with_cache() {
467        let mut store = Store::open("").expect("opened");
468        let mut in_progress = store.begin_transaction().expect("began");
469        in_progress
470            .import(fixture_path("cities.schema"))
471            .expect("transacted schema");
472        in_progress
473            .import(fixture_path("all_seattle.edn"))
474            .expect("transacted data");
475        in_progress
476            .cache(
477                &kw!(:neighborhood/district),
478                CacheDirection::Forward,
479                CacheAction::Register,
480            )
481            .expect("cache done");
482        in_progress
483            .cache(
484                &kw!(:district/name),
485                CacheDirection::Forward,
486                CacheAction::Register,
487            )
488            .expect("cache done");
489        in_progress
490            .cache(
491                &kw!(:neighborhood/name),
492                CacheDirection::Reverse,
493                CacheAction::Register,
494            )
495            .expect("cache done");
496
497        let query = r#"[:find ?district
498                        :in ?hood
499                        :where
500                        [?neighborhood :neighborhood/name ?hood]
501                        [?neighborhood :neighborhood/district ?d]
502                        [?d :district/name ?district]]"#;
503        let hood = "Beacon Hill";
504        let inputs =
505            QueryInputs::with_value_sequence(vec![(var!(?hood), TypedValue::typed_string(hood))]);
506        let mut prepared = in_progress.q_prepare(query, inputs).expect("prepared");
507        match prepared {
508            PreparedQuery::Constant {
509                select: ref _select,
510            } => {}
511            _ => panic!(),
512        };
513
514        let start = std::time::Instant::now();
515        let results = prepared.run(None).expect("results");
516        let end = std::time::Instant::now();
517        println!(
518            "Prepared cache execution took {}µs",
519            (end - start).as_micros()
520        );
521        assert_eq!(
522            results.into_rel().expect("result"),
523            vec![vec![TypedValue::typed_string("Greater Duwamish")]].into()
524        );
525    }
526
527    trait StoreCache {
528        fn get_entid_for_value(&self, attr: Entid, val: &TypedValue) -> Option<Entid>;
529        fn is_attribute_cached_reverse(&self, attr: Entid) -> bool;
530        fn is_attribute_cached_forward(&self, attr: Entid) -> bool;
531    }
532
533    impl StoreCache for Store {
534        fn get_entid_for_value(&self, attr: Entid, val: &TypedValue) -> Option<Entid> {
535            let cache = self.conn.current_cache();
536            cache.get_entid_for_value(attr, val)
537        }
538
539        fn is_attribute_cached_forward(&self, attr: Entid) -> bool {
540            self.conn.current_cache().is_attribute_cached_forward(attr)
541        }
542
543        fn is_attribute_cached_reverse(&self, attr: Entid) -> bool {
544            self.conn.current_cache().is_attribute_cached_reverse(attr)
545        }
546    }
547
548    #[test]
549    fn test_cache_mutation() {
550        let mut store = Store::open("").expect("opened");
551
552        {
553            let mut in_progress = store.begin_transaction().expect("begun");
554            in_progress
555                .transact(
556                    r#"[
557                {  :db/ident       :foo/bar
558                   :db/cardinality :db.cardinality/one
559                   :db/index       true
560                   :db/unique      :db.unique/identity
561                   :db/valueType   :db.type/long },
562                {  :db/ident       :foo/baz
563                   :db/cardinality :db.cardinality/one
564                   :db/valueType   :db.type/boolean }
565                {  :db/ident       :foo/x
566                   :db/cardinality :db.cardinality/many
567                   :db/valueType   :db.type/long }]"#,
568                )
569                .expect("transact");
570
571            // Cache one….
572            in_progress
573                .cache(
574                    &kw!(:foo/bar),
575                    CacheDirection::Reverse,
576                    CacheAction::Register,
577                )
578                .expect("cache done");
579            in_progress.commit().expect("commit");
580        }
581
582        let foo_bar = store
583            .conn
584            .current_schema()
585            .get_entid(&kw!(:foo/bar))
586            .expect("foo/bar")
587            .0;
588        let foo_baz = store
589            .conn
590            .current_schema()
591            .get_entid(&kw!(:foo/baz))
592            .expect("foo/baz")
593            .0;
594        let foo_x = store
595            .conn
596            .current_schema()
597            .get_entid(&kw!(:foo/x))
598            .expect("foo/x")
599            .0;
600
601        // … and cache the others via the store.
602        store
603            .cache(&kw!(:foo/baz), CacheDirection::Both)
604            .expect("cache done");
605        store
606            .cache(&kw!(:foo/x), CacheDirection::Forward)
607            .expect("cache done");
608        {
609            assert!(store.is_attribute_cached_reverse(foo_bar));
610            assert!(store.is_attribute_cached_forward(foo_baz));
611            assert!(store.is_attribute_cached_reverse(foo_baz));
612            assert!(store.is_attribute_cached_forward(foo_x));
613        }
614
615        // Add some data.
616        {
617            let mut in_progress = store.begin_transaction().expect("begun");
618
619            {
620                assert!(in_progress.cache.is_attribute_cached_reverse(foo_bar));
621                assert!(in_progress.cache.is_attribute_cached_forward(foo_baz));
622                assert!(in_progress.cache.is_attribute_cached_reverse(foo_baz));
623                assert!(in_progress.cache.is_attribute_cached_forward(foo_x));
624
625                assert!(in_progress
626                    .cache
627                    .overlay
628                    .is_attribute_cached_reverse(foo_bar));
629                assert!(in_progress
630                    .cache
631                    .overlay
632                    .is_attribute_cached_forward(foo_baz));
633                assert!(in_progress
634                    .cache
635                    .overlay
636                    .is_attribute_cached_reverse(foo_baz));
637                assert!(in_progress.cache.overlay.is_attribute_cached_forward(foo_x));
638            }
639
640            in_progress
641                .transact(
642                    r#"[
643                {:foo/bar 15, :foo/baz false, :foo/x [1, 2, 3]}
644                {:foo/bar 99, :foo/baz true}
645                {:foo/bar -2, :foo/baz true}
646                ]"#,
647                )
648                .expect("transact");
649
650            // Data is in the cache.
651            let first = in_progress
652                .cache
653                .get_entid_for_value(foo_bar, &TypedValue::Long(15))
654                .expect("id");
655            assert_eq!(
656                in_progress
657                    .cache
658                    .get_value_for_entid(&in_progress.schema, foo_baz, first)
659                    .expect("val"),
660                &TypedValue::Boolean(false)
661            );
662
663            // All three values for :foo/x.
664            let all_three: BTreeSet<TypedValue> = in_progress
665                .cache
666                .get_values_for_entid(&in_progress.schema, foo_x, first)
667                .expect("val")
668                .iter()
669                .cloned()
670                .collect();
671            assert_eq!(
672                all_three,
673                vec![1, 2, 3].into_iter().map(TypedValue::Long).collect()
674            );
675
676            in_progress.commit().expect("commit");
677        }
678
679        // Data is still in the cache.
680        {
681            let first = store
682                .get_entid_for_value(foo_bar, &TypedValue::Long(15))
683                .expect("id");
684            let cache: SQLiteAttributeCache = store.conn.current_cache();
685            assert_eq!(
686                cache
687                    .get_value_for_entid(&store.conn.current_schema(), foo_baz, first)
688                    .expect("val"),
689                &TypedValue::Boolean(false)
690            );
691
692            let all_three: BTreeSet<TypedValue> = cache
693                .get_values_for_entid(&store.conn.current_schema(), foo_x, first)
694                .expect("val")
695                .iter()
696                .cloned()
697                .collect();
698            assert_eq!(
699                all_three,
700                vec![1, 2, 3].into_iter().map(TypedValue::Long).collect()
701            );
702        }
703
704        // We can remove data and the cache reflects it, immediately and after commit.
705        {
706            let mut in_progress = store.begin_transaction().expect("began");
707            let first = in_progress
708                .cache
709                .get_entid_for_value(foo_bar, &TypedValue::Long(15))
710                .expect("id");
711            in_progress
712                .transact(format!("[[:db/retract {} :foo/x 2]]", first).as_str())
713                .expect("transact");
714
715            let only_two: BTreeSet<TypedValue> = in_progress
716                .cache
717                .get_values_for_entid(&in_progress.schema, foo_x, first)
718                .expect("val")
719                .iter()
720                .cloned()
721                .collect();
722            assert_eq!(
723                only_two,
724                vec![1, 3].into_iter().map(TypedValue::Long).collect()
725            );
726
727            // Rollback: unchanged.
728        }
729        {
730            let first = store
731                .get_entid_for_value(foo_bar, &TypedValue::Long(15))
732                .expect("id");
733            let cache: SQLiteAttributeCache = store.conn.current_cache();
734            assert_eq!(
735                cache
736                    .get_value_for_entid(&store.conn.current_schema(), foo_baz, first)
737                    .expect("val"),
738                &TypedValue::Boolean(false)
739            );
740
741            let all_three: BTreeSet<TypedValue> = cache
742                .get_values_for_entid(&store.conn.current_schema(), foo_x, first)
743                .expect("val")
744                .iter()
745                .cloned()
746                .collect();
747            assert_eq!(
748                all_three,
749                vec![1, 2, 3].into_iter().map(TypedValue::Long).collect()
750            );
751        }
752
753        // Try again, but this time commit.
754        {
755            let mut in_progress = store.begin_transaction().expect("began");
756            let first = in_progress
757                .cache
758                .get_entid_for_value(foo_bar, &TypedValue::Long(15))
759                .expect("id");
760            in_progress
761                .transact(format!("[[:db/retract {} :foo/x 2]]", first).as_str())
762                .expect("transact");
763            in_progress.commit().expect("committed");
764        }
765        {
766            let first = store
767                .get_entid_for_value(foo_bar, &TypedValue::Long(15))
768                .expect("id");
769            let cache: SQLiteAttributeCache = store.conn.current_cache();
770            assert_eq!(
771                cache
772                    .get_value_for_entid(&store.conn.current_schema(), foo_baz, first)
773                    .expect("val"),
774                &TypedValue::Boolean(false)
775            );
776
777            let only_two: BTreeSet<TypedValue> = cache
778                .get_values_for_entid(&store.conn.current_schema(), foo_x, first)
779                .expect("val")
780                .iter()
781                .cloned()
782                .collect();
783            assert_eq!(
784                only_two,
785                vec![1, 3].into_iter().map(TypedValue::Long).collect()
786            );
787        }
788    }
789
790    fn test_register_observer() {
791        let mut conn = Store::open("").unwrap();
792
793        let key = "Test Observer".to_string();
794        let tx_observer = TxObserver::new(BTreeSet::new(), move |_obs_key, _batch| {});
795
796        conn.register_observer(key.clone(), Arc::new(tx_observer));
797        assert!(conn.is_registered_as_observer(&key));
798    }
799
800    #[test]
801    fn test_deregister_observer() {
802        let mut conn = Store::open("").unwrap();
803
804        let key = "Test Observer".to_string();
805
806        let tx_observer = TxObserver::new(BTreeSet::new(), move |_obs_key, _batch| {});
807
808        conn.register_observer(key.clone(), Arc::new(tx_observer));
809        assert!(conn.is_registered_as_observer(&key));
810
811        conn.unregister_observer(&key);
812
813        assert!(!conn.is_registered_as_observer(&key));
814    }
815
816    fn add_schema(conn: &mut Store) {
817        // transact some schema
818        let mut in_progress = conn.begin_transaction().expect("expected in progress");
819        in_progress
820            .ensure_vocabulary(&Definition::new(
821                kw!(:todo/items),
822                1,
823                vec![
824                    (
825                        kw!(:todo/uuid),
826                        AttributeBuilder::helpful()
827                            .value_type(ValueType::Uuid)
828                            .multival(false)
829                            .unique(Unique::Value)
830                            .index(true)
831                            .build(),
832                    ),
833                    (
834                        kw!(:todo/name),
835                        AttributeBuilder::helpful()
836                            .value_type(ValueType::String)
837                            .multival(false)
838                            .fulltext(true)
839                            .build(),
840                    ),
841                    (
842                        kw!(:todo/completion_date),
843                        AttributeBuilder::helpful()
844                            .value_type(ValueType::Instant)
845                            .multival(false)
846                            .build(),
847                    ),
848                    (
849                        kw!(:label/name),
850                        AttributeBuilder::helpful()
851                            .value_type(ValueType::String)
852                            .multival(false)
853                            .unique(Unique::Value)
854                            .fulltext(true)
855                            .index(true)
856                            .build(),
857                    ),
858                    (
859                        kw!(:label/color),
860                        AttributeBuilder::helpful()
861                            .value_type(ValueType::String)
862                            .multival(false)
863                            .build(),
864                    ),
865                ],
866            ))
867            .expect("expected vocubulary");
868        in_progress.commit().expect("Expected vocabulary committed");
869    }
870
871    #[derive(Default, Debug)]
872    struct ObserverOutput {
873        txids: Vec<i64>,
874        changes: Vec<BTreeSet<i64>>,
875        called_key: Option<String>,
876    }
877
878    #[test]
879    fn test_observer_notified_on_registered_change() {
880        let mut conn = Store::open("").unwrap();
881        add_schema(&mut conn);
882
883        let name_entid: Entid = conn
884            .conn()
885            .current_schema()
886            .get_entid(&kw!(:todo/name))
887            .expect("entid to exist for name")
888            .into();
889        let date_entid: Entid = conn
890            .conn()
891            .current_schema()
892            .get_entid(&kw!(:todo/completion_date))
893            .expect("entid to exist for completion_date")
894            .into();
895        let mut registered_attrs = BTreeSet::new();
896        registered_attrs.insert(name_entid);
897        registered_attrs.insert(date_entid);
898
899        let key = "Test Observing".to_string();
900
901        let output = Arc::new(Mutex::new(ObserverOutput::default()));
902
903        let mut_output = Arc::downgrade(&output);
904        let (tx, rx): (mpsc::Sender<()>, mpsc::Receiver<()>) = mpsc::channel();
905        // because the TxObserver is in an Arc and is therefore Sync, we have to wrap the Sender in a Mutex to also
906        // make it Sync.
907        let thread_tx = Mutex::new(tx);
908        let tx_observer = Arc::new(TxObserver::new(registered_attrs, move |obs_key, batch| {
909            if let Some(out) = mut_output.upgrade() {
910                let mut o = out.lock().unwrap();
911                o.called_key = Some(obs_key.to_string());
912                for (tx_id, changes) in batch.into_iter() {
913                    o.txids.push(*tx_id);
914                    o.changes.push(changes.clone());
915                }
916                o.txids.sort_unstable();
917            }
918            thread_tx.lock().unwrap().send(()).unwrap();
919        }));
920
921        conn.register_observer(key.clone(), Arc::clone(&tx_observer));
922        assert!(conn.is_registered_as_observer(&key));
923
924        let mut tx_ids = Vec::new();
925        let mut changesets = Vec::new();
926        let db_tx_instant_entid: Entid = conn
927            .conn()
928            .current_schema()
929            .get_entid(&kw!(:db/txInstant))
930            .expect("entid to exist for :db/txInstant")
931            .into();
932        let uuid_entid: Entid = conn
933            .conn()
934            .current_schema()
935            .get_entid(&kw!(:todo/uuid))
936            .expect("entid to exist for name")
937            .into();
938        {
939            let mut in_progress = conn.begin_transaction().expect("expected transaction");
940            for i in 0..3 {
941                let mut changeset = BTreeSet::new();
942                changeset.insert(db_tx_instant_entid);
943                let name = format!("todo{}", i);
944                let uuid = Uuid::new_v4();
945                let mut builder = in_progress.builder().describe_tempid(&name);
946                builder
947                    .add(kw!(:todo/uuid), TypedValue::Uuid(uuid))
948                    .expect("Expected added uuid");
949                changeset.insert(uuid_entid);
950                builder
951                    .add(kw!(:todo/name), TypedValue::typed_string(name))
952                    .expect("Expected added name");
953                changeset.insert(name_entid);
954                if i % 2 == 0 {
955                    builder
956                        .add(kw!(:todo/completion_date), TypedValue::current_instant())
957                        .expect("Expected added date");
958                    changeset.insert(date_entid);
959                }
960                let (ip, r) = builder.transact();
961                let report = r.expect("expected a report");
962                tx_ids.push(report.tx_id);
963                changesets.push(changeset);
964                in_progress = ip;
965            }
966            let mut builder = in_progress.builder().describe_tempid("Label");
967            builder
968                .add(kw!(:label/name), TypedValue::typed_string("Label 1"))
969                .expect("Expected added name");
970            builder
971                .add(kw!(:label/color), TypedValue::typed_string("blue"))
972                .expect("Expected added color");
973            builder.commit().expect("expect transaction to occur");
974        }
975
976        let delay = Duration::from_millis(100);
977        let _ = rx.recv_timeout(delay);
978
979        let out = Arc::try_unwrap(output).expect("unwrapped");
980        let o = out.into_inner().expect("Expected an Output");
981        assert_eq!(o.called_key, Some(key));
982        assert_eq!(o.txids, tx_ids);
983        assert_eq!(o.changes, changesets);
984    }
985
986    #[test]
987    fn test_observer_not_notified_on_unregistered_change() {
988        let mut conn = Store::open("").unwrap();
989        add_schema(&mut conn);
990
991        let name_entid: Entid = conn
992            .conn()
993            .current_schema()
994            .get_entid(&kw!(:todo/name))
995            .expect("entid to exist for name")
996            .into();
997        let date_entid: Entid = conn
998            .conn()
999            .current_schema()
1000            .get_entid(&kw!(:todo/completion_date))
1001            .expect("entid to exist for completion_date")
1002            .into();
1003        let mut registered_attrs = BTreeSet::new();
1004        registered_attrs.insert(name_entid);
1005        registered_attrs.insert(date_entid);
1006
1007        let key = "Test Observing".to_string();
1008
1009        let output = Arc::new(Mutex::new(ObserverOutput::default()));
1010
1011        let mut_output = Arc::downgrade(&output);
1012        let (tx, rx): (mpsc::Sender<()>, mpsc::Receiver<()>) = mpsc::channel();
1013        let thread_tx = Mutex::new(tx);
1014        let tx_observer = Arc::new(TxObserver::new(registered_attrs, move |obs_key, batch| {
1015            if let Some(out) = mut_output.upgrade() {
1016                let mut o = out.lock().unwrap();
1017                o.called_key = Some(obs_key.to_string());
1018                for (tx_id, changes) in batch.into_iter() {
1019                    o.txids.push(*tx_id);
1020                    o.changes.push(changes.clone());
1021                }
1022                o.txids.sort_unstable();
1023            }
1024            thread_tx.lock().unwrap().send(()).unwrap();
1025        }));
1026
1027        conn.register_observer(key.clone(), Arc::clone(&tx_observer));
1028        assert!(conn.is_registered_as_observer(&key));
1029
1030        let tx_ids = Vec::<Entid>::new();
1031        let changesets = Vec::<BTreeSet<Entid>>::new();
1032        {
1033            let mut in_progress = conn.begin_transaction().expect("expected transaction");
1034            for i in 0..3 {
1035                let name = format!("label{}", i);
1036                let mut builder = in_progress.builder().describe_tempid(&name);
1037                builder
1038                    .add(kw!(:label/name), TypedValue::typed_string(name))
1039                    .expect("Expected added name");
1040                builder
1041                    .add(kw!(:label/color), TypedValue::typed_string("blue"))
1042                    .expect("Expected added color");
1043                let (ip, _) = builder.transact();
1044                in_progress = ip;
1045            }
1046        }
1047
1048        let delay = Duration::from_millis(100);
1049        let _ = rx.recv_timeout(delay);
1050
1051        let out = Arc::try_unwrap(output).expect("unwrapped");
1052        let o = out.into_inner().expect("Expected an Output");
1053        assert_eq!(o.called_key, None);
1054        assert_eq!(o.txids, tx_ids);
1055        assert_eq!(o.changes, changesets);
1056    }
1057}