ordinary-storage 0.11.1

// Copyright (C) 2026 The Ordinary Authors.
//
// SPDX-License-Identifier: BSD-3-Clause

pub mod delete;
pub mod get;
pub mod insert;
pub mod list;
pub mod query;
pub mod update;

use crate::stores::database::delete::{DeleteError, delete};
use crate::stores::database::get::get;
use crate::stores::database::insert::insert;
use crate::stores::database::list::list;
use crate::stores::database::query::query;
use crate::stores::database::update::update;
use crate::{Accessor, RefDepth, Transaction, field_to_bytes, push_field_from_bytes};
use anyhow::bail;
use bytes::{BufMut, Bytes, BytesMut};
use chacha20poly1305::{
    KeyInit, XChaCha20Poly1305, XNonce,
    aead::{Aead, Generate},
};
use flexbuffers::Reader;
use ordinary_config::{DatabaseLimits, DatabaseModelConfig, UuidVersion};
use ordinary_types::{Field, Kind};
use parking_lot::Mutex;
use rustc_hash::FxHashMap;
use saferlmdb::{
    self as lmdb, Database, DatabaseOptions, Environment, ReadTransaction, WriteTransaction,
};
use std::sync::Arc;
use uuid::Uuid;

pub use crate::stores::database::{
    get::GetError, insert::InsertError, query::QueryError, update::UpdateError,
};

pub struct DatabaseStore {
    pub limits: DatabaseLimits,

    env: Arc<Environment>,

    /// Vec of auto incrementing ints per-model
    auto_ints: Vec<Arc<Mutex<usize>>>,

    /// key([`model_byte`][int_id]) -> value(flexbuffer encoded)
    item_db: Arc<Database<'static>>,
    /// key([`model_idx`][field_idx][`indexed_value`]) -> value([kind][int_id])
    index_db: Arc<Database<'static>>,
    /// key([`model_idx`][field_idx][`indexed_value`]) -> value(Vec<[kind][int_id]>)
    query_db: Arc<Database<'static>>,

    /// used for server-side field encryption
    encryption_key: [u8; 32],

    model_configs: Vec<DatabaseModelConfig>,
    model_config_map: FxHashMap<String, DatabaseModelConfig>,

    /// only for refs that are many. only a single many reference
    /// to another model can exist on a given model.
    ///
    /// (model, field) -> (model idx, field idx, fields, `is_many`)
    #[allow(clippy::type_complexity)]
    cross_refs: FxHashMap<(u8, u8), (u8, u8, Vec<Field>, bool)>,

    log_sizes: bool,
}

impl DatabaseStore {
    #[allow(clippy::too_many_lines)]
    pub fn new(
        limits: DatabaseLimits,
        mut model_configs: Vec<DatabaseModelConfig>,
        encryption_key: [u8; 32],
        env: &Arc<Environment>,
        log_sizes: bool,
    ) -> anyhow::Result<Self> {
        let item_db = Arc::new(Database::open(
            env.clone(),
            Some("item"),
            &DatabaseOptions::new(lmdb::db::Flags::CREATE),
        )?);
        let index_db = Arc::new(Database::open(
            env.clone(),
            Some("index"),
            &DatabaseOptions::new(lmdb::db::Flags::CREATE),
        )?);
        let query_db = Arc::new(Database::open(
            env.clone(),
            Some("query"),
            &DatabaseOptions::new(lmdb::db::Flags::CREATE),
        )?);

        let mut auto_ints = vec![];
        let mut cross_refs = FxHashMap::default();

        model_configs.sort_by_key(|a| a.idx);
        for config in &mut model_configs {
            config.fields.sort_by_key(|a| a.idx);
        }

        let mut model_config_map = FxHashMap::default();

        for config in &model_configs {
            model_config_map.insert(config.name.clone(), config.clone());

            let txn = ReadTransaction::new(env.clone())?;
            let access = txn.access();

            let mut item_cursor = txn.cursor(item_db.clone())?;

            match item_cursor.seek_range_k::<[u8], [u8]>(&access, &[config.idx]) {
                Ok((key, _)) => {
                    let model_idx = key[0];

                    if model_idx == config.idx {
                        let mut curr = &key[1..];

                        loop {
                            if let Ok((key, _)) = item_cursor.next::<[u8], [u8]>(&access) {
                                if key[0] == config.idx {
                                    curr = &key[1..];
                                } else {
                                    match curr.len() {
                                        1 => auto_ints.push(Arc::new(Mutex::new(curr[0] as usize))),
                                        2 => auto_ints.push(Arc::new(Mutex::new(
                                            u16::from_be_bytes(curr.try_into()?) as usize,
                                        ))),
                                        4 => auto_ints.push(Arc::new(Mutex::new(
                                            u32::from_be_bytes(curr.try_into()?) as usize,
                                        ))),
                                        _ => auto_ints.push(Arc::new(Mutex::new(usize::try_from(
                                            u64::from_be_bytes(curr.try_into()?),
                                        )?))),
                                    }
                                    break;
                                }
                            } else {
                                match curr.len() {
                                    1 => auto_ints.push(Arc::new(Mutex::new(curr[0] as usize))),
                                    2 => auto_ints.push(Arc::new(Mutex::new(u16::from_be_bytes(
                                        curr.try_into()?,
                                    )
                                        as usize))),
                                    4 => auto_ints.push(Arc::new(Mutex::new(u32::from_be_bytes(
                                        curr.try_into()?,
                                    )
                                        as usize))),
                                    _ => auto_ints.push(Arc::new(Mutex::new(usize::try_from(
                                        u64::from_be_bytes(curr.try_into()?),
                                    )?))),
                                }
                                break;
                            }
                        }
                    } else {
                        auto_ints.push(Arc::new(Mutex::new(0)));
                    }
                }
                Err(_) => {
                    auto_ints.push(Arc::new(Mutex::new(0)));
                }
            }

            for field in &config.fields {
                if let Kind::Ref {
                    model,
                    field: _ref_field,
                    many: _many,
                } = &field.kind
                {
                    for ref_config in &model_configs {
                        if ref_config.name == *model {
                            for ref_field in &ref_config.fields {
                                if let Kind::Ref {
                                    model: back_ref_model,
                                    many: back_ref_many,
                                    field: _back_ref_field,
                                } = &ref_field.kind
                                    && *back_ref_model == config.name
                                {
                                    cross_refs.insert(
                                        (config.idx, field.idx),
                                        (
                                            ref_config.idx,
                                            ref_field.idx,
                                            ref_config.fields.clone(),
                                            back_ref_many == &Some(true),
                                        ),
                                    );
                                }
                            }
                        }
                    }
                }
            }
        }

        Ok(Self {
            limits,
            env: env.clone(),

            auto_ints,
            encryption_key,

            item_db,
            index_db,
            query_db,

            model_configs,
            model_config_map,

            cross_refs,

            log_sizes,
        })
    }

    fn encrypt_compress_copy_to(
        &self,
        field: &Field,
        reader: &Reader<&[u8]>,
        item: &mut flexbuffers::VectorBuilder,
    ) -> anyhow::Result<()> {
        if field.encrypted == Some(true) {
            if field.indexed == Some(true) || field.queryable == Some(true) {
                bail!("cannot encrypt indexed/queryable fields");
            }

            let cipher = XChaCha20Poly1305::new(&self.encryption_key.into());
            let nonce = XNonce::generate();

            let plaintext = if field.compressed == Some(true) {
                let bytes = field_to_bytes(field, reader);
                Bytes::copy_from_slice(
                    zstd::stream::encode_all(std::io::Cursor::new(bytes.as_ref()), 4)?.as_slice(),
                )
            } else {
                field_to_bytes(field, reader)
            };

            match cipher.encrypt(&nonce, plaintext.as_ref()) {
                Ok(mut v) => {
                    if v.len() as u64 > self.limits.max_field_size {
                        if field.compressed == Some(true) {
                            bail!(
                                "i = {}, nm = {}, encrypted + compressed field exceeds size limits",
                                field.idx,
                                field.name
                            )
                        }

                        bail!(
                            "i = {}, nm = {}, encrypted field exceeds size limits",
                            field.idx,
                            field.name
                        )
                    }

                    v.extend_from_slice(&nonce);
                    item.push(flexbuffers::Blob(v.as_ref()));
                }
                Err(err) => {
                    tracing::error!(%err);
                }
            }
        } else if field.compressed == Some(true) {
            let bytes = field_to_bytes(field, reader);
            let compressed = zstd::stream::encode_all(std::io::Cursor::new(bytes.as_ref()), 4)?;

            if compressed.len() as u64 > self.limits.max_field_size {
                bail!(
                    "i = {}, nm = {}, compressed field exceeds size limits",
                    field.idx,
                    field.name
                );
            }

            item.push(flexbuffers::Blob(compressed.as_slice()));
        } else {
            let mut size = 0;
            field.kind.copy_to(reader, item, Some(&mut size))?;

            if size as u64 > self.limits.max_field_size {
                bail!(
                    "i = {}, nm = {}, field exceeds size limits",
                    field.idx,
                    field.name
                );
            }
        }

        Ok(())
    }

    pub fn insert(
        &self,
        model_idx: u8,
        payload: &[u8],
        txn: Option<WriteTransaction<'_>>,
    ) -> anyhow::Result<[u8; 16], InsertError> {
        let Some(model_config) = self.model_configs.get(model_idx as usize) else {
            return Err(InsertError::ModelNotFound(model_idx));
        };

        let uuid = if let Some(uuid_version) = &model_config.uuid {
            match uuid_version {
                UuidVersion::V4 => Uuid::new_v4(),
                UuidVersion::V7 => Uuid::now_v7(),
            }
        } else {
            Uuid::new_v4()
        };

        insert(self, model_config, uuid, payload, txn)
    }

    pub fn update(
        &self,
        model_idx: u8,
        uuid: [u8; 16],
        payload: &[u8],
        txn: Option<WriteTransaction<'_>>,
    ) -> anyhow::Result<(), UpdateError> {
        let Some(model_config) = self.model_configs.get(model_idx as usize) else {
            return Err(UpdateError::ModelNotFound(model_idx));
        };

        update(self, model_config, uuid, payload, txn)
    }

    pub fn delete(
        &self,
        model_idx: u8,
        uuid: [u8; 16],
        txn: Option<WriteTransaction<'_>>,
    ) -> anyhow::Result<(), DeleteError> {
        let Some(model_config) = self.model_configs.get(model_idx as usize) else {
            return Err(DeleteError::ModelNotFound(model_idx));
        };

        delete(self, model_config, uuid, txn)
    }

    pub fn get(
        &self,
        model_idx: u8,
        field_idx: u8,
        index_value: &[u8],
        ref_depth: &RefDepth,
        txn: Option<Transaction<'_>>,
    ) -> anyhow::Result<Bytes, GetError> {
        let Some(model_config) = self.model_configs.get(model_idx as usize) else {
            return Err(GetError::ModelNotFound(model_idx));
        };

        get(self, model_config, field_idx, index_value, ref_depth, txn)
    }

    pub fn query(
        &self,
        model_idx: u8,
        field_idx: u8,
        expression: u8,
        query_value: &[u8],
        ref_depth: &RefDepth,
        txn: Option<Transaction<'_>>,
    ) -> anyhow::Result<Bytes, QueryError> {
        let Some(model_config) = self.model_configs.get(model_idx as usize) else {
            return Err(QueryError::ModelNotFound(model_idx));
        };

        query(
            self,
            model_config,
            field_idx,
            expression,
            query_value,
            ref_depth,
            txn,
        )
    }

    /// *inefficient and primarily useful for admin purposes*
    pub fn list(&self, model_idx: u8, cursor: Option<&[u8]>) -> anyhow::Result<Bytes> {
        let Some(model_config) = self.model_configs.get(model_idx as usize) else {
            bail!("no model for idx {model_idx}")
        };

        list(self, model_config, cursor)
    }

    #[allow(clippy::too_many_lines)]
    fn build_item(
        &self,
        model_config: &DatabaseModelConfig,
        access: &Accessor,
        item_key: &[u8],
        item_vec: &mut flexbuffers::VectorBuilder,
        ref_depth: &RefDepth,
        mut count: usize,
    ) -> anyhow::Result<usize> {
        let item = match access {
            Accessor::Write(access) => access.get::<[u8], [u8]>(&self.item_db, item_key)?,
            Accessor::Const(access) => access.get::<[u8], [u8]>(&self.item_db, item_key)?,
        };

        let root_vec = Reader::get_root(item)?.as_vector();

        let item_uuid = root_vec.idx(0).as_blob();
        item_vec.push(item_uuid);

        for field in &model_config.fields {
            match &field.kind {
                Kind::Ref {
                    model,
                    field: _,
                    many,
                } => {
                    if let Some(((field_idx, limit, cursor), ref_depth)) =
                        ref_depth.0.iter().find(|v| v.0.0 == field.idx)
                    {
                        if let Some(model_config) = self.model_config_map.get(model) {
                            if many == &Some(true) {
                                let mut skip_key = match cursor {
                                    Some(cursor) => {
                                        let mut index_key = BytesMut::new();

                                        index_key.put_u8(model_config.idx);
                                        index_key.put_u8(0);
                                        index_key.put(cursor.as_ref());

                                        let item_key = match access {
                                            Accessor::Write(access) => access.get::<[u8], [u8]>(
                                                &self.index_db,
                                                index_key.as_ref(),
                                            )?,
                                            Accessor::Const(access) => access.get::<[u8], [u8]>(
                                                &self.index_db,
                                                index_key.as_ref(),
                                            )?,
                                        };
                                        Some(item_key)
                                    }
                                    None => None,
                                };

                                let mut limit = *limit as usize;
                                let mut ref_list_vec = item_vec.start_vector();

                                for ref_id in &root_vec.idx(*field_idx as usize).as_vector() {
                                    if limit == 0 {
                                        break;
                                    }

                                    let ref_id_bytes = ref_id.as_blob().0;

                                    if let Some(sk) = skip_key {
                                        if sk == ref_id_bytes {
                                            skip_key = None;
                                        }

                                        continue;
                                    }

                                    let mut ref_vec = ref_list_vec.start_vector();

                                    count = self.build_item(
                                        model_config,
                                        access,
                                        ref_id_bytes,
                                        &mut ref_vec,
                                        ref_depth,
                                        count,
                                    )?;

                                    ref_vec.end_vector();

                                    limit -= 1;
                                }

                                ref_list_vec.end_vector();
                            } else if *limit > 0 {
                                let ref_id_bytes = root_vec.idx(*field_idx as usize).as_blob().0;

                                let mut ref_vec = item_vec.start_vector();

                                count = self.build_item(
                                    model_config,
                                    access,
                                    ref_id_bytes,
                                    &mut ref_vec,
                                    ref_depth,
                                    count,
                                )?;

                                ref_vec.end_vector();
                            } else {
                                item_vec.push(());
                            }
                        }
                    } else {
                        item_vec.push(());
                    }
                }
                _ => {
                    if field.encrypted == Some(true) {
                        let cipher = XChaCha20Poly1305::new(&self.encryption_key.into());
                        let bytes = root_vec.idx(field.idx as usize).as_blob().0;
                        let len = bytes.len() - 24;

                        let ciphertext = &bytes[..len - 24];
                        let nonce: [u8; 24] = bytes[len - 24..].try_into()?;
                        let nonce = XNonce::try_from(nonce.as_slice())?;

                        match cipher.decrypt(&nonce, ciphertext) {
                            Ok(plaintext) => {
                                if field.compressed == Some(true) {
                                    let decompressed =
                                        zstd::stream::decode_all(plaintext.as_slice())?;
                                    push_field_from_bytes(field, &decompressed, item_vec)?;
                                } else {
                                    push_field_from_bytes(field, &plaintext, item_vec)?;
                                }
                            }
                            Err(err) => tracing::error!(%err),
                        }
                    } else if field.compressed == Some(true) {
                        let bytes = root_vec.idx(field.idx as usize).as_blob().0;

                        let decompressed = zstd::stream::decode_all(bytes)?;

                        push_field_from_bytes(field, &decompressed, item_vec)?;
                    } else {
                        field
                            .kind
                            .copy_to(&root_vec.idx(field.idx as usize), item_vec, None)?;
                    }
                }
            }
        }

        Ok(count + 1)
    }
}