Skip to main content

howler_core/
db.rs

1use crate::models::{Sighting, Source, SpeciesStatus};
2use anyhow::{Context, Result};
3use chrono::{DateTime, Utc};
4use rusqlite::{params, Connection};
5use std::sync::Mutex;
6
7pub struct Database {
8    conn: Mutex<Connection>,
9}
10
11impl Database {
12    pub fn new(path: &str) -> Result<Self> {
13        let conn = Connection::open(path).context("Failed to open database")?;
14
15        let db = Database {
16            conn: Mutex::new(conn),
17        };
18        db.init_schema()?;
19        Ok(db)
20    }
21
22    fn init_schema(&self) -> Result<()> {
23        let conn = self.conn.lock().unwrap();
24        conn.execute(
25            "CREATE TABLE IF NOT EXISTS sightings (
26                id INTEGER PRIMARY KEY AUTOINCREMENT,
27                species TEXT NOT NULL,
28                scientific_name TEXT,
29                latitude REAL NOT NULL,
30                longitude REAL NOT NULL,
31                observed_on TEXT NOT NULL,
32                source TEXT NOT NULL,
33                source_id TEXT NOT NULL,
34                details TEXT,
35                UNIQUE(source, source_id)
36            )",
37            [],
38        )
39        .context("Failed to create sightings table")?;
40
41        conn.execute(
42            "CREATE TABLE IF NOT EXISTS species_status (
43                id INTEGER PRIMARY KEY AUTOINCREMENT,
44                scientific_name TEXT NOT NULL UNIQUE,
45                common_name TEXT,
46                red_list_category TEXT,
47                population_trend TEXT,
48                threats TEXT
49            )",
50            [],
51        )
52        .context("Failed to create species_status table")?;
53
54        // New tables for advanced features
55        conn.execute(
56            "CREATE TABLE IF NOT EXISTS packs (
57                pack_id INTEGER PRIMARY KEY AUTOINCREMENT,
58                territory_geometry TEXT,
59                estimated_size INTEGER,
60                center_latitude REAL,
61                center_longitude REAL,
62                area_km2 REAL
63            )",
64            [],
65        )
66        .context("Failed to create packs table")?;
67
68        conn.execute(
69            "CREATE TABLE IF NOT EXISTS individuals (
70                individual_id INTEGER PRIMARY KEY AUTOINCREMENT,
71                individual_identifier TEXT UNIQUE,
72                species TEXT,
73                sex TEXT,
74                age_class TEXT,
75                pack_id INTEGER,
76                FOREIGN KEY (pack_id) REFERENCES packs(pack_id)
77            )",
78            [],
79        )
80        .context("Failed to create individuals table")?;
81
82        conn.execute(
83            "CREATE TABLE IF NOT EXISTS movements (
84                movement_id INTEGER PRIMARY KEY AUTOINCREMENT,
85                from_sighting_id INTEGER,
86                to_sighting_id INTEGER,
87                distance_km REAL,
88                bearing_degrees REAL,
89                duration_seconds INTEGER,
90                speed_kmh REAL,
91                FOREIGN KEY (from_sighting_id) REFERENCES sightings(id),
92                FOREIGN KEY (to_sighting_id) REFERENCES sightings(id)
93            )",
94            [],
95        )
96        .context("Failed to create movements table")?;
97
98        // Indexes for performance
99        conn.execute(
100            "CREATE INDEX IF NOT EXISTS idx_sightings_species ON sightings(species)",
101            [],
102        )
103        .context("Failed to create species index")?;
104
105        conn.execute(
106            "CREATE INDEX IF NOT EXISTS idx_sightings_source ON sightings(source)",
107            [],
108        )
109        .context("Failed to create source index")?;
110
111        conn.execute(
112            "CREATE INDEX IF NOT EXISTS idx_sightings_observed_on ON sightings(observed_on)",
113            [],
114        )
115        .context("Failed to create observed_on index")?;
116
117        conn.execute(
118            "CREATE INDEX IF NOT EXISTS idx_sightings_location ON sightings(latitude, longitude)",
119            [],
120        )
121        .context("Failed to create location index")?;
122
123        conn.execute(
124            "CREATE INDEX IF NOT EXISTS idx_individuals_pack ON individuals(pack_id)",
125            [],
126        )
127        .context("Failed to create individuals pack index")?;
128
129        // Performance optimizations
130        let _ = conn.query_row("PRAGMA journal_mode = WAL", [], |row| {
131            row.get::<_, String>(0)
132        });
133        let _ = conn.execute("PRAGMA synchronous = NORMAL", []);
134        let _ = conn.execute("PRAGMA cache_size = -64000", []);
135
136        Ok(())
137    }
138
139    pub fn insert_sighting(&self, sighting: &Sighting) -> Result<i64> {
140        let conn = self.conn.lock().unwrap();
141        conn.execute(
142            "INSERT OR REPLACE INTO sightings 
143             (species, scientific_name, latitude, longitude, observed_on, source, source_id, details)
144             VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8)",
145            params![
146                sighting.species,
147                sighting.scientific_name,
148                sighting.latitude,
149                sighting.longitude,
150                sighting.observed_on.to_rfc3339(),
151                sighting.source.to_string(),
152                sighting.source_id,
153                sighting.details,
154            ],
155        )
156        .context("Failed to insert sighting")?;
157
158        Ok(conn.last_insert_rowid())
159    }
160
161    /// Batch insert sightings for better performance
162    pub fn insert_sightings_batch(&self, sightings: &[Sighting]) -> Result<usize> {
163        let mut conn = self.conn.lock().unwrap();
164        let tx = conn.transaction().context("Failed to begin transaction")?;
165
166        let mut count = 0;
167        for sighting in sightings {
168            tx.execute(
169                "INSERT OR REPLACE INTO sightings 
170                 (species, scientific_name, latitude, longitude, observed_on, source, source_id, details)
171                 VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8)",
172                params![
173                    sighting.species,
174                    sighting.scientific_name,
175                    sighting.latitude,
176                    sighting.longitude,
177                    sighting.observed_on.to_rfc3339(),
178                    sighting.source.to_string(),
179                    sighting.source_id,
180                    sighting.details,
181                ],
182            )
183            .context("Failed to insert sighting in batch")?;
184            count += 1;
185        }
186
187        tx.commit().context("Failed to commit transaction")?;
188        Ok(count)
189    }
190
191    pub fn get_all_sightings(&self) -> Result<Vec<Sighting>> {
192        let conn = self.conn.lock().unwrap();
193        let mut stmt = conn
194            .prepare(
195                "SELECT id, species, scientific_name, latitude, longitude, observed_on, source, source_id, details
196                 FROM sightings",
197            )
198            .context("Failed to prepare query")?;
199
200        let sightings = stmt
201            .query_map([], |row| {
202                let observed_on_str: String = row.get(5)?;
203                let observed_on = DateTime::parse_from_rfc3339(&observed_on_str)
204                    .map(|dt| dt.with_timezone(&Utc))
205                    .unwrap_or_else(|_| Utc::now());
206
207                let source_str: String = row.get(6)?;
208                let source = match source_str.as_str() {
209                    "GBIF" => Source::GBIF,
210                    "Movebank" => Source::Movebank,
211                    "iNaturalist" => Source::INaturalist,
212                    "IUCN" => Source::IUCN,
213                    _ => Source::GBIF,
214                };
215
216                Ok(Sighting {
217                    id: Some(row.get(0)?),
218                    species: row.get(1)?,
219                    scientific_name: row.get(2)?,
220                    latitude: row.get(3)?,
221                    longitude: row.get(4)?,
222                    observed_on,
223                    source,
224                    source_id: row.get(7)?,
225                    details: row.get(8)?,
226                })
227            })
228            .context("Failed to execute query")?;
229
230        let mut result = Vec::new();
231        for sighting in sightings {
232            result.push(sighting.context("Failed to parse sighting")?);
233        }
234        Ok(result)
235    }
236
237    pub fn insert_species_status(&self, status: &SpeciesStatus) -> Result<i64> {
238        let conn = self.conn.lock().unwrap();
239        conn.execute(
240            "INSERT OR REPLACE INTO species_status 
241             (scientific_name, common_name, red_list_category, population_trend, threats)
242             VALUES (?1, ?2, ?3, ?4, ?5)",
243            params![
244                status.scientific_name,
245                status.common_name,
246                status.red_list_category,
247                status.population_trend,
248                status.threats,
249            ],
250        )
251        .context("Failed to insert species status")?;
252
253        Ok(conn.last_insert_rowid())
254    }
255
256    pub fn get_species_status(&self, scientific_name: &str) -> Result<Option<SpeciesStatus>> {
257        let conn = self.conn.lock().unwrap();
258        let mut stmt = conn
259            .prepare(
260                "SELECT id, scientific_name, common_name, red_list_category, population_trend, threats
261                 FROM species_status WHERE scientific_name = ?1",
262            )
263            .context("Failed to prepare query")?;
264
265        let mut result = None;
266        let mut rows = stmt.query(params![scientific_name])?;
267
268        if let Some(row) = rows.next()? {
269            result = Some(SpeciesStatus {
270                id: Some(row.get(0)?),
271                scientific_name: row.get(1)?,
272                common_name: row.get(2)?,
273                red_list_category: row.get(3)?,
274                population_trend: row.get(4)?,
275                threats: row.get(5)?,
276            });
277        }
278
279        Ok(result)
280    }
281
282    /// Get a reference to the underlying connection for auth/annotation modules
283    pub fn connection(&self) -> std::sync::MutexGuard<'_, Connection> {
284        self.conn.lock().unwrap()
285    }
286}