use std::{io::BufRead, sync::Arc};
use clap::Parser;
use prost::Message;
use crate::{
clinvar_minimal,
common::{self, keys},
pbs::annonars::clinvar::v1::minimal::ReferenceAssertion,
};
#[derive(Parser, Debug, Clone)]
#[command(about = "import minimal ClinVar data into RocksDB", long_about = None)]
pub struct Args {
#[arg(long, value_enum)]
pub genome_release: common::cli::GenomeRelease,
#[arg(long, required = true)]
pub path_in_jsonl: String,
#[arg(long)]
pub path_out_rocksdb: String,
#[arg(long, default_value = "clinvar")]
pub cf_name: String,
#[arg(long, default_value = "clinvar_by_accession")]
pub cf_name_by_accession: String,
#[arg(long)]
pub path_wal_dir: Option<String>,
}
fn jsonl_import(
db: &rocksdb::DBWithThreadMode<rocksdb::MultiThreaded>,
args: &Args,
) -> Result<(), anyhow::Error> {
let cf_data = db.cf_handle(&args.cf_name).unwrap();
let cf_by_accession = db.cf_handle(&args.cf_name_by_accession).unwrap();
let reader: Box<dyn std::io::Read> = if args.path_in_jsonl.ends_with(".gz") {
Box::new(flate2::read::GzDecoder::new(std::fs::File::open(
&args.path_in_jsonl,
)?))
} else {
Box::new(std::fs::File::open(&args.path_in_jsonl)?)
};
let reader = std::io::BufReader::new(reader);
for line in reader.lines() {
let line = line?;
let record = match serde_json::from_str::<clinvar_minimal::cli::reading::Record>(&line) {
Ok(record) => record,
Err(e) => {
tracing::warn!("skipping line because of error: {}", e);
continue;
}
};
let clinvar_minimal::cli::reading::Record {
rcv,
vcv,
title,
clinical_significance,
review_status,
sequence_location,
..
} = record;
let clinical_significance: crate::pbs::annonars::clinvar::v1::minimal::ClinicalSignificance =
clinical_significance.into();
let review_status: crate::pbs::annonars::clinvar::v1::minimal::ReviewStatus =
review_status.into();
let clinvar_minimal::cli::reading::SequenceLocation {
assembly,
chr,
start,
stop,
reference_allele_vcf,
alternate_allele_vcf,
..
} = sequence_location;
if let (Some(start), Some(stop), Some(reference_allele_vcf), Some(alternate_allele_vcf)) =
(start, stop, reference_allele_vcf, alternate_allele_vcf)
{
let var = keys::Var::from(
&chr,
start as i32,
&reference_allele_vcf,
&alternate_allele_vcf,
);
let key: Vec<u8> = var.into();
let data = db
.get_cf(&cf_data, key.clone())
.map_err(|e| anyhow::anyhow!("problem querying database: {}", e));
match data {
Err(e) => {
tracing::warn!("skipping line because of error: {}", e);
continue;
}
Ok(data) => {
db.put_cf(&cf_by_accession, rcv.as_bytes(), &key)?;
db.put_cf(&cf_by_accession, vcv.as_bytes(), &key)?;
let record = if let Some(data) = data {
let mut record =
crate::pbs::annonars::clinvar::v1::minimal::Record::decode(&data[..])?;
record.reference_assertions.push(
crate::pbs::annonars::clinvar::v1::minimal::ReferenceAssertion {
rcv,
title,
clinical_significance: clinical_significance.into(),
review_status: review_status.into(),
},
);
record
.reference_assertions
.sort_by_key(|a| (a.clinical_significance, a.review_status));
record
} else {
crate::pbs::annonars::clinvar::v1::minimal::Record {
release: assembly,
chromosome: chr,
start,
stop,
reference: reference_allele_vcf,
alternative: alternate_allele_vcf,
vcv,
reference_assertions: vec![ReferenceAssertion {
rcv,
title,
clinical_significance: clinical_significance.into(),
review_status: review_status.into(),
}],
}
};
let buf = record.encode_to_vec();
db.put_cf(&cf_data, &key, &buf)?;
}
}
}
}
Ok(())
}
pub fn run(common: &common::cli::Args, args: &Args) -> Result<(), anyhow::Error> {
tracing::info!("Starting 'clinvar-minimal import' command");
tracing::info!("common = {:#?}", &common);
tracing::info!("args = {:#?}", &args);
tracing::info!("Opening RocksDB for writing ...");
let before_opening_rocksdb = std::time::Instant::now();
let options = rocksdb_utils_lookup::tune_options(
rocksdb::Options::default(),
args.path_wal_dir.as_ref().map(|s| s.as_ref()),
);
let cf_names = &["meta", &args.cf_name, &args.cf_name_by_accession];
let db = Arc::new(rocksdb::DB::open_cf_with_opts(
&options,
common::readlink_f(&args.path_out_rocksdb)?,
cf_names
.iter()
.map(|name| (name.to_string(), options.clone()))
.collect::<Vec<_>>(),
)?);
tracing::info!(" writing meta information");
let cf_meta = db.cf_handle("meta").unwrap();
db.put_cf(&cf_meta, "annonars-version", crate::VERSION)?;
db.put_cf(
&cf_meta,
"genome-release",
format!("{}", args.genome_release),
)?;
db.put_cf(&cf_meta, "db-name", "clinvar-minimal")?;
tracing::info!(
"... done opening RocksDB for writing in {:?}",
before_opening_rocksdb.elapsed()
);
tracing::info!("Importing JSONL file ...");
let before_import = std::time::Instant::now();
jsonl_import(&db, args)?;
tracing::info!(
"... done importing JSONL file in {:?}",
before_import.elapsed()
);
tracing::info!("Running RocksDB compaction ...");
let before_compaction = std::time::Instant::now();
rocksdb_utils_lookup::force_compaction_cf(&db, cf_names, Some(" "), true)?;
tracing::info!(
"... done compacting RocksDB in {:?}",
before_compaction.elapsed()
);
tracing::info!("All done. Have a nice day!");
Ok(())
}
#[cfg(test)]
mod test {
use super::*;
use clap_verbosity_flag::Verbosity;
use temp_testdir::TempDir;
#[test]
fn smoke_test_import_jsonl() {
let tmp_dir = TempDir::default();
let common = common::cli::Args {
verbose: Verbosity::new(1, 0),
};
let args = Args {
genome_release: common::cli::GenomeRelease::Grch37,
path_in_jsonl: String::from("tests/clinvar-minimal/clinvar-seqvars-grch37-tgds.jsonl"),
path_out_rocksdb: format!("{}", tmp_dir.join("out-rocksdb").display()),
cf_name: String::from("clinvar"),
cf_name_by_accession: String::from("clinvar_by_accession"),
path_wal_dir: None,
};
run(&common, &args).unwrap();
}
}