mr-ability 0.6.0

//! # RocksDB 底层封装
//!
//! 对 RocksDB 的列族（Column Family）操作进行封装，
//! 为上层 [`MemoryStore`]、[`FacetStore`]、[`EdgeStore`] 和 [`RocksDBVectorStore`] 提供基础读写能力。
//!
//! ## 列族划分
//!
//! | 列族 | 用途 |
//! |------|------|
//! | `memories` | 记忆主存储（UUID → JSON） |
//! | `by_tag` | 标签索引（`tag:id` → UUID） |
//! | `deleted` | 已删除记忆索引 |
//! | `importance` | 重要性索引（UUID → f64 字符串） |
//! | `projects` | 项目信息 |
//! | `config` | 配置键值对 |
//! | `facets` | Facet 主存储（UUID → JSON） |
//! | `memory_facets` | 记忆→Facet 索引（`memory_id:facet_id` → UUID） |
//! | `edges` | 边主存储（UUID → JSON） |
//! | `node_edges_out` | 出边索引（`source_id:edge_type` → `[edge_ids]`） |
//! | `node_edges_in` | 入边索引（`target_id:edge_type` → `[edge_ids]`） |
//! | `propagation_deltas` | 传播增量主存储（UUID → JSON） |
//! | `pending_propagations` | 待应用增量队列（`target_id:priority:delta_id` → UUID） |
//! | `applied_propagations` | 已应用增量索引（`source_id:target_id` → UUID） |
//! | `rules` | 检索规则主存储（UUID → JSON） |
//! | `rule_stats` | 规则统计（UUID → JSON） |
//! | `memory_summaries` | 记忆摘要（UUID → String） |
//! | `memory_keywords` | 记忆关键词（UUID → JSON） |

use anyhow::{Context, Result};
use rocksdb::{ColumnFamilyDescriptor, Options, DB};
use std::path::Path;

/// 记忆主存储列族
const CF_MEMORIES: &str = "memories";
/// 标签索引列族
const CF_BY_TAG: &str = "by_tag";
/// 已删除记忆索引列族
const CF_DELETED: &str = "deleted";
/// 重要性索引列族
const CF_IMPORTANCE: &str = "importance";
/// 项目信息列族
const CF_PROJECTS: &str = "projects";
/// 配置键值对列族
const CF_CONFIG: &str = "config";
/// Facet 主存储列族
const CF_FACETS: &str = "facets";
/// 记忆→Facet 索引列族
const CF_MEMORY_FACETS: &str = "memory_facets";
/// 边主存储列族
const CF_EDGES: &str = "edges";
/// 出边索引列族
const CF_NODE_EDGES_OUT: &str = "node_edges_out";
/// 入边索引列族
const CF_NODE_EDGES_IN: &str = "node_edges_in";
/// 传播增量主存储列族
const CF_PROPAGATION_DELTAS: &str = "propagation_deltas";
/// 待应用增量队列列族
const CF_PENDING_PROPAGATIONS: &str = "pending_propagations";
/// 已应用增量索引列族
const CF_APPLIED_PROPAGATIONS: &str = "applied_propagations";
/// 检索规则主存储列族
const CF_RULES: &str = "rules";
/// 规则统计列族
const CF_RULE_STATS: &str = "rule_stats";
/// 记忆摘要列族
const CF_MEMORY_SUMMARIES: &str = "memory_summaries";
/// 记忆关键词列族
const CF_MEMORY_KEYWORDS: &str = "memory_keywords";

/// RocksDB 存储封装，提供列族访问和基础读写操作。
pub struct RocksDBStore {
    db: DB,
}

impl RocksDBStore {
    /// 打开 RocksDB 数据库，自动创建所有列族。
    pub fn open(path: &Path) -> Result<Self> {
        let mut options = Options::default();
        options.create_if_missing(true);
        options.create_missing_column_families(true);

        let cfs = vec![
            ColumnFamilyDescriptor::new(CF_MEMORIES, Options::default()),
            ColumnFamilyDescriptor::new(CF_BY_TAG, Options::default()),
            ColumnFamilyDescriptor::new(CF_DELETED, Options::default()),
            ColumnFamilyDescriptor::new(CF_IMPORTANCE, Options::default()),
            ColumnFamilyDescriptor::new(CF_PROJECTS, Options::default()),
            ColumnFamilyDescriptor::new(CF_CONFIG, Options::default()),
            ColumnFamilyDescriptor::new(CF_FACETS, Options::default()),
            ColumnFamilyDescriptor::new(CF_MEMORY_FACETS, Options::default()),
            ColumnFamilyDescriptor::new(CF_EDGES, Options::default()),
            ColumnFamilyDescriptor::new(CF_NODE_EDGES_OUT, Options::default()),
            ColumnFamilyDescriptor::new(CF_NODE_EDGES_IN, Options::default()),
            ColumnFamilyDescriptor::new(CF_PROPAGATION_DELTAS, Options::default()),
            ColumnFamilyDescriptor::new(CF_PENDING_PROPAGATIONS, Options::default()),
            ColumnFamilyDescriptor::new(CF_APPLIED_PROPAGATIONS, Options::default()),
            ColumnFamilyDescriptor::new(CF_RULES, Options::default()),
            ColumnFamilyDescriptor::new(CF_RULE_STATS, Options::default()),
            ColumnFamilyDescriptor::new(CF_MEMORY_SUMMARIES, Options::default()),
            ColumnFamilyDescriptor::new(CF_MEMORY_KEYWORDS, Options::default()),
        ];

        let db = DB::open_cf_descriptors(&options, path, cfs).context("Failed to open RocksDB")?;

        Ok(Self { db })
    }

    /// 获取 `memories` 列族句柄。
    pub fn cf_memories(&self) -> Result<&rocksdb::ColumnFamily> {
        self.db
            .cf_handle(CF_MEMORIES)
            .context("Column family 'memories' not found")
    }

    /// 获取 `by_tag` 列族句柄。
    pub fn cf_by_tag(&self) -> Result<&rocksdb::ColumnFamily> {
        self.db
            .cf_handle(CF_BY_TAG)
            .context("Column family 'by_tag' not found")
    }

    /// 获取 `deleted` 列族句柄。
    pub fn cf_deleted(&self) -> Result<&rocksdb::ColumnFamily> {
        self.db
            .cf_handle(CF_DELETED)
            .context("Column family 'deleted' not found")
    }

    /// 获取 `importance` 列族句柄。
    pub fn cf_importance(&self) -> Result<&rocksdb::ColumnFamily> {
        self.db
            .cf_handle(CF_IMPORTANCE)
            .context("Column family 'importance' not found")
    }

    /// 获取 `projects` 列族句柄。
    pub fn cf_projects(&self) -> Result<&rocksdb::ColumnFamily> {
        self.db
            .cf_handle(CF_PROJECTS)
            .context("Column family 'projects' not found")
    }

    /// 获取 `config` 列族句柄。
    pub fn cf_config(&self) -> Result<&rocksdb::ColumnFamily> {
        self.db
            .cf_handle(CF_CONFIG)
            .context("Column family 'config' not found")
    }

    /// 获取 `facets` 列族句柄。
    pub fn cf_facets(&self) -> Result<&rocksdb::ColumnFamily> {
        self.db
            .cf_handle(CF_FACETS)
            .context("Column family 'facets' not found")
    }

    /// 获取 `memory_facets` 列族句柄。
    pub fn cf_memory_facets(&self) -> Result<&rocksdb::ColumnFamily> {
        self.db
            .cf_handle(CF_MEMORY_FACETS)
            .context("Column family 'memory_facets' not found")
    }

    /// 获取 `edges` 列族句柄。
    pub fn cf_edges(&self) -> Result<&rocksdb::ColumnFamily> {
        self.db
            .cf_handle(CF_EDGES)
            .context("Column family 'edges' not found")
    }

    /// 获取 `node_edges_out` 列族句柄。
    pub fn cf_node_edges_out(&self) -> Result<&rocksdb::ColumnFamily> {
        self.db
            .cf_handle(CF_NODE_EDGES_OUT)
            .context("Column family 'node_edges_out' not found")
    }

    /// 获取 `node_edges_in` 列族句柄。
    pub fn cf_node_edges_in(&self) -> Result<&rocksdb::ColumnFamily> {
        self.db
            .cf_handle(CF_NODE_EDGES_IN)
            .context("Column family 'node_edges_in' not found")
    }

    /// 获取 `propagation_deltas` 列族句柄。
    pub fn cf_propagation_deltas(&self) -> Result<&rocksdb::ColumnFamily> {
        self.db
            .cf_handle(CF_PROPAGATION_DELTAS)
            .context("Column family 'propagation_deltas' not found")
    }

    /// 获取 `pending_propagations` 列族句柄。
    pub fn cf_pending_propagations(&self) -> Result<&rocksdb::ColumnFamily> {
        self.db
            .cf_handle(CF_PENDING_PROPAGATIONS)
            .context("Column family 'pending_propagations' not found")
    }

    /// 获取 `applied_propagations` 列族句柄。
    pub fn cf_applied_propagations(&self) -> Result<&rocksdb::ColumnFamily> {
        self.db
            .cf_handle(CF_APPLIED_PROPAGATIONS)
            .context("Column family 'applied_propagations' not found")
    }

    /// 获取 `rules` 列族句柄。
    pub fn cf_rules(&self) -> Result<&rocksdb::ColumnFamily> {
        self.db
            .cf_handle(CF_RULES)
            .context("Column family 'rules' not found")
    }

    /// 获取 `rule_stats` 列族句柄。
    pub fn cf_rule_stats(&self) -> Result<&rocksdb::ColumnFamily> {
        self.db
            .cf_handle(CF_RULE_STATS)
            .context("Column family 'rule_stats' not found")
    }

    /// 获取 `memory_summaries` 列族句柄。
    pub fn cf_memory_summaries(&self) -> Result<&rocksdb::ColumnFamily> {
        self.db
            .cf_handle(CF_MEMORY_SUMMARIES)
            .context("Column family 'memory_summaries' not found")
    }

    /// 获取 `memory_keywords` 列族句柄。
    pub fn cf_memory_keywords(&self) -> Result<&rocksdb::ColumnFamily> {
        self.db
            .cf_handle(CF_MEMORY_KEYWORDS)
            .context("Column family 'memory_keywords' not found")
    }

    /// 向指定列族写入键值对。
    pub fn put_cf(&self, cf: &rocksdb::ColumnFamily, key: &[u8], value: &[u8]) -> Result<()> {
        self.db
            .put_cf(cf, key, value)
            .context("Failed to put value")
    }

    /// 从指定列族读取键对应的值。
    pub fn get_cf(&self, cf: &rocksdb::ColumnFamily, key: &[u8]) -> Result<Option<Vec<u8>>> {
        self.db.get_cf(cf, key).context("Failed to get value")
    }

    /// 从指定列族删除键。
    pub fn delete_cf(&self, cf: &rocksdb::ColumnFamily, key: &[u8]) -> Result<()> {
        self.db.delete_cf(cf, key).context("Failed to delete value")
    }

    /// 获取指定列族的原始迭代器，用于全量遍历。
    pub fn iter_cf(&self, cf: &rocksdb::ColumnFamily) -> rocksdb::DBRawIterator<'_> {
        self.db.raw_iterator_cf(cf)
    }
}

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

    #[test]
    fn test_rocksdb_open() {
        let dir = tempdir().unwrap();
        let store = RocksDBStore::open(dir.path()).unwrap();
        assert!(store.cf_memories().is_ok());
    }

    #[test]
    fn test_rocksdb_put_get() {
        let dir = tempdir().unwrap();
        let store = RocksDBStore::open(dir.path()).unwrap();

        let cf = store.cf_memories().unwrap();
        store.put_cf(cf, b"test_key", b"test_value").unwrap();

        let value = store.get_cf(cf, b"test_key").unwrap();
        assert_eq!(value, Some(b"test_value".to_vec()));
    }
}