Skip to main content

alice_open_data/
lib.rs

1#[cfg(not(target_arch = "wasm32"))]
2use std::{
3    fs::{DirBuilder, File},
4    io::Write,
5    path::PathBuf,
6};
7
8use failure::{format_err, Error};
9use reqwest::{Client, Url};
10
11#[cfg(test)]
12mod tests;
13
14fn root_url() -> Url {
15    if cfg!(target_arch = "wasm32") {
16        // Proxy with CORS properly set
17        "http://127.0.0.1:3030/opendata/".parse().unwrap()
18    } else {
19        "http://opendata.web.cern.ch".parse().unwrap()
20    }
21}
22
23/// Download the given file to the local collection
24#[cfg(not(target_arch = "wasm32"))]
25pub async fn download(base_dir: PathBuf, url: Url) -> Result<usize, Error> {
26    let mut dest = base_dir;
27    let mut sub_dir = url.path().to_owned();
28    // Remove the leading "\" from the original path
29    sub_dir.remove(0);
30    dest.push(sub_dir);
31    // Do not re-download if the file already exists
32    if dest.exists() {
33        return Ok(0);
34    }
35    // Make sure the dir exists
36    if let Some(dir) = dest.parent() {
37        DirBuilder::new().recursive(true).create(dir)?;
38    }
39    // Try downloading with re-tries
40    let mut n_retries = 3;
41    let bytes = loop {
42        let result = try_download(url.clone()).await;
43        if result.is_ok() || n_retries <= 0 {
44            break result;
45        } else {
46            n_retries -= 1;
47        }
48    }?;
49    let mut f = File::create(dest)?;
50    Ok(f.write(&bytes)?)
51}
52
53async fn try_download(url: Url) -> Result<Vec<u8>, Error> {
54    let resp = Client::new().get(url).send().await?;
55    Ok(resp
56        .error_for_status()?
57        .bytes()
58        .await?
59        .into_iter()
60        .collect())
61}
62
63/// Base path to the local ALICE open data directory
64#[cfg(not(target_arch = "wasm32"))]
65pub fn data_dir() -> Result<PathBuf, Error> {
66    let mut dir = dirs::home_dir().ok_or_else(|| format_err!("No home directory"))?;
67    dir.push("lhc_open_data");
68    Ok(dir)
69}
70
71/// Hardcoded path to a specific file. Useful for testing.
72/// That file should be the the first to be downloaded automatically.
73#[cfg(not(target_arch = "wasm32"))]
74pub fn test_file() -> Result<PathBuf, Error> {
75    let mut dir = data_dir()?;
76    dir.push("eos/opendata/alice/2010/LHC10h/000139038/ESD/0001/AliESDs.root");
77    Ok(dir)
78}
79
80/// Path to all files of `LHC10h`
81#[cfg(not(target_arch = "wasm32"))]
82pub fn all_files_10h() -> Result<Vec<PathBuf>, Error> {
83    let mut search_dir = data_dir()?;
84    search_dir.push("**/AliESDs.root");
85    let files: Vec<_> = glob::glob(search_dir.to_str().unwrap())
86        .expect("Can't resolve glob")
87        .map(|path| path.unwrap())
88        .collect();
89    Ok(files)
90}
91
92pub async fn get_file_list(run: u32) -> Result<Vec<Url>, Error> {
93    // Due to CORS we have to change the urls based on the target for now
94    let uri = root_url().join(match run {
95        139_038 => "record/1102/files/ALICE_LHC10h_PbPb_ESD_139038_file_index.txt",
96        139_173 => "record/1103/files/ALICE_LHC10h_PbPb_ESD_139173_file_index.txt",
97        139_437 => "record/1104/files/ALICE_LHC10h_PbPb_ESD_139437_file_index.txt",
98        139_438 => "record/1105/files/ALICE_LHC10h_PbPb_ESD_139438_file_index.txt",
99        139_465 => "record/1106/files/ALICE_LHC10h_PbPb_ESD_139465_file_index.txt",
100        _ => return Err(format_err!("Invalid run number")),
101    })?;
102
103    let req = Client::new().get(uri);
104    let resp = req.send().await?;
105    if resp.status().is_success() {
106        let content = resp.text().await?;
107        content
108            .lines()
109            .map(|l| root_url().join(&l[26..]).map_err(Into::into))
110            .collect()
111    } else {
112        Err(format_err!("Could not download list of files"))
113    }
114}