heisenberg_data_processing/raw/
fetch.rs

1use std::{fs, path::Path};
2
3use futures::StreamExt;
4use indicatif::{ProgressBar, ProgressStyle};
5use reqwest::Client;
6use tempfile::NamedTempFile;
7use tokio::io::AsyncWriteExt;
8use tracing::{info, instrument};
9use zip::ZipArchive;
10
11use super::Result;
12use crate::{
13    DataError, DataSource,
14    raw::{
15        admin::{Admin1CodeRawData, Admin2CodeRawData},
16        country_info::CountryInfoRawData,
17        feature_codes::FeatureCodesRawData,
18        places::PlacesRawData,
19    },
20};
21
22pub struct TempData {
23    pub places: PlacesRawData,
24    pub country_info: CountryInfoRawData,
25    pub admin1: Admin1CodeRawData,
26    pub admin2: Admin2CodeRawData,
27    pub feature_codes: FeatureCodesRawData,
28}
29impl TempData {
30    #[must_use]
31    pub fn from_temp_files(
32        places: NamedTempFile,
33        country_info: NamedTempFile,
34        admin1: NamedTempFile,
35        admin2: NamedTempFile,
36        feature_codes: NamedTempFile,
37    ) -> Self {
38        Self {
39            places: PlacesRawData::new(places),
40            country_info: CountryInfoRawData::new(country_info),
41            admin1: Admin1CodeRawData::new(admin1),
42            admin2: Admin2CodeRawData::new(admin2),
43            feature_codes: FeatureCodesRawData::new(feature_codes),
44        }
45    }
46}
47
48#[instrument(name = "Download data", skip_all, level = "info")]
49pub fn download_data(data_source: &DataSource) -> Result<TempData> {
50    let data_source_url = data_source
51        .places_url()
52        .ok_or_else(|| DataError::NoDataDirProvided)?;
53
54    tokio::runtime::Runtime::new()?.block_on(async {
55        let client = Client::new();
56
57        let (places, country_info, admin1, admin2, feature_codes) = tokio::try_join!(
58            download_places_data(&client, &data_source_url),
59            download_country_info(&client),
60            download_admin1_codes(&client),
61            download_admin2_codes(&client),
62            download_feature_codes(&client),
63        )?;
64
65        Ok(TempData::from_temp_files(
66            places,
67            country_info,
68            admin1,
69            admin2,
70            feature_codes,
71        ))
72    })
73}
74
75/*
76
77cities15000.zip: Generated dataset with 13176 admin rows, 17960 place rows (~12s)
78warning: heisenberg@0.1.0:   Admin: 1657242 bytes (1618.4 KB)
79warning: heisenberg@0.1.0:   Place: 1880950 bytes (1836.9 KB)
80warning: heisenberg@0.1.0:   Total: 3538192 bytes (3455.3 KB)
81
82cities5000.zip: Generated dataset with 24021 admin rows, 38525 place rows (~14s)
83warning: heisenberg@0.1.0:   Admin: 2381696 bytes (2325.9 KB)
84warning: heisenberg@0.1.0:   Place: 3581318 bytes (3497.4 KB)
85warning: heisenberg@0.1.0:   Total: 5963014 bytes (5823.3 KB)
86
87cities1000.zip: Generated dataset with 62733 admin rows, 94526 place rows (~16s)
88warning: heisenberg@0.1.0:   Admin: 4464743 bytes (4360.1 KB)
89warning: heisenberg@0.1.0:   Place: 7256275 bytes (7086.2 KB)
90warning: heisenberg@0.1.0:   Total: 11721018 bytes (11446.3 KB)
91
92cities500.zip: Generated dataset with 81545 admin rows, 136847 place rows (~1m 10s)
93warning: heisenberg@0.1.0:   Admin: 5293457 bytes (5169.4 KB)
94warning: heisenberg@0.1.0:   Place: 9504661 bytes (9281.9 KB)
95warning: heisenberg@0.1.0:   Total: 14798118 bytes (14451.3 KB)
96*/
97
98async fn download_places_data(client: &Client, url: &str) -> Result<NamedTempFile> {
99    download_zip_and_extract_first_entry_to_temp_file(client, url).await
100}
101
102const COUNTRY_INFO_URL: &str = "https://download.geonames.org/export/dump/countryInfo.txt";
103/// Downloads the country info file from `GeoNames` and extracts it to a temporary file.
104async fn download_country_info(client: &Client) -> Result<NamedTempFile> {
105    download_to_temp_file(client, COUNTRY_INFO_URL).await
106}
107
108const FEATURE_CODES_URL: &str = "https://download.geonames.org/export/dump/featureCodes_en.txt";
109/// Downloads the feature codes file from `GeoNames` and extracts it to a temporary file.
110async fn download_feature_codes(client: &Client) -> Result<NamedTempFile> {
111    download_to_temp_file(client, FEATURE_CODES_URL).await
112}
113
114const ADMIN1_CODES_URL: &str = "https://download.geonames.org/export/dump/admin1CodesASCII.txt";
115/// Downloads the admin1 codes file from `GeoNames` and extracts it to a temporary file.
116async fn download_admin1_codes(client: &Client) -> Result<NamedTempFile> {
117    download_to_temp_file(client, ADMIN1_CODES_URL).await
118}
119
120const ADMIN2_CODES_URL: &str = "https://download.geonames.org/export/dump/admin2Codes.txt";
121/// Downloads the admin2 codes file from `GeoNames` and extracts it to a temporary file.
122async fn download_admin2_codes(client: &Client) -> Result<NamedTempFile> {
123    download_to_temp_file(client, ADMIN2_CODES_URL).await
124}
125
126async fn download_to_temp_file(client: &Client, url: &str) -> Result<NamedTempFile> {
127    info!(url, "Starting download");
128    let response = client.get(url).send().await?.error_for_status()?;
129
130    let total_size = response.content_length().unwrap_or(0);
131
132    let pb = ProgressBar::new(total_size);
133    pb.set_style(ProgressStyle::default_bar()
134        .template("{msg}\n{spinner:.green} [{elapsed_precise}] [{wide_bar:.cyan/blue}] {bytes}/{total_bytes} ({bytes_per_sec}, {eta})").expect("Progress bar template")
135        .progress_chars("█░"));
136    pb.set_message(format!(
137        "Downloading {}",
138        url.split('/').next_back().unwrap_or(url)
139    ));
140
141    let temp_file = NamedTempFile::new()?;
142    let mut dest_file = tokio::fs::File::create(temp_file.path()).await?;
143    //let mut dest_file = fs::File::create(temp_file.path())?;
144
145    let mut stream = response.bytes_stream();
146    while let Some(item) = stream.next().await {
147        let chunk = item?;
148        dest_file.write_all(&chunk).await?;
149        pb.inc(chunk.len() as u64);
150    }
151    dest_file.flush().await?; // Ensure all bytes are written
152    pb.finish_and_clear();
153    println!();
154    Ok(temp_file)
155}
156
157async fn download_zip_and_extract_first_entry_to_temp_file(
158    client: &Client,
159    zip_url: &str,
160) -> Result<NamedTempFile> {
161    info!(zip_url, "Starting ZIP download");
162    // Download the zip file itself (will show its own progress bar)
163    let zip_temp_file = download_to_temp_file(client, zip_url).await?;
164    info!(path = ?zip_temp_file.path(), "ZIP download complete");
165
166    let zip_file_path = zip_temp_file.path().to_path_buf();
167
168    let extracted_content_temp_file =
169        tokio::task::spawn_blocking(move || extract_first_entry_from_zip(zip_file_path)).await??;
170
171    Ok(extracted_content_temp_file)
172}
173
174fn extract_first_entry_from_zip(zip_file_path: impl AsRef<Path>) -> Result<NamedTempFile> {
175    let zip_fs_file = fs::File::open(&zip_file_path)?;
176    let mut archive = ZipArchive::new(zip_fs_file)?;
177
178    if archive.is_empty() {
179        return Err(zip::result::ZipError::FileNotFound.into());
180    }
181
182    let mut file_in_zip = archive.by_index(0)?;
183
184    let extracted_content_temp_file = NamedTempFile::with_suffix(".txt")?;
185    let mut extracted_fs_file = fs::File::create(extracted_content_temp_file.path())?;
186
187    std::io::copy(&mut file_in_zip, &mut extracted_fs_file)?;
188    info!(path = ?extracted_content_temp_file.path(), "File extracted successfully");
189
190    Ok(extracted_content_temp_file)
191}
heisenberg_data_processing/raw/fetch.rs

heisenberg_data_processing/raw/
fetch.rs
