diff --git a/core/store/src/db/rocksdb.rs b/core/store/src/db/rocksdb.rs index 0f435398acb..aff7d91b506 100644 --- a/core/store/src/db/rocksdb.rs +++ b/core/store/src/db/rocksdb.rs @@ -297,6 +297,13 @@ impl RocksDB { _ => unreachable!(), } } + + pub fn compact_column(&self, col: DBCol) -> io::Result<()> { + let none = Option::<&[u8]>::None; + tracing::info!(target: "db", column = %col, "Compact column"); + self.db.compact_range_cf(self.cf_handle(col)?, none, none); + Ok(()) + } } impl Database for RocksDB { @@ -375,10 +382,8 @@ impl Database for RocksDB { } fn compact(&self) -> io::Result<()> { - let none = Option::<&[u8]>::None; for col in DBCol::iter() { - tracing::info!(target: "db", column = %col, "Compacted column"); - self.db.compact_range_cf(self.cf_handle(col)?, none, none); + self.compact_column(col)?; } Ok(()) } diff --git a/tools/database/src/analyse_data_size_distribution.rs b/tools/database/src/analyse_data_size_distribution.rs index b41e5319aa5..b6142394b6c 100644 --- a/tools/database/src/analyse_data_size_distribution.rs +++ b/tools/database/src/analyse_data_size_distribution.rs @@ -8,7 +8,7 @@ use std::sync::{Arc, Mutex}; use std::{panic, println}; use strum::IntoEnumIterator; -use crate::utils::open_rocksdb; +use crate::utils::{open_rocksdb, resolve_column}; #[derive(Parser)] pub(crate) struct AnalyseDataSizeDistributionCommand { @@ -21,10 +21,6 @@ pub(crate) struct AnalyseDataSizeDistributionCommand { top_k: usize, } -fn resolve_db_col(col: &str) -> Option { - DBCol::iter().filter(|db_col| <&str>::from(db_col) == col).next() -} - #[derive(Clone)] struct ColumnFamilyCountAndSize { number_of_pairs: usize, @@ -186,19 +182,17 @@ fn read_all_pairs(db: &RocksDB, col_families: &Vec) -> DataSizeDistributi DataSizeDistribution::new(key_sizes, value_sizes, column_families) } -fn get_column_families(input_col: &Option) -> Vec { +fn get_column_families(input_col: &Option) -> anyhow::Result> { match input_col { - Some(column_name) => { - vec![resolve_db_col(&column_name).unwrap()] - } - None => DBCol::iter().collect(), + Some(column_name) => Ok(vec![resolve_column(column_name)?]), + None => Ok(DBCol::iter().collect()), } } impl AnalyseDataSizeDistributionCommand { pub(crate) fn run(&self, home: &PathBuf) -> anyhow::Result<()> { let db = open_rocksdb(home, near_store::Mode::ReadOnly)?; - let column_families = get_column_families(&self.column); + let column_families = get_column_families(&self.column)?; let results = read_all_pairs(&db, &column_families); results.print_results(self.top_k); diff --git a/tools/database/src/compact.rs b/tools/database/src/compact.rs index 026172d2aba..d70f29bf574 100644 --- a/tools/database/src/compact.rs +++ b/tools/database/src/compact.rs @@ -1,15 +1,23 @@ -use crate::utils::open_rocksdb; +use crate::utils::{open_rocksdb, resolve_column}; use clap::Parser; use near_store::db::Database; use std::path::PathBuf; #[derive(Parser)] -pub(crate) struct RunCompactionCommand {} +pub(crate) struct RunCompactionCommand { + /// If specified only this column will compacted + #[arg(short, long)] + column: Option, +} impl RunCompactionCommand { pub(crate) fn run(&self, home: &PathBuf) -> anyhow::Result<()> { let db = open_rocksdb(home, near_store::Mode::ReadWrite)?; - db.compact()?; + if let Some(col_name) = &self.column { + db.compact_column(resolve_column(col_name)?)?; + } else { + db.compact()?; + } eprintln!("Compaction is finished!"); Ok(()) } diff --git a/tools/database/src/utils.rs b/tools/database/src/utils.rs index f84c66aea0d..85eb6de28bf 100644 --- a/tools/database/src/utils.rs +++ b/tools/database/src/utils.rs @@ -1,5 +1,9 @@ use std::path::Path; +use anyhow::anyhow; +use near_store::DBCol; +use strum::IntoEnumIterator; + pub(crate) fn open_rocksdb( home: &Path, mode: near_store::Mode, @@ -13,3 +17,10 @@ pub(crate) fn open_rocksdb( near_store::db::RocksDB::open(&db_path, store_config, mode, near_store::Temperature::Hot)?; Ok(rocksdb) } + +pub(crate) fn resolve_column(col_name: &str) -> anyhow::Result { + DBCol::iter() + .filter(|db_col| <&str>::from(db_col) == col_name) + .next() + .ok_or_else(|| anyhow!("column {col_name} does not exist")) +}