bufkit_data/archive/
query.rs

1use rusqlite::OptionalExtension;
2use std::{collections::HashSet, io::Read, iter::FromIterator, str::FromStr};
3
4use crate::{
5    errors::BufkitDataErr,
6    models::Model,
7    site::{SiteInfo, StateProv, StationNumber},
8};
9
10mod station_summary;
11pub use station_summary::StationSummary;
12
13impl crate::Archive {
14    /// Retrieve a list of sites in the archive.
15    pub fn sites(&self) -> Result<Vec<SiteInfo>, BufkitDataErr> {
16        let mut stmt = self
17            .db_conn
18            .prepare(include_str!("query/retrieve_sites.sql"))?;
19
20        let vals: Result<Vec<SiteInfo>, BufkitDataErr> = stmt
21            .query_and_then([], Self::parse_row_to_site)?
22            .map(|res| res.map_err(BufkitDataErr::Database))
23            .collect();
24
25        vals
26    }
27
28    fn parse_row_to_site(row: &rusqlite::Row) -> Result<SiteInfo, rusqlite::Error> {
29        let station_num: u32 = row.get(0)?;
30        let station_num = StationNumber::from(station_num);
31
32        let name: Option<String> = row.get(1)?;
33        let notes: Option<String> = row.get(3)?;
34        let state: Option<StateProv> = row
35            .get::<_, String>(2)
36            .ok()
37            .and_then(|a_string| StateProv::from_str(&a_string).ok());
38
39        let time_zone: Option<chrono::FixedOffset> =
40            row.get::<_, i32>(4).ok().map(|offset: i32| {
41                if offset < 0 {
42                    chrono::FixedOffset::west_opt(offset.abs()).unwrap()
43                } else {
44                    chrono::FixedOffset::east_opt(offset).unwrap()
45                }
46            });
47
48        Ok(SiteInfo {
49            station_num,
50            name,
51            notes,
52            state,
53            time_zone,
54        })
55    }
56
57    /// Retrieve the sites with their most recent station id for the given model.
58    pub fn sites_and_ids_for(
59        &self,
60        model: Model,
61    ) -> Result<Vec<(SiteInfo, String)>, BufkitDataErr> {
62        self.db_conn
63            .execute("DROP TABLE IF EXISTS temp_ids", [])?;
64        self.db_conn.execute(
65            "
66                CREATE TEMP TABLE temp_ids AS
67                SELECT files.id, files.station_num
68                FROM files JOIN (
69                    SELECT files.station_num, MAX(files.init_time) as maxtime
70                    FROM files 
71                    WHERE files.model = ?1
72                    GROUP BY files.station_num) as maxs
73                ON maxs.station_num = files.station_num AND files.init_time = maxs.maxtime
74                WHERE files.model = ?1
75            ",
76            &[&model.as_static_str()],
77        )?;
78
79        let mut stmt = self.db_conn.prepare(
80            "
81                SELECT 
82                    sites.station_num,
83                    sites.name, 
84                    sites.state, 
85                    sites.notes, 
86                    sites.tz_offset_sec, 
87                    temp_ids.id
88                FROM sites JOIN temp_ids ON temp_ids.station_num = sites.station_num
89            ",
90        )?;
91
92        let parse_row = |row: &rusqlite::Row| -> Result<(SiteInfo, String), rusqlite::Error> {
93            let site_info = Self::parse_row_to_site(row)?;
94            let site_id: String = row.get(5)?;
95            Ok((site_info, site_id))
96        };
97
98        let vals: Result<Vec<(SiteInfo, String)>, BufkitDataErr> = stmt
99            .query_and_then([], parse_row)?
100            .map(|res| res.map_err(BufkitDataErr::Database))
101            .collect();
102
103        vals
104    }
105
106    /// Retrieve the information about a single site id
107    pub fn site(&self, station_num: StationNumber) -> Option<SiteInfo> {
108        self.db_conn
109            .query_row_and_then(
110                "
111                    SELECT
112                         station_num,
113                         name,
114                         state,
115                         notes,
116                         tz_offset_sec
117                    FROM sites 
118                    WHERE station_num = ?1
119                ",
120                &[&Into::<u32>::into(station_num)],
121                Self::parse_row_to_site,
122            )
123            .ok()
124    }
125
126    /// Get a list of models in the archive for this site.
127    pub fn models(&self, station_num: StationNumber) -> Result<Vec<Model>, BufkitDataErr> {
128        let station_num: u32 = Into::<u32>::into(station_num);
129
130        let mut stmt = self
131            .db_conn
132            .prepare("SELECT DISTINCT model FROM files WHERE station_num = ?1")?;
133
134        let vals: Result<Vec<Model>, BufkitDataErr> = stmt
135            .query_map(&[&station_num], |row| row.get::<_, String>(0))?
136            .map(|res| res.map_err(BufkitDataErr::Database))
137            .map(|res| {
138                res.and_then(|name| Model::from_str(&name).map_err(BufkitDataErr::StrumError))
139            })
140            .collect();
141
142        vals
143    }
144
145    /// Retrieve a file from the archive.
146    pub fn retrieve(
147        &self,
148        station_num: StationNumber,
149        model: Model,
150        init_time: chrono::NaiveDateTime,
151    ) -> Result<String, BufkitDataErr> {
152        let station_num: u32 = Into::<u32>::into(station_num);
153
154        let file_name: Result<String, _> = self.db_conn.query_row(
155            "SELECT file_name FROM files WHERE station_num = ?1 AND model = ?2 AND init_time = ?3",
156            &[
157                &station_num as &dyn rusqlite::types::ToSql,
158                &model.as_static_str() as &dyn rusqlite::types::ToSql,
159                &init_time as &dyn rusqlite::types::ToSql,
160            ],
161            |row| row.get(0),
162        );
163
164        let file_name = match file_name {
165            Ok(fname) => fname,
166            Err(rusqlite::Error::QueryReturnedNoRows) => return Err(BufkitDataErr::NotInIndex),
167            Err(x) => return Err(BufkitDataErr::Database(x)),
168        };
169
170        let file = std::fs::File::open(self.data_root().join(file_name))?;
171        let mut decoder = flate2::read::GzDecoder::new(file);
172        let mut s = String::new();
173        decoder.read_to_string(&mut s)?;
174        Ok(s)
175    }
176
177    /// Retrieve the  most recent file.
178    pub fn retrieve_most_recent(
179        &self,
180        station_num: StationNumber,
181        model: Model,
182    ) -> Result<String, BufkitDataErr> {
183        let station_num: u32 = Into::<u32>::into(station_num);
184
185        let file_name: Result<String, _> = self.db_conn.query_row(
186            "
187                SELECT file_name 
188                FROM files 
189                WHERE station_num = ?1 AND model = ?2 
190                ORDER BY init_time DESC 
191                LIMIT 1
192            ",
193            &[
194                &station_num as &dyn rusqlite::types::ToSql,
195                &model.as_static_str() as &dyn rusqlite::types::ToSql,
196            ],
197            |row| row.get(0),
198        );
199
200        let file_name = match file_name {
201            Ok(fname) => fname,
202            Err(rusqlite::Error::QueryReturnedNoRows) => return Err(BufkitDataErr::NotInIndex),
203            Err(x) => return Err(BufkitDataErr::Database(x)),
204        };
205
206        let file = std::fs::File::open(self.data_root().join(file_name))?;
207        let mut decoder = flate2::read::GzDecoder::new(file);
208        let mut s = String::new();
209        decoder.read_to_string(&mut s)?;
210        Ok(s)
211    }
212
213    /// Retrieve all the soundings with any data valid between the start and end times.
214    pub fn retrieve_all_valid_in(
215        &self,
216        station_num: StationNumber,
217        model: Model,
218        start: chrono::NaiveDateTime,
219        end: chrono::NaiveDateTime,
220    ) -> Result<impl Iterator<Item = String>, BufkitDataErr> {
221        let station_num: u32 = Into::<u32>::into(station_num);
222
223        let mut stmt = self.db_conn.prepare(
224            "
225                    SELECT file_name 
226                    FROM files 
227                    WHERE station_num = ?1 AND model = ?2 AND 
228                        (
229                            (init_time <= ?3 AND end_time >= ?4) OR 
230                            (init_time >= ?3 AND init_time < ?4) OR 
231                            (end_time > ?3 AND end_time <= ?4)
232                        )
233                    ORDER BY init_time ASC 
234                ",
235        )?;
236
237        let file_names: Vec<String> = stmt
238            .query_map(
239                &[
240                    &station_num as &dyn rusqlite::types::ToSql,
241                    &model.as_static_str() as &dyn rusqlite::types::ToSql,
242                    &start as &dyn rusqlite::types::ToSql,
243                    &end as &dyn rusqlite::types::ToSql,
244                ],
245                |row| row.get(0),
246            )?
247            .filter_map(|res| res.ok())
248            .collect();
249
250        if file_names.is_empty() {
251            return Err(BufkitDataErr::NotInIndex);
252        }
253
254        let root = self.data_root();
255        Ok(file_names.into_iter().filter_map(move |fname| {
256            std::fs::File::open(root.join(fname)).ok().and_then(|f| {
257                let mut decoder = flate2::read::GzDecoder::new(f);
258                let mut s = String::new();
259                match decoder.read_to_string(&mut s) {
260                    Ok(_) => Some(s),
261                    Err(_) => None,
262                }
263            })
264        }))
265    }
266
267    /// Check to see if a file is present in the archive and it is retrieveable.
268    pub fn file_exists(
269        &self,
270        site: StationNumber,
271        model: Model,
272        init_time: chrono::NaiveDateTime,
273    ) -> Result<bool, BufkitDataErr> {
274        let num_records: i32 = self.db_conn.query_row(
275            "SELECT COUNT(*) FROM files WHERE station_num = ?1 AND model = ?2 AND init_time = ?3",
276            &[
277                &Into::<i64>::into(site) as &dyn rusqlite::types::ToSql,
278                &model.as_static_str() as &dyn rusqlite::types::ToSql,
279                &init_time as &dyn rusqlite::types::ToSql,
280            ],
281            |row| row.get(0),
282        )?;
283
284        Ok(num_records == 1)
285    }
286
287    /// Retrieve the most recent station number used with this ID and model.
288    pub fn station_num_for_id(
289        &self,
290        id: &str,
291        model: Model,
292    ) -> Result<StationNumber, BufkitDataErr> {
293        let station_num: Result<u32, _> = self.db_conn.query_row(
294            include_str!("query/station_num_for_id_and_model.sql"),
295            &[
296                &id.to_uppercase() as &dyn rusqlite::types::ToSql,
297                &model.as_static_str(),
298            ],
299            |row| row.get(0),
300        );
301
302        let station_num = match station_num {
303            Ok(num) => StationNumber::from(num),
304            Err(rusqlite::Error::QueryReturnedNoRows) => return Err(BufkitDataErr::NotInIndex),
305            Err(x) => return Err(BufkitDataErr::Database(x)),
306        };
307
308        Ok(station_num)
309    }
310
311    /// Retrieve a list of site ids use with the station number.
312    pub fn ids(
313        &self,
314        station_num: StationNumber,
315        model: Model,
316    ) -> Result<Vec<String>, BufkitDataErr> {
317        let station_num: u32 = Into::<u32>::into(station_num);
318
319        let mut stmt = self.db_conn.prepare(
320            "
321                SELECT DISTINCT id 
322                FROM files
323                WHERE station_num = ?1 AND model = ?2
324            ",
325        )?;
326
327        let sites: Result<Vec<String>, _> = stmt
328            .query_map(
329                &[
330                    &station_num as &dyn rusqlite::types::ToSql,
331                    &model.as_static_str() as &dyn rusqlite::types::ToSql,
332                ],
333                |row| row.get(0),
334            )?
335            .collect();
336
337        sites.map_err(BufkitDataErr::Database)
338    }
339
340    /// Retrieve the most recently used ID with a site.
341    pub fn most_recent_id(
342        &self,
343        station_num: StationNumber,
344        model: Model,
345    ) -> Result<Option<String>, BufkitDataErr> {
346        let station_num_raw: u32 = Into::<u32>::into(station_num);
347
348        let mut stmt = self.db_conn.prepare(
349            "
350                SELECT id, init_time 
351                FROM files
352                WHERE station_num = ?1 AND model = ?2
353                ORDER BY init_time DESC
354                LIMIT 1
355            ",
356        )?;
357
358        let most_recent_site: String = match stmt
359            .query_row(
360                &[
361                    &station_num_raw as &dyn rusqlite::types::ToSql,
362                    &model.as_static_str() as &dyn rusqlite::types::ToSql,
363                ],
364                |row| row.get(0),
365            )
366            .optional()?
367        {
368            Some(id) => id,
369            None => return Ok(None),
370        };
371
372        let most_recent_station_num = self.station_num_for_id(&most_recent_site, model)?;
373
374        if most_recent_station_num == station_num {
375            Ok(Some(most_recent_site))
376        } else {
377            Ok(None)
378        }
379    }
380
381    /// Get an inventory of soundings for a site & model.
382    pub fn inventory(
383        &self,
384        station_num: StationNumber,
385        model: Model,
386    ) -> Result<Vec<chrono::NaiveDateTime>, BufkitDataErr> {
387        let station_num: u32 = Into::<u32>::into(station_num);
388
389        let mut stmt = self.db_conn.prepare(
390            "
391                SELECT init_time
392                FROM files
393                WHERE station_num = ?1 AND model = ?2
394                ORDER BY init_time ASC
395            ",
396        )?;
397
398        let inv: Result<Vec<chrono::NaiveDateTime>, _> = stmt
399            .query_map(
400                &[
401                    &station_num as &dyn rusqlite::types::ToSql,
402                    &model.as_static_str() as &dyn rusqlite::types::ToSql,
403                ],
404                |row| row.get(0),
405            )?
406            .collect();
407
408        inv.map_err(BufkitDataErr::Database)
409    }
410
411    /// Get list of missing init times.
412    ///
413    /// If time_range is `None`, this will find the first and last entries and then look for any
414    /// gaps. If time_range is specified, then the end times are inclusive.
415    pub fn missing_inventory(
416        &self,
417        station_num: StationNumber,
418        model: Model,
419        time_range: Option<(chrono::NaiveDateTime, chrono::NaiveDateTime)>,
420    ) -> Result<Vec<chrono::NaiveDateTime>, BufkitDataErr> {
421        let (start, end) = if let Some((start, end)) = time_range {
422            (start, end)
423        } else {
424            self.first_and_last_dates(station_num, model)?
425        };
426
427        let inv = self.inventory(station_num, model)?;
428        let inv: HashSet<chrono::NaiveDateTime> = HashSet::from_iter(inv.into_iter());
429
430        let mut to_ret = vec![];
431        for curr_time in model.all_runs(&start, &end) {
432            if !inv.contains(&curr_time) {
433                to_ret.push(curr_time);
434            }
435        }
436
437        Ok(to_ret)
438    }
439
440    /// Get the number of files in the archive for the given station and model.
441    pub fn count(&self, station_num: StationNumber, model: Model) -> Result<u32, BufkitDataErr> {
442        let station_num: u32 = Into::<u32>::into(station_num);
443        self.db_conn
444            .query_row(
445                "
446                SELECT COUNT(*)
447                FROM files
448                WHERE station_num = ?1 AND model = ?2
449            ",
450                &[
451                    &station_num as &dyn rusqlite::types::ToSql,
452                    &model.as_static_str(),
453                ],
454                |row| row.get(0),
455            )
456            .map_err(BufkitDataErr::Database)
457    }
458
459    fn first_and_last_dates(
460        &self,
461        station_num: StationNumber,
462        model: Model,
463    ) -> Result<(chrono::NaiveDateTime, chrono::NaiveDateTime), BufkitDataErr> {
464        let station_num: u32 = Into::<u32>::into(station_num);
465
466        let start = self.db_conn.query_row(
467            "
468                    SELECT init_time
469                    FROM files
470                    WHERE station_num = ?1 AND model = ?2
471                    ORDER BY init_time ASC
472                    LIMIT 1
473                ",
474            &[
475                &station_num as &dyn rusqlite::types::ToSql,
476                &model.as_static_str() as &dyn rusqlite::types::ToSql,
477            ],
478            |row| row.get(0),
479        )?;
480
481        let end = self.db_conn.query_row(
482            "
483                    SELECT init_time
484                    FROM files
485                    WHERE station_num = ?1 AND model = ?2
486                    ORDER BY init_time DESC
487                    LIMIT 1
488                ",
489            &[
490                &station_num as &dyn rusqlite::types::ToSql,
491                &model.as_static_str() as &dyn rusqlite::types::ToSql,
492            ],
493            |row| row.get(0),
494        )?;
495
496        Ok((start, end))
497    }
498}
499
500#[cfg(test)]
501mod unit {
502    use super::*;
503    use crate::archive::unit::*; // test helpers.
504
505    use chrono::NaiveDate;
506
507    #[test]
508    fn test_sites_and_ids_for() {
509        let TestArchive {
510            tmp: _tmp,
511            mut arch,
512        } = create_test_archive().expect("Failed to create test archive.");
513
514        let test_sites = &get_test_sites();
515
516        for site in test_sites {
517            arch.add_site(site).expect("Error adding site.");
518        }
519
520        fill_test_archive(&mut arch);
521
522        let sites_and_ids = arch
523            .sites_and_ids_for(Model::GFS)
524            .expect("error getting sites & ids");
525        println!("There are {} entries in the variable.", sites_and_ids.len());
526        for (site, id) in &sites_and_ids {
527            println!("{:?} - {}", site, id);
528            if site.station_num == StationNumber::from(727730) {
529                assert_eq!(&id, &"KMSO");
530            } else {
531                panic!("Test data must have changed, not recognized.");
532            }
533        }
534    }
535
536    #[test]
537    fn test_site_info() {
538        let TestArchive { tmp: _tmp, arch } =
539            create_test_archive().expect("Failed to create test archive.");
540
541        let test_sites = &get_test_sites();
542
543        for site in test_sites {
544            arch.add_site(site).expect("Error adding site.");
545        }
546
547        let si = arch
548            .site(StationNumber::from(1))
549            .expect("Error retrieving site.");
550        assert_eq!(si.name, Some("Chicago/O'Hare".to_owned()));
551        assert_eq!(si.notes, Some("Major air travel hub.".to_owned()));
552        assert_eq!(si.state, Some(StateProv::IL));
553        assert_eq!(si.time_zone, None);
554
555        let si = arch
556            .site(StationNumber::from(2))
557            .expect("Error retrieving site.");
558        assert_eq!(si.name, Some("Seattle".to_owned()));
559        assert_eq!(
560            si.notes,
561            Some("A coastal city with coffe and rain".to_owned())
562        );
563        assert_eq!(si.state, Some(StateProv::WA));
564        assert_eq!(si.time_zone, Some(chrono::FixedOffset::west(8 * 3600)));
565
566        let si = arch
567            .site(StationNumber::from(3))
568            .expect("Error retrieving site.");
569        assert_eq!(si.name, Some("Missoula".to_owned()));
570        assert_eq!(si.notes, Some("In a valley.".to_owned()));
571        assert_eq!(si.state, None);
572        assert_eq!(si.time_zone, Some(chrono::FixedOffset::west(7 * 3600)));
573
574        assert!(arch.site(StationNumber::from(0)).is_none());
575        assert!(arch.site(StationNumber::from(100)).is_none());
576    }
577
578    #[test]
579    fn test_models_for_site() {
580        let TestArchive {
581            tmp: _tmp,
582            mut arch,
583        } = create_test_archive().expect("Failed to create test archive.");
584
585        fill_test_archive(&mut arch);
586
587        let kmso = StationNumber::from(727730); // Station number for KMSO
588        let models = arch.models(kmso).expect("Error querying archive.");
589
590        assert!(models.contains(&Model::GFS));
591        assert!(models.contains(&Model::NAM));
592        assert!(!models.contains(&Model::NAM4KM));
593    }
594
595    #[test]
596    fn test_retrieve() {
597        let TestArchive {
598            tmp: _tmp,
599            mut arch,
600        } = create_test_archive().expect("Failed to create test archive.");
601
602        fill_test_archive(&mut arch);
603
604        let kmso = StationNumber::from(727730); // Station number for KMSO
605        let init_time = NaiveDate::from_ymd(2017, 4, 1).and_hms(18, 0, 0);
606        let model = Model::GFS;
607
608        let res = arch.retrieve(kmso, model, init_time);
609        assert!(res.is_ok());
610
611        let init_time = NaiveDate::from_ymd(2117, 4, 1).and_hms(18, 0, 0);
612        let res = arch.retrieve(kmso, model, init_time);
613        match res {
614            Err(BufkitDataErr::NotInIndex) => {}
615            Err(_) => panic!("Wrong error type returned."),
616            Ok(_) => panic!("This should not exist in the database."),
617        }
618    }
619
620    #[test]
621    fn test_retrieve_most_recent() {
622        let TestArchive {
623            tmp: _tmp,
624            mut arch,
625        } = create_test_archive().expect("Failed to create test archive.");
626
627        fill_test_archive(&mut arch);
628
629        let kmso = StationNumber::from(727730); // Station number for KMSO
630        let init_time = NaiveDate::from_ymd(2017, 4, 1).and_hms(18, 0, 0);
631        let model = Model::GFS;
632
633        let res = arch.retrieve_most_recent(kmso, model);
634
635        if let Ok(str_data) = res {
636            let retrieved_init_time = sounding_bufkit::BufkitData::init(&str_data, "")
637                .expect("Failure parsing.")
638                .into_iter()
639                .next()
640                .expect("No data in file?")
641                .0
642                .valid_time()
643                .expect("No valid time with sounding?");
644
645            assert_eq!(retrieved_init_time, init_time);
646        } else {
647            panic!("Nothing found!");
648        }
649    }
650
651    #[test]
652    fn test_file_exists() {
653        let TestArchive {
654            tmp: _tmp,
655            mut arch,
656        } = create_test_archive().expect("Failed to create test archive.");
657
658        fill_test_archive(&mut arch);
659
660        let kmso_station_num = StationNumber::from(727730); // Station number for KMSO
661        let model = Model::NAM;
662
663        let first = NaiveDate::from_ymd(2017, 4, 1).and_hms(0, 0, 0);
664        let second = NaiveDate::from_ymd(2017, 4, 1).and_hms(12, 0, 0);
665        let last = NaiveDate::from_ymd(2017, 4, 1).and_hms(18, 0, 0);
666        let missing = NaiveDate::from_ymd(2017, 4, 1).and_hms(6, 0, 0);
667        assert!(arch.file_exists(kmso_station_num, model, first).unwrap());
668        assert!(arch.file_exists(kmso_station_num, model, second).unwrap());
669        assert!(arch.file_exists(kmso_station_num, model, last).unwrap());
670        assert!(!arch.file_exists(kmso_station_num, model, missing).unwrap());
671    }
672
673    #[test]
674    fn test_station_num_for_id() {
675        let TestArchive {
676            tmp: _tmp,
677            mut arch,
678        } = create_test_archive().expect("Failed to create test archive.");
679
680        fill_test_archive(&mut arch);
681
682        let kmso_station_num = StationNumber::from(727730); // Station number for KMSO
683
684        if let Ok(retrieved_station_num) = arch.station_num_for_id("kmso", Model::GFS) {
685            assert_eq!(retrieved_station_num, kmso_station_num);
686        } else {
687            panic!("Could not find station number!");
688        }
689
690        if let Ok(retrieved_station_num) = arch.station_num_for_id("KMSO", Model::GFS) {
691            assert_eq!(retrieved_station_num, kmso_station_num);
692        } else {
693            panic!("Could not find station number!");
694        }
695
696        if let Ok(retrieved_station_num) = arch.station_num_for_id("KmSo", Model::NAM) {
697            assert_eq!(retrieved_station_num, kmso_station_num);
698        } else {
699            panic!("Could not find station number!");
700        }
701
702        match arch.station_num_for_id("xyz", Model::GFS) {
703            Err(BufkitDataErr::NotInIndex) => {}
704            Ok(num) => panic!("Found station that does not exists! station_num = {}", num),
705            Err(err) => panic!("Other error: {}", err),
706        }
707    }
708
709    #[test]
710    fn test_ids() {
711        let TestArchive {
712            tmp: _tmp,
713            mut arch,
714        } = create_test_archive().expect("Failed to create test archive.");
715
716        fill_test_archive(&mut arch);
717
718        let kmso_station_num = StationNumber::from(727730); // Station number for KMSO
719        let ids = arch
720            .ids(kmso_station_num, Model::GFS)
721            .expect("Database error.");
722        assert!(ids.contains(&"KMSO".to_owned()));
723        assert_eq!(ids.len(), 1);
724
725        let ids = arch
726            .ids(kmso_station_num, Model::NAM)
727            .expect("Database error.");
728        assert!(ids.contains(&"KMSO".to_owned()));
729        assert_eq!(ids.len(), 1);
730
731        let ids = arch
732            .ids(kmso_station_num, Model::NAM4KM)
733            .expect("Database error.");
734        assert_eq!(ids.len(), 0);
735
736        let fake_station_num = StationNumber::from(5);
737        let ids = arch
738            .ids(fake_station_num, Model::GFS)
739            .expect("Database error.");
740        assert_eq!(ids.len(), 0);
741    }
742
743    #[test]
744    fn test_most_recent_id() {
745        let TestArchive {
746            tmp: _tmp,
747            mut arch,
748        } = create_test_archive().expect("Failed to create test archive.");
749
750        fill_test_archive(&mut arch);
751
752        let kmso_station_num = StationNumber::from(727730); // Station number for KMSO
753        let id = arch
754            .most_recent_id(kmso_station_num, Model::GFS)
755            .expect("Database error.");
756        assert_eq!(id.unwrap(), "KMSO".to_owned());
757
758        let id = arch
759            .most_recent_id(kmso_station_num, Model::NAM)
760            .expect("Database error.");
761        assert_eq!(id.unwrap(), "KMSO".to_owned());
762    }
763
764    #[test]
765    fn test_inventory() {
766        let TestArchive {
767            tmp: _tmp,
768            mut arch,
769        } = create_test_archive().expect("Failed to create test archive.");
770
771        fill_test_archive(&mut arch);
772
773        let kmso = StationNumber::from(727730); // Station number for KMSO
774        let first = NaiveDate::from_ymd(2017, 4, 1).and_hms(0, 0, 0);
775        let second = NaiveDate::from_ymd(2017, 4, 1).and_hms(12, 0, 0);
776        let last = NaiveDate::from_ymd(2017, 4, 1).and_hms(18, 0, 0);
777        let missing = NaiveDate::from_ymd(2017, 4, 1).and_hms(6, 0, 0);
778
779        let inv = arch.inventory(kmso, Model::NAM).expect("Data base error?");
780        assert!(inv.contains(&first));
781        assert!(inv.contains(&second));
782        assert!(inv.contains(&last));
783        assert!(!inv.contains(&missing));
784    }
785
786    #[test]
787    fn test_missing_inventory() {
788        let TestArchive {
789            tmp: _tmp,
790            mut arch,
791        } = create_test_archive().expect("Failed to create test archive.");
792
793        fill_test_archive(&mut arch);
794
795        let kmso = StationNumber::from(727730); // Station number for KMSO
796        let first = NaiveDate::from_ymd(2017, 4, 1).and_hms(0, 0, 0);
797        let second = NaiveDate::from_ymd(2017, 4, 1).and_hms(12, 0, 0);
798        let last = NaiveDate::from_ymd(2017, 4, 1).and_hms(18, 0, 0);
799        let missing = NaiveDate::from_ymd(2017, 4, 1).and_hms(6, 0, 0);
800
801        let missing_times = arch
802            .missing_inventory(kmso, Model::NAM, None)
803            .expect("Data base error?");
804        assert!(!missing_times.contains(&first));
805        assert!(!missing_times.contains(&second));
806        assert!(!missing_times.contains(&last));
807        assert!(missing_times.contains(&missing));
808
809        let larger_range = (
810            NaiveDate::from_ymd(2017, 3, 31).and_hms(0, 0, 0),
811            NaiveDate::from_ymd(2017, 4, 2).and_hms(12, 0, 0),
812        );
813        let missing_times = arch
814            .missing_inventory(kmso, Model::NAM, Some(larger_range))
815            .expect("Data base error?");
816        assert!(missing_times.contains(&NaiveDate::from_ymd(2017, 3, 31).and_hms(0, 0, 0)));
817        assert!(missing_times.contains(&NaiveDate::from_ymd(2017, 3, 31).and_hms(6, 0, 0)));
818        assert!(missing_times.contains(&NaiveDate::from_ymd(2017, 3, 31).and_hms(12, 0, 0)));
819        assert!(missing_times.contains(&NaiveDate::from_ymd(2017, 3, 31).and_hms(18, 0, 0)));
820        assert!(!missing_times.contains(&first));
821        assert!(!missing_times.contains(&second));
822        assert!(!missing_times.contains(&last));
823        assert!(missing_times.contains(&missing));
824        assert!(missing_times.contains(&NaiveDate::from_ymd(2017, 4, 2).and_hms(0, 0, 0)));
825        assert!(missing_times.contains(&NaiveDate::from_ymd(2017, 4, 2).and_hms(6, 0, 0)));
826        assert!(missing_times.contains(&NaiveDate::from_ymd(2017, 4, 2).and_hms(12, 0, 0)));
827    }
828
829    #[test]
830    fn test_retrieve_all_valid_in() {
831        let TestArchive {
832            tmp: _tmp,
833            mut arch,
834        } = create_test_archive().expect("Failed to create test archive.");
835
836        fill_test_archive(&mut arch);
837
838        let kmso = StationNumber::from(727730); // Station number for KMSO
839        let start = NaiveDate::from_ymd(2017, 4, 1).and_hms(0, 0, 0);
840        let end = NaiveDate::from_ymd(2017, 4, 1).and_hms(12, 0, 0);
841        assert_eq!(
842            arch.retrieve_all_valid_in(kmso, Model::GFS, start, end)
843                .unwrap()
844                .into_iter()
845                .count(),
846            1
847        );
848
849        let end = NaiveDate::from_ymd(2017, 4, 2).and_hms(0, 0, 0);
850        assert_eq!(
851            arch.retrieve_all_valid_in(kmso, Model::GFS, start, end)
852                .unwrap()
853                .into_iter()
854                .count(),
855            3
856        );
857
858        let start = NaiveDate::from_ymd(2017, 4, 1).and_hms(18, 0, 0);
859        assert_eq!(
860            arch.retrieve_all_valid_in(kmso, Model::GFS, start, end)
861                .unwrap()
862                .into_iter()
863                .count(),
864            3
865        );
866
867        let start = NaiveDate::from_ymd(2017, 4, 1).and_hms(0, 0, 0);
868        let end = NaiveDate::from_ymd(2017, 4, 1).and_hms(12, 0, 0);
869        assert_eq!(
870            arch.retrieve_all_valid_in(kmso, Model::NAM, start, end)
871                .unwrap()
872                .into_iter()
873                .count(),
874            1
875        );
876
877        let end = NaiveDate::from_ymd(2017, 4, 2).and_hms(0, 0, 0);
878        assert_eq!(
879            arch.retrieve_all_valid_in(kmso, Model::NAM, start, end)
880                .unwrap()
881                .into_iter()
882                .count(),
883            3
884        );
885
886        let start = NaiveDate::from_ymd(2017, 4, 1).and_hms(18, 0, 0);
887        assert_eq!(
888            arch.retrieve_all_valid_in(kmso, Model::NAM, start, end)
889                .unwrap()
890                .into_iter()
891                .count(),
892            3
893        );
894    }
895
896    #[test]
897    fn test_count() {
898        let TestArchive {
899            tmp: _tmp,
900            mut arch,
901        } = create_test_archive().expect("Failed to create test archive.");
902
903        fill_test_archive(&mut arch);
904
905        let kmso = StationNumber::from(727730); // Station number for KMSO
906        let model = Model::GFS;
907        assert_eq!(arch.count(kmso, model).unwrap(), 3);
908
909        let model = Model::NAM;
910        assert_eq!(arch.count(kmso, model).unwrap(), 3);
911    }
912}
bufkit_data/archive/query.rs

bufkit_data/archive/
query.rs
