Skip to main content

smriti/db/
geonames.rs

1//! GeoNames database access helpers.
2
3use std::fs::File;
4use std::io::{BufRead, BufReader};
5use std::path::Path;
6use std::path::PathBuf;
7
8use rusqlite::{Connection, OpenFlags};
9
10/// Resolve GeoNames DB path.
11///
12/// Searches the same candidate roots as the rest of bootstrap (asset
13/// install dir, $SMRITI_ASSET_DIR / legacy $PHOTOVAULT_ASSET_DIR,
14/// executable dir, project root, /usr/lib/smriti, /usr/lib/photovault)
15/// plus the CWD-relative `data/`.
16pub fn geonames_db_path() -> PathBuf {
17    let mut first_existing = None;
18    // 1. Walk every bootstrap-known asset root.
19    for root in candidate_geonames_roots() {
20        let p = root.join("data").join("geonames.db");
21        if p.exists() {
22            if geonames_db_is_current(&p) {
23                return p;
24            }
25            first_existing.get_or_insert(p);
26        }
27    }
28    // 2. Last-resort: CWD-relative literal "data/geonames.db".
29    let cwd = PathBuf::from("data").join("geonames.db");
30    if cwd.exists() {
31        if geonames_db_is_current(&cwd) {
32            return cwd;
33        }
34        first_existing.get_or_insert(cwd);
35    }
36    if let Some(existing) = first_existing {
37        return existing;
38    }
39    // 3. Fallback path that may not exist; caller checks.
40    crate::bootstrap::default_asset_install_dir()
41        .join("data")
42        .join("geonames.db")
43}
44
45/// Candidate root directories under which `data/geonames.db` may live.
46/// Mirrors `bootstrap::candidate_asset_roots` but inlined here to avoid
47/// pulling that helper into a public surface.
48fn candidate_geonames_roots() -> Vec<PathBuf> {
49    let mut roots = Vec::new();
50    roots.push(crate::bootstrap::default_asset_install_dir());
51    if let Ok(from_env) = std::env::var("SMRITI_ASSET_DIR") {
52        roots.push(PathBuf::from(from_env));
53    }
54    if let Ok(from_env) = std::env::var("PHOTOVAULT_ASSET_DIR") {
55        roots.push(PathBuf::from(from_env));
56    }
57    if let Ok(exe) = std::env::current_exe() {
58        if let Some(dir) = exe.parent() {
59            roots.push(dir.to_path_buf());
60            // Debian install layout (smriti primary, legacy photovault fallback)
61            roots.push(dir.join("..").join("lib").join("smriti"));
62            roots.push(dir.join("..").join("lib").join("photovault"));
63            // target/debug/<binary> → walk up to workspace root
64            roots.push(dir.join("..").join(".."));
65            roots.push(dir.join("..").join("..").join(".."));
66        }
67    }
68    let cwd = crate::bootstrap::project_root();
69    // `cargo tauri dev` typically sets CWD to src-tauri/. Walk up so the
70    // dev tree's data/ directory is reachable.
71    if let Some(parent) = cwd.parent() {
72        roots.push(parent.to_path_buf());
73        if let Some(grand) = parent.parent() {
74            roots.push(grand.to_path_buf());
75        }
76    }
77    roots.push(cwd);
78    roots.push(PathBuf::from("/usr/lib/smriti"));
79    roots.push(PathBuf::from("/usr/lib/photovault"));
80    roots
81}
82
83/// Check if bundled GeoNames DB exists.
84pub fn geonames_db_exists() -> bool {
85    geonames_schema_is_current()
86}
87
88/// Schema version stamp. Bump whenever the column layout changes so
89/// `geonames_schema_is_current` can detect a stale on-disk DB and
90/// trigger a rebuild.
91///   v2: added `feature_code` (PPLC / PPLA / PPLA2 / PPL / …). Lets us
92///       prefer real cities over GeoNames "PPL" entries that inherit
93///       a metro-wide population — e.g. Rasapudipalem (PPL, pop
94///       1,728,128 in upstream data) shouldn't outrank Visakhapatnam
95///       (PPLA2, pop 1,063,178).
96pub const GEONAMES_SCHEMA_VERSION: i64 = 2;
97
98/// Returns `true` when the on-disk geonames.db matches the current
99/// schema. False if the file is missing OR built against an older
100/// layout. Callers (asset setup, the geocoder) should rebuild when
101/// this returns false.
102pub fn geonames_schema_is_current() -> bool {
103    let path = geonames_db_path();
104    geonames_db_is_current(&path)
105}
106
107/// Validate a specific GeoNames database without creating or mutating it.
108pub fn geonames_db_is_current(path: &Path) -> bool {
109    if !path.exists() {
110        return false;
111    }
112    let Ok(conn) = Connection::open_with_flags(path, OpenFlags::SQLITE_OPEN_READ_ONLY) else {
113        return false;
114    };
115    // Probe both the v2 column and the data. A zero-byte, interrupted,
116    // or schema-only database must not make asset setup report success.
117    conn.query_row(
118        "SELECT COUNT(*) FROM cities WHERE feature_code IS NOT NULL",
119        [],
120        |row| row.get::<_, i64>(0),
121    )
122    .is_ok_and(|count| count > 1_000)
123}
124
125pub fn build_geonames_db(project_root: &Path) -> Result<(), String> {
126    let data_dir = project_root.join("data");
127    let countries_path = data_dir.join("country_codes.txt");
128    let cities_path = data_dir.join("cities1000.txt");
129    // Building and lookup have different semantics. Always put the output
130    // beside the supplied sources; setup scripts and release packaging use
131    // this exact contract.
132    let db_path = data_dir.join("geonames.db");
133
134    std::fs::create_dir_all(&data_dir).map_err(|e| format!("Failed to create db dir: {}", e))?;
135
136    // Build beside the destination and only replace the live DB after a
137    // successful commit. Interrupted setup leaves the previous DB intact.
138    let temp_path = data_dir.join(format!("geonames.db.tmp-{}", std::process::id()));
139    if temp_path.exists() {
140        std::fs::remove_file(&temp_path)
141            .map_err(|e| format!("Failed removing stale GeoNames temp DB: {}", e))?;
142    }
143
144    if let Err(error) = build_geonames_db_file(&countries_path, &cities_path, &temp_path) {
145        let _ = std::fs::remove_file(&temp_path);
146        return Err(error);
147    }
148    if !geonames_db_is_current(&temp_path) {
149        let _ = std::fs::remove_file(&temp_path);
150        return Err(
151            "Built GeoNames database failed validation; source data may be incomplete".to_string(),
152        );
153    }
154
155    let backup_path = data_dir.join(format!("geonames.db.backup-{}", std::process::id()));
156    if backup_path.exists() {
157        std::fs::remove_file(&backup_path)
158            .map_err(|e| format!("Failed removing stale GeoNames backup: {}", e))?;
159    }
160    let had_existing = db_path.exists();
161    if had_existing {
162        std::fs::rename(&db_path, &backup_path)
163            .map_err(|e| format!("Failed backing up old GeoNames DB: {}", e))?;
164    }
165    if let Err(error) = std::fs::rename(&temp_path, &db_path) {
166        if had_existing {
167            let _ = std::fs::rename(&backup_path, &db_path);
168        }
169        return Err(format!("Failed installing rebuilt GeoNames DB: {}", error));
170    }
171    if had_existing {
172        let _ = std::fs::remove_file(&backup_path);
173    }
174
175    Ok(())
176}
177
178fn build_geonames_db_file(
179    countries_path: &Path,
180    cities_path: &Path,
181    db_path: &Path,
182) -> Result<(), String> {
183    let mut conn = Connection::open(db_path).map_err(|e| format!("Failed to open DB: {}", e))?;
184
185    conn.execute_batch(
186        r#"
187        CREATE TABLE IF NOT EXISTS cities (
188            id INTEGER PRIMARY KEY,
189            name TEXT NOT NULL,
190            ascii_name TEXT NOT NULL,
191            latitude REAL NOT NULL,
192            longitude REAL NOT NULL,
193            country_code TEXT NOT NULL,
194            country_name TEXT NOT NULL,
195            population INTEGER,
196            feature_code TEXT,
197            timezone TEXT
198        );
199
200        CREATE TABLE IF NOT EXISTS countries (
201            code TEXT PRIMARY KEY,
202            name TEXT NOT NULL
203        );
204
205        CREATE INDEX IF NOT EXISTS idx_cities_coords ON cities(latitude, longitude);
206        CREATE INDEX IF NOT EXISTS idx_cities_feature ON cities(feature_code);
207        "#,
208    )
209    .map_err(|e| format!("Failed creating schema: {}", e))?;
210
211    let tx = conn
212        .transaction()
213        .map_err(|e| format!("Failed starting transaction: {}", e))?;
214
215    let countries = std::fs::read_to_string(countries_path)
216        .map_err(|e| format!("Failed reading {}: {}", countries_path.display(), e))?;
217    for line in countries.lines() {
218        let parts: Vec<&str> = line.split('\t').collect();
219        if parts.len() >= 2 {
220            // Windows PowerShell 5's `Set-Content -Encoding UTF8` prefixes
221            // the first code with a BOM. Strip it or all Andorran city rows
222            // fail their country lookup during the Windows setup path.
223            let country_code = parts[0].trim_start_matches('\u{feff}');
224            if country_code.starts_with('#') {
225                continue;
226            }
227            // Accept compact `code<TAB>name` input and raw GeoNames
228            // countryInfo.txt input produced by older packaging jobs.
229            let country_name = if parts.len() >= 5 { parts[4] } else { parts[1] };
230            tx.execute(
231                "INSERT OR IGNORE INTO countries (code, name) VALUES (?1, ?2)",
232                [country_code, country_name],
233            )
234            .map_err(|e| format!("Failed inserting country: {}", e))?;
235        }
236    }
237
238    let file = File::open(cities_path)
239        .map_err(|e| format!("Failed opening {}: {}", cities_path.display(), e))?;
240    let reader = BufReader::new(file);
241    let mut stmt = tx
242        .prepare(
243            r#"
244            INSERT OR IGNORE INTO cities
245                (id, name, ascii_name, latitude, longitude, country_code, country_name, population, feature_code, timezone)
246            VALUES
247                (?1, ?2, ?3, ?4, ?5, ?6, COALESCE((SELECT name FROM countries WHERE code = ?6), ?6), ?7, ?8, ?9)
248            "#,
249        )
250        .map_err(|e| format!("Failed preparing statement: {}", e))?;
251
252    for line in reader.lines() {
253        let line = line.map_err(|e| format!("Failed reading cities line: {}", e))?;
254        let parts: Vec<&str> = line.split('\t').collect();
255        // cities1000.txt columns (tab-separated):
256        //   0=id 1=name 2=ascii_name 3=alt_names 4=lat 5=lng
257        //   6=feature_class 7=feature_code 8=country_code 9=cc2
258        //   10..13=admin codes 14=population 15=elevation 16=dem
259        //   17=timezone 18=mod_date
260        if parts.len() >= 18 {
261            let id: i64 = parts[0].parse().unwrap_or(0);
262            let name = parts[1];
263            let ascii_name = parts[2];
264            let lat: f64 = parts[4].parse().unwrap_or(0.0);
265            let lon: f64 = parts[5].parse().unwrap_or(0.0);
266            let country_code = parts[8];
267            let population: i64 = parts[14].parse().unwrap_or(0);
268            let feature_code = parts[7];
269            let timezone = parts[17];
270
271            stmt.execute(rusqlite::params![
272                id,
273                name,
274                ascii_name,
275                lat,
276                lon,
277                country_code,
278                population,
279                feature_code,
280                timezone
281            ])
282            .map_err(|e| format!("Failed inserting city row: {}", e))?;
283        }
284    }
285
286    drop(stmt);
287    tx.commit()
288        .map_err(|e| format!("Failed committing GeoNames DB: {}", e))?;
289
290    Ok(())
291}
292
293#[cfg(test)]
294mod tests {
295    use super::*;
296
297    #[test]
298    fn builder_writes_beside_sources_and_accepts_raw_country_info() {
299        let temp = tempfile::tempdir().expect("tempdir");
300        let data = temp.path().join("data");
301        std::fs::create_dir_all(&data).expect("data dir");
302        std::fs::write(
303            data.join("country_codes.txt"),
304            "\u{feff}ZZ\tZZZ\t999\tZZ\tTest Country\n",
305        )
306        .expect("countries");
307
308        let mut cities = String::new();
309        for id in 1..=1_001 {
310            cities.push_str(&format!(
311                "{id}\tCity {id}\tCity {id}\t\t1.0\t2.0\tP\tPPL\tZZ\t\t\t\t\t\t{id}\t\t\tUTC\t2026-01-01\n"
312            ));
313        }
314        std::fs::write(data.join("cities1000.txt"), cities).expect("cities");
315
316        build_geonames_db(temp.path()).expect("build");
317
318        let db = data.join("geonames.db");
319        assert!(db.exists());
320        assert!(!temp.path().join("geonames.db").exists());
321        let conn = Connection::open(db).expect("open built db");
322        let country: String = conn
323            .query_row("SELECT country_name FROM cities WHERE id = 1", [], |row| {
324                row.get(0)
325            })
326            .expect("country name");
327        assert_eq!(country, "Test Country");
328    }
329}