use crate::search::Search;
use crate::{Format, HtsGetError};
use crate::{
bam_search::BamSearch,
bcf_search::BcfSearch,
cram_search::CramSearch,
vcf_search::VcfSearch,
{HtsGet, Query, Response, Result},
};
use async_trait::async_trait;
use htsget_config::config::location::Locations;
use htsget_config::resolver::{ResolveResponse, StorageResolver};
use htsget_config::storage;
use htsget_storage::Storage;
use tracing::debug;
use tracing::instrument;
#[derive(Debug, Clone)]
pub struct HtsGetFromStorage {
storage: Storage,
}
#[async_trait]
impl HtsGet for Locations {
async fn search(self, mut query: Query) -> Result<Response> {
self
.resolve_request::<HtsGetFromStorage>(&mut query)
.await
.ok_or_else(|| HtsGetError::not_found("failed to match query with storage"))?
}
}
#[async_trait]
impl HtsGet for HtsGetFromStorage {
#[instrument(level = "debug", skip(self))]
async fn search(self, query: Query) -> Result<Response> {
debug!(format = ?query.format(), ?query, "searching {:?}, with query {:?}", query.format(), query);
match query.format() {
Format::Bam => BamSearch::new(self.into_inner()).search(query).await,
Format::Cram => CramSearch::new(self.into_inner()).search(query).await,
Format::Vcf => VcfSearch::new(self.into_inner()).search(query).await,
Format::Bcf => BcfSearch::new(self.into_inner()).search(query).await,
}
}
}
#[async_trait]
impl ResolveResponse for HtsGetFromStorage {
async fn from_file(file_storage: &storage::file::File, query: &Query) -> Result<Response> {
let storage = Storage::from_file(file_storage, query).await?;
let searcher = HtsGetFromStorage::new(storage);
searcher.search(query.clone()).await
}
#[cfg(feature = "aws")]
async fn from_s3(s3_storage: &storage::s3::S3, query: &Query) -> Result<Response> {
let storage = Storage::from_s3(s3_storage, query).await;
let searcher = HtsGetFromStorage::new(storage?);
searcher.search(query.clone()).await
}
#[cfg(feature = "url")]
async fn from_url(url_storage_config: storage::url::Url, query: &Query) -> Result<Response> {
let storage = Storage::from_url(url_storage_config, query).await;
let searcher = HtsGetFromStorage::new(storage?);
searcher.search(query.clone()).await
}
#[cfg(feature = "url")]
async fn from_json_path(
json_path_storage: storage::json_path::JsonPath,
query: &Query,
) -> Result<Response> {
let storage = Storage::from_json_path(json_path_storage, query).await;
let searcher = HtsGetFromStorage::new(storage?);
searcher.search(query.clone()).await
}
}
impl HtsGetFromStorage {
pub fn new(storage: Storage) -> Self {
Self { storage }
}
pub fn storage(&self) -> &Storage {
&self.storage
}
pub fn into_inner(self) -> Storage {
self.storage
}
}
#[cfg(test)]
pub(crate) mod tests {
use std::fs;
use std::future::Future;
use std::path::{Path, PathBuf};
#[cfg(feature = "aws")]
use {
htsget_storage::s3::S3Storage, htsget_test::aws_mocks::with_s3_test_server, std::fs::create_dir,
};
use htsget_config::config::location::{Location, PrefixOrId, SimpleLocation};
use htsget_config::storage;
use htsget_config::storage::Backend;
use htsget_config::types::Class::Body;
use htsget_config::types::Scheme::Http;
use htsget_storage::local::FileStorage;
#[cfg(feature = "experimental")]
use htsget_test::c4gh::decrypt_data;
use htsget_test::http::concat::ConcatResponse;
use http::uri::Authority;
use tempfile::TempDir;
use crate::bam_search::tests::{
BAM_FILE_NAME, expected_url as bam_expected_url, with_local_storage as with_bam_local_storage,
};
use crate::vcf_search::tests::{
VCF_FILE_NAME_SPEC, expected_url as vcf_expected_url,
with_local_storage as with_vcf_local_storage,
};
use crate::{Headers, Url};
use super::*;
#[tokio::test]
async fn search_bam() {
with_bam_local_storage(|storage| async move {
let htsget = HtsGetFromStorage::new(storage);
let query = Query::new_with_default_request("htsnexus_test_NA12878", Format::Bam);
let response = htsget.search(query).await;
println!("{response:#?}");
let expected_response = Ok(Response::new(
Format::Bam,
vec![
Url::new(bam_expected_url())
.with_headers(Headers::default().with_header("Range", "bytes=0-2596798")),
],
));
assert_eq!(response, expected_response);
Some((BAM_FILE_NAME.to_string(), (response.unwrap(), Body).into()))
})
.await;
}
#[tokio::test]
async fn search_vcf() {
with_vcf_local_storage(|storage| async move {
let htsget = HtsGetFromStorage::new(storage);
let filename = "spec-v4.3";
let query = Query::new_with_default_request(filename, Format::Vcf);
let response = htsget.search(query).await;
println!("{response:#?}");
assert_eq!(response, expected_vcf_response(filename));
Some((
VCF_FILE_NAME_SPEC.to_string(),
(response.unwrap(), Body).into(),
))
})
.await;
}
#[tokio::test]
async fn from_local_storage() {
with_config_local_storage(
|_, local_storage| async move {
let filename = "spec-v4.3";
let query = Query::new_with_default_request(filename, Format::Vcf);
let response = HtsGetFromStorage::from_file(&local_storage, &query).await;
assert_eq!(response, expected_vcf_response(filename));
Some((
VCF_FILE_NAME_SPEC.to_string(),
(response.unwrap(), Body).into(),
))
},
"data/vcf",
&[],
)
.await;
}
#[tokio::test]
async fn search_resolvers() {
with_config_local_storage(
|_, local_storage| async {
let locations = Locations::new(vec![Location::Simple(
SimpleLocation::new(
Backend::File(local_storage),
"".to_string(),
Some(PrefixOrId::Prefix("".to_string())),
)
.into(),
)]);
let filename = "spec-v4.3";
let query = Query::new_with_default_request(filename, Format::Vcf);
let response = locations.search(query).await;
assert_eq!(response, expected_vcf_response(filename));
Some((
VCF_FILE_NAME_SPEC.to_string(),
(response.unwrap(), Body).into(),
))
},
"data/vcf",
&[],
)
.await;
}
fn expected_vcf_response(filename: &str) -> Result<Response> {
Ok(Response::new(
Format::Vcf,
vec![
Url::new(vcf_expected_url(filename))
.with_headers(Headers::default().with_header("Range", "bytes=0-850")),
],
))
}
async fn copy_files_from(from_path: &str, to_path: &Path, copy_files: &[&str]) -> PathBuf {
let mut base_path = std::env::current_dir()
.unwrap()
.parent()
.unwrap()
.join(from_path);
for file_name in copy_files {
fs::copy(base_path.join(file_name), to_path.join(file_name)).unwrap();
}
if !copy_files.is_empty() {
base_path = PathBuf::from(to_path);
}
base_path
}
async fn with_config_local_storage_map<M, F, Fut>(
test: F,
path: &str,
copy_files: &[&str],
map: M,
) where
F: FnOnce(PathBuf, storage::file::File) -> Fut,
Fut: Future<Output = Option<(String, ConcatResponse)>>,
M: FnOnce(&[u8]) -> Vec<u8>,
{
let tmp_dir = TempDir::new().unwrap();
let base_path = copy_files_from(path, tmp_dir.path(), copy_files).await;
println!("{base_path:#?}");
let response = test(
base_path.clone(),
storage::file::File::new(
Http,
Authority::from_static("127.0.0.1:8081"),
base_path.to_str().unwrap().to_string(),
),
)
.await;
read_records(response, &base_path, map).await;
}
async fn with_config_local_storage<F, Fut>(test: F, path: &str, copy_files: &[&str])
where
F: FnOnce(PathBuf, storage::file::File) -> Fut,
Fut: Future<Output = Option<(String, ConcatResponse)>>,
{
with_config_local_storage_map(test, path, copy_files, |b| b.to_vec()).await;
}
async fn read_records<F>(response: Option<(String, ConcatResponse)>, base_path: &Path, map: F)
where
F: FnOnce(&[u8]) -> Vec<u8>,
{
if let Some((target_file, response)) = response {
let records = response
.concat_from_file_path(&base_path.join(target_file))
.await
.unwrap();
let bytes = map(records.merged_bytes());
records.set_bytes(bytes).read_records().await.unwrap();
}
}
pub(crate) async fn with_local_storage_fn<F, Fut>(test: F, path: &str, copy_files: &[&str])
where
F: FnOnce(Storage) -> Fut,
Fut: Future<Output = Option<(String, ConcatResponse)>>,
{
with_config_local_storage(
|base_path, local_storage| async {
test(Storage::new(
FileStorage::new(base_path, local_storage, vec![]).unwrap(),
))
.await
},
path,
copy_files,
)
.await;
}
#[cfg(feature = "experimental")]
pub(crate) async fn with_local_storage_c4gh<F, Fut>(test: F)
where
F: FnOnce(Storage) -> Fut,
Fut: Future<Output = Option<(String, ConcatResponse)>>,
{
with_config_local_storage_map(
|base_path, local_storage| async {
test(Storage::new(
FileStorage::new(base_path, local_storage, vec![]).unwrap(),
))
.await
},
"data/c4gh",
&[],
decrypt_data,
)
.await;
}
#[cfg(feature = "aws")]
pub(crate) async fn with_aws_storage_fn<F, Fut>(test: F, path: &str, copy_files: &[&str])
where
F: FnOnce(Storage) -> Fut,
Fut: Future<Output = Option<(String, ConcatResponse)>>,
{
let tmp_dir = TempDir::new().unwrap();
let to_path = tmp_dir.path().join("folder");
create_dir(&to_path).unwrap();
let base_path = copy_files_from(path, &to_path, copy_files).await;
with_aws_s3_storage_fn(
|storage| async {
let response = test(storage).await;
read_records(response, &base_path, |b| b.to_vec()).await;
},
"folder".to_string(),
base_path.parent().unwrap(),
)
.await;
}
#[cfg(feature = "aws")]
pub(crate) async fn with_aws_s3_storage_fn<F, Fut>(test: F, folder_name: String, base_path: &Path)
where
F: FnOnce(Storage) -> Fut,
Fut: Future<Output = ()>,
{
with_s3_test_server(base_path, |client| async move {
test(Storage::new(S3Storage::new(client, folder_name))).await;
})
.await;
}
}