gecko-dev/third_party/rust/rkv/examples/simple-store.rs

// Any copyright is dedicated to the Public Domain.
// http://creativecommons.org/publicdomain/zero/1.0/

//! A simple rkv demo that showcases the basic usage (put/get/delete) of rkv.
//!
//! You can test this out by running:
//!
//!     cargo run --example simple-store

use std::fs;

use tempfile::Builder;

use rkv::{
    backend::{SafeMode, SafeModeDatabase, SafeModeEnvironment, SafeModeRwTransaction},
    Manager, Rkv, StoreOptions, Value,
};

type MultiStore = rkv::MultiStore<SafeModeDatabase>;
type Writer<'w> = rkv::Writer<SafeModeRwTransaction<'w>>;

fn getput<'w, 's>(store: MultiStore, writer: &'w mut Writer, ids: &'s mut Vec<String>) {
    let keys = vec!["str1", "str2", "str3"];
    // we convert the writer into a cursor so that we can safely read
    for k in keys.iter() {
        // this is a multi-valued database, so get returns an iterator
        let mut iter = store.get(writer, k).unwrap();
        while let Some(Ok((_key, val))) = iter.next() {
            if let Value::Str(s) = val {
                ids.push(s.to_owned());
            } else {
                panic!("didn't get a string back!");
            }
        }
    }
    for id in ids {
        store.put(writer, &id, &Value::Blob(b"weeeeeee")).unwrap();
    }
}

fn delete(store: MultiStore, writer: &mut Writer) {
    let keys = vec!["str1", "str2", "str3"];
    let vals = vec!["string uno", "string quatro", "string siete"];
    // we convert the writer into a cursor so that we can safely read
    for i in 0..keys.len() {
        store
            .delete(writer, &keys[i], &Value::Str(vals[i]))
            .unwrap();
    }
}

fn main() {
    let root = Builder::new().prefix("simple-db").tempdir().unwrap();
    fs::create_dir_all(root.path()).unwrap();
    let p = root.path();

    // The manager enforces that each process opens the same lmdb environment at most once
    let mut manager = Manager::<SafeModeEnvironment>::singleton().write().unwrap();
    let created_arc = manager.get_or_create(p, Rkv::new::<SafeMode>).unwrap();
    let k = created_arc.read().unwrap();

    // Creates a store called "store"
    let store = k.open_single("store", StoreOptions::create()).unwrap();
    let multistore = k.open_multi("multistore", StoreOptions::create()).unwrap();

    println!("Inserting data...");
    {
        // Use a writer to mutate the store
        let mut writer = k.write().unwrap();
        store.put(&mut writer, "int", &Value::I64(1234)).unwrap();
        store
            .put(&mut writer, "uint", &Value::U64(1234_u64))
            .unwrap();
        store
            .put(&mut writer, "float", &Value::F64(1234.0.into()))
            .unwrap();
        store
            .put(&mut writer, "instant", &Value::Instant(1_528_318_073_700))
            .unwrap();
        store
            .put(&mut writer, "boolean", &Value::Bool(true))
            .unwrap();
        store
            .put(&mut writer, "string", &Value::Str("héllo, yöu"))
            .unwrap();
        store
            .put(
                &mut writer,
                "json",
                &Value::Json(r#"{"foo":"bar", "number": 1}"#),
            )
            .unwrap();
        store
            .put(&mut writer, "blob", &Value::Blob(b"blob"))
            .unwrap();
        writer.commit().unwrap();
    }

    println!("Testing getput");
    {
        let mut ids = Vec::new();
        let mut writer = k.write().unwrap();
        multistore
            .put(&mut writer, "str1", &Value::Str("string uno"))
            .unwrap();
        multistore
            .put(&mut writer, "str1", &Value::Str("string dos"))
            .unwrap();
        multistore
            .put(&mut writer, "str1", &Value::Str("string tres"))
            .unwrap();
        multistore
            .put(&mut writer, "str2", &Value::Str("string quatro"))
            .unwrap();
        multistore
            .put(&mut writer, "str2", &Value::Str("string cinco"))
            .unwrap();
        multistore
            .put(&mut writer, "str2", &Value::Str("string seis"))
            .unwrap();
        multistore
            .put(&mut writer, "str3", &Value::Str("string siete"))
            .unwrap();
        multistore
            .put(&mut writer, "str3", &Value::Str("string ocho"))
            .unwrap();
        multistore
            .put(&mut writer, "str3", &Value::Str("string nueve"))
            .unwrap();
        getput(multistore, &mut writer, &mut ids);
        writer.commit().unwrap();
        let mut writer = k.write().unwrap();
        delete(multistore, &mut writer);
        writer.commit().unwrap();
    }

    println!("Looking up keys...");
    {
        // Use a reader to query the store
        let reader = k.read().unwrap();
        println!("Get int {:?}", store.get(&reader, "int").unwrap());
        println!("Get uint {:?}", store.get(&reader, "uint").unwrap());
        println!("Get float {:?}", store.get(&reader, "float").unwrap());
        println!("Get instant {:?}", store.get(&reader, "instant").unwrap());
        println!("Get boolean {:?}", store.get(&reader, "boolean").unwrap());
        println!("Get string {:?}", store.get(&reader, "string").unwrap());
        println!("Get json {:?}", store.get(&reader, "json").unwrap());
        println!("Get blob {:?}", store.get(&reader, "blob").unwrap());
        println!(
            "Get non-existent {:?}",
            store.get(&reader, "non-existent").unwrap()
        );
    }

    println!("Looking up keys via Writer.get()...");
    {
        let mut writer = k.write().unwrap();
        store.put(&mut writer, "foo", &Value::Str("bar")).unwrap();
        store.put(&mut writer, "bar", &Value::Str("baz")).unwrap();
        store.delete(&mut writer, "foo").unwrap();
        println!(
            "It should be None! ({:?})",
            store.get(&writer, "foo").unwrap()
        );
        println!("Get bar ({:?})", store.get(&writer, "bar").unwrap());
        writer.commit().unwrap();
        let reader = k.read().expect("reader");
        println!(
            "It should be None! ({:?})",
            store.get(&reader, "foo").unwrap()
        );
        println!("Get bar {:?}", store.get(&reader, "bar").unwrap());
    }

    println!("Aborting transaction...");
    {
        // Aborting a write transaction rollbacks the change(s)
        let mut writer = k.write().unwrap();
        store.put(&mut writer, "foo", &Value::Str("bar")).unwrap();
        writer.abort();

        let reader = k.read().expect("reader");
        println!(
            "It should be None! ({:?})",
            store.get(&reader, "foo").unwrap()
        );
        // Explicitly aborting a transaction is not required unless an early
        // abort is desired, since both read and write transactions will
        // implicitly be aborted once they go out of scope.
    }

    println!("Deleting keys...");
    {
        // Deleting a key/value also requires a write transaction
        let mut writer = k.write().unwrap();
        store.put(&mut writer, "foo", &Value::Str("bar")).unwrap();
        store.delete(&mut writer, "foo").unwrap();
        println!(
            "It should be None! ({:?})",
            store.get(&writer, "foo").unwrap()
        );
        writer.commit().unwrap();

        // Committing a transaction consumes the writer, preventing you
        // from reusing it by failing and reporting a compile-time error.
        // This line would report error[E0382]: use of moved value: `writer`.
        // store.put(&mut writer, "baz", &Value::Str("buz")).unwrap();
    }

    println!("Clearing store...");
    {
        // Clearing a store deletes all the entries in that store
        let mut writer = k.write().unwrap();
        store.put(&mut writer, "foo", &Value::Str("bar")).unwrap();
        store.put(&mut writer, "bar", &Value::Str("baz")).unwrap();
        store.clear(&mut writer).unwrap();
        writer.commit().unwrap();

        let reader = k.read().expect("reader");
        println!(
            "It should be None! ({:?})",
            store.get(&reader, "foo").unwrap()
        );
        println!(
            "It should be None! ({:?})",
            store.get(&reader, "bar").unwrap()
        );
    }

    println!("Write and read on multiple stores...");
    {
        let another_store = k
            .open_single("another_store", StoreOptions::create())
            .unwrap();
        let mut writer = k.write().unwrap();
        store.put(&mut writer, "foo", &Value::Str("bar")).unwrap();
        another_store
            .put(&mut writer, "foo", &Value::Str("baz"))
            .unwrap();
        writer.commit().unwrap();

        let reader = k.read().unwrap();
        println!(
            "Get from store value: {:?}",
            store.get(&reader, "foo").unwrap()
        );
        println!(
            "Get from another store value: {:?}",
            another_store.get(&reader, "foo").unwrap()
        );
    }
}