feat: support column_family_metadata, column_family_metadata_cf (#810)

2 files changed, 84 insertions, 0 deletions

diff --git a/src/db.rs b/src/db.rs
index c6f043d..6511307 100644
--- a/src/db.rs
+++ b/src/db.rs
@@ -2007,6 +2007,47 @@ impl<T: ThreadMode, D: DBInner> DBCommon<T, D> {
         }
     }
 
+    /// Obtains the LSM-tree meta data of the default column family of the DB
+    pub fn get_column_family_metadata(&self) -> ColumnFamilyMetaData {
+        unsafe {
+            let ptr = ffi::rocksdb_get_column_family_metadata(self.inner.inner());
+
+            let metadata = ColumnFamilyMetaData {
+                size: ffi::rocksdb_column_family_metadata_get_size(ptr),
+                name: from_cstr(ffi::rocksdb_column_family_metadata_get_name(ptr)),
+                file_count: ffi::rocksdb_column_family_metadata_get_file_count(ptr),
+            };
+
+            // destroy
+            ffi::rocksdb_column_family_metadata_destroy(ptr);
+
+            // return
+            metadata
+        }
+    }
+
+    /// Obtains the LSM-tree meta data of the specified column family of the DB
+    pub fn get_column_family_metadata_cf(
+        &self,
+        cf: &impl AsColumnFamilyRef,
+    ) -> ColumnFamilyMetaData {
+        unsafe {
+            let ptr = ffi::rocksdb_get_column_family_metadata_cf(self.inner.inner(), cf.inner());
+
+            let metadata = ColumnFamilyMetaData {
+                size: ffi::rocksdb_column_family_metadata_get_size(ptr),
+                name: from_cstr(ffi::rocksdb_column_family_metadata_get_name(ptr)),
+                file_count: ffi::rocksdb_column_family_metadata_get_file_count(ptr),
+            };
+
+            // destroy
+            ffi::rocksdb_column_family_metadata_destroy(ptr);
+
+            // return
+            metadata
+        }
+    }
+
     /// Returns a list of all table files with their level, start key
     /// and end key
     pub fn live_files(&self) -> Result<Vec<LiveFile>, Error> {
@@ -2197,6 +2238,18 @@ impl<T: ThreadMode, I: DBInner> fmt::Debug for DBCommon<T, I> {
     }
 }
 
+/// The metadata that describes a column family.
+#[derive(Debug, Clone)]
+pub struct ColumnFamilyMetaData {
+    // The size of this column family in bytes, which is equal to the sum of
+    // the file size of its "levels".
+    pub size: u64,
+    // The name of the column family.
+    pub name: String,
+    // The number of files in this column family.
+    pub file_count: usize,
+}
+
 /// The metadata that describes a SST file
 #[derive(Debug, Clone)]
 pub struct LiveFile {
diff --git a/tests/test_db.rs b/tests/test_db.rs
index a4563e3..1f073f2 100644
--- a/tests/test_db.rs
+++ b/tests/test_db.rs
@@ -25,6 +25,7 @@ use rocksdb::{
     DBWithThreadMode, Env, Error, ErrorKind, FifoCompactOptions, IteratorMode, MultiThreaded,
     Options, PerfContext, PerfMetric, ReadOptions, SingleThreaded, SliceTransform, Snapshot,
     UniversalCompactOptions, UniversalCompactionStopStyle, WriteBatch, DB,
+    DEFAULT_COLUMN_FAMILY_NAME,
 };
 use util::{assert_iter, pair, DBPath};
 
@@ -376,6 +377,36 @@ fn set_option_cf_test() {
 }
 
 #[test]
+fn set_column_family_metadata_test() {
+    let path = DBPath::new("_set_column_family_metadata_test");
+    {
+        let mut opts = Options::default();
+        opts.create_if_missing(true);
+        opts.create_missing_column_families(true);
+        let db = DB::open_cf(&opts, &path, vec![DEFAULT_COLUMN_FAMILY_NAME, "cf2"]).unwrap();
+
+        let cf1 = db.cf_handle(DEFAULT_COLUMN_FAMILY_NAME).unwrap();
+        db.put_cf(&cf1, b"key1", b"value").unwrap();
+
+        let cf2 = db.cf_handle("cf2").unwrap();
+        db.put_cf(&cf2, b"key1", b"value").unwrap();
+        db.put_cf(&cf2, b"key2", b"value").unwrap();
+        db.put_cf(&cf2, b"key3", b"value").unwrap();
+
+        db.flush_cf(&cf1).unwrap();
+        db.flush_cf(&cf2).unwrap();
+
+        let default_cf_metadata = db.get_column_family_metadata();
+        assert_eq!(default_cf_metadata.size > 150, true);
+        assert_eq!(default_cf_metadata.file_count, 1);
+
+        let cf2_metadata = db.get_column_family_metadata_cf(&cf2);
+        assert_eq!(cf2_metadata.size > default_cf_metadata.size, true);
+        assert_eq!(cf2_metadata.file_count, 1);
+    }
+}
+
+#[test]
 fn test_sequence_number() {
     let path = DBPath::new("_rust_rocksdb_test_sequence_number");
     {