popgetter_core/
lib.rs

1use std::path::Path;
2
3use anyhow::Result;
4#[cfg(feature = "cache")]
5use anyhow::{anyhow, Context};
6use data_request_spec::DataRequestSpec;
7use log::{debug, error};
8use metadata::Metadata;
9use polars::frame::DataFrame;
10use search::{Params, SearchParams, SearchResults};
11
12use crate::config::Config;
13
14// Re-exports
15pub use column_names as COL;
16
17// Modules
18pub mod column_names;
19pub mod config;
20pub mod data_request_spec;
21pub mod error;
22#[cfg(feature = "formatters")]
23pub mod formatters;
24pub mod geo;
25pub mod metadata;
26pub mod parquet;
27pub mod search;
28
29/// Type for popgetter metadata, config and API
30#[derive(Debug, PartialEq)]
31pub struct Popgetter {
32    pub metadata: Metadata,
33    pub config: Config,
34}
35
36impl Popgetter {
37    /// Setup the Popgetter object with default configuration
38    pub async fn new() -> Result<Self> {
39        Self::new_with_config(Config::default()).await
40    }
41
42    /// Setup the Popgetter object with custom configuration
43    pub async fn new_with_config(config: Config) -> Result<Self> {
44        debug!("config: {config:?}");
45        let metadata = metadata::load_all(&config).await?;
46        Ok(Self { metadata, config })
47    }
48
49    // Only include method with "cache" feature since it requires a filesystem
50    #[cfg(feature = "cache")]
51    /// Setup the Popgetter object with custom configuration from cache
52    pub async fn new_with_config_and_cache(config: Config) -> Result<Self> {
53        // On macOS: ~/Library/Caches
54        let path = dirs::cache_dir()
55            .ok_or(anyhow!("Failed to get cache directory"))?
56            .join("popgetter");
57        Popgetter::new_with_config_and_cache_path(config, path).await
58    }
59
60    // Only include method with "cache" feature since it requires a filesystem
61    #[cfg(feature = "cache")]
62    async fn new_with_config_and_cache_path<P: AsRef<Path>>(
63        config: Config,
64        path: P,
65    ) -> Result<Self> {
66        // Try to read metadata from cache
67        if path.as_ref().exists() {
68            match Popgetter::new_from_cache_path(config.clone(), &path) {
69                Ok(popgetter) => return Ok(popgetter),
70                Err(err) => {
71                    // Log error, continue without cache and attempt to create one
72                    error!("Failed to read metadata from cache with error: {err}");
73                }
74            }
75        }
76        // If no metadata cache, get metadata and try to cache
77        std::fs::create_dir_all(&path)?;
78        let popgetter = Popgetter::new_with_config(config).await?;
79        // If error creating cache, remove cache path
80        if let Err(err) = popgetter.metadata.write_cache(&path) {
81            std::fs::remove_dir_all(&path).with_context(|| {
82                "Failed to remove cache dir following error writing cache: {err}"
83            })?;
84            Err(err)?
85        }
86        Ok(popgetter)
87    }
88
89    // Only include method with "cache" feature since it requires a filesystem
90    #[cfg(feature = "cache")]
91    fn new_from_cache_path<P: AsRef<Path>>(config: Config, path: P) -> Result<Self> {
92        let metadata = Metadata::from_cache(path)?;
93        Ok(Self { metadata, config })
94    }
95
96    /// Generates `SearchResults` using popgetter given `SearchParams`
97    // TODO: consider reverting to an API where `SearchParams` are moved, add benches
98    pub fn search(&self, search_params: &SearchParams) -> SearchResults {
99        search_params
100            .clone()
101            .search(&self.metadata.combined_metric_source_geometry())
102    }
103
104    /// Downloads data using popgetter given a `DataRequestSpec`
105    pub async fn download_data_request_spec(
106        &self,
107        data_request_spec: &DataRequestSpec,
108    ) -> Result<DataFrame> {
109        let params: Params = data_request_spec.clone().try_into()?;
110        let search_results = self.search(&params.search);
111        search_results
112            .download(&self.config, &params.download)
113            .await
114    }
115
116    /// Downloads data using popgetter given `Params`
117    pub async fn download_params(&self, params: &Params) -> Result<DataFrame> {
118        self.search(&params.search)
119            .download(&self.config, &params.download)
120            .await
121    }
122}
123
124#[cfg(test)]
125#[cfg(feature = "cache")]
126mod tests {
127
128    use tempfile::TempDir;
129
130    use super::*;
131
132    #[tokio::test]
133    async fn test_popgetter_cache() -> anyhow::Result<()> {
134        let tempdir = TempDir::new()?;
135        let config = Config::default();
136        let popgetter = Popgetter::new_with_config(config.clone()).await?;
137        popgetter.metadata.write_cache(&tempdir)?;
138        let popgetter_from_cache =
139            Popgetter::new_with_config_and_cache_path(config, tempdir).await?;
140        assert_eq!(popgetter, popgetter_from_cache);
141        Ok(())
142    }
143}
popgetter_core/lib.rs

popgetter_core/
lib.rs
