KeyValueTable.hpp

Go to the documentation of this file.

#pragma once
#ifndef _MDBX_CONTAINERS_KEY_VALUE_TABLE_HPP_INCLUDED
#define _MDBX_CONTAINERS_KEY_VALUE_TABLE_HPP_INCLUDED

 
#include "common.hpp"
#include <map>
#include <unordered_set>
 
namespace mdbxc {

    template<class KeyT, class ValueT>
    class KeyValueTable final : public BaseTable {
    public:

        KeyValueTable(std::shared_ptr<Connection> connection,
                      std::string name = "kv_store",
                      MDBX_db_flags_t flags = MDBX_DB_DEFAULTS | MDBX_CREATE)
            : BaseTable(std::move(connection),
                        std::move(name),
                        flags | get_mdbx_flags<KeyT>()) {}

        explicit KeyValueTable(const Config& config,
                               std::string name = "kv_store",
                               MDBX_db_flags_t flags = MDBX_DB_DEFAULTS | MDBX_CREATE)
            : BaseTable(Connection::create(config),
                        std::move(name),
                        flags | get_mdbx_flags<KeyT>()) {}

        ~KeyValueTable() override final = default;
 
        // --- Operators ---

        template<template <class...> class ContainerT>
        KeyValueTable& operator=(const ContainerT<KeyT, ValueT>& container) {
            with_transaction([this, &container](MDBX_txn* txn) {
                db_reconcile(container, txn);
            }, TransactionMode::WRITABLE);
            return *this;
        }
        
        KeyValueTable& operator=(const std::vector<std::pair<KeyT, ValueT>>& container) {
            with_transaction([this, &container](MDBX_txn* txn) {
                db_reconcile(container, txn);
            }, TransactionMode::WRITABLE);
            return *this;
        }

        template<template<class...> class ContainerT = std::map>
        auto operator()() {
            using ReturnT = std::conditional_t<
                std::is_same_v<ContainerT<KeyT, ValueT>, std::vector<std::pair<KeyT, ValueT>>>,
                std::vector<std::pair<KeyT, ValueT>>,
                ContainerT<KeyT, ValueT>
            >;
 
            ReturnT container;
            with_transaction([this, &container](MDBX_txn* txn) {
                db_load(container, txn);
            }, TransactionMode::READ_ONLY);
            return container;
        }

        class AssignmentProxy {
        public:
            AssignmentProxy(KeyValueTable& db, KeyT key)
                : m_db(db), m_key(std::move(key)) {}

            AssignmentProxy& operator=(const ValueT& value) {
                m_db.insert_or_assign(m_key, value);
                return *this;
            }

            operator ValueT() const {
                auto val = m_db.find(m_key);
                if (val) return *val;
                ValueT def{}; // default construct if missing
                m_db.insert_or_assign(m_key, def); // persist the default value
                return def;
            }
 
        private:
            KeyValueTable& m_db; 
            KeyT m_key;          
        };
        
        AssignmentProxy operator[](const KeyT& key) {
            return AssignmentProxy(*this, key);
        }
 
        // --- Existing methods ---

        template<template <class...> class ContainerT>
        void load(ContainerT<KeyT, ValueT>& container, MDBX_txn* txn = nullptr) {
            with_transaction([this, &container](MDBX_txn* txn) {
                db_load(container, txn);
            }, TransactionMode::READ_ONLY, txn);
        }
        
        template<template <class...> class ContainerT>
        void load(ContainerT<KeyT, ValueT>& container, const Transaction& txn) {
            load(container, txn.handle());
        }
        
        void load(std::vector<std::pair<KeyT, ValueT>>& container, MDBX_txn* txn = nullptr) {
            with_transaction([this, &container](MDBX_txn* txn) {
                db_load(container, txn);
            }, TransactionMode::READ_ONLY, txn);
        }
        
        void load(std::vector<std::pair<KeyT, ValueT>>& container, const Transaction& txn) {
            load(container, txn.handle());
        }

        template<template<class...> class ContainerT = std::map>
        auto retrieve_all(MDBX_txn* txn = nullptr) {
            using ReturnT = std::conditional_t<
                std::is_same_v<ContainerT<KeyT, ValueT>, std::vector<std::pair<KeyT, ValueT>>>,
                std::vector<std::pair<KeyT, ValueT>>,
                ContainerT<KeyT, ValueT>
            >;
 
            ReturnT container;
            with_transaction([this, &container](MDBX_txn* txn) {
                db_load(container, txn);
            }, TransactionMode::READ_ONLY, txn);
            return container;
        }
        
        template<template<class...> class ContainerT = std::map>
        auto retrieve_all(const Transaction& txn) {
            return retrieve_all<ContainerT>(txn.handle());
        }

        template<template <class...> class ContainerT>
        void append(const ContainerT<KeyT, ValueT>& container, MDBX_txn* txn = nullptr) {
            with_transaction([this, &container](MDBX_txn* txn) {
                db_append(container, txn);
            }, TransactionMode::WRITABLE, txn);
        }
        
        template<template <class...> class ContainerT>
        void append(const ContainerT<KeyT, ValueT>& container, const Transaction& txn) {
            append(container, txn.handle());
        }
        
        
        void append(const std::vector<std::pair<KeyT, ValueT>>& container, MDBX_txn* txn = nullptr) {
            with_transaction([this, &container](MDBX_txn* txn) {
                db_append(container, txn);
            }, TransactionMode::WRITABLE, txn);
        }

        void append(const std::vector<std::pair<KeyT, ValueT>>& container, const Transaction& txn) {
            append(container, txn.handle());
        }

        template<template <class...> class ContainerT>
        void reconcile(const ContainerT<KeyT, ValueT>& container, MDBX_txn* txn = nullptr) {
            with_transaction([this, &container](MDBX_txn* txn) {
                db_reconcile(container, txn);
            }, TransactionMode::WRITABLE, txn);
        }
        
        template<template <class...> class ContainerT>
        void reconcile(const ContainerT<KeyT, ValueT>& container, const Transaction& txn) {
            reconcile(container, txn.handle());
        }
        
        void reconcile(const std::vector<std::pair<KeyT, ValueT>>& container, MDBX_txn* txn = nullptr) {
            with_transaction([this, &container](MDBX_txn* txn) {
                db_reconcile(container, txn);
            }, TransactionMode::WRITABLE, txn);
        }

        void reconcile(const std::vector<std::pair<KeyT, ValueT>>& container, const Transaction& txn) {
            reconcile(container, txn.handle());
        }

        bool insert(const KeyT &key, const ValueT &value, MDBX_txn* txn = nullptr) {
            bool res;
            with_transaction([this, &key, &value, &res](MDBX_txn* txn) {
                res = db_insert_if_absent(key, value, txn);
            }, TransactionMode::WRITABLE, txn);
            return res;
        }
        
        bool insert(const KeyT &key, const ValueT &value, const Transaction& txn) {
            return insert(key, value, txn.handle());
        }

        bool insert(const std::pair<KeyT, ValueT> &pair, MDBX_txn* txn = nullptr) {
            bool res;
            with_transaction([this, &pair, &res](MDBX_txn* txn) {
                res = db_insert_if_absent(pair.first, pair.second, txn);
            }, TransactionMode::WRITABLE, txn);
            return res;
        }
        
        bool insert(const std::pair<KeyT, ValueT> &pair, const Transaction& txn) {
            return insert(pair, txn.handle());
        }
        
        void insert_or_assign(const KeyT &key, const ValueT &value, MDBX_txn* txn = nullptr) {
            with_transaction([this, &key, &value](MDBX_txn* txn) {
                db_insert_or_assign(key, value, txn);
            }, TransactionMode::WRITABLE, txn);
        }
        
        void insert_or_assign(const KeyT &key, const ValueT &value, const Transaction& txn) {
            insert_or_assign(key, value, txn.handle());
        }

        void insert_or_assign(const std::pair<KeyT, ValueT> &pair, MDBX_txn* txn = nullptr) {
            insert_or_assign([&pair](MDBX_txn* txn) {
                db_insert_or_assign(pair.first, pair.second, txn);
            }, TransactionMode::WRITABLE, txn);
        }
        
        void insert_or_assign(const std::pair<KeyT, ValueT> &pair, const Transaction& txn) {
            insert_or_assign(pair, txn.handle());
        }
        
        ValueT at(const KeyT& key, MDBX_txn* txn = nullptr) const {
            ValueT value;
            with_transaction([this, &key, &value](MDBX_txn* txn) {
                if (!db_get(key, value, txn)) {
                    throw std::out_of_range("Key not found in database");
                }
            }, TransactionMode::READ_ONLY, txn);
            return value;
        }
        
        ValueT at(const KeyT& key, const Transaction& txn) const {
            return at(key, txn.handle());
        }
        
        bool try_get(const KeyT& key, ValueT& out, MDBX_txn* txn) const {
            bool res;
            with_transaction([this, &key, &out, &res](MDBX_txn* txn) {
                res = db_get(key, out, txn);
            }, TransactionMode::READ_ONLY, txn);
            return res;
        }
        
        bool try_get(const KeyT& key, ValueT& out, const Transaction& txn) const {
            return try_get(key, out, txn.handle());
        }
 
#if __cplusplus >= 201703L   
        std::optional<ValueT> find(const KeyT& key, MDBX_txn* txn = nullptr) const {
            std::optional<ValueT> result;
            with_transaction([this, &key, &result](MDBX_txn* txn) {
                ValueT tmp;
                if (db_get(key, tmp, txn)) {
                    result = std::move(tmp);
                }
            }, TransactionMode::READ_ONLY, txn);
            return result;
        }
        
        std::optional<ValueT> find(const KeyT& key, const Transaction& txn) const {
            return find(key, txn.handle());
        }
#endif

        std::pair<bool, ValueT> find_compat(const KeyT& key, MDBX_txn* txn = nullptr) const {
            std::pair<bool, ValueT> result{false, ValueT{}};
            with_transaction([this, &key, &result](MDBX_txn* txn) {
                if (db_get(key, result.second, txn)) {
                    result.first = true;
                }
            }, TransactionMode::READ_ONLY, txn);
            return result;
        }
        
        std::pair<bool, ValueT> find_compat(const KeyT& key, const Transaction& txn) const {
            return find_compat(key, txn.handle());
        }

        bool contains(const KeyT& key, MDBX_txn* txn = nullptr) const {
            bool res;
            with_transaction([this, &key, &res](MDBX_txn* txn) {
                res = db_contains(key, txn);
            }, TransactionMode::READ_ONLY, txn);
            return res;
        }
        
        bool contains(const KeyT& key, const Transaction& txn) const {
            return contains(key, txn.handle());
        }

        std::size_t count(MDBX_txn* txn = nullptr) const {
            std::size_t res = 0;
            with_transaction([this, &res](MDBX_txn* txn) {
                res = db_count(txn);
            }, TransactionMode::READ_ONLY, txn);
            return res;
        }
        
        std::size_t count(const Transaction& txn) const {
            return count(txn.handle());
        }

        bool empty(MDBX_txn* txn = nullptr) const {
            std::size_t res = 0;
            with_transaction([this, &res](MDBX_txn* txn) {
                res = db_count(txn);
            }, TransactionMode::READ_ONLY, txn);
            return (res == 0);
        }
        
        bool empty(const Transaction& txn) const {
            return empty(txn.handle());
        }

        bool erase(const KeyT &key, MDBX_txn* txn = nullptr) {
            bool res;
            with_transaction([this, &key, &res](MDBX_txn* txn) {
                res = db_erase(key, txn);
            }, TransactionMode::WRITABLE, txn);
            return res;
        }
        
        bool erase(const KeyT &key, const Transaction& txn) {
            return erase(key, txn.handle());
        }

        void clear(MDBX_txn* txn = nullptr) {
            with_transaction([this](MDBX_txn* txn) {
                db_clear(txn);
            }, TransactionMode::WRITABLE, txn);
        }
        
        void clear(const Transaction& txn) {
            clear(txn.handle());
        }
 
    private:
    
        template<typename F>
        void with_transaction(F&& action, TransactionMode mode = TransactionMode::WRITABLE, MDBX_txn* txn = nullptr) const {
            if (txn) {
                action(txn);
                return;
            }
            txn = thread_txn(); // reuse transaction bound to this thread if any
            if (txn) {
                action(txn);
                return;
            }
            
            auto txn_guard = m_connection->transaction(mode);
            try {
                action(txn_guard.handle());
                txn_guard.commit();
            } catch(...) {
                // Ensure rollback on any exception and rethrow to the caller
                try {
                    txn_guard.rollback();
                } catch(...) {
                    // Ignore rollback errors here
                }
                throw;
            }
        }
    
        template<template <class...> class ContainerT>
        void db_load(ContainerT<KeyT, ValueT>& container, MDBX_txn* txn_handle) {
            MDBX_cursor* cursor = nullptr;
            check_mdbx(mdbx_cursor_open(txn_handle, m_dbi, &cursor), "Failed to open MDBX cursor");
 
            MDBX_val db_key, db_val;
            while (mdbx_cursor_get(cursor, &db_key, &db_val, MDBX_NEXT) == MDBX_SUCCESS) {
                auto&& key = deserialize_value<KeyT>(db_key);
                auto&& value = deserialize_value<ValueT>(db_val);
                container.emplace(std::move(key), std::move(value));
            }
 
            mdbx_cursor_close(cursor);
        }
        
        void db_load(std::vector<std::pair<KeyT, ValueT>>& out_vector, MDBX_txn* txn) {
            MDBX_cursor* cursor = nullptr;
            check_mdbx(mdbx_cursor_open(txn, m_dbi, &cursor), "Failed to open MDBX cursor");
 
            MDBX_val db_key, db_val;
            while (mdbx_cursor_get(cursor, &db_key, &db_val, MDBX_NEXT) == MDBX_SUCCESS) {
                auto&& key = deserialize_value<KeyT>(db_key);
                auto&& value = deserialize_value<ValueT>(db_val);
                out_vector.emplace_back(std::move(key), std::move(value));
            }
 
            mdbx_cursor_close(cursor);
        }

        bool db_get(const KeyT& key, ValueT& value, MDBX_txn* txn_handle) const {
            MDBX_val db_key = serialize_key(key);
            MDBX_val db_val;
            int rc = mdbx_get(txn_handle, m_dbi, &db_key, &db_val);
            if (rc == MDBX_NOTFOUND) return false;
            check_mdbx(rc, "Failed to retrieve value");
            value = deserialize_value<ValueT>(db_val);
            return true;
        }

        bool db_contains(const KeyT& key, MDBX_txn* txn_handle) const {
            MDBX_val db_key = serialize_key(key);
            MDBX_val db_val; // dummy
            int rc = mdbx_get(txn_handle, m_dbi, &db_key, &db_val);
            if (rc == MDBX_NOTFOUND) return false;
            check_mdbx(rc, "Failed to check key presence");
            return true;
        }

        std::size_t db_count(MDBX_txn* txn_handle) const {
            MDBX_stat stat;
            check_mdbx(mdbx_dbi_stat(txn_handle, m_dbi, &stat, sizeof(stat)), "Failed to query database statistics");
            return stat.ms_entries;
        }

        template<template <class...> class ContainerT>
        void db_append(const ContainerT<KeyT, ValueT>& container, MDBX_txn* txn_handle) {
            for (const auto& pair : container) {
                MDBX_val db_key = serialize_key(pair.first);
                MDBX_val db_val = serialize_value(pair.second);
                check_mdbx(
                    mdbx_put(txn_handle, m_dbi, &db_key, &db_val, MDBX_UPSERT),
                    "Failed to write record"
                );
            }
        }
        
        void db_append(const std::vector<std::pair<KeyT, ValueT>>& container, MDBX_txn* txn_handle) {
            for (const auto& [key, value] : container) {
                MDBX_val db_key = serialize_key(key);
                MDBX_val db_val = serialize_value(value);
                check_mdbx(
                    mdbx_put(txn_handle, m_dbi, &db_key, &db_val, MDBX_UPSERT),
                    "Failed to write record"
                );
            }
        }

        template<template <class...> class ContainerT>
        void db_reconcile(const ContainerT<KeyT, ValueT>& container, MDBX_txn* txn_handle) {
            // 1. Collect all keys from the container
            std::unordered_set<KeyT> new_keys;
            for (const auto& pair : container) {
                new_keys.insert(pair.first);
                MDBX_val db_key = serialize_key(pair.first);
                MDBX_val db_val = serialize_value(pair.second);
                check_mdbx(
                    mdbx_put(txn_handle, m_dbi, &db_key, &db_val, MDBX_UPSERT),
                    "Failed to write record"
                );
            }
 
            // 2. Iterate over existing keys in the DB and remove the extras
            MDBX_cursor* cursor = nullptr;
            check_mdbx(mdbx_cursor_open(txn_handle, m_dbi, &cursor), "Failed to open MDBX cursor");
 
            MDBX_val db_key, db_val;
            while (mdbx_cursor_get(cursor, &db_key, &db_val, MDBX_NEXT) == MDBX_SUCCESS) {
                auto&& key = deserialize_value<KeyT>(db_key);
                if (new_keys.find(key) == new_keys.end()) {
                    check_mdbx(mdbx_cursor_del(cursor, MDBX_CURRENT), "Failed to delete record using cursor");
                }
            }
 
            mdbx_cursor_close(cursor);
        }
        
        void db_reconcile(const std::vector<std::pair<KeyT, ValueT>>& container, MDBX_txn* txn_handle) {
            // 1. Collect keys and upsert values
            std::unordered_set<KeyT> new_keys;
            for (size_t i = 0; i < container.size(); ++i) {
                const KeyT& key = container[i].first;
                const ValueT& value = container[i].second;
 
                new_keys.insert(key);
 
                MDBX_val db_key = serialize_key(key);
                MDBX_val db_val = serialize_value(value);
                check_mdbx(
                    mdbx_put(txn_handle, m_dbi, &db_key, &db_val, MDBX_UPSERT),
                    "Failed to write record"
                );
            }
 
            // 2. Delete stale keys from DB
            MDBX_cursor* cursor = nullptr;
            check_mdbx(mdbx_cursor_open(txn_handle, m_dbi, &cursor), "Failed to open MDBX cursor");
 
            MDBX_val db_key, db_val;
            while (mdbx_cursor_get(cursor, &db_key, &db_val, MDBX_NEXT) == MDBX_SUCCESS) {
                KeyT key = deserialize_value<KeyT>(db_key);
                if (new_keys.find(key) == new_keys.end()) {
                    check_mdbx(mdbx_cursor_del(cursor, MDBX_CURRENT), "Failed to delete record using cursor");
                }
            }
 
            mdbx_cursor_close(cursor);
        }
        
        bool db_insert_if_absent(const KeyT& key, const ValueT& value, MDBX_txn* txn_handle) {
            MDBX_val db_key = serialize_key(key);
            MDBX_val db_val = serialize_value(value);
            int rc = mdbx_put(txn_handle, m_dbi, &db_key, &db_val, MDBX_NOOVERWRITE);
 
            if (rc == MDBX_SUCCESS)
                return true;
            if (rc == MDBX_KEYEXIST)
                return false;
 
            check_mdbx(rc, "Failed to insert key-value pair");
            return false;
        }
        
        void db_insert_or_assign(const KeyT& key, const ValueT& value, MDBX_txn* txn_handle) {
            MDBX_val db_key = serialize_key(key);
            MDBX_val db_val = serialize_value(value);
            check_mdbx(
                mdbx_put(txn_handle, m_dbi, &db_key, &db_val, MDBX_UPSERT),  // or 0
                "Failed to insert or assign key-value pair"
            );
        }

        bool db_erase(const KeyT& key, MDBX_txn* txn_handle) {
            MDBX_val db_key = serialize_key(key);
            int rc = mdbx_del(txn_handle, m_dbi, &db_key, nullptr);
            if (rc == MDBX_SUCCESS) return true;
            if (rc == MDBX_NOTFOUND) return false;
            check_mdbx(rc, "Failed to erase key");
            return false;
        }

        void db_clear(MDBX_txn* txn_handle) {
            check_mdbx(mdbx_drop(txn_handle, m_dbi, 0), "Failed to clear table");
        }
 
    }; // KeyValueTable
 
}; // namespace mdbxc
 
#endif // _MDBX_CONTAINERS_KEY_VALUE_TABLE_HPP_INCLUDED