Skip to main content

mentat/
conn.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::borrow::Borrow;
14
15use std::collections::BTreeMap;
16
17use std::sync::{Arc, Mutex};
18
19use rusqlite::TransactionBehavior;
20
21pub use core_traits::{Attribute, Entid, KnownEntid, StructuredMap, TypedValue, ValueType};
22
23use mentat_core::{HasSchema, Keyword, Schema, TxReport, ValueRc};
24
25use mentat_db::cache::{InProgressSQLiteAttributeCache, SQLiteAttributeCache};
26
27use mentat_db::db;
28use mentat_db::{
29    InProgressObserverTransactWatcher, PartitionMap, TxObservationService, TxObserver,
30};
31
32use mentat_query_pull::{
33    pull_attributes_for_entities, pull_attributes_for_entity, pull_pattern_for_entity,
34};
35
36use mentat_transaction::{CacheAction, CacheDirection, InProgress, InProgressRead, Metadata};
37
38use public_traits::errors::{MentatError, Result};
39
40use mentat_transaction::query::{
41    lookup_value_for_attribute, lookup_values_for_attribute, q_explain, q_once, q_prepare,
42    q_uncached, Known, PreparedResult, QueryExplanation, QueryInputs, QueryOutput,
43};
44
45/// A mutable, safe reference to the current Mentat store.
46pub struct Conn {
47    /// `Mutex` since all reads and writes need to be exclusive.  Internally, owned data for the
48    /// volatile parts (generation and partition map), and `Arc` for the infrequently changing parts
49    /// (schema, cache) that we want to share across threads.  A consuming thread may use a shared
50    /// reference after the `Conn`'s `Metadata` has moved on.
51    ///
52    /// The motivating case is multiple query threads taking references to the current schema to
53    /// perform long-running queries while a single writer thread moves the metadata -- partition
54    /// map and schema -- forward.
55    ///
56    /// We want the attribute cache to be isolated across transactions, updated within
57    /// `InProgress` writes, and updated in the `Conn` on commit. To achieve this we
58    /// store the cache itself in an `Arc` inside `SQLiteAttributeCache`, so that `.get_mut()`
59    /// gives us copy-on-write semantics.
60    /// We store that cached `Arc` here in a `Mutex`, so that the main copy can be carefully
61    /// replaced on commit.
62    metadata: Mutex<Metadata>,
63
64    // TODO: maintain set of change listeners or handles to transaction report queues. #298.
65
66    // TODO: maintain cache of query plans that could be shared across threads and invalidated when
67    // the schema changes. #315.
68    pub(crate) tx_observer_service: Mutex<TxObservationService>,
69}
70
71impl Conn {
72    // Intentionally not public.
73    fn new(partition_map: PartitionMap, schema: Schema) -> Conn {
74        Conn {
75            metadata: Mutex::new(Metadata::new(
76                0,
77                partition_map,
78                Arc::new(schema),
79                Default::default(),
80            )),
81            tx_observer_service: Mutex::new(TxObservationService::new()),
82        }
83    }
84
85    /// Take the metadata lock, recovering it if an earlier panic poisoned it.
86    ///
87    /// Nothing that can fail (SQL, projection) runs while the lock is held: readers
88    /// take a [`Conn::snapshot`] and release it, and `InProgress::commit` only swaps
89    /// in already-computed values. So a poisoned lock still guards consistent data.
90    fn lock_metadata(&self) -> std::sync::MutexGuard<'_, Metadata> {
91        self.metadata
92            .lock()
93            .unwrap_or_else(std::sync::PoisonError::into_inner)
94    }
95
96    /// The current schema and attribute cache, both cheap `Arc` clones, taken
97    /// under the lock and released before any query runs. An interrupted or
98    /// failing query therefore never holds (or poisons) the metadata lock.
99    fn snapshot(&self) -> (Arc<Schema>, SQLiteAttributeCache) {
100        let m = self.lock_metadata();
101        (m.schema.clone(), m.attribute_cache.clone())
102    }
103
104    pub fn connect(sqlite: &mut rusqlite::Connection) -> Result<Conn> {
105        let db = db::ensure_current_version(sqlite)?;
106        Ok(Conn::new(db.partition_map, db.schema))
107    }
108
109    /// Yield a clone of the current `Schema` instance.
110    pub fn current_schema(&self) -> Arc<Schema> {
111        // We always unwrap the mutex lock: if it's poisoned, this will propogate panics to all
112        // accessing threads.  This is perhaps not reasonable; we expect the mutex to be held for
113        // very short intervals, but a panic during a critical update section is possible, since the
114        // lock encapsulates committing a SQL transaction.
115        //
116        // That being said, in the future we will provide an interface to take the mutex, providing
117        // maximum flexibility for Mentat consumers.
118        //
119        // This approach might need to change when we support interrupting query threads (#297), and
120        // will definitely need to change if we support interrupting transactor threads.
121        //
122        // Improving this is tracked by https://github.com/mozilla/mentat/issues/356.
123        self.lock_metadata().schema.clone()
124    }
125
126    pub fn current_cache(&self) -> SQLiteAttributeCache {
127        self.lock_metadata().attribute_cache.clone()
128    }
129
130    pub fn last_tx_id(&self) -> Entid {
131        // The mutex is taken during this entire method.
132        let metadata = self.lock_metadata();
133
134        metadata.partition_map[":db.part/tx"].next_entid() - 1
135    }
136
137    /// Query the Mentat store, using the given connection and the current metadata.
138    pub fn q_once<T>(
139        &self,
140        sqlite: &rusqlite::Connection,
141        query: &str,
142        inputs: T,
143    ) -> Result<QueryOutput>
144    where
145        T: Into<Option<QueryInputs>>,
146    {
147        // Doesn't clone, unlike `current_schema`.
148        let (schema, cache) = self.snapshot();
149        let known = Known::new(&schema, Some(&cache));
150        q_once(sqlite, known, query, inputs)
151    }
152
153    /// `q_once`, plus the attributes it filtered by value (see `auto_index`).
154    pub fn q_once_value_filtered<T>(
155        &self,
156        sqlite: &rusqlite::Connection,
157        query: &str,
158        inputs: T,
159    ) -> Result<(QueryOutput, std::collections::BTreeSet<Entid>)>
160    where
161        T: Into<Option<QueryInputs>>,
162    {
163        let (schema, cache) = self.snapshot();
164        let known = Known::new(&schema, Some(&cache));
165        mentat_transaction::query::q_once_value_filtered(sqlite, known, query, inputs)
166    }
167
168    /// Query the store against a historical basis reconstructed as of transaction
169    /// `tx` (inclusive): the value that was current at `tx`.
170    pub fn q_once_as_of<T>(
171        &self,
172        sqlite: &rusqlite::Connection,
173        query: &str,
174        inputs: T,
175        tx: Entid,
176    ) -> Result<QueryOutput>
177    where
178        T: Into<Option<QueryInputs>>,
179    {
180        let basis = mentat_query_algebrizer::TemporalBound::AsOf(tx);
181        self.q_once_temporal(sqlite, query, inputs, basis)
182    }
183
184    /// Query the store seeing only the current datoms transacted after
185    /// transaction `tx` (Datomic's `since`).
186    pub fn q_once_since<T>(
187        &self,
188        sqlite: &rusqlite::Connection,
189        query: &str,
190        inputs: T,
191        tx: Entid,
192    ) -> Result<QueryOutput>
193    where
194        T: Into<Option<QueryInputs>>,
195    {
196        let basis = mentat_query_algebrizer::TemporalBound::Since(tx);
197        self.q_once_temporal(sqlite, query, inputs, basis)
198    }
199
200    /// Query the store on a historical basis (as-of, since, or both).
201    pub fn q_once_temporal<T>(
202        &self,
203        sqlite: &rusqlite::Connection,
204        query: &str,
205        inputs: T,
206        basis: mentat_query_algebrizer::TemporalBound,
207    ) -> Result<QueryOutput>
208    where
209        T: Into<Option<QueryInputs>>,
210    {
211        let (schema, cache) = self.snapshot();
212        let known = Known::new(&schema, Some(&cache));
213        mentat_transaction::query::q_once_temporal(sqlite, known, query, inputs, basis)
214    }
215
216    /// Query the Mentat store, using the given connection and the current metadata,
217    /// but without using the cache.
218    pub fn q_uncached<T>(
219        &self,
220        sqlite: &rusqlite::Connection,
221        query: &str,
222        inputs: T,
223    ) -> Result<QueryOutput>
224    where
225        T: Into<Option<QueryInputs>>,
226    {
227        let (schema, _) = self.snapshot();
228        q_uncached(sqlite, &schema, query, inputs)
229    }
230
231    pub fn q_prepare<'sqlite, T>(
232        &self,
233        sqlite: &'sqlite rusqlite::Connection,
234        query: &str,
235        inputs: T,
236    ) -> PreparedResult<'sqlite>
237    where
238        T: Into<Option<QueryInputs>>,
239    {
240        let (schema, cache) = self.snapshot();
241        let known = Known::new(&schema, Some(&cache));
242        q_prepare(sqlite, known, query, inputs)
243    }
244
245    pub fn q_explain<T>(
246        &self,
247        sqlite: &rusqlite::Connection,
248        query: &str,
249        inputs: T,
250    ) -> Result<QueryExplanation>
251    where
252        T: Into<Option<QueryInputs>>,
253    {
254        let (schema, cache) = self.snapshot();
255        let known = Known::new(&schema, Some(&cache));
256        q_explain(sqlite, known, query, inputs)
257    }
258
259    /// `q_explain` against a historical basis (`None` = current).
260    pub fn q_explain_temporal<T>(
261        &self,
262        sqlite: &rusqlite::Connection,
263        query: &str,
264        inputs: T,
265        temporal: Option<mentat_query_algebrizer::TemporalBound>,
266    ) -> Result<QueryExplanation>
267    where
268        T: Into<Option<QueryInputs>>,
269    {
270        let (schema, cache) = self.snapshot();
271        let known = Known::new(&schema, Some(&cache));
272        mentat_transaction::query::q_explain_temporal(sqlite, known, query, inputs, temporal)
273    }
274
275    pub fn pull_attributes_for_entities<E, A>(
276        &self,
277        sqlite: &rusqlite::Connection,
278        entities: E,
279        attributes: A,
280    ) -> Result<BTreeMap<Entid, ValueRc<StructuredMap>>>
281    where
282        E: IntoIterator<Item = Entid>,
283        A: IntoIterator<Item = Entid>,
284    {
285        let (schema, _) = self.snapshot();
286        pull_attributes_for_entities(&schema, &mentat_sql::sqlite(sqlite), entities, attributes)
287            .map_err(|e| e.into())
288    }
289
290    pub fn pull_attributes_for_entity<A>(
291        &self,
292        sqlite: &rusqlite::Connection,
293        entity: Entid,
294        attributes: A,
295    ) -> Result<StructuredMap>
296    where
297        A: IntoIterator<Item = Entid>,
298    {
299        let (schema, _) = self.snapshot();
300        pull_attributes_for_entity(&schema, &mentat_sql::sqlite(sqlite), entity, attributes)
301            .map_err(|e| e.into())
302    }
303
304    /// Pull `pattern` (EDN, e.g. `[:person/name {:person/friend [*]}]`) for
305    /// `entity` on `basis` (`None` = current, else as-of / since).
306    pub fn pull_on(
307        &self,
308        sqlite: &rusqlite::Connection,
309        entity: Entid,
310        pattern: &str,
311        basis: Option<mentat_query_algebrizer::TemporalBound>,
312    ) -> Result<StructuredMap> {
313        let (schema, _) = self.snapshot();
314        let specs = edn::parse::pull_pattern(pattern)?;
315        Ok(pull_pattern_for_entity(
316            &schema,
317            &mentat_sql::sqlite(sqlite),
318            entity,
319            specs,
320            basis,
321        )?)
322    }
323
324    pub fn lookup_values_for_attribute(
325        &self,
326        sqlite: &rusqlite::Connection,
327        entity: Entid,
328        attribute: &edn::Keyword,
329    ) -> Result<Vec<TypedValue>> {
330        let (schema, cache) = self.snapshot();
331        let known = Known::new(&schema, Some(&cache));
332        lookup_values_for_attribute(sqlite, known, entity, attribute)
333    }
334
335    pub fn lookup_value_for_attribute(
336        &self,
337        sqlite: &rusqlite::Connection,
338        entity: Entid,
339        attribute: &edn::Keyword,
340    ) -> Result<Option<TypedValue>> {
341        let (schema, cache) = self.snapshot();
342        let known = Known::new(&schema, Some(&cache));
343        lookup_value_for_attribute(sqlite, known, entity, attribute)
344    }
345
346    /// Take a SQLite transaction.
347    fn begin_transaction_with_behavior<'m, 'conn>(
348        &'m mut self,
349        sqlite: &'conn mut rusqlite::Connection,
350        behavior: TransactionBehavior,
351    ) -> Result<InProgress<'m, 'conn>> {
352        let tx = sqlite.transaction_with_behavior(behavior)?;
353        let (current_generation, current_partition_map, current_schema, cache_cow) = {
354            // The mutex is taken during this block.
355            let current: &Metadata = &self.lock_metadata();
356            (
357                current.generation,
358                // Expensive, but the partition map is updated after every committed transaction.
359                current.partition_map.clone(),
360                // Cheap.
361                current.schema.clone(),
362                current.attribute_cache.clone(),
363            )
364        };
365
366        Ok(InProgress {
367            mutex: &self.metadata,
368            transaction: tx,
369            generation: current_generation,
370            partition_map: current_partition_map,
371            schema: (*current_schema).clone(),
372            cache: InProgressSQLiteAttributeCache::from_cache(cache_cow),
373            use_caching: true,
374            tx_observer: &self.tx_observer_service,
375            tx_observer_watcher: InProgressObserverTransactWatcher::new(),
376        })
377    }
378
379    // Helper to avoid passing connections around.
380    // Make both args mutable so that we can't have parallel access.
381    pub fn begin_read<'m, 'conn>(
382        &'m mut self,
383        sqlite: &'conn mut rusqlite::Connection,
384    ) -> Result<InProgressRead<'m, 'conn>> {
385        self.begin_transaction_with_behavior(sqlite, TransactionBehavior::Deferred)
386            .map(|ip| InProgressRead { in_progress: ip })
387    }
388
389    pub fn begin_uncached_read<'m, 'conn>(
390        &'m mut self,
391        sqlite: &'conn mut rusqlite::Connection,
392    ) -> Result<InProgressRead<'m, 'conn>> {
393        self.begin_transaction_with_behavior(sqlite, TransactionBehavior::Deferred)
394            .map(|mut ip| {
395                ip.use_caching(false);
396                InProgressRead { in_progress: ip }
397            })
398    }
399
400    /// IMMEDIATE means 'start the transaction now, but don't exclude readers'. It prevents other
401    /// connections from taking immediate or exclusive transactions. This is appropriate for our
402    /// writes and `InProgress`: it means we are ready to write whenever we want to, and nobody else
403    /// can start a transaction that's not `DEFERRED`, but we don't need exclusivity yet.
404    pub fn begin_transaction<'m, 'conn>(
405        &'m mut self,
406        sqlite: &'conn mut rusqlite::Connection,
407    ) -> Result<InProgress<'m, 'conn>> {
408        self.begin_transaction_with_behavior(sqlite, TransactionBehavior::Immediate)
409    }
410
411    /// Transact entities against the Mentat store, using the given connection and the current
412    /// metadata.
413    pub fn transact<B>(
414        &mut self,
415        sqlite: &mut rusqlite::Connection,
416        transaction: B,
417    ) -> Result<TxReport>
418    where
419        B: Borrow<str>,
420    {
421        // Parse outside the SQL transaction. This is a tradeoff: we are limiting the scope of the
422        // transaction, and indeed we don't even create a SQL transaction if the provided input is
423        // invalid, but it means SQLite errors won't be found until the parse is complete, and if
424        // there's a race for the database (don't do that!) we are less likely to win it.
425        let entities = edn::parse::entities(transaction.borrow())?;
426
427        let mut in_progress = self.begin_transaction(sqlite)?;
428        let report = in_progress.transact_entities(entities)?;
429        in_progress.commit()?;
430
431        Ok(report)
432    }
433
434    /// Speculatively transact `transaction` and return its `TxReport` without
435    /// persisting anything (Datomic `d/with` / pg_mentat `mentat.with`). The
436    /// enclosing SQLite transaction is rolled back on drop, so no changes are
437    /// committed even though a real IMMEDIATE transaction was opened.
438    pub fn transact_speculative<B>(
439        &mut self,
440        sqlite: &mut rusqlite::Connection,
441        transaction: B,
442    ) -> Result<TxReport>
443    where
444        B: Borrow<str>,
445    {
446        let entities = edn::parse::entities(transaction.borrow())?;
447
448        let mut in_progress = self.begin_transaction(sqlite)?;
449        let report = in_progress.transact_entities_speculative(entities)?;
450        // Deliberately do NOT commit: drop rolls back the outer transaction.
451        drop(in_progress);
452
453        Ok(report)
454    }
455
456    /// Adds or removes the values of a given attribute to an in-memory cache.
457    /// The attribute should be a namespaced string: e.g., `:foo/bar`.
458    /// `cache_action` determines if the attribute should be added or removed from the cache.
459    /// CacheAction::Add is idempotent - each attribute is only added once.
460    /// CacheAction::Remove throws an error if the attribute does not currently exist in the cache.
461    pub fn cache(
462        &mut self,
463        sqlite: &mut rusqlite::Connection,
464        schema: &Schema,
465        attribute: &Keyword,
466        cache_direction: CacheDirection,
467        cache_action: CacheAction,
468    ) -> Result<()> {
469        let mut metadata = self.lock_metadata();
470        let attribute_entid: Entid;
471
472        // Immutable borrow of metadata.
473        {
474            attribute_entid = metadata
475                .schema
476                .attribute_for_ident(attribute)
477                .ok_or_else(|| MentatError::UnknownAttribute(attribute.to_string()))?
478                .1
479                .into();
480        }
481
482        let cache = &mut metadata.attribute_cache;
483        match cache_action {
484            CacheAction::Register => match cache_direction {
485                CacheDirection::Both => cache.register(schema, sqlite, attribute_entid),
486                CacheDirection::Forward => cache.register_forward(schema, sqlite, attribute_entid),
487                CacheDirection::Reverse => cache.register_reverse(schema, sqlite, attribute_entid),
488            }
489            .map_err(|e| e.into()),
490            CacheAction::Deregister => {
491                cache.unregister(attribute_entid);
492                Ok(())
493            }
494        }
495    }
496
497    pub fn register_observer(&mut self, key: String, observer: Arc<TxObserver>) {
498        self.tx_observer_service
499            .lock()
500            .unwrap()
501            .register(key, observer);
502    }
503
504    pub fn unregister_observer(&mut self, key: &str) {
505        self.tx_observer_service.lock().unwrap().deregister(key);
506    }
507}
508
509#[cfg(test)]
510mod tests {
511    use super::*;
512
513    use std::time::Instant;
514
515    use core_traits::{Binding, TypedValue};
516
517    use mentat_core::CachedAttributes;
518
519    use mentat_transaction::query::Variable;
520
521    use mentat_transaction::query::{IntoResult, QueryInputs, QueryResults};
522
523    use mentat_db::USER0;
524
525    use mentat_transaction::Queryable;
526
527    #[test]
528    fn test_transact_does_not_collide_existing_entids() {
529        let mut sqlite = db::new_connection("").unwrap();
530        let mut conn = Conn::connect(&mut sqlite).unwrap();
531
532        // Let's find out the next ID that'll be allocated. We're going to try to collide with it
533        // a bit later.
534        let next =
535            conn.metadata.lock().expect("metadata").partition_map[":db.part/user"].next_entid();
536        let t = format!("[[:db/add {} :db.schema/attribute \"tempid\"]]", next + 1);
537
538        match conn.transact(&mut sqlite, t.as_str()) {
539            Err(MentatError::DbError(e)) => {
540                assert_eq!(
541                    e.kind(),
542                    ::db_traits::errors::DbErrorKind::UnallocatedEntid(next + 1)
543                );
544            }
545            x => panic!("expected db error, got {:?}", x),
546        }
547
548        // Transact two more tempids.
549        let t = "[[:db/add \"one\" :db.schema/attribute \"more\"]]";
550        let report = conn.transact(&mut sqlite, t).expect("transact succeeded");
551        assert_eq!(report.tempids["more"], next);
552        assert_eq!(report.tempids["one"], next + 1);
553    }
554
555    #[test]
556    fn test_transact_does_not_collide_new_entids() {
557        let mut sqlite = db::new_connection("").unwrap();
558        let mut conn = Conn::connect(&mut sqlite).unwrap();
559
560        // Let's find out the next ID that'll be allocated. We're going to try to collide with it.
561        let next =
562            conn.metadata.lock().expect("metadata").partition_map[":db.part/user"].next_entid();
563
564        // If this were to be resolved, we'd get [:db/add 65537 :db.schema/attribute 65537], but
565        // we should reject this, because the first ID was provided by the user!
566        let t = format!("[[:db/add {} :db.schema/attribute \"tempid\"]]", next);
567
568        match conn.transact(&mut sqlite, t.as_str()) {
569            Err(MentatError::DbError(e)) => {
570                // All this, despite this being the ID we were about to allocate!
571                assert_eq!(
572                    e.kind(),
573                    ::db_traits::errors::DbErrorKind::UnallocatedEntid(next)
574                );
575            }
576            x => panic!("expected db error, got {:?}", x),
577        }
578
579        // And if we subsequently transact in a way that allocates one ID, we _will_ use that one.
580        // Note that `10` is a bootstrapped entid; we use it here as a known-good value.
581        let t = "[[:db/add 10 :db.schema/attribute \"temp\"]]";
582        let report = conn.transact(&mut sqlite, t).expect("transact succeeded");
583        assert_eq!(report.tempids["temp"], next);
584    }
585
586    /// Return the entid that will be allocated to the next transacted tempid.
587    fn get_next_entid(conn: &Conn) -> Entid {
588        let partition_map = &conn.metadata.lock().unwrap().partition_map;
589        partition_map.get(":db.part/user").unwrap().next_entid()
590    }
591
592    #[test]
593    fn test_compound_transact() {
594        let mut sqlite = db::new_connection("").unwrap();
595        let mut conn = Conn::connect(&mut sqlite).unwrap();
596
597        let tempid_offset = get_next_entid(&conn);
598
599        let t = "[[:db/add \"one\" :db/ident :a/keyword1] \
600                 [:db/add \"two\" :db/ident :a/keyword2]]";
601
602        // This can refer to `t`, 'cos they occur in separate txes.
603        let t2 = "[{:db.schema/attribute \"three\", :db/ident :a/keyword1}]";
604
605        // Scoped borrow of `conn`.
606        {
607            let mut in_progress = conn
608                .begin_transaction(&mut sqlite)
609                .expect("begun successfully");
610            let report = in_progress.transact(t).expect("transacted successfully");
611            let one = *report.tempids.get("one").expect("found one");
612            let two = *report.tempids.get("two").expect("found two");
613            assert!(one != two);
614            assert!(one == tempid_offset || one == tempid_offset + 1);
615            assert!(two == tempid_offset || two == tempid_offset + 1);
616
617            println!(
618                "RES: {:?}",
619                in_progress
620                    .q_once("[:find ?v :where [?x :db/ident ?v]]", None)
621                    .unwrap()
622            );
623
624            let during = in_progress
625                .q_once("[:find ?x . :where [?x :db/ident :a/keyword1]]", None)
626                .expect("query succeeded");
627            assert_eq!(
628                during.results,
629                QueryResults::Scalar(Some(TypedValue::Ref(one).into()))
630            );
631
632            let report = in_progress.transact(t2).expect("t2 succeeded");
633            in_progress.commit().expect("commit succeeded");
634            let three = *report.tempids.get("three").expect("found three");
635            assert!(one != three);
636            assert!(two != three);
637        }
638
639        // The DB part table changed.
640        let tempid_offset_after = get_next_entid(&conn);
641        assert_eq!(tempid_offset + 3, tempid_offset_after);
642    }
643
644    #[test]
645    fn test_transact_speculative() {
646        let mut sqlite = db::new_connection("").unwrap();
647        let mut conn = Conn::connect(&mut sqlite).unwrap();
648
649        let before = get_next_entid(&conn);
650
651        // Speculatively add an ident. The report should reflect the change...
652        let t = "[[:db/add \"e\" :db/ident :spec/keyword]]";
653        let report = conn
654            .transact_speculative(&mut sqlite, t)
655            .expect("speculative transact succeeded");
656        let e = *report.tempids.get("e").expect("tempid resolved");
657        assert_eq!(e, before, "speculative report allocates the next entid");
658
659        // ...but nothing is persisted: the ident is not in the store.
660        let found = conn
661            .q_once(
662                &mut sqlite,
663                "[:find ?x . :where [?x :db/ident :spec/keyword]]",
664                None,
665            )
666            .expect("query succeeded");
667        assert_eq!(
668            found.results,
669            QueryResults::Scalar(None),
670            "speculative datom must not persist"
671        );
672
673        // ...and the partition counter was rolled back, so a *real* transact
674        // reuses the same entid the speculative run reported.
675        assert_eq!(get_next_entid(&conn), before, "entid not consumed");
676        let real = conn
677            .transact(&mut sqlite, t)
678            .expect("real transact succeeded");
679        assert_eq!(*real.tempids.get("e").expect("tempid"), before);
680        assert_eq!(get_next_entid(&conn), before + 1);
681    }
682
683    #[test]
684    fn test_simple_prepared_query() {
685        let mut c = db::new_connection("").expect("Couldn't open conn.");
686        let mut conn = Conn::connect(&mut c).expect("Couldn't open DB.");
687        conn.transact(
688            &mut c,
689            r#"[
690            [:db/add "s" :db/ident :foo/boolean]
691            [:db/add "s" :db/valueType :db.type/boolean]
692            [:db/add "s" :db/cardinality :db.cardinality/one]
693        ]"#,
694        )
695        .expect("successful transaction");
696
697        let report = conn
698            .transact(
699                &mut c,
700                r#"[
701            [:db/add "u" :foo/boolean true]
702            [:db/add "p" :foo/boolean false]
703        ]"#,
704            )
705            .expect("successful transaction");
706        let yes = *report.tempids.get("u").expect("found it");
707
708        let vv = Variable::from_valid_name("?v");
709
710        let values = QueryInputs::with_value_sequence(vec![(vv, true.into())]);
711
712        let read = conn.begin_read(&mut c).expect("read");
713
714        // N.B., you might choose to algebrize _without_ validating that the
715        // types are known. In this query we know that `?v` must be a boolean,
716        // and so we can kinda generate our own required input types!
717        let mut prepared = read
718            .q_prepare(
719                r#"[:find [?x ...]
720                                              :in ?v
721                                              :where [?x :foo/boolean ?v]]"#,
722                values,
723            )
724            .expect("prepare succeeded");
725
726        let yeses = prepared.run(None).expect("result");
727        assert_eq!(
728            yeses.results,
729            QueryResults::Coll(vec![TypedValue::Ref(yes).into()])
730        );
731
732        let yeses_again = prepared.run(None).expect("result");
733        assert_eq!(
734            yeses_again.results,
735            QueryResults::Coll(vec![TypedValue::Ref(yes).into()])
736        );
737    }
738
739    #[test]
740    fn test_compound_rollback() {
741        let mut sqlite = db::new_connection("").unwrap();
742        let mut conn = Conn::connect(&mut sqlite).unwrap();
743
744        let tempid_offset = get_next_entid(&conn);
745
746        // Nothing in the store => USER0 should be our starting point.
747        assert_eq!(tempid_offset, USER0);
748
749        let t = "[[:db/add \"one\" :db/ident :a/keyword1] \
750                 [:db/add \"two\" :db/ident :a/keyword2]]";
751
752        // Scoped borrow of `sqlite`.
753        {
754            let mut in_progress = conn
755                .begin_transaction(&mut sqlite)
756                .expect("begun successfully");
757            let report = in_progress.transact(t).expect("transacted successfully");
758
759            let one = *report.tempids.get("one").expect("found it");
760            let two = *report.tempids.get("two").expect("found it");
761
762            // The IDs are contiguous, starting at the previous part index.
763            assert!(one != two);
764            assert!(one == tempid_offset || one == tempid_offset + 1);
765            assert!(two == tempid_offset || two == tempid_offset + 1);
766
767            // Inside the InProgress we can see our changes.
768            let during = in_progress
769                .q_once("[:find ?x . :where [?x :db/ident :a/keyword1]]", None)
770                .expect("query succeeded");
771
772            assert_eq!(
773                during.results,
774                QueryResults::Scalar(Some(TypedValue::Ref(one).into()))
775            );
776
777            // And we can do direct lookup, too.
778            let kw = in_progress
779                .lookup_value_for_attribute(one, &edn::Keyword::namespaced("db", "ident"))
780                .expect("lookup succeeded");
781            assert_eq!(
782                kw,
783                Some(TypedValue::Keyword(
784                    edn::Keyword::namespaced("a", "keyword1").into()
785                ))
786            );
787
788            in_progress.rollback().expect("rollback succeeded");
789        }
790
791        let after = conn
792            .q_once(
793                &sqlite,
794                "[:find ?x . :where [?x :db/ident :a/keyword1]]",
795                None,
796            )
797            .expect("query succeeded");
798        assert_eq!(after.results, QueryResults::Scalar(None));
799
800        // The DB part table is unchanged.
801        let tempid_offset_after = get_next_entid(&conn);
802        assert_eq!(tempid_offset, tempid_offset_after);
803    }
804
805    #[test]
806    fn test_transact_errors() {
807        let mut sqlite = db::new_connection("").unwrap();
808        let mut conn = Conn::connect(&mut sqlite).unwrap();
809
810        // Good: empty transaction.
811        let report = conn.transact(&mut sqlite, "[]").unwrap();
812        assert_eq!(report.tx_id, 0x10000000 + 1);
813
814        // Bad EDN: missing closing ']'.
815        let report = conn.transact(&mut sqlite, "[[:db/add \"t\" :db/ident :a/keyword]");
816        match report.expect_err("expected transact to fail for bad edn") {
817            MentatError::EdnParseError(_) => {}
818            x => panic!("expected EDN parse error, got {:?}", x),
819        }
820
821        // Good EDN.
822        let report = conn
823            .transact(&mut sqlite, "[[:db/add \"t\" :db/ident :a/keyword]]")
824            .unwrap();
825        assert_eq!(report.tx_id, 0x10000000 + 2);
826
827        // Bad transaction data: missing leading :db/add.
828        let report = conn.transact(&mut sqlite, "[[\"t\" :db/ident :b/keyword]]");
829        match report.expect_err("expected transact error") {
830            MentatError::EdnParseError(_) => {}
831            x => panic!("expected EDN parse error, got {:?}", x),
832        }
833
834        // Good transaction data.
835        let report = conn
836            .transact(&mut sqlite, "[[:db/add \"u\" :db/ident :b/keyword]]")
837            .unwrap();
838        assert_eq!(report.tx_id, 0x10000000 + 3);
839
840        // Bad transaction based on state of store: conflicting upsert.
841        let report = conn.transact(
842            &mut sqlite,
843            "[[:db/add \"u\" :db/ident :a/keyword]
844                                                  [:db/add \"u\" :db/ident :b/keyword]]",
845        );
846        match report.expect_err("expected transact error") {
847            MentatError::DbError(e) => match e.kind() {
848                ::db_traits::errors::DbErrorKind::SchemaConstraintViolation(_) => {}
849                _ => panic!("expected SchemaConstraintViolation"),
850            },
851            x => panic!("expected db error, got {:?}", x),
852        }
853    }
854
855    #[test]
856    fn test_add_to_cache_failure_no_attribute() {
857        let mut sqlite = db::new_connection("").unwrap();
858        let mut conn = Conn::connect(&mut sqlite).unwrap();
859        let _report = conn
860            .transact(
861                &mut sqlite,
862                r#"[
863            {  :db/ident       :foo/bar
864               :db/valueType   :db.type/long
865               :db/cardinality :db.cardinality/one },
866            {  :db/ident       :foo/baz
867               :db/valueType   :db.type/boolean
868               :db/cardinality :db.cardinality/one }]"#,
869            )
870            .unwrap();
871
872        let kw = kw!(:foo/bat);
873        let schema = conn.current_schema();
874        let res = conn.cache(
875            &mut sqlite,
876            &schema,
877            &kw,
878            CacheDirection::Forward,
879            CacheAction::Register,
880        );
881        match res.expect_err("expected cache to fail") {
882            MentatError::UnknownAttribute(msg) => assert_eq!(msg, ":foo/bat"),
883            x => panic!("expected UnknownAttribute error, got {:?}", x),
884        }
885    }
886
887    // TODO expand tests to cover lookup_value_for_attribute comparing with and without caching
888    #[test]
889    fn test_lookup_attribute_with_caching() {
890        let mut sqlite = db::new_connection("").unwrap();
891        let mut conn = Conn::connect(&mut sqlite).unwrap();
892        let _report = conn
893            .transact(
894                &mut sqlite,
895                r#"[
896            {  :db/ident       :foo/bar
897               :db/valueType   :db.type/long
898               :db/cardinality :db.cardinality/one },
899            {  :db/ident       :foo/baz
900               :db/valueType   :db.type/boolean
901               :db/cardinality :db.cardinality/one }]"#,
902            )
903            .expect("transaction expected to succeed");
904
905        {
906            let mut in_progress = conn.begin_transaction(&mut sqlite).expect("transaction");
907            for _ in 1..100 {
908                let _report = in_progress
909                    .transact(
910                        r#"[
911            {  :foo/bar        100
912               :foo/baz        false },
913            {  :foo/bar        200
914               :foo/baz        true },
915            {  :foo/bar        100
916               :foo/baz        false },
917            {  :foo/bar        300
918               :foo/baz        true },
919            {  :foo/bar        400
920               :foo/baz        false },
921            {  :foo/bar        500
922               :foo/baz        true }]"#,
923                    )
924                    .expect("transaction expected to succeed");
925            }
926            in_progress.commit().expect("Committed");
927        }
928
929        let entities = conn
930            .q_once(&sqlite, r#"[:find ?e . :where [?e :foo/bar 400]]"#, None)
931            .expect("Expected query to work")
932            .into_scalar()
933            .expect("expected rel results");
934        let first = entities.expect("expected a result");
935        let entid = match first {
936            Binding::Scalar(TypedValue::Ref(entid)) => entid,
937            x => panic!("expected Some(Ref), got {:?}", x),
938        };
939
940        let kw = kw!(:foo/bar);
941        let start = Instant::now();
942        let uncached_val = conn
943            .lookup_value_for_attribute(&sqlite, entid, &kw)
944            .expect("Expected value on lookup");
945        let finish = Instant::now();
946        let uncached_elapsed_time = finish.duration_since(start);
947        println!("Uncached time: {:?}", uncached_elapsed_time);
948
949        let schema = conn.current_schema();
950        conn.cache(
951            &mut sqlite,
952            &schema,
953            &kw,
954            CacheDirection::Forward,
955            CacheAction::Register,
956        )
957        .expect("expected caching to work");
958
959        for _ in 1..5 {
960            let start = Instant::now();
961            let cached_val = conn
962                .lookup_value_for_attribute(&sqlite, entid, &kw)
963                .expect("Expected value on lookup");
964            let finish = Instant::now();
965            let cached_elapsed_time = finish.duration_since(start);
966            assert_eq!(cached_val, uncached_val);
967
968            println!("Cached time: {:?}", cached_elapsed_time);
969            assert!(cached_elapsed_time < uncached_elapsed_time);
970        }
971    }
972
973    #[test]
974    fn test_cache_usage() {
975        let mut sqlite = db::new_connection("").unwrap();
976        let mut conn = Conn::connect(&mut sqlite).unwrap();
977
978        let db_ident = (*conn.current_schema())
979            .get_entid(&kw!(:db/ident))
980            .expect("db_ident")
981            .0;
982        let db_type = (*conn.current_schema())
983            .get_entid(&kw!(:db/valueType))
984            .expect("db_ident")
985            .0;
986        println!("db/ident is {}", db_ident);
987        println!("db/type is {}", db_type);
988        let query = format!(
989            "[:find ?ident . :where [?e {} :db/doc][?e {} ?type][?type {} ?ident]]",
990            db_ident, db_type, db_ident
991        );
992
993        println!("Query is {}", query);
994
995        assert!(!conn.current_cache().is_attribute_cached_forward(db_ident));
996
997        {
998            let mut ip = conn.begin_transaction(&mut sqlite).expect("began");
999
1000            let ident = ip
1001                .q_once(query.as_str(), None)
1002                .into_scalar_result()
1003                .expect("query");
1004            assert_eq!(
1005                ident,
1006                Some(TypedValue::typed_ns_keyword("db.type", "string").into())
1007            );
1008
1009            let start = std::time::Instant::now();
1010            ip.q_once(query.as_str(), None)
1011                .into_scalar_result()
1012                .expect("query");
1013            let end = std::time::Instant::now();
1014            println!("Uncached took {}µs", (end - start).as_micros());
1015
1016            ip.cache(
1017                &kw!(:db/ident),
1018                CacheDirection::Forward,
1019                CacheAction::Register,
1020            )
1021            .expect("registered");
1022            ip.cache(
1023                &kw!(:db/valueType),
1024                CacheDirection::Forward,
1025                CacheAction::Register,
1026            )
1027            .expect("registered");
1028
1029            assert!(ip.cache.is_attribute_cached_forward(db_ident));
1030
1031            let ident = ip
1032                .q_once(query.as_str(), None)
1033                .into_scalar_result()
1034                .expect("query");
1035            assert_eq!(
1036                ident,
1037                Some(TypedValue::typed_ns_keyword("db.type", "string").into())
1038            );
1039
1040            let start = std::time::Instant::now();
1041            ip.q_once(query.as_str(), None)
1042                .into_scalar_result()
1043                .expect("query");
1044            let end = std::time::Instant::now();
1045            println!("Cached took {}µs", (end - start).as_micros());
1046
1047            // If we roll back the change, our caching operations are also rolled back.
1048            ip.rollback().expect("rolled back");
1049        }
1050
1051        assert!(!conn.current_cache().is_attribute_cached_forward(db_ident));
1052
1053        {
1054            let mut ip = conn.begin_transaction(&mut sqlite).expect("began");
1055
1056            let ident = ip
1057                .q_once(query.as_str(), None)
1058                .into_scalar_result()
1059                .expect("query");
1060            assert_eq!(
1061                ident,
1062                Some(TypedValue::typed_ns_keyword("db.type", "string").into())
1063            );
1064            ip.cache(
1065                &kw!(:db/ident),
1066                CacheDirection::Forward,
1067                CacheAction::Register,
1068            )
1069            .expect("registered");
1070            ip.cache(
1071                &kw!(:db/valueType),
1072                CacheDirection::Forward,
1073                CacheAction::Register,
1074            )
1075            .expect("registered");
1076
1077            assert!(ip.cache.is_attribute_cached_forward(db_ident));
1078
1079            ip.commit().expect("rolled back");
1080        }
1081
1082        assert!(conn.current_cache().is_attribute_cached_forward(db_ident));
1083        assert!(conn.current_cache().is_attribute_cached_forward(db_type));
1084    }
1085}