mentat/

store.rs

// Copyright 2016-2018 Mozilla
//
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use
// this file except in compliance with the License. You may obtain a copy of the
// License at http://www.apache.org/licenses/LICENSE-2.0
// Unless required by applicable law or agreed to in writing, software distributed
// under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
// CONDITIONS OF ANY KIND, either express or implied. See the License for the
// specific language governing permissions and limitations under the License.

#![allow(dead_code)]

use std::collections::BTreeMap;

use std::sync::Arc;

use core_traits::{Entid, StructuredMap, TypedValue};

use mentat_core::{HasSchema, Keyword, TxReport, ValueRc};
use mentat_db::TxObserver;

use mentat_transaction::{
    CacheAction, CacheDirection, InProgress, InProgressRead, Pullable, Queryable,
};

use crate::conn::Conn;

use public_traits::errors::Result;

use mentat_transaction::query::{PreparedResult, QueryExplanation, QueryInputs, QueryOutput};

/// A convenience wrapper around a single SQLite connection and a Conn. This is suitable
/// for applications that don't require complex connection management.
/// The `:db/fn` attribute: the mino source of a user transaction function.
const DB_FN_ATTRIBUTE: &str =
    "[{:db/ident :db/fn :db/valueType :db.type/string :db/cardinality :db.cardinality/one}]";

/// The live `:db/fn` source on the entity whose `:db/ident` is `ident`, if any.
#[cfg(feature = "mino")]
pub(crate) fn db_fn_source(
    store: &Store,
    ident: &str,
) -> std::result::Result<Option<String>, String> {
    let db_fn = Keyword::namespaced("db", "fn");
    let schema = store.conn.current_schema();
    if schema.get_entid(&db_fn).is_none() {
        return Ok(None);
    }
    let Ok(edn::Value::Keyword(kw)) = edn::parse::value(ident).map(|v| v.without_spans()) else {
        return Ok(None);
    };
    let Some(e) = schema.get_entid(&kw) else {
        return Ok(None); // not an ident at all
    };
    match store
        .lookup_value_for_attribute(e.0, &db_fn)
        .map_err(|err| err.to_string())?
    {
        Some(TypedValue::String(s)) => Ok(Some(s.as_ref().clone())),
        _ => Ok(None),
    }
}

pub struct Store {
    conn: Conn,
    sqlite: rusqlite::Connection,
    advisor: std::cell::RefCell<crate::auto_index::Advisor>,
}

impl Store {
    /// Open a store at the supplied path, ensuring that it includes the bootstrap schema.
    pub fn open(path: &str) -> Result<Store> {
        let mut connection = crate::new_connection(path)?;
        let conn = Conn::connect(&mut connection)?;
        Ok(Store {
            conn,
            sqlite: connection,
            advisor: Default::default(),
        })
    }

    pub fn transact(&mut self, transaction: &str) -> Result<TxReport> {
        let transaction = self.expand_tx_fns(transaction)?;
        let mut ip = self.begin_transaction()?;
        let report = ip.transact(transaction.as_str())?;
        ip.commit()?;
        Ok(report)
    }

    /// Prepare `transaction` for the transactor: replace every invocation of a
    /// user-defined transaction function (`[:my/fn arg ...]` on an ident that
    /// carries `:db/fn`) with the tx-data it returns, and define the `:db/fn`
    /// attribute itself the first time a transaction installs a function.
    /// [`Store::transact`] and [`Store::with_speculative`] call this; a host
    /// that drives [`InProgress::transact`] directly calls it first.
    ///
    /// Functions run in the sandboxed mino interpreter (`mino` feature), see
    /// the database as of the start of the transaction, and nest at most
    /// `MENTAT_TXFN_MAX_DEPTH` (default 16) deep. Without the `mino` feature
    /// only the `:db/fn` attribute bootstrap applies.
    pub fn expand_tx_fns(&mut self, transaction: &str) -> Result<String> {
        let has_db_fn_attr = self
            .conn
            .current_schema()
            .get_entid(&Keyword::namespaced("db", "fn"))
            .is_some();
        if !has_db_fn_attr && transaction.contains(":db/fn") {
            // The embedded core schema predates transaction functions; define
            // `:db/fn` alongside the first function installed, rather than
            // changing every store's bootstrap.
            self.transact_raw(DB_FN_ATTRIBUTE)?;
        }
        // Only a store that has (or is about to have) `:db/fn` can hold a
        // transaction function, so other stores skip the interpreter entirely.
        #[cfg(feature = "mino")]
        if has_db_fn_attr || transaction.contains(":db/fn") {
            // ponytail: swaps in a throwaway in-memory Store (~ms) to move this
            // one into the script backend; only paid by transactions that could
            // invoke a function. Borrowing it instead needs a non-'static
            // ScriptBackend.
            let store = std::mem::replace(self, Store::open("")?);
            let (store, expanded) = crate::script::expand_tx_fns(store, transaction);
            *self = store;
            return expanded.map_err(public_traits::errors::MentatError::TransactionFunction);
        }
        Ok(transaction.to_string())
    }

    /// Transact without transaction-function expansion.
    fn transact_raw(&mut self, transaction: &str) -> Result<TxReport> {
        let mut ip = self.begin_transaction()?;
        let report = ip.transact(transaction)?;
        ip.commit()?;
        Ok(report)
    }

    /// Speculatively transact `transaction` and return its `TxReport` WITHOUT
    /// committing (Datomic `d/with`). Delegates to [`Conn::transact_speculative`],
    /// which opens a real SQLite transaction and rolls it back on drop, so the
    /// store is left unchanged. This is the thin `&mut Store` wrapper the
    /// scripting layer needs, since `transact_speculative` lives on `Conn` and
    /// wants both the `Conn` and the `rusqlite::Connection` (both private here).
    pub fn with_speculative(&mut self, transaction: &str) -> Result<TxReport> {
        let transaction = self.expand_tx_fns(transaction)?;
        self.conn
            .transact_speculative(&mut self.sqlite, transaction.as_str())
    }
}

#[cfg(feature = "sqlcipher")]
impl Store {
    /// Variant of `open` that allows a key (for encryption/decryption) to be
    /// supplied. Fails unless linked against sqlcipher (or something else that
    /// supports the Sqlite Encryption Extension).
    pub fn open_with_key(path: &str, encryption_key: &str) -> Result<Store> {
        let mut connection = crate::new_connection_with_key(path, encryption_key)?;
        let conn = Conn::connect(&mut connection)?;
        Ok(Store {
            conn,
            sqlite: connection,
            advisor: Default::default(),
        })
    }

    /// Change the key for a database that was opened using `open_with_key` (using `PRAGMA
    /// rekey`). Fails unless linked against sqlcipher (or something else that supports the Sqlite
    /// Encryption Extension).
    pub fn change_encryption_key(&mut self, new_encryption_key: &str) -> Result<()> {
        crate::change_encryption_key(&self.sqlite, new_encryption_key)?;
        Ok(())
    }
}

impl Store {
    /// Intended for use from tests.
    pub fn sqlite_mut(&mut self) -> &mut rusqlite::Connection {
        &mut self.sqlite
    }

    /// A shared borrow of the underlying SQLite connection, for read-only
    /// helpers (e.g. `mentat_db::debug::datoms`) that only need `&Connection`.
    pub fn sqlite_ref(&self) -> &rusqlite::Connection {
        &self.sqlite
    }

    #[cfg(test)]
    pub fn is_registered_as_observer(&self, key: &str) -> bool {
        self.conn
            .tx_observer_service
            .lock()
            .unwrap()
            .is_registered(key)
    }
}

impl Store {
    pub fn dismantle(self) -> (rusqlite::Connection, Conn) {
        (self.sqlite, self.conn)
    }

    pub fn conn(&self) -> &Conn {
        &self.conn
    }

    pub fn begin_read(&mut self) -> Result<InProgressRead<'_, '_>> {
        self.conn.begin_read(&mut self.sqlite)
    }

    pub fn begin_transaction(&mut self) -> Result<InProgress<'_, '_>> {
        self.conn.begin_transaction(&mut self.sqlite)
    }

    pub fn cache(&mut self, attr: &Keyword, direction: CacheDirection) -> Result<()> {
        let schema = &self.conn.current_schema();
        self.conn.cache(
            &mut self.sqlite,
            schema,
            attr,
            direction,
            CacheAction::Register,
        )
    }

    pub fn register_observer(&mut self, key: String, observer: Arc<TxObserver>) {
        self.conn.register_observer(key, observer);
    }

    pub fn unregister_observer(&mut self, key: &str) {
        self.conn.unregister_observer(key);
    }

    pub fn last_tx_id(&self) -> Entid {
        self.conn.last_tx_id()
    }

    /// How mentat manages secondary indexes for this `Store` (see [`AutoIndex`]).
    /// The default comes from `MENTAT_AUTO_INDEX` (`off`/`schema`/`adaptive`),
    /// else `Schema`.
    ///
    /// [`AutoIndex`]: crate::AutoIndex
    pub fn set_auto_index(&mut self, mode: crate::AutoIndex) {
        self.advisor.get_mut().mode = mode;
    }

    /// Adaptive thresholds: create an index once an attribute is value-filtered
    /// by `min_uses` queries within one tuning period; drop a mentat-created
    /// index unused for a whole period and for at least `idle`.
    pub fn set_index_tuning(&mut self, min_uses: u32, idle: std::time::Duration) {
        let a = self.advisor.get_mut();
        a.min_uses = min_uses.max(1);
        a.idle = idle;
    }

    /// Apply (or, with `dry_run`, only report) index changes per the current
    /// [`AutoIndex`](crate::AutoIndex) mode. Only ever drops indexes mentat
    /// created (listed in `mentat_managed_indexes`).
    pub fn tune_indexes(&mut self, dry_run: bool) -> Result<Vec<crate::IndexAction>> {
        self.tune(dry_run)
    }

    fn tune(&self, dry_run: bool) -> Result<Vec<crate::IndexAction>> {
        let schema = self.conn.current_schema();
        self.advisor
            .borrow_mut()
            .tune(&self.sqlite, &schema, dry_run)
    }

    /// Run a Datalog query against a historical basis reconstructed as of
    /// transaction `tx` (inclusive) -- the state that was current at `tx`.
    pub fn q_once_as_of<T>(&self, query: &str, inputs: T, tx: Entid) -> Result<QueryOutput>
    where
        T: Into<Option<QueryInputs>>,
    {
        self.conn.q_once_as_of(&self.sqlite, query, inputs, tx)
    }

    /// `EXPLAIN QUERY PLAN` for a query, optionally against a historical basis.
    pub fn q_explain_temporal<T>(
        &self,
        query: &str,
        inputs: T,
        temporal: Option<crate::TemporalBound>,
    ) -> Result<QueryExplanation>
    where
        T: Into<Option<QueryInputs>>,
    {
        self.conn
            .q_explain_temporal(&self.sqlite, query, inputs, temporal)
    }

    /// Run a Datalog query seeing only datoms transacted after transaction `tx`
    /// (`since`). Typically used with a history pattern `[?e ?a ?v ?tx ?added]`.
    pub fn q_once_since<T>(&self, query: &str, inputs: T, tx: Entid) -> Result<QueryOutput>
    where
        T: Into<Option<QueryInputs>>,
    {
        self.conn.q_once_since(&self.sqlite, query, inputs, tx)
    }
}

impl Queryable for Store {
    fn q_once<T>(&self, query: &str, inputs: T) -> Result<QueryOutput>
    where
        T: Into<Option<QueryInputs>>,
    {
        if self.advisor.borrow().mode != crate::AutoIndex::Adaptive {
            return self.conn.q_once(&self.sqlite, query, inputs);
        }
        let (out, attrs) = self
            .conn
            .q_once_value_filtered(&self.sqlite, query, inputs)?;
        if self.advisor.borrow_mut().record(&attrs) {
            // Best effort: e.g. another connection may hold the write lock.
            if let Err(e) = self.tune(false) {
                log::warn!("automatic index tuning failed: {e}");
            }
        }
        Ok(out)
    }

    fn q_prepare<T>(&self, query: &str, inputs: T) -> PreparedResult<'_>
    where
        T: Into<Option<QueryInputs>>,
    {
        self.conn.q_prepare(&self.sqlite, query, inputs)
    }

    fn q_explain<T>(&self, query: &str, inputs: T) -> Result<QueryExplanation>
    where
        T: Into<Option<QueryInputs>>,
    {
        self.conn.q_explain(&self.sqlite, query, inputs)
    }

    fn lookup_values_for_attribute<E>(
        &self,
        entity: E,
        attribute: &edn::Keyword,
    ) -> Result<Vec<TypedValue>>
    where
        E: Into<Entid>,
    {
        self.conn
            .lookup_values_for_attribute(&self.sqlite, entity.into(), attribute)
    }

    fn lookup_value_for_attribute<E>(
        &self,
        entity: E,
        attribute: &edn::Keyword,
    ) -> Result<Option<TypedValue>>
    where
        E: Into<Entid>,
    {
        self.conn
            .lookup_value_for_attribute(&self.sqlite, entity.into(), attribute)
    }
}

impl Pullable for Store {
    fn pull_attributes_for_entities<E, A>(
        &self,
        entities: E,
        attributes: A,
    ) -> Result<BTreeMap<Entid, ValueRc<StructuredMap>>>
    where
        E: IntoIterator<Item = Entid>,
        A: IntoIterator<Item = Entid>,
    {
        self.conn
            .pull_attributes_for_entities(&self.sqlite, entities, attributes)
    }

    fn pull_attributes_for_entity<A>(&self, entity: Entid, attributes: A) -> Result<StructuredMap>
    where
        A: IntoIterator<Item = Entid>,
    {
        self.conn
            .pull_attributes_for_entity(&self.sqlite, entity, attributes)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    use uuid::Uuid;

    use std::collections::BTreeSet;
    use std::path::{Path, PathBuf};
    use std::sync::mpsc;
    use std::sync::Mutex;
    use std::time::Duration;

    use mentat_db::cache::SQLiteAttributeCache;

    use core_traits::{TypedValue, ValueType};

    use mentat_core::{CachedAttributes, HasSchema};

    use mentat_transaction::entity_builder::BuildTerms;

    use mentat_transaction::query::PreparedQuery;

    use mentat_query_algebrizer::QueryInputs;

    use crate::vocabulary::{AttributeBuilder, Definition, VersionedStore};

    use core_traits::attribute::Unique;

    fn fixture_path(rest: &str) -> PathBuf {
        let fixtures = Path::new("../fixtures/");
        fixtures.join(Path::new(rest))
    }

    #[test]
    fn test_prepared_query_with_cache() {
        let mut store = Store::open("").expect("opened");
        let mut in_progress = store.begin_transaction().expect("began");
        in_progress
            .import(fixture_path("cities.schema"))
            .expect("transacted schema");
        in_progress
            .import(fixture_path("all_seattle.edn"))
            .expect("transacted data");
        in_progress
            .cache(
                &kw!(:neighborhood/district),
                CacheDirection::Forward,
                CacheAction::Register,
            )
            .expect("cache done");
        in_progress
            .cache(
                &kw!(:district/name),
                CacheDirection::Forward,
                CacheAction::Register,
            )
            .expect("cache done");
        in_progress
            .cache(
                &kw!(:neighborhood/name),
                CacheDirection::Reverse,
                CacheAction::Register,
            )
            .expect("cache done");

        let query = r#"[:find ?district
                        :in ?hood
                        :where
                        [?neighborhood :neighborhood/name ?hood]
                        [?neighborhood :neighborhood/district ?d]
                        [?d :district/name ?district]]"#;
        let hood = "Beacon Hill";
        let inputs =
            QueryInputs::with_value_sequence(vec![(var!(?hood), TypedValue::typed_string(hood))]);
        let mut prepared = in_progress.q_prepare(query, inputs).expect("prepared");
        match prepared {
            PreparedQuery::Constant {
                select: ref _select,
            } => {}
            _ => panic!(),
        };

        let start = std::time::Instant::now();
        let results = prepared.run(None).expect("results");
        let end = std::time::Instant::now();
        println!(
            "Prepared cache execution took {}µs",
            (end - start).as_micros()
        );
        assert_eq!(
            results.into_rel().expect("result"),
            vec![vec![TypedValue::typed_string("Greater Duwamish")]].into()
        );
    }

    trait StoreCache {
        fn get_entid_for_value(&self, attr: Entid, val: &TypedValue) -> Option<Entid>;
        fn is_attribute_cached_reverse(&self, attr: Entid) -> bool;
        fn is_attribute_cached_forward(&self, attr: Entid) -> bool;
    }

    impl StoreCache for Store {
        fn get_entid_for_value(&self, attr: Entid, val: &TypedValue) -> Option<Entid> {
            let cache = self.conn.current_cache();
            cache.get_entid_for_value(attr, val)
        }

        fn is_attribute_cached_forward(&self, attr: Entid) -> bool {
            self.conn.current_cache().is_attribute_cached_forward(attr)
        }

        fn is_attribute_cached_reverse(&self, attr: Entid) -> bool {
            self.conn.current_cache().is_attribute_cached_reverse(attr)
        }
    }

    #[test]
    fn test_cache_mutation() {
        let mut store = Store::open("").expect("opened");

        {
            let mut in_progress = store.begin_transaction().expect("begun");
            in_progress
                .transact(
                    r#"[
                {  :db/ident       :foo/bar
                   :db/cardinality :db.cardinality/one
                   :db/index       true
                   :db/unique      :db.unique/identity
                   :db/valueType   :db.type/long },
                {  :db/ident       :foo/baz
                   :db/cardinality :db.cardinality/one
                   :db/valueType   :db.type/boolean }
                {  :db/ident       :foo/x
                   :db/cardinality :db.cardinality/many
                   :db/valueType   :db.type/long }]"#,
                )
                .expect("transact");

            // Cache one….
            in_progress
                .cache(
                    &kw!(:foo/bar),
                    CacheDirection::Reverse,
                    CacheAction::Register,
                )
                .expect("cache done");
            in_progress.commit().expect("commit");
        }

        let foo_bar = store
            .conn
            .current_schema()
            .get_entid(&kw!(:foo/bar))
            .expect("foo/bar")
            .0;
        let foo_baz = store
            .conn
            .current_schema()
            .get_entid(&kw!(:foo/baz))
            .expect("foo/baz")
            .0;
        let foo_x = store
            .conn
            .current_schema()
            .get_entid(&kw!(:foo/x))
            .expect("foo/x")
            .0;

        // … and cache the others via the store.
        store
            .cache(&kw!(:foo/baz), CacheDirection::Both)
            .expect("cache done");
        store
            .cache(&kw!(:foo/x), CacheDirection::Forward)
            .expect("cache done");
        {
            assert!(store.is_attribute_cached_reverse(foo_bar));
            assert!(store.is_attribute_cached_forward(foo_baz));
            assert!(store.is_attribute_cached_reverse(foo_baz));
            assert!(store.is_attribute_cached_forward(foo_x));
        }

        // Add some data.
        {
            let mut in_progress = store.begin_transaction().expect("begun");

            {
                assert!(in_progress.cache.is_attribute_cached_reverse(foo_bar));
                assert!(in_progress.cache.is_attribute_cached_forward(foo_baz));
                assert!(in_progress.cache.is_attribute_cached_reverse(foo_baz));
                assert!(in_progress.cache.is_attribute_cached_forward(foo_x));

                assert!(in_progress
                    .cache
                    .overlay
                    .is_attribute_cached_reverse(foo_bar));
                assert!(in_progress
                    .cache
                    .overlay
                    .is_attribute_cached_forward(foo_baz));
                assert!(in_progress
                    .cache
                    .overlay
                    .is_attribute_cached_reverse(foo_baz));
                assert!(in_progress.cache.overlay.is_attribute_cached_forward(foo_x));
            }

            in_progress
                .transact(
                    r#"[
                {:foo/bar 15, :foo/baz false, :foo/x [1, 2, 3]}
                {:foo/bar 99, :foo/baz true}
                {:foo/bar -2, :foo/baz true}
                ]"#,
                )
                .expect("transact");

            // Data is in the cache.
            let first = in_progress
                .cache
                .get_entid_for_value(foo_bar, &TypedValue::Long(15))
                .expect("id");
            assert_eq!(
                in_progress
                    .cache
                    .get_value_for_entid(&in_progress.schema, foo_baz, first)
                    .expect("val"),
                &TypedValue::Boolean(false)
            );

            // All three values for :foo/x.
            let all_three: BTreeSet<TypedValue> = in_progress
                .cache
                .get_values_for_entid(&in_progress.schema, foo_x, first)
                .expect("val")
                .iter()
                .cloned()
                .collect();
            assert_eq!(
                all_three,
                vec![1, 2, 3].into_iter().map(TypedValue::Long).collect()
            );

            in_progress.commit().expect("commit");
        }

        // Data is still in the cache.
        {
            let first = store
                .get_entid_for_value(foo_bar, &TypedValue::Long(15))
                .expect("id");
            let cache: SQLiteAttributeCache = store.conn.current_cache();
            assert_eq!(
                cache
                    .get_value_for_entid(&store.conn.current_schema(), foo_baz, first)
                    .expect("val"),
                &TypedValue::Boolean(false)
            );

            let all_three: BTreeSet<TypedValue> = cache
                .get_values_for_entid(&store.conn.current_schema(), foo_x, first)
                .expect("val")
                .iter()
                .cloned()
                .collect();
            assert_eq!(
                all_three,
                vec![1, 2, 3].into_iter().map(TypedValue::Long).collect()
            );
        }

        // We can remove data and the cache reflects it, immediately and after commit.
        {
            let mut in_progress = store.begin_transaction().expect("began");
            let first = in_progress
                .cache
                .get_entid_for_value(foo_bar, &TypedValue::Long(15))
                .expect("id");
            in_progress
                .transact(format!("[[:db/retract {} :foo/x 2]]", first).as_str())
                .expect("transact");

            let only_two: BTreeSet<TypedValue> = in_progress
                .cache
                .get_values_for_entid(&in_progress.schema, foo_x, first)
                .expect("val")
                .iter()
                .cloned()
                .collect();
            assert_eq!(
                only_two,
                vec![1, 3].into_iter().map(TypedValue::Long).collect()
            );

            // Rollback: unchanged.
        }
        {
            let first = store
                .get_entid_for_value(foo_bar, &TypedValue::Long(15))
                .expect("id");
            let cache: SQLiteAttributeCache = store.conn.current_cache();
            assert_eq!(
                cache
                    .get_value_for_entid(&store.conn.current_schema(), foo_baz, first)
                    .expect("val"),
                &TypedValue::Boolean(false)
            );

            let all_three: BTreeSet<TypedValue> = cache
                .get_values_for_entid(&store.conn.current_schema(), foo_x, first)
                .expect("val")
                .iter()
                .cloned()
                .collect();
            assert_eq!(
                all_three,
                vec![1, 2, 3].into_iter().map(TypedValue::Long).collect()
            );
        }

        // Try again, but this time commit.
        {
            let mut in_progress = store.begin_transaction().expect("began");
            let first = in_progress
                .cache
                .get_entid_for_value(foo_bar, &TypedValue::Long(15))
                .expect("id");
            in_progress
                .transact(format!("[[:db/retract {} :foo/x 2]]", first).as_str())
                .expect("transact");
            in_progress.commit().expect("committed");
        }
        {
            let first = store
                .get_entid_for_value(foo_bar, &TypedValue::Long(15))
                .expect("id");
            let cache: SQLiteAttributeCache = store.conn.current_cache();
            assert_eq!(
                cache
                    .get_value_for_entid(&store.conn.current_schema(), foo_baz, first)
                    .expect("val"),
                &TypedValue::Boolean(false)
            );

            let only_two: BTreeSet<TypedValue> = cache
                .get_values_for_entid(&store.conn.current_schema(), foo_x, first)
                .expect("val")
                .iter()
                .cloned()
                .collect();
            assert_eq!(
                only_two,
                vec![1, 3].into_iter().map(TypedValue::Long).collect()
            );
        }
    }

    fn test_register_observer() {
        let mut conn = Store::open("").unwrap();

        let key = "Test Observer".to_string();
        let tx_observer = TxObserver::new(BTreeSet::new(), move |_obs_key, _batch| {});

        conn.register_observer(key.clone(), Arc::new(tx_observer));
        assert!(conn.is_registered_as_observer(&key));
    }

    #[test]
    fn test_deregister_observer() {
        let mut conn = Store::open("").unwrap();

        let key = "Test Observer".to_string();

        let tx_observer = TxObserver::new(BTreeSet::new(), move |_obs_key, _batch| {});

        conn.register_observer(key.clone(), Arc::new(tx_observer));
        assert!(conn.is_registered_as_observer(&key));

        conn.unregister_observer(&key);

        assert!(!conn.is_registered_as_observer(&key));
    }

    fn add_schema(conn: &mut Store) {
        // transact some schema
        let mut in_progress = conn.begin_transaction().expect("expected in progress");
        in_progress
            .ensure_vocabulary(&Definition::new(
                kw!(:todo/items),
                1,
                vec![
                    (
                        kw!(:todo/uuid),
                        AttributeBuilder::helpful()
                            .value_type(ValueType::Uuid)
                            .multival(false)
                            .unique(Unique::Value)
                            .index(true)
                            .build(),
                    ),
                    (
                        kw!(:todo/name),
                        AttributeBuilder::helpful()
                            .value_type(ValueType::String)
                            .multival(false)
                            .fulltext(true)
                            .build(),
                    ),
                    (
                        kw!(:todo/completion_date),
                        AttributeBuilder::helpful()
                            .value_type(ValueType::Instant)
                            .multival(false)
                            .build(),
                    ),
                    (
                        kw!(:label/name),
                        AttributeBuilder::helpful()
                            .value_type(ValueType::String)
                            .multival(false)
                            .unique(Unique::Value)
                            .fulltext(true)
                            .index(true)
                            .build(),
                    ),
                    (
                        kw!(:label/color),
                        AttributeBuilder::helpful()
                            .value_type(ValueType::String)
                            .multival(false)
                            .build(),
                    ),
                ],
            ))
            .expect("expected vocubulary");
        in_progress.commit().expect("Expected vocabulary committed");
    }

    #[derive(Default, Debug)]
    struct ObserverOutput {
        txids: Vec<i64>,
        changes: Vec<BTreeSet<i64>>,
        called_key: Option<String>,
    }

    #[test]
    fn test_observer_notified_on_registered_change() {
        let mut conn = Store::open("").unwrap();
        add_schema(&mut conn);

        let name_entid: Entid = conn
            .conn()
            .current_schema()
            .get_entid(&kw!(:todo/name))
            .expect("entid to exist for name")
            .into();
        let date_entid: Entid = conn
            .conn()
            .current_schema()
            .get_entid(&kw!(:todo/completion_date))
            .expect("entid to exist for completion_date")
            .into();
        let mut registered_attrs = BTreeSet::new();
        registered_attrs.insert(name_entid);
        registered_attrs.insert(date_entid);

        let key = "Test Observing".to_string();

        let output = Arc::new(Mutex::new(ObserverOutput::default()));

        let mut_output = Arc::downgrade(&output);
        let (tx, rx): (mpsc::Sender<()>, mpsc::Receiver<()>) = mpsc::channel();
        // because the TxObserver is in an Arc and is therefore Sync, we have to wrap the Sender in a Mutex to also
        // make it Sync.
        let thread_tx = Mutex::new(tx);
        let tx_observer = Arc::new(TxObserver::new(registered_attrs, move |obs_key, batch| {
            if let Some(out) = mut_output.upgrade() {
                let mut o = out.lock().unwrap();
                o.called_key = Some(obs_key.to_string());
                for (tx_id, changes) in batch.into_iter() {
                    o.txids.push(*tx_id);
                    o.changes.push(changes.clone());
                }
                o.txids.sort_unstable();
            }
            thread_tx.lock().unwrap().send(()).unwrap();
        }));

        conn.register_observer(key.clone(), Arc::clone(&tx_observer));
        assert!(conn.is_registered_as_observer(&key));

        let mut tx_ids = Vec::new();
        let mut changesets = Vec::new();
        let db_tx_instant_entid: Entid = conn
            .conn()
            .current_schema()
            .get_entid(&kw!(:db/txInstant))
            .expect("entid to exist for :db/txInstant")
            .into();
        let uuid_entid: Entid = conn
            .conn()
            .current_schema()
            .get_entid(&kw!(:todo/uuid))
            .expect("entid to exist for name")
            .into();
        {
            let mut in_progress = conn.begin_transaction().expect("expected transaction");
            for i in 0..3 {
                let mut changeset = BTreeSet::new();
                changeset.insert(db_tx_instant_entid);
                let name = format!("todo{}", i);
                let uuid = Uuid::new_v4();
                let mut builder = in_progress.builder().describe_tempid(&name);
                builder
                    .add(kw!(:todo/uuid), TypedValue::Uuid(uuid))
                    .expect("Expected added uuid");
                changeset.insert(uuid_entid);
                builder
                    .add(kw!(:todo/name), TypedValue::typed_string(name))
                    .expect("Expected added name");
                changeset.insert(name_entid);
                if i % 2 == 0 {
                    builder
                        .add(kw!(:todo/completion_date), TypedValue::current_instant())
                        .expect("Expected added date");
                    changeset.insert(date_entid);
                }
                let (ip, r) = builder.transact();
                let report = r.expect("expected a report");
                tx_ids.push(report.tx_id);
                changesets.push(changeset);
                in_progress = ip;
            }
            let mut builder = in_progress.builder().describe_tempid("Label");
            builder
                .add(kw!(:label/name), TypedValue::typed_string("Label 1"))
                .expect("Expected added name");
            builder
                .add(kw!(:label/color), TypedValue::typed_string("blue"))
                .expect("Expected added color");
            builder.commit().expect("expect transaction to occur");
        }

        let delay = Duration::from_millis(100);
        let _ = rx.recv_timeout(delay);

        let out = Arc::try_unwrap(output).expect("unwrapped");
        let o = out.into_inner().expect("Expected an Output");
        assert_eq!(o.called_key, Some(key));
        assert_eq!(o.txids, tx_ids);
        assert_eq!(o.changes, changesets);
    }

    #[test]
    fn test_observer_not_notified_on_unregistered_change() {
        let mut conn = Store::open("").unwrap();
        add_schema(&mut conn);

        let name_entid: Entid = conn
            .conn()
            .current_schema()
            .get_entid(&kw!(:todo/name))
            .expect("entid to exist for name")
            .into();
        let date_entid: Entid = conn
            .conn()
            .current_schema()
            .get_entid(&kw!(:todo/completion_date))
            .expect("entid to exist for completion_date")
            .into();
        let mut registered_attrs = BTreeSet::new();
        registered_attrs.insert(name_entid);
        registered_attrs.insert(date_entid);

        let key = "Test Observing".to_string();

        let output = Arc::new(Mutex::new(ObserverOutput::default()));

        let mut_output = Arc::downgrade(&output);
        let (tx, rx): (mpsc::Sender<()>, mpsc::Receiver<()>) = mpsc::channel();
        let thread_tx = Mutex::new(tx);
        let tx_observer = Arc::new(TxObserver::new(registered_attrs, move |obs_key, batch| {
            if let Some(out) = mut_output.upgrade() {
                let mut o = out.lock().unwrap();
                o.called_key = Some(obs_key.to_string());
                for (tx_id, changes) in batch.into_iter() {
                    o.txids.push(*tx_id);
                    o.changes.push(changes.clone());
                }
                o.txids.sort_unstable();
            }
            thread_tx.lock().unwrap().send(()).unwrap();
        }));

        conn.register_observer(key.clone(), Arc::clone(&tx_observer));
        assert!(conn.is_registered_as_observer(&key));

        let tx_ids = Vec::<Entid>::new();
        let changesets = Vec::<BTreeSet<Entid>>::new();
        {
            let mut in_progress = conn.begin_transaction().expect("expected transaction");
            for i in 0..3 {
                let name = format!("label{}", i);
                let mut builder = in_progress.builder().describe_tempid(&name);
                builder
                    .add(kw!(:label/name), TypedValue::typed_string(name))
                    .expect("Expected added name");
                builder
                    .add(kw!(:label/color), TypedValue::typed_string("blue"))
                    .expect("Expected added color");
                let (ip, _) = builder.transact();
                in_progress = ip;
            }
        }

        let delay = Duration::from_millis(100);
        let _ = rx.recv_timeout(delay);

        let out = Arc::try_unwrap(output).expect("unwrapped");
        let o = out.into_inner().expect("Expected an Output");
        assert_eq!(o.called_key, None);
        assert_eq!(o.txids, tx_ids);
        assert_eq!(o.changes, changesets);
    }
}