1use crate::location::{Location, State};
2use crate::radar::{
3 Radar, RadarId, RadarImageDataLayer, RadarImageFeature, RadarImageFeatureLayer,
4 RadarImageLegend, RadarType,
5};
6use crate::station::WeatherStation;
7use anyhow::{anyhow, Result};
8use chrono::{TimeZone, Utc};
9use rusqlite::{named_params, params, Connection, Row};
10use std::fs;
11use std::path::PathBuf;
12use std::str::FromStr;
13use tracing::{debug, info};
14
15pub struct Database {
16 path: PathBuf,
17 pub conn: Connection,
18}
19impl Database {
20 pub fn from_path(path: PathBuf) -> Result<Database> {
21 if let Some(parent) = path.parent() {
22 fs::create_dir_all(parent)?;
23 }
24 let connection = Connection::open(&path)?;
25 Ok(Self {
26 path,
27 conn: connection,
28 })
29 }
30
31 pub fn init(&self) -> Result<()> {
32 info!("creating database at {}", self.path.display());
33 self.conn.execute_batch(include_str!("../sql/schema.sql"))?;
34 Ok(())
35 }
36
37 pub fn insert_stations(
38 &mut self,
39 stations: impl Iterator<Item = WeatherStation>,
40 ) -> Result<()> {
41 let tx = self.conn.transaction()?;
42 let mut stmt = tx.prepare(include_str!("../sql/insert_station.sql"))?;
43
44 for station in stations {
45 stmt.execute(named_params! {
46 ":id": station.id,
47 ":district_id": station.district_id,
48 ":name": station.name,
49 ":start": station.start,
50 ":end": station.end,
51 ":latitude": station.latitude,
52 ":longitude": station.longitude,
53 ":source": station.source,
54 ":state": station.state,
55 ":height": station.height,
56 ":barometric_height": station.barometric_height,
57 ":wmo_id": station.wmo_id,
58 })?;
59 }
60
61 stmt.finalize()?;
62 tx.commit()?;
63 Ok(())
64 }
65
66 pub fn get_station(&self, bom_id: u32) -> Result<WeatherStation> {
67 let mut stmt = self.conn.prepare("SELECT * FROM station WHERE id = (?)")?;
68 let mut binding = stmt.query(params![bom_id])?;
69 let row = binding.next()?.unwrap();
70
71 Ok(WeatherStation {
72 id: row.get(0)?,
73 district_id: row.get(1)?,
74 name: row.get(2)?,
75 start: row.get(3)?,
76 end: row.get(4)?,
77 latitude: row.get(5)?,
78 longitude: row.get(6)?,
79 source: row.get(7)?,
80 state: row.get(8)?,
81 height: row.get(9)?,
82 barometric_height: row.get(10)?,
83 wmo_id: row.get(11)?,
84 })
85 }
86
87 pub fn insert_radars(&mut self, radars: &[Radar], legends: &[RadarImageLegend]) -> Result<()> {
88 let tx = self.conn.transaction()?;
89 let mut stmt = tx.prepare(include_str!("../sql/insert_radar.sql"))?;
90
91 for radar in radars {
92 stmt.execute(named_params! {
93 ":id": radar.id,
94 ":name": radar.name,
95 ":full_name": radar.full_name,
96 ":latitude": radar.latitude,
97 ":longitude": radar.longitude,
98 ":state": radar.state,
99 ":type_": radar.r#type,
100 ":group_": radar.group,
101 })?;
102 }
103
104 stmt.finalize()?;
105
106 let mut stmt = tx.prepare("INSERT INTO radar_legend (id, image) VALUES (?, ?)")?;
107 for legend in legends {
108 stmt.execute(params![legend.r#type.id(), legend.png_buf])?;
109 }
110 stmt.finalize()?;
111 tx.commit()?;
112 Ok(())
113 }
114
115 pub fn load_radars(&self) -> Result<Vec<Radar>> {
116 let mut stmt = self.conn.prepare("SELECT * FROM radar")?;
117 let mut radars = Vec::new();
118 let mut rows = stmt.query([])?;
119 while let Some(row) = rows.next()? {
120 radars.push(self.row_to_radar(row)?);
121 }
122 Ok(radars)
123 }
124
125 fn row_to_radar(&self, row: &Row) -> Result<Radar> {
126 Ok(Radar {
127 id: row.get(0)?,
128 name: row.get(1)?,
129 full_name: row.get(2)?,
130 latitude: row.get(3)?,
131 longitude: row.get(4)?,
132 state: row.get(5)?,
133 r#type: row.get(6)?,
134 group: row.get(7)?,
135 })
136 }
137
138 pub fn get_radar(&self, id: u32) -> Result<Radar> {
139 let mut stmt = self.conn.prepare("SELECT * FROM radar WHERE id = (?)")?;
140 let mut binding = stmt.query(params![id])?;
141 let row = binding.next()?.unwrap();
142 self.row_to_radar(row)
143 }
144
145 pub fn get_radar_data_layers(
146 &self,
147 id: RadarId,
148 type_: &RadarType,
149 max_frames: Option<u64>,
150 ) -> Result<Vec<RadarImageDataLayer>> {
151 let max_frames = max_frames.map(|i| i as i32).unwrap_or(-1);
152 let params = params![id, type_.id() as u8, max_frames];
153 let sql = include_str!("../sql/get_radar_data_layers.sql");
154 let mut stmt = self.conn.prepare(sql)?;
155 let mut layers = Vec::new();
156 let mut rows = stmt.query(params)?;
157 while let Some(row) = rows.next()? {
158 let radar_type_id: u8 = row.get(1)?;
159 let image = RadarImageDataLayer {
160 radar_id: row.get(0)?,
161 radar_type: RadarType::from_id(radar_type_id as char)?,
162 png_buf: row.get(2)?,
163 datetime: Utc.timestamp_opt(row.get(3)?, 0).unwrap(),
164 filename: row.get(4)?,
165 };
166 layers.push(image);
167 }
168 if layers.is_empty() {
169 return Err(anyhow!(
170 "No data layers found for radar ID {id} type {type_}"
171 ));
172 }
173 layers.reverse();
174 Ok(layers)
175 }
176
177 pub fn get_radar_data_layer_names(&mut self, radar_type: &RadarType) -> Result<Vec<String>> {
178 debug!(
179 "Loading existing radar data layer names from {}",
180 self.path.display()
181 );
182 let params = params![radar_type.id() as u8];
183 let sql = "SELECT filename FROM radar_data_layer WHERE radar_type_id = (?)";
184 let mut stmt = self.conn.prepare(sql)?;
185 let mut names = Vec::new();
186 let mut rows = stmt.query(params)?;
187 while let Some(row) = rows.next()? {
188 names.push(row.get(0)?);
189 }
190 Ok(names)
191 }
192 pub fn get_radar_feature_layers(
193 &mut self,
194 id: RadarId,
195 type_: &RadarType,
196 ) -> Result<Vec<RadarImageFeatureLayer>> {
197 let params = params![id, type_.size().id() as u8];
198 let sql = include_str!("../sql/get_radar_feature_layers.sql");
199 let mut stmt = self.conn.prepare(sql)?;
200 let mut layers = Vec::new();
201 let mut rows = stmt.query(params)?;
202
203 while let Some(row) = rows.next()? {
204 let layer_type: String = row.get(1)?;
205 let radar_type_id: u8 = row.get(2)?;
206 let layer = RadarImageFeatureLayer {
207 radar_id: row.get(0)?,
208 feature: RadarImageFeature::from_str(&layer_type)?,
209 size: RadarType::from_id(radar_type_id as char)?,
210 png_buf: row.get(3)?,
211 filename: row.get(4)?,
212 };
213 layers.push(layer);
214 }
215 if layers.is_empty() {
216 return Err(anyhow!(
217 "No feature layers found for radar ID {id} type {type_}"
218 ));
219 }
220 Ok(layers)
221 }
222
223 pub fn get_radar_legend(&mut self, radar_type: &RadarType) -> Result<RadarImageLegend> {
224 let legend_type = radar_type.legend_type();
225 let sql = "SELECT image FROM radar_legend WHERE id = (?)";
226 let mut stmt = self.conn.prepare(sql)?;
227 let mut binding = stmt.query(params![legend_type.id()])?;
228 let row = binding.next()?.unwrap();
229 let legend = RadarImageLegend {
230 r#type: legend_type,
231 png_buf: row.get_unwrap(0),
232 };
233 Ok(legend)
234 }
235
236 pub fn insert_radar_data_layers(&mut self, layers: &[RadarImageDataLayer]) -> Result<()> {
237 let tx = self.conn.transaction()?;
238 let mut stmt = tx.prepare(include_str!("../sql/insert_radar_data_layer.sql"))?;
239 for layer in layers {
240 debug!(
241 "Inserting data layer {} into {}",
242 layer.filename,
243 self.path.display()
244 );
245 stmt.execute(named_params! {
246 ":image": layer.png_buf,
247 ":radar_id": layer.radar_id,
248 ":radar_type_id": layer.radar_type.id() as u8,
249 ":timestamp": layer.datetime.timestamp(),
250 ":filename": layer.filename,
251 })?;
252 }
253 stmt.finalize()?;
254 tx.commit()?;
255 Ok(())
256 }
257
258 pub fn insert_radar_feature_layers(&mut self, layers: &[RadarImageFeatureLayer]) -> Result<()> {
259 let tx = self.conn.transaction()?;
260 let mut stmt = tx.prepare(include_str!("../sql/insert_radar_feature_layer.sql"))?;
261 for layer in layers {
262 debug!(
263 "Inserting feature layer {} into {}",
264 layer.filename,
265 self.path.display()
266 );
267 stmt.execute(named_params! {
268 ":image": layer.png_buf,
269 ":radar_id": layer.radar_id,
270 ":radar_type_id": layer.size.id() as u8,
271 ":feature": layer.feature.to_string(),
272 ":filename": layer.filename,
273 })?;
274 }
275 stmt.finalize()?;
276 tx.commit()?;
277 Ok(())
278 }
279
280 pub fn delete_radar_data_layers(&mut self, layers: &[RadarImageDataLayer]) -> Result<()> {
281 let tx = self.conn.transaction()?;
282 let sql = "DELETE FROM radar_data_layer WHERE filename = (?)";
283 let mut stmt = tx.prepare(sql)?;
284 for layer in layers {
285 debug!(
286 "Deleting data layer {} from {}",
287 layer.filename,
288 self.path.display()
289 );
290 stmt.execute(params![layer.filename])?;
291 }
292 stmt.finalize()?;
293 tx.commit()?;
294 Ok(())
295 }
296
297 pub fn insert_location(&self, location: &Location) -> Result<()> {
298 debug!(
299 "Inserting location {} into {}",
300 location.id,
301 self.path.display()
302 );
303 let mut stmt = self
304 .conn
305 .prepare_cached(include_str!("../sql/insert_location.sql"))?;
306
307 stmt.execute(named_params! {
308 ":id": location.id,
309 ":geohash": location.geohash,
310 ":station_id": location.station.as_ref().map(|s| s.id),
311 ":has_wave": location.has_wave,
312 ":latitude": location.latitude,
313 ":longitude": location.longitude,
314 ":marine_area_id": location.marine_area_id,
315 ":name": location.name,
316 ":state": location.state.to_string(),
317 ":postcode": location.postcode,
318 ":tidal_point": location.tidal_point,
319 ":timezone": location.timezone,
320 ":weather": serde_json::to_string(&location.weather)?,
321 })?;
322
323 Ok(())
324 }
325
326 pub fn update_weather(&self, location: &Location) -> Result<()> {
327 debug!(
328 "Updating {}'s weather in {}",
329 location.id,
330 self.path.display()
331 );
332 let mut stmt = self
333 .conn
334 .prepare("UPDATE location SET weather = (?) WHERE id = (?)")?;
335 let weather = serde_json::to_string(&location.weather)?;
336 stmt.execute(params![weather, location.id])?;
337 Ok(())
338 }
339
340 pub fn get_location(&self, id: &str) -> Result<Location> {
341 let mut stmt = self.conn.prepare("SELECT * FROM location WHERE id = (?)")?;
342 let mut binding = stmt.query(params![id])?;
343 let Some(row) = binding.next()? else {
344 return Err(anyhow!(
345 "No record of Location {} in {}",
346 id,
347 self.path.display()
348 ));
349 };
350
351 let station = if let Some(station_id) = row.get(2)? {
352 Some(self.get_station(station_id)?)
353 } else {
354 None
355 };
356 let state_name: String = row.get(8)?;
357 let state = State::from_str(&state_name).unwrap();
358 let weather_json: String = row.get(12)?;
359 let weather = serde_json::from_str(&weather_json).unwrap();
360
361 Ok(Location {
362 id: row.get(0)?,
363 geohash: row.get(1)?,
364 station,
365 has_wave: row.get(3)?,
366 latitude: row.get(4)?,
367 longitude: row.get(5)?,
368 marine_area_id: row.get(6)?,
369 name: row.get(7)?,
370 state,
371 postcode: row.get(9)?,
372 tidal_point: row.get(10)?,
373 timezone: row.get(11)?,
374 weather,
375 })
376 }
377
378 pub fn get_locations(&self, location_ids: &Vec<String>) -> Result<Vec<Location>> {
379 let mut locations = Vec::new();
380 for id in location_ids {
381 let location = self.get_location(id)?;
382 locations.push(location);
383 }
384 Ok(locations)
385 }
386}