1use std::path::Path;
2use std::time::{SystemTime, UNIX_EPOCH};
3
4use rusqlite::{params, Connection};
5
6use crate::error::Result;
7use crate::models::{
8 CertificateInfo, NewProbeResult, ProbeResult, ProbeStatus, TimingReport, TlsReport, TlsStatus,
9};
10
11#[derive(Debug)]
12pub struct SqliteStore {
13 conn: Connection,
14}
15
16impl SqliteStore {
17 pub fn open<P: AsRef<Path>>(path: P) -> Result<Self> {
18 let conn = Connection::open(path)?;
19 let store = Self { conn };
20 store.init()?;
21
22 Ok(store)
23 }
24
25 pub fn in_memory() -> Result<Self> {
26 let conn = Connection::open_in_memory()?;
27 let store = Self { conn };
28 store.init()?;
29
30 Ok(store)
31 }
32
33 pub fn init(&self) -> Result<()> {
34 self.conn.execute(
35 "CREATE TABLE IF NOT EXISTS runs (
36 id INTEGER PRIMARY KEY AUTOINCREMENT,
37 target TEXT NOT NULL,
38 request_path TEXT NOT NULL DEFAULT '/',
39 dns_ms INTEGER NOT NULL,
40 tcp_ms INTEGER NOT NULL,
41 tls_ms INTEGER NOT NULL,
42 ttfb_ms INTEGER NOT NULL,
43 download_ms INTEGER NOT NULL DEFAULT 0,
44 total_ms INTEGER NOT NULL,
45 status TEXT NOT NULL,
46 response_status_code INTEGER,
47 response_status_text TEXT,
48 response_bytes INTEGER,
49 error_message TEXT,
50 tls_enabled INTEGER,
51 tls_status TEXT,
52 tls_version TEXT,
53 tls_cipher_suite TEXT,
54 tls_certificate_subject TEXT,
55 tls_certificate_issuer TEXT,
56 tls_certificate_not_before TEXT,
57 tls_certificate_not_after TEXT,
58 tls_error_message TEXT,
59 created_at_ms INTEGER NOT NULL
60 )",
61 [],
62 )?;
63
64 self.ensure_column("runs", "request_path", "TEXT NOT NULL DEFAULT '/'")?;
65 self.ensure_column("runs", "response_status_code", "INTEGER")?;
66 self.ensure_column("runs", "response_status_text", "TEXT")?;
67 self.ensure_column("runs", "response_bytes", "INTEGER")?;
68 self.ensure_column("runs", "download_ms", "INTEGER NOT NULL DEFAULT 0")?;
69
70 self.conn.execute(
71 "CREATE INDEX IF NOT EXISTS idx_runs_target_created_at
72 ON runs(target, created_at_ms DESC)",
73 [],
74 )?;
75
76 Ok(())
77 }
78
79 pub fn save_result(&self, result: &NewProbeResult) -> Result<ProbeResult> {
80 let created_at_ms = current_time_ms()?;
81
82 self.conn.execute(
83 "INSERT INTO runs (
84 target,
85 request_path,
86 dns_ms,
87 tcp_ms,
88 tls_ms,
89 ttfb_ms,
90 download_ms,
91 total_ms,
92 status,
93 response_status_code,
94 response_status_text,
95 response_bytes,
96 error_message,
97 tls_enabled,
98 tls_status,
99 tls_version,
100 tls_cipher_suite,
101 tls_certificate_subject,
102 tls_certificate_issuer,
103 tls_certificate_not_before,
104 tls_certificate_not_after,
105 tls_error_message,
106 created_at_ms
107 ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14, ?15, ?16, ?17, ?18, ?19, ?20, ?21, ?22, ?23)",
108 params![
109 &result.target,
110 &result.request_path,
111 result.report.dns_ms as i64,
112 result.report.tcp_ms as i64,
113 result.report.tls_ms as i64,
114 result.report.ttfb_ms as i64,
115 result.report.download_ms as i64,
116 result.report.total_ms as i64,
117 probe_status_as_str(result.status),
118 result.response_status_code.map(i64::from),
119 &result.response_status_text,
120 result.response_bytes.map(|value| value as i64),
121 &result.error_message,
122 result.tls.as_ref().map(|tls| tls.enabled as i64),
123 result.tls.as_ref().map(|tls| tls_status_as_str(tls.status)),
124 result.tls.as_ref().and_then(|tls| tls.version.as_deref()),
125 result
126 .tls
127 .as_ref()
128 .and_then(|tls| tls.cipher_suite.as_deref()),
129 result
130 .tls
131 .as_ref()
132 .and_then(|tls| tls.certificate.as_ref().map(|cert| cert.subject.as_str())),
133 result
134 .tls
135 .as_ref()
136 .and_then(|tls| tls.certificate.as_ref().map(|cert| cert.issuer.as_str())),
137 result.tls.as_ref().and_then(|tls| {
138 tls.certificate
139 .as_ref()
140 .and_then(|cert| cert.not_before.as_deref())
141 }),
142 result.tls.as_ref().and_then(|tls| {
143 tls.certificate
144 .as_ref()
145 .and_then(|cert| cert.not_after.as_deref())
146 }),
147 result
148 .tls
149 .as_ref()
150 .and_then(|tls| tls.error_message.as_deref()),
151 created_at_ms,
152 ],
153 )?;
154
155 Ok(ProbeResult {
156 id: self.conn.last_insert_rowid(),
157 target: result.target.clone(),
158 request_path: result.request_path.clone(),
159 report: result.report.clone(),
160 status: result.status,
161 response_status_code: result.response_status_code,
162 response_status_text: result.response_status_text.clone(),
163 response_bytes: result.response_bytes,
164 error_message: result.error_message.clone(),
165 tls: result.tls.clone(),
166 created_at_ms,
167 })
168 }
169
170 pub fn save_report(&self, target: &str, report: &TimingReport) -> Result<ProbeResult> {
171 let result = NewProbeResult {
172 target: target.to_string(),
173 request_path: request_path_from_target(target),
174 report: report.clone(),
175 status: ProbeStatus::Succeeded,
176 response_status_code: None,
177 response_status_text: None,
178 response_bytes: None,
179 error_message: None,
180 tls: None,
181 };
182
183 self.save_result(&result)
184 }
185
186 pub fn recent_results(&self, limit: usize) -> Result<Vec<ProbeResult>> {
187 let mut statement = self.conn.prepare(
188 "SELECT id, target, request_path, dns_ms, tcp_ms, tls_ms, ttfb_ms, download_ms, total_ms, status,
189 response_status_code, response_status_text, response_bytes, error_message,
190 tls_enabled, tls_status, tls_version, tls_cipher_suite,
191 tls_certificate_subject, tls_certificate_issuer,
192 tls_certificate_not_before, tls_certificate_not_after,
193 tls_error_message, created_at_ms
194 FROM runs
195 ORDER BY created_at_ms DESC, id DESC
196 LIMIT ?1",
197 )?;
198
199 let rows = statement.query_map([limit as i64], read_probe_result)?;
200
201 rows.collect::<std::result::Result<Vec<_>, _>>()
202 .map_err(Into::into)
203 }
204
205 pub fn recent_results_for_target(&self, target: &str, limit: usize) -> Result<Vec<ProbeResult>> {
206 let mut statement = self.conn.prepare(
207 "SELECT id, target, request_path, dns_ms, tcp_ms, tls_ms, ttfb_ms, download_ms, total_ms, status,
208 response_status_code, response_status_text, response_bytes, error_message,
209 tls_enabled, tls_status, tls_version, tls_cipher_suite,
210 tls_certificate_subject, tls_certificate_issuer,
211 tls_certificate_not_before, tls_certificate_not_after,
212 tls_error_message, created_at_ms
213 FROM runs
214 WHERE target = ?1
215 ORDER BY created_at_ms DESC, id DESC
216 LIMIT ?2",
217 )?;
218
219 let rows = statement.query_map(params![target, limit as i64], read_probe_result)?;
220
221 rows.collect::<std::result::Result<Vec<_>, _>>()
222 .map_err(Into::into)
223 }
224
225 pub fn recent_runs(&self, limit: usize) -> Result<Vec<ProbeResult>> {
226 self.recent_results(limit)
227 }
228
229 pub fn delete_results_for_target(&self, target: &str) -> Result<usize> {
230 let deleted = self
231 .conn
232 .execute("DELETE FROM runs WHERE target = ?1", params![target])?;
233
234 Ok(deleted)
235 }
236
237 pub fn count_results_for_target(&self, target: &str) -> Result<usize> {
238 let mut statement = self
239 .conn
240 .prepare("SELECT COUNT(*) FROM runs WHERE target = ?1")?;
241 let count = statement.query_row(params![target], |row| row.get::<_, i64>(0))?;
242
243 Ok(count.max(0) as usize)
244 }
245
246 pub fn latest_result_for_target(&self, target: &str) -> Result<Option<ProbeResult>> {
247 let mut statement = self.conn.prepare(
248 "SELECT id, target, request_path, dns_ms, tcp_ms, tls_ms, ttfb_ms, download_ms, total_ms, status,
249 response_status_code, response_status_text, response_bytes, error_message,
250 tls_enabled, tls_status, tls_version, tls_cipher_suite,
251 tls_certificate_subject, tls_certificate_issuer,
252 tls_certificate_not_before, tls_certificate_not_after,
253 tls_error_message, created_at_ms
254 FROM runs
255 WHERE target = ?1
256 ORDER BY created_at_ms DESC, id DESC
257 LIMIT 1",
258 )?;
259
260 let mut rows = statement.query([target])?;
261
262 match rows.next()? {
263 Some(row) => Ok(Some(read_probe_result(row)?)),
264 None => Ok(None),
265 }
266 }
267
268 fn ensure_column(&self, table: &str, column: &str, definition: &str) -> Result<()> {
269 let pragma = format!("PRAGMA table_info({table})");
270 let mut statement = self.conn.prepare(&pragma)?;
271 let rows = statement.query_map([], |row| row.get::<_, String>(1))?;
272 let existing_columns = rows.collect::<std::result::Result<Vec<_>, _>>()?;
273
274 if !existing_columns.iter().any(|existing| existing == column) {
275 let alter = format!("ALTER TABLE {table} ADD COLUMN {column} {definition}");
276 self.conn.execute(&alter, [])?;
277 }
278
279 Ok(())
280 }
281
282 pub fn latest_run_for_target(&self, target: &str) -> Result<Option<ProbeResult>> {
283 self.latest_result_for_target(target)
284 }
285}
286
287fn current_time_ms() -> Result<i64> {
288 let now = SystemTime::now().duration_since(UNIX_EPOCH)?;
289
290 Ok(now.as_millis() as i64)
291}
292
293fn read_probe_result(row: &rusqlite::Row<'_>) -> rusqlite::Result<ProbeResult> {
294 Ok(ProbeResult {
295 id: row.get(0)?,
296 target: row.get(1)?,
297 request_path: row.get(2)?,
298 report: TimingReport {
299 dns_ms: row.get::<_, i64>(3)? as u64,
300 tcp_ms: row.get::<_, i64>(4)? as u64,
301 tls_ms: row.get::<_, i64>(5)? as u64,
302 ttfb_ms: row.get::<_, i64>(6)? as u64,
303 download_ms: row.get::<_, i64>(7)? as u64,
304 total_ms: row.get::<_, i64>(8)? as u64,
305 },
306 status: probe_status_from_str(&row.get::<_, String>(9)?)?,
307 response_status_code: row.get::<_, Option<i64>>(10)?.map(|value| value as u16),
308 response_status_text: row.get(11)?,
309 response_bytes: row.get::<_, Option<i64>>(12)?.map(|value| value as u64),
310 error_message: row.get(13)?,
311 tls: read_tls_report(row)?,
312 created_at_ms: row.get(23)?,
313 })
314}
315
316fn probe_status_as_str(status: ProbeStatus) -> &'static str {
317 match status {
318 ProbeStatus::Succeeded => "succeeded",
319 ProbeStatus::Failed => "failed",
320 }
321}
322
323fn probe_status_from_str(value: &str) -> rusqlite::Result<ProbeStatus> {
324 match value {
325 "succeeded" => Ok(ProbeStatus::Succeeded),
326 "failed" => Ok(ProbeStatus::Failed),
327 _ => Err(rusqlite::Error::InvalidQuery),
328 }
329}
330
331fn tls_status_as_str(status: TlsStatus) -> &'static str {
332 match status {
333 TlsStatus::NotUsed => "not_used",
334 TlsStatus::HandshakeSucceeded => "handshake_succeeded",
335 TlsStatus::HandshakeFailed => "handshake_failed",
336 }
337}
338
339fn tls_status_from_str(value: &str) -> rusqlite::Result<TlsStatus> {
340 match value {
341 "not_used" => Ok(TlsStatus::NotUsed),
342 "handshake_succeeded" => Ok(TlsStatus::HandshakeSucceeded),
343 "handshake_failed" => Ok(TlsStatus::HandshakeFailed),
344 _ => Err(rusqlite::Error::InvalidQuery),
345 }
346}
347
348fn read_tls_report(row: &rusqlite::Row<'_>) -> rusqlite::Result<Option<TlsReport>> {
349 let enabled = row.get::<_, Option<i64>>(14)?;
350
351 let Some(enabled) = enabled else {
352 return Ok(None);
353 };
354
355 let subject = row.get::<_, Option<String>>(18)?;
356 let issuer = row.get::<_, Option<String>>(19)?;
357 let not_before = row.get::<_, Option<String>>(20)?;
358 let not_after = row.get::<_, Option<String>>(21)?;
359
360 let certificate = match (subject, issuer, not_before, not_after) {
361 (Some(subject), Some(issuer), not_before, not_after) => Some(CertificateInfo {
362 subject,
363 issuer,
364 not_before,
365 not_after,
366 }),
367 _ => None,
368 };
369
370 Ok(Some(TlsReport {
371 enabled: enabled != 0,
372 status: tls_status_from_str(&row.get::<_, String>(15)?)?,
373 version: row.get(16)?,
374 cipher_suite: row.get(17)?,
375 certificate,
376 error_message: row.get(22)?,
377 }))
378}
379
380fn request_path_from_target(target: &str) -> String {
381 url::Url::parse(target)
382 .ok()
383 .map(|url| {
384 let mut path = url.path().to_string();
385 if path.is_empty() {
386 path.push('/');
387 }
388 if let Some(query) = url.query() {
389 path.push('?');
390 path.push_str(query);
391 }
392 path
393 })
394 .unwrap_or_else(|| "/".to_string())
395}
396
397#[cfg(test)]
398mod tests {
399 use crate::engine::timing::build_report;
400 use crate::models::{NewProbeResult, ProbeStatus, TlsReport, TlsStatus};
401
402 use super::SqliteStore;
403
404 #[test]
405 fn save_report_and_read_recent_results() {
406 let store = SqliteStore::in_memory().expect("store should initialize");
407 let report = build_report(10, 20, 30, 40, 5);
408
409 let saved_result = store
410 .save_report("http://localhost:3000", &report)
411 .expect("report should save");
412
413 let runs = store
414 .recent_results(10)
415 .expect("recent results should load");
416
417 assert_eq!(runs.len(), 1);
418 assert_eq!(runs[0].id, saved_result.id);
419 assert_eq!(runs[0].target, "http://localhost:3000");
420 assert_eq!(runs[0].request_path, "/");
421 assert_eq!(runs[0].report, report);
422 assert_eq!(runs[0].status, ProbeStatus::Succeeded);
423 assert_eq!(runs[0].response_status_code, None);
424 assert_eq!(runs[0].error_message, None);
425 assert_eq!(runs[0].tls, None);
426 }
427
428 #[test]
429 fn latest_result_for_target_returns_most_recent_match() {
430 let store = SqliteStore::in_memory().expect("store should initialize");
431
432 let first_report = build_report(1, 2, 3, 4, 5);
433 let second_report = build_report(4, 5, 6, 7, 8);
434
435 store
436 .save_report("http://localhost:3000", &first_report)
437 .expect("first report should save");
438 let latest_result = store
439 .save_report("http://localhost:3000", &second_report)
440 .expect("second report should save");
441 store
442 .save_report("http://localhost:4000", &first_report)
443 .expect("other target report should save");
444
445 let run = store
446 .latest_result_for_target("http://localhost:3000")
447 .expect("latest result should load")
448 .expect("target should have a run");
449
450 assert_eq!(run.id, latest_result.id);
451 assert_eq!(run.target, "http://localhost:3000");
452 assert_eq!(run.request_path, "/");
453 assert_eq!(run.report, second_report);
454 assert_eq!(run.status, ProbeStatus::Succeeded);
455 }
456
457 #[test]
458 fn recent_results_for_target_filters_by_target() {
459 let store = SqliteStore::in_memory().expect("store should initialize");
460 let report = build_report(1, 2, 3, 4, 5);
461
462 store
463 .save_report("http://localhost:3000", &report)
464 .expect("first target report should save");
465 store
466 .save_report("http://localhost:4000", &report)
467 .expect("second target report should save");
468 store
469 .save_report("http://localhost:3000", &report)
470 .expect("third target report should save");
471
472 let runs = store
473 .recent_results_for_target("http://localhost:3000", 10)
474 .expect("filtered results should load");
475
476 assert_eq!(runs.len(), 2);
477 assert!(runs.iter().all(|run| run.target == "http://localhost:3000"));
478 }
479
480 #[test]
481 fn delete_results_for_target_removes_matching_rows() {
482 let store = SqliteStore::in_memory().expect("store should initialize");
483 let report = build_report(1, 2, 3, 4, 5);
484
485 store
486 .save_report("http://localhost:3000", &report)
487 .expect("first target report should save");
488 store
489 .save_report("http://localhost:4000", &report)
490 .expect("second target report should save");
491 store
492 .save_report("http://localhost:3000", &report)
493 .expect("third target report should save");
494
495 let deleted = store
496 .delete_results_for_target("http://localhost:3000")
497 .expect("deletion should succeed");
498
499 assert_eq!(deleted, 2);
500 assert_eq!(
501 store
502 .recent_results_for_target("http://localhost:3000", 10)
503 .expect("filtered results should load")
504 .len(),
505 0
506 );
507 assert_eq!(
508 store
509 .recent_results_for_target("http://localhost:4000", 10)
510 .expect("remaining results should load")
511 .len(),
512 1
513 );
514 assert_eq!(
515 store
516 .count_results_for_target("http://localhost:3000")
517 .expect("count should load"),
518 0
519 );
520 assert_eq!(
521 store
522 .count_results_for_target("http://localhost:4000")
523 .expect("remaining count should load"),
524 1
525 );
526 }
527
528 #[test]
529 fn save_result_persists_failure_details() {
530 let store = SqliteStore::in_memory().expect("store should initialize");
531 let result = NewProbeResult {
532 target: "http://localhost:3000".to_string(),
533 request_path: "/".to_string(),
534 report: build_report(0, 0, 0, 0, 0),
535 status: ProbeStatus::Failed,
536 response_status_code: None,
537 response_status_text: None,
538 response_bytes: None,
539 error_message: Some("request timed out".to_string()),
540 tls: Some(TlsReport {
541 enabled: true,
542 status: TlsStatus::HandshakeFailed,
543 version: None,
544 cipher_suite: None,
545 certificate: None,
546 error_message: Some("request timed out".to_string()),
547 }),
548 };
549
550 let saved = store
551 .save_result(&result)
552 .expect("result should save");
553
554 let loaded = store
555 .latest_result_for_target("http://localhost:3000")
556 .expect("latest result should load")
557 .expect("target should have a result");
558
559 assert_eq!(loaded.id, saved.id);
560 assert_eq!(loaded.status, ProbeStatus::Failed);
561 assert_eq!(loaded.request_path, "/");
562 assert_eq!(loaded.error_message.as_deref(), Some("request timed out"));
563 assert_eq!(
564 loaded.tls.expect("tls should be stored").status,
565 TlsStatus::HandshakeFailed
566 );
567 }
568}