feat: expose open_rocksdb() functions (#240)

varfish-org · Oct 13, 2023 · 9f9fd2d · 9f9fd2d
1 parent a4e5283
commit 9f9fd2d
Show file tree

Hide file tree

Showing 10 changed files with 172 additions and 74 deletions.
diff --git a/src/clinvar_genes/cli/query.rs b/src/clinvar_genes/cli/query.rs
@@ -28,16 +28,18 @@ pub struct Args {
     pub hgnc_id: String,
 }
 
-/// Open RocksDB database.
-fn open_rocksdb(
-    args: &Args,
+/// Open RocksDb given path and column family name for data and metadata.
+pub fn open_rocksdb<P: AsRef<std::path::Path>>(
+    path_rocksdb: P,
+    cf_data: &str,
+    cf_meta: &str,
 ) -> Result<Arc<rocksdb::DBWithThreadMode<rocksdb::MultiThreaded>>, anyhow::Error> {
     tracing::info!("Opening RocksDB database ...");
     let before_open = std::time::Instant::now();
-    let cf_names = &["meta", &args.cf_name];
+    let cf_names = &[cf_meta, cf_data];
     let db = Arc::new(rocksdb::DB::open_cf_for_read_only(
         &rocksdb::Options::default(),
-        &args.path_rocksdb,
+        &path_rocksdb,
         cf_names,
         true,
     )?);
@@ -50,6 +52,13 @@ fn open_rocksdb(
     Ok(db)
 }
 
+/// Open RocksDB database from command line arguments.
+pub fn open_rocksdb_from_args(
+    args: &Args,
+) -> Result<Arc<rocksdb::DBWithThreadMode<rocksdb::MultiThreaded>>, anyhow::Error> {
+    open_rocksdb(&args.path_rocksdb, &args.cf_name, "meta")
+}
+
 /// Print values to `out_writer`.
 fn print_record(
     out_writer: &mut Box<dyn std::io::Write>,
@@ -72,7 +81,7 @@ pub fn run(common: &common::cli::Args, args: &Args) -> Result<(), anyhow::Error>
     tracing::info!("args = {:#?}", &args);
 
     // Open the RocksDB database.
-    let db = open_rocksdb(args)?;
+    let db = open_rocksdb_from_args(args)?;
     let cf_data = db.cf_handle(&args.cf_name).unwrap();
 
     // Obtain writer to output.

diff --git a/src/clinvar_minimal/cli/query.rs b/src/clinvar_minimal/cli/query.rs
@@ -34,27 +34,29 @@ pub struct Args {
 
 /// Meta information as read from database.
 #[derive(Debug)]
-struct Meta {
+pub struct Meta {
     /// Genome release of data in database.
     pub genome_release: String,
 }
 
-/// Open RocksDB database.
-fn open_rocksdb(
-    args: &Args,
+/// Open RocksDb given path and column family name for data and metadata.
+pub fn open_rocksdb<P: AsRef<std::path::Path>>(
+    path_rocksdb: P,
+    cf_data: &str,
+    cf_meta: &str,
 ) -> Result<(Arc<rocksdb::DBWithThreadMode<rocksdb::MultiThreaded>>, Meta), anyhow::Error> {
     tracing::info!("Opening RocksDB database ...");
     let before_open = std::time::Instant::now();
-    let cf_names = &["meta", &args.cf_name];
+    let cf_names = &[cf_meta, cf_data];
     let db = Arc::new(rocksdb::DB::open_cf_for_read_only(
         &rocksdb::Options::default(),
-        &args.path_rocksdb,
+        &path_rocksdb,
         cf_names,
         true,
     )?);
     tracing::info!("  reading meta information");
     let meta = {
-        let cf_meta = db.cf_handle("meta").unwrap();
+        let cf_meta = db.cf_handle(cf_meta).unwrap();
         let meta_genome_release = String::from_utf8(
             db.get_cf(&cf_meta, "genome-release")?
                 .ok_or_else(|| anyhow::anyhow!("missing value meta:genome-release"))?,
@@ -73,6 +75,13 @@ fn open_rocksdb(
     Ok((db, meta))
 }
 
+/// Open RocksDB database from command line arguments.
+pub fn open_rocksdb_from_args(
+    args: &Args,
+) -> Result<(Arc<rocksdb::DBWithThreadMode<rocksdb::MultiThreaded>>, Meta), anyhow::Error> {
+    open_rocksdb(&args.path_rocksdb, &args.cf_name, "meta")
+}
+
 fn print_record(
     out_writer: &mut Box<dyn std::io::Write>,
     output_format: common::cli::OutputFormat,
@@ -116,7 +125,7 @@ pub fn run(common: &common::cli::Args, args: &Args) -> Result<(), anyhow::Error>
     tracing::info!("common = {:#?}", &common);
     tracing::info!("args = {:#?}", &args);
 
-    let (db, meta) = open_rocksdb(args)?;
+    let (db, meta) = open_rocksdb_from_args(args)?;
     let cf_data = db.cf_handle(&args.cf_name).unwrap();
 
     // Obtain writer to output.

diff --git a/src/cons/cli/query.rs b/src/cons/cli/query.rs
@@ -48,27 +48,29 @@ pub struct ArgsQuery {
 
 /// Meta information as read from database.
 #[derive(Debug)]
-struct Meta {
+pub struct Meta {
     /// Genome release of data in database.
     pub genome_release: String,
 }
 
-/// Open RocksDB database.
-fn open_rocksdb(
-    args: &Args,
+/// Open RocksDb given path and column family name for data and metadata.
+pub fn open_rocksdb<P: AsRef<std::path::Path>>(
+    path_rocksdb: P,
+    cf_data: &str,
+    cf_meta: &str,
 ) -> Result<(Arc<rocksdb::DBWithThreadMode<rocksdb::MultiThreaded>>, Meta), anyhow::Error> {
     tracing::info!("Opening RocksDB database ...");
     let before_open = std::time::Instant::now();
-    let cf_names = &["meta", &args.cf_name];
+    let cf_names = &[cf_meta, cf_data];
     let db = Arc::new(rocksdb::DB::open_cf_for_read_only(
         &rocksdb::Options::default(),
-        &args.path_rocksdb,
+        &path_rocksdb,
         cf_names,
         true,
     )?);
     tracing::info!("  reading meta information");
     let meta = {
-        let cf_meta = db.cf_handle("meta").unwrap();
+        let cf_meta = db.cf_handle(cf_meta).unwrap();
         let meta_genome_release = String::from_utf8(
             db.get_cf(&cf_meta, "genome-release")?
                 .ok_or_else(|| anyhow::anyhow!("missing value meta:genome-release"))?,
@@ -87,6 +89,13 @@ fn open_rocksdb(
     Ok((db, meta))
 }
 
+/// Open RocksDB database from command line arguments.
+pub fn open_rocksdb_from_args(
+    args: &Args,
+) -> Result<(Arc<rocksdb::DBWithThreadMode<rocksdb::MultiThreaded>>, Meta), anyhow::Error> {
+    open_rocksdb(&args.path_rocksdb, &args.cf_name, "meta")
+}
+
 /// Print values to `out_writer`.
 fn print_values(
     out_writer: &mut Box<dyn std::io::Write>,
@@ -109,7 +118,7 @@ pub fn run(common: &common::cli::Args, args: &Args) -> Result<(), anyhow::Error>
     tracing::info!("args = {:#?}", &args);
 
     // Open the RocksDB database.
-    let (db, meta) = open_rocksdb(args)?;
+    let (db, meta) = open_rocksdb_from_args(args)?;
     let cf_data = db.cf_handle(&args.cf_name).unwrap();
 
     // Obtain writer to output.

diff --git a/src/dbsnp/cli/query.rs b/src/dbsnp/cli/query.rs
@@ -34,7 +34,7 @@ pub struct Args {
 
 /// Meta information as read from database.
 #[derive(Debug)]
-struct Meta {
+pub struct Meta {
     /// Genome release of data in database.
     pub genome_release: String,
     /// Name of the database.
@@ -43,22 +43,24 @@ struct Meta {
     pub db_version: String,
 }
 
-/// Open RocksDB database.
-fn open_rocksdb(
-    args: &Args,
+/// Open RocksDb given path and column family name for data and metadata.
+pub fn open_rocksdb<P: AsRef<std::path::Path>>(
+    path_rocksdb: P,
+    cf_data: &str,
+    cf_meta: &str,
 ) -> Result<(Arc<rocksdb::DBWithThreadMode<rocksdb::MultiThreaded>>, Meta), anyhow::Error> {
     tracing::info!("Opening RocksDB database ...");
     let before_open = std::time::Instant::now();
-    let cf_names = &["meta", &args.cf_name];
+    let cf_names = &[cf_meta, cf_data];
     let db = Arc::new(rocksdb::DB::open_cf_for_read_only(
         &rocksdb::Options::default(),
-        &args.path_rocksdb,
+        &path_rocksdb,
         cf_names,
         true,
     )?);
     tracing::info!("  reading meta information");
     let meta = {
-        let cf_meta = db.cf_handle("meta").unwrap();
+        let cf_meta = db.cf_handle(cf_meta).unwrap();
         let meta_db_name = String::from_utf8(
             db.get_cf(&cf_meta, "db-name")?
                 .ok_or_else(|| anyhow::anyhow!("missing value meta:db-schema"))?,
@@ -89,6 +91,13 @@ fn open_rocksdb(
     Ok((db, meta))
 }
 
+/// Open RocksDB database.
+pub fn open_rocksdb_from_args(
+    args: &Args,
+) -> Result<(Arc<rocksdb::DBWithThreadMode<rocksdb::MultiThreaded>>, Meta), anyhow::Error> {
+    open_rocksdb(&args.path_rocksdb, &args.cf_name, "meta")
+}
+
 fn print_record(
     out_writer: &mut Box<dyn std::io::Write>,
     output_format: common::cli::OutputFormat,
@@ -132,7 +141,7 @@ pub fn run(common: &common::cli::Args, args: &Args) -> Result<(), anyhow::Error>
     tracing::info!("common = {:#?}", &common);
     tracing::info!("args = {:#?}", &args);
 
-    let (db, meta) = open_rocksdb(args)?;
+    let (db, meta) = open_rocksdb_from_args(args)?;
     let cf_data = db.cf_handle(&args.cf_name).unwrap();
 
     // Obtain writer to output.

diff --git a/src/freqs/cli/query.rs b/src/freqs/cli/query.rs
@@ -34,22 +34,26 @@ pub struct Meta {
     pub genome_release: String,
 }
 
-/// Open RocksDB database.
-fn open_rocksdb(
-    args: &Args,
+/// Open RocksDb given path and column family name for data and metadata.
+pub fn open_rocksdb<P: AsRef<std::path::Path>>(
+    path_rocksdb: P,
+    cf_auto: &str,
+    cf_gono: &str,
+    cf_mtdna: &str,
+    cf_meta: &str,
 ) -> Result<(Arc<rocksdb::DBWithThreadMode<rocksdb::MultiThreaded>>, Meta), anyhow::Error> {
     tracing::info!("Opening RocksDB database ...");
     let before_open = std::time::Instant::now();
-    let cf_names = &["meta", "autosomal", "gonosomal", "mitochondrial"];
+    let cf_names = &[cf_meta, cf_auto, cf_gono, cf_mtdna];
     let db = Arc::new(rocksdb::DB::open_cf_for_read_only(
         &rocksdb::Options::default(),
-        &args.path_rocksdb,
+        &path_rocksdb,
         cf_names,
         true,
     )?);
     tracing::info!("  reading meta information");
     let meta = {
-        let cf_meta = db.cf_handle("meta").unwrap();
+        let cf_meta = db.cf_handle(cf_meta).unwrap();
         let meta_genome_release = String::from_utf8(
             db.get_cf(&cf_meta, "genome-release")?
                 .ok_or_else(|| anyhow::anyhow!("missing value meta:genome-release"))?,
@@ -68,6 +72,19 @@ fn open_rocksdb(
     Ok((db, meta))
 }
 
+/// Open RocksDB database from command line arguments.
+pub fn open_rocksdb_from_args(
+    args: &Args,
+) -> Result<(Arc<rocksdb::DBWithThreadMode<rocksdb::MultiThreaded>>, Meta), anyhow::Error> {
+    open_rocksdb(
+        &args.path_rocksdb,
+        "autosomal",
+        "gonosomal",
+        "mitochondrial",
+        "meta",
+    )
+}
+
 fn query_for_variant(
     variant: &spdi::Var,
     db: &rocksdb::DBWithThreadMode<rocksdb::MultiThreaded>,
@@ -124,7 +141,7 @@ pub fn run(common: &common::cli::Args, args: &Args) -> Result<(), anyhow::Error>
         }
     };
 
-    let (db, _meta) = open_rocksdb(args)?;
+    let (db, _meta) = open_rocksdb_from_args(args)?;
 
     tracing::info!("Running query...");
     let before_query = std::time::Instant::now();

diff --git a/src/genes/cli/query.rs b/src/genes/cli/query.rs
@@ -28,16 +28,18 @@ pub struct Args {
     pub hgnc_id: String,
 }
 
-/// Open RocksDB database.
-fn open_rocksdb(
-    args: &Args,
+/// Open RocksDb given path and column family name for data and metadata.
+pub fn open_rocksdb<P: AsRef<std::path::Path>>(
+    path_rocksdb: P,
+    cf_data: &str,
+    cf_meta: &str,
 ) -> Result<Arc<rocksdb::DBWithThreadMode<rocksdb::MultiThreaded>>, anyhow::Error> {
     tracing::info!("Opening RocksDB database ...");
     let before_open = std::time::Instant::now();
-    let cf_names = &["meta", &args.cf_name];
+    let cf_names: &[&str; 2] = &[cf_meta, cf_data];
     let db = Arc::new(rocksdb::DB::open_cf_for_read_only(
         &rocksdb::Options::default(),
-        &args.path_rocksdb,
+        &path_rocksdb,
         cf_names,
         true,
     )?);
@@ -50,6 +52,13 @@ fn open_rocksdb(
     Ok(db)
 }
 
+/// Open RocksDB database from command line arguments.
+pub fn open_rocksdb_from_args(
+    args: &Args,
+) -> Result<Arc<rocksdb::DBWithThreadMode<rocksdb::MultiThreaded>>, anyhow::Error> {
+    open_rocksdb(&args.path_rocksdb, &args.cf_name, "meta")
+}
+
 /// Print values to `out_writer`.
 fn print_record(
     out_writer: &mut Box<dyn std::io::Write>,
@@ -72,7 +81,7 @@ pub fn run(common: &common::cli::Args, args: &Args) -> Result<(), anyhow::Error>
     tracing::info!("args = {:#?}", &args);
 
     // Open the RocksDB database.
-    let db = open_rocksdb(args)?;
+    let db = open_rocksdb_from_args(args)?;
     let cf_data = db.cf_handle(&args.cf_name).unwrap();
 
     // Obtain writer to output.

diff --git a/src/gnomad_mtdna/cli/query.rs b/src/gnomad_mtdna/cli/query.rs
@@ -34,27 +34,29 @@ pub struct Args {
 
 /// Meta information as read from database.
 #[derive(Debug)]
-struct Meta {
+pub struct Meta {
     /// Genome release of data in database.
     pub genome_release: String,
 }
 
-/// Open RocksDB database.
-fn open_rocksdb(
-    args: &Args,
+/// Open RocksDb given path and column family name for data and metadata.
+pub fn open_rocksdb<P: AsRef<std::path::Path>>(
+    path_rocksdb: P,
+    cf_data: &str,
+    cf_meta: &str,
 ) -> Result<(Arc<rocksdb::DBWithThreadMode<rocksdb::MultiThreaded>>, Meta), anyhow::Error> {
     tracing::info!("Opening RocksDB database ...");
     let before_open = std::time::Instant::now();
-    let cf_names = &["meta", &args.cf_name];
+    let cf_names: &[&str; 2] = &[cf_meta, cf_data];
     let db = Arc::new(rocksdb::DB::open_cf_for_read_only(
         &rocksdb::Options::default(),
-        &args.path_rocksdb,
+        &path_rocksdb,
         cf_names,
         true,
     )?);
     tracing::info!("  reading meta information");
     let meta = {
-        let cf_meta = db.cf_handle("meta").unwrap();
+        let cf_meta = db.cf_handle(cf_meta).unwrap();
         let meta_genome_release = String::from_utf8(
             db.get_cf(&cf_meta, "genome-release")?
                 .ok_or_else(|| anyhow::anyhow!("missing value meta:genome-release"))?,
@@ -73,6 +75,13 @@ fn open_rocksdb(
     Ok((db, meta))
 }
 
+/// Open RocksDB database from command line arguments.
+pub fn open_rocksdb_from_args(
+    args: &Args,
+) -> Result<(Arc<rocksdb::DBWithThreadMode<rocksdb::MultiThreaded>>, Meta), anyhow::Error> {
+    open_rocksdb(&args.path_rocksdb, &args.cf_name, "meta")
+}
+
 fn print_record(
     out_writer: &mut Box<dyn std::io::Write>,
     output_format: common::cli::OutputFormat,
@@ -116,7 +125,7 @@ pub fn run(common: &common::cli::Args, args: &Args) -> Result<(), anyhow::Error>
     tracing::info!("common = {:#?}", &common);
     tracing::info!("args = {:#?}", &args);
 
-    let (db, meta) = open_rocksdb(args)?;
+    let (db, meta) = open_rocksdb_from_args(args)?;
     let cf_data = db.cf_handle(&args.cf_name).unwrap();
 
     // Obtain writer to output.