path: root/base/src/database.rs

/*
    This file is part of jellything (https://codeberg.org/metamuffin/jellything)
    which is licensed under the GNU Affero General Public License (version 3); see /COPYING.
    Copyright (C) 2024 metamuffin <metamuffin.org>
*/
use anyhow::Context;
use bincode::{Decode, Encode};
use jellycommon::{
    user::{NodeUserData, User},
    ExtendedNode, Node,
};
use log::info;
use serde::{Deserialize, Serialize};
use std::{borrow::Borrow, fs::create_dir_all, ops::Deref, path::Path, sync::RwLock};
use tantivy::{
    directory::MmapDirectory,
    schema::{Field, Schema, FAST, INDEXED, STORED, TEXT},
    DateOptions, Index, IndexReader, IndexWriter, ReloadPolicy,
};

pub use redb::*;
pub use tantivy;

pub const T_USER: TableDefinition<&str, Ser<User>> = TableDefinition::new("user");
pub const T_USER_NODE: TableDefinition<(&str, &str), Ser<NodeUserData>> =
    TableDefinition::new("user_node");
pub const T_INVITE: TableDefinition<&str, Ser<()>> = TableDefinition::new("invite");
pub const T_NODE: TableDefinition<&str, Ser<Node>> = TableDefinition::new("node");
pub const T_NODE_EXTENDED: TableDefinition<&str, Ser<ExtendedNode>> =
    TableDefinition::new("node-ext");
pub const T_NODE_IMPORT: TableDefinition<&str, Ser<Vec<(Vec<usize>, Node)>>> =
    TableDefinition::new("node-import");

pub struct DataAcid {
    pub inner: redb::Database,
    pub node_index: NodeIndex,
}

impl DataAcid {
    pub fn open(path: &Path) -> Result<Self, anyhow::Error> {
        create_dir_all(path).context("creating database directory")?;
        info!("opening kv store...");
        let db = redb::Database::create(path.join("data")).context("opening kv store")?;
        info!("opening node index...");
        let ft_node = NodeIndex::new(path).context("in node index")?;
        let r = Self {
            inner: db,
            node_index: ft_node,
        };

        {
            // this creates all tables such that read operations on them do not fail.
            let txn = r.begin_write()?;
            drop(txn.open_table(T_INVITE)?);
            drop(txn.open_table(T_USER)?);
            drop(txn.open_table(T_USER_NODE)?);
            drop(txn.open_table(T_NODE)?);
            drop(txn.open_table(T_NODE_IMPORT)?);
            drop(txn.open_table(T_NODE_EXTENDED)?);
            txn.commit()?;
        }

        info!("ready");
        Ok(r)
    }
}

impl Deref for DataAcid {
    type Target = Database;
    fn deref(&self) -> &Self::Target {
        &self.inner
    }
}

pub struct NodeIndex {
    pub schema: Schema,
    pub reader: IndexReader,
    pub writer: RwLock<IndexWriter>,
    pub index: Index,
    pub id: Field,
    pub title: Field,
    pub releasedate: Field,
    pub description: Field,
    pub f_index: Field,
}
impl NodeIndex {
    fn new(path: &Path) -> anyhow::Result<Self> {
        let mut schema = Schema::builder();
        let id = schema.add_text_field("id", TEXT | STORED | FAST);
        let title = schema.add_text_field("title", TEXT);
        let description = schema.add_text_field("description", TEXT);
        let f_index = schema.add_u64_field("index", FAST);
        let releasedate = schema.add_date_field(
            "releasedate",
            DateOptions::from(INDEXED)
                .set_fast()
                .set_precision(tantivy::DateTimePrecision::Seconds),
        );
        let schema = schema.build();
        create_dir_all(path.join("node_index"))?;
        let directory =
            MmapDirectory::open(path.join("node_index")).context("opening index directory")?;
        let index = Index::open_or_create(directory, schema.clone()).context("creating index")?;
        let reader = index
            .reader_builder()
            .reload_policy(ReloadPolicy::OnCommitWithDelay)
            .try_into()
            .context("creating reader")?;
        let writer = index.writer(30_000_000).context("creating writer")?;
        Ok(Self {
            index,
            writer: writer.into(),
            reader,
            schema,
            f_index,
            releasedate,
            id,
            description,
            title,
        })
    }
}

pub trait TableExt<Key, KeyRef, Value> {
    fn get(self, db: &DataAcid, key: KeyRef) -> anyhow::Result<Option<Value>>;
    fn insert(self, db: &DataAcid, key: KeyRef, value: Value) -> anyhow::Result<()>;
    fn remove(self, db: &DataAcid, key: KeyRef) -> anyhow::Result<Option<Value>>;
}
impl<'a, 'b, 'c, Key, Value, KeyRef> TableExt<Key, KeyRef, Value>
    for TableDefinition<'a, Key, Ser<Value>>
where
    Key: Borrow<<Key as RedbValue>::SelfType<'b>> + redb::RedbKey,
    Value: Encode + Decode + std::fmt::Debug + Serialize + for<'x> Deserialize<'x>,
    KeyRef: Borrow<<Key as redb::RedbValue>::SelfType<'c>>,
{
    fn get(self, db: &DataAcid, key: KeyRef) -> anyhow::Result<Option<Value>> {
        let txn = db.inner.begin_read()?;
        let table = txn.open_table(self)?;
        let user = table.get(key)?.map(|v| v.value().0);
        drop(table);
        Ok(user)
    }
    fn insert(self, db: &DataAcid, key: KeyRef, value: Value) -> anyhow::Result<()> {
        let txn = db.inner.begin_write()?;
        let mut table = txn.open_table(self)?;
        table.insert(key, Ser(value))?;
        drop(table);
        txn.commit()?;
        Ok(())
    }
    fn remove(self, db: &DataAcid, key: KeyRef) -> anyhow::Result<Option<Value>> {
        let txn = db.inner.begin_write()?;
        let mut table = txn.open_table(self)?;
        let prev = table.remove(key)?.map(|v| v.value().0);
        drop(table);
        txn.commit()?;
        Ok(prev)
    }
}

// pub trait TableIterExt<
//     'a,
//     Key: redb::RedbKey + 'static,
//     Value: redb::RedbValue + 'static,
//     F: FnOnce(&redb::Range<'a, Key, Value>) -> anyhow::Result<T>,
//     T: 'static,
// >
// {
//     fn iter(self, db: &'a DataAcid, f: F) -> anyhow::Result<T>;
// }
// impl<'a, Key, Value, F, T> TableIterExt<'a, Key, Value, F, T>
//     for TableDefinition<'static, Key, Value>
// where
//     Key: redb::RedbKey,
//     Value: redb::RedbValue,
//     F: FnOnce(&redb::Range<'a, Key, Value>) -> anyhow::Result<T>,
//     T: 'static,
// {
//     fn iter(self, db: &DataAcid, f: F) -> anyhow::Result<T> {
//         let txn = db.begin_read()?;
//         let table = txn.open_table(self)?;
//         let iter = table.iter()?;
//         let ret = f(&iter)?;
//         drop(iter);
//         drop(table);
//         drop(txn);
//         Ok(ret)
//     }
// }

#[derive(Debug)]
#[cfg(not(feature = "db_json"))]
pub struct Ser<T>(pub T);
#[cfg(not(feature = "db_json"))]
impl<T: Encode + Decode + std::fmt::Debug> RedbValue for Ser<T> {
    type SelfType<'a> = Ser<T>
    where
        Self: 'a;
    type AsBytes<'a> = Vec<u8>
    where
        Self: 'a;

    fn fixed_width() -> Option<usize> {
        None
    }

    fn from_bytes<'a>(data: &'a [u8]) -> Self::SelfType<'a>
    where
        Self: 'a,
    {
        Ser(bincode::decode_from_slice(data, bincode::config::legacy())
            .unwrap()
            .0)
    }

    fn as_bytes<'a, 'b: 'a>(value: &'a Self::SelfType<'b>) -> Self::AsBytes<'a>
    where
        Self: 'a,
        Self: 'b,
    {
        bincode::encode_to_vec(&value.0, bincode::config::legacy()).unwrap()
    }

    fn type_name() -> redb::TypeName {
        TypeName::new("bincode")
    }
}

#[derive(Debug)]
#[cfg(feature = "db_json")]
pub struct Ser<T>(pub T);
#[cfg(feature = "db_json")]
impl<T: Serialize + for<'a> Deserialize<'a> + std::fmt::Debug> RedbValue for Ser<T> {
    type SelfType<'a> = Ser<T>
    where
        Self: 'a;
    type AsBytes<'a> = Vec<u8>
    where
        Self: 'a;

    fn fixed_width() -> Option<usize> {
        None
    }

    fn from_bytes<'a>(data: &'a [u8]) -> Self::SelfType<'a>
    where
        Self: 'a,
    {
        Ser(serde_json::from_slice(data).unwrap())
    }

    fn as_bytes<'a, 'b: 'a>(value: &'a Self::SelfType<'b>) -> Self::AsBytes<'a>
    where
        Self: 'a,
        Self: 'b,
    {
        serde_json::to_vec(&value.0).unwrap()
    }

    fn type_name() -> redb::TypeName {
        TypeName::new("json")
    }
}