bufkit_data/archive/
root.rs

1use crate::{errors::BufkitDataErr, models::Model, site::StationNumber, Archive};
2use rusqlite::ToSql;
3
4impl Archive {
5    const DATA_DIR: &'static str = "data";
6    const DB_FILE: &'static str = "index.db";
7
8    /// Initialize a new archive.
9    pub fn create(root: &dyn AsRef<std::path::Path>) -> Result<Self, BufkitDataErr> {
10        let data_root = root.as_ref().join(Archive::DATA_DIR);
11        let db_file = root.as_ref().join(Archive::DB_FILE);
12        let root = root.as_ref().to_path_buf();
13
14        std::fs::create_dir_all(&data_root)?; // The folder to store the sounding files.
15
16        // Create and set up the archive
17        let db_conn = rusqlite::Connection::open_with_flags(
18            db_file,
19            rusqlite::OpenFlags::SQLITE_OPEN_READ_WRITE | rusqlite::OpenFlags::SQLITE_OPEN_CREATE,
20        )?;
21
22        db_conn.execute_batch(include_str!("root/create_index.sql"))?;
23
24        Ok(Archive { root, db_conn })
25    }
26
27    /// Open an existing archive.
28    pub fn connect(root: &dyn AsRef<std::path::Path>) -> Result<Self, BufkitDataErr> {
29        let db_file = root.as_ref().join(Archive::DB_FILE);
30        let root = root.as_ref().to_path_buf();
31
32        // Create and set up the archive
33        let db_conn = rusqlite::Connection::open_with_flags(
34            db_file,
35            rusqlite::OpenFlags::SQLITE_OPEN_READ_WRITE,
36        )?;
37
38        Ok(Archive { root, db_conn })
39    }
40
41    /// Retrieve a path to the root. Allows caller to store files in the archive.
42    pub fn root(&self) -> &std::path::Path {
43        &self.root
44    }
45
46    /// Get the directory the data files are stored in.
47    pub(crate) fn data_root(&self) -> std::path::PathBuf {
48        self.root.join(Archive::DATA_DIR)
49    }
50
51    /// Export part of the archive.
52    pub fn export(
53        &self,
54        stations: &[StationNumber],
55        models: &[Model],
56        start: chrono::NaiveDateTime,
57        end: chrono::NaiveDateTime,
58        dest: &std::path::Path,
59    ) -> Result<(), BufkitDataErr> {
60        let new_db = Self::create(&dest)?;
61        let db_file = new_db.root.join(Archive::DB_FILE);
62
63        let statement = &format!(
64            "ATTACH '{}' AS ex;",
65            db_file.to_str().ok_or_else(|| BufkitDataErr::GeneralError(
66                "Unable to convert path to string".to_owned()
67            ))?
68        );
69        self.db_conn.execute(statement, [])?;
70
71        let mut sites_stmt = self.db_conn.prepare(
72            "
73                INSERT INTO ex.sites 
74                SELECT * FROM main.sites 
75                WHERE main.sites.station_num = ?1
76            ",
77        )?;
78
79        let mut files_stmt = self.db_conn.prepare(
80            "
81                INSERT INTO ex.files
82                SELECT * FROM main.files 
83                WHERE main.files.station_num = ?1 AND main.files.model = ?2 
84                    AND main.files.init_time >= ?3 AND main.files.init_time <= ?4
85            ",
86        )?;
87
88        let source_dir = self.root.join(Archive::DATA_DIR);
89        let dest_dir = dest.join(Archive::DATA_DIR);
90        let mut file_names_stmt = self.db_conn.prepare(
91            "
92                SELECT ex.files.file_name FROM ex.files
93                WHERE ex.files.station_num = ?1 AND ex.files.model = ?2
94                    AND ex.files.init_time >= ?3 AND ex.files.init_time <= ?4
95            ",
96        )?;
97
98        for &stn in stations {
99            let stn_num: u32 = stn.into();
100            sites_stmt.execute([stn_num])?;
101
102            for &model in models {
103                files_stmt.execute(&[
104                    &stn_num as &dyn ToSql,
105                    &model.as_static_str(),
106                    &start,
107                    &end,
108                ])?;
109
110                let fnames = file_names_stmt.query_and_then(
111                    &[&stn_num as &dyn ToSql, &model.as_static_str(), &start, &end],
112                    |row| -> Result<String, _> { row.get(0) },
113                )?;
114
115                for fname in fnames {
116                    let fname = fname?;
117                    let src = source_dir.join(&fname);
118                    let dest = dest_dir.join(fname);
119                    std::fs::copy(src, dest)?;
120                }
121            }
122        }
123
124        Ok(())
125    }
126}
127
128#[cfg(test)]
129mod unit {
130    use super::*;
131    use crate::archive::unit::*; // Test setup and tear down.
132
133    #[test]
134    fn test_archive_create_new() {
135        assert!(create_test_archive().is_ok());
136    }
137
138    #[test]
139    fn test_archive_connect() {
140        let TestArchive { tmp, arch } =
141            create_test_archive().expect("Failed to create test archive.");
142        drop(arch);
143
144        assert!(Archive::connect(&tmp.path()).is_ok());
145        assert!(Archive::connect(&"unlikely_directory_in_my_project").is_err());
146    }
147
148    #[test]
149    fn test_get_root() {
150        let TestArchive { tmp, arch } =
151            create_test_archive().expect("Failed to create test archive.");
152
153        let root = arch.root();
154        assert_eq!(root, tmp.path());
155    }
156}
bufkit_data/archive/root.rs

bufkit_data/archive/
root.rs
