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 proxy_overhead_us: 0,
307 },
308 status: probe_status_from_str(&row.get::<_, String>(9)?)?,
309 response_status_code: row.get::<_, Option<i64>>(10)?.map(|value| value as u16),
310 response_status_text: row.get(11)?,
311 response_bytes: row.get::<_, Option<i64>>(12)?.map(|value| value as u64),
312 error_message: row.get(13)?,
313 tls: read_tls_report(row)?,
314 created_at_ms: row.get(23)?,
315 })
316}
317
318fn probe_status_as_str(status: ProbeStatus) -> &'static str {
319 match status {
320 ProbeStatus::Succeeded => "succeeded",
321 ProbeStatus::Failed => "failed",
322 }
323}
324
325fn probe_status_from_str(value: &str) -> rusqlite::Result<ProbeStatus> {
326 match value {
327 "succeeded" => Ok(ProbeStatus::Succeeded),
328 "failed" => Ok(ProbeStatus::Failed),
329 _ => Err(rusqlite::Error::InvalidQuery),
330 }
331}
332
333fn tls_status_as_str(status: TlsStatus) -> &'static str {
334 match status {
335 TlsStatus::NotUsed => "not_used",
336 TlsStatus::HandshakeSucceeded => "handshake_succeeded",
337 TlsStatus::HandshakeFailed => "handshake_failed",
338 }
339}
340
341fn tls_status_from_str(value: &str) -> rusqlite::Result<TlsStatus> {
342 match value {
343 "not_used" => Ok(TlsStatus::NotUsed),
344 "handshake_succeeded" => Ok(TlsStatus::HandshakeSucceeded),
345 "handshake_failed" => Ok(TlsStatus::HandshakeFailed),
346 _ => Err(rusqlite::Error::InvalidQuery),
347 }
348}
349
350fn read_tls_report(row: &rusqlite::Row<'_>) -> rusqlite::Result<Option<TlsReport>> {
351 let enabled = row.get::<_, Option<i64>>(14)?;
352
353 let Some(enabled) = enabled else {
354 return Ok(None);
355 };
356
357 let subject = row.get::<_, Option<String>>(18)?;
358 let issuer = row.get::<_, Option<String>>(19)?;
359 let not_before = row.get::<_, Option<String>>(20)?;
360 let not_after = row.get::<_, Option<String>>(21)?;
361
362 let certificate = match (subject, issuer, not_before, not_after) {
363 (Some(subject), Some(issuer), not_before, not_after) => Some(CertificateInfo {
364 subject,
365 issuer,
366 not_before,
367 not_after,
368 }),
369 _ => None,
370 };
371
372 Ok(Some(TlsReport {
373 enabled: enabled != 0,
374 status: tls_status_from_str(&row.get::<_, String>(15)?)?,
375 version: row.get(16)?,
376 cipher_suite: row.get(17)?,
377 certificate,
378 error_message: row.get(22)?,
379 }))
380}
381
382fn request_path_from_target(target: &str) -> String {
383 url::Url::parse(target)
384 .ok()
385 .map(|url| {
386 let mut path = url.path().to_string();
387 if path.is_empty() {
388 path.push('/');
389 }
390 if let Some(query) = url.query() {
391 path.push('?');
392 path.push_str(query);
393 }
394 path
395 })
396 .unwrap_or_else(|| "/".to_string())
397}
398
399#[cfg(test)]
400mod tests {
401 use crate::engine::timing::build_report;
402 use crate::models::{NewProbeResult, ProbeStatus, TlsReport, TlsStatus};
403
404 use super::SqliteStore;
405
406 #[test]
407 fn save_report_and_read_recent_results() {
408 let store = SqliteStore::in_memory().expect("store should initialize");
409 let report = build_report(10, 20, 30, 40, 5);
410
411 let saved_result = store
412 .save_report("http://localhost:3000", &report)
413 .expect("report should save");
414
415 let runs = store
416 .recent_results(10)
417 .expect("recent results should load");
418
419 assert_eq!(runs.len(), 1);
420 assert_eq!(runs[0].id, saved_result.id);
421 assert_eq!(runs[0].target, "http://localhost:3000");
422 assert_eq!(runs[0].request_path, "/");
423 assert_eq!(runs[0].report, report);
424 assert_eq!(runs[0].status, ProbeStatus::Succeeded);
425 assert_eq!(runs[0].response_status_code, None);
426 assert_eq!(runs[0].error_message, None);
427 assert_eq!(runs[0].tls, None);
428 }
429
430 #[test]
431 fn latest_result_for_target_returns_most_recent_match() {
432 let store = SqliteStore::in_memory().expect("store should initialize");
433
434 let first_report = build_report(1, 2, 3, 4, 5);
435 let second_report = build_report(4, 5, 6, 7, 8);
436
437 store
438 .save_report("http://localhost:3000", &first_report)
439 .expect("first report should save");
440 let latest_result = store
441 .save_report("http://localhost:3000", &second_report)
442 .expect("second report should save");
443 store
444 .save_report("http://localhost:4000", &first_report)
445 .expect("other target report should save");
446
447 let run = store
448 .latest_result_for_target("http://localhost:3000")
449 .expect("latest result should load")
450 .expect("target should have a run");
451
452 assert_eq!(run.id, latest_result.id);
453 assert_eq!(run.target, "http://localhost:3000");
454 assert_eq!(run.request_path, "/");
455 assert_eq!(run.report, second_report);
456 assert_eq!(run.status, ProbeStatus::Succeeded);
457 }
458
459 #[test]
460 fn recent_results_for_target_filters_by_target() {
461 let store = SqliteStore::in_memory().expect("store should initialize");
462 let report = build_report(1, 2, 3, 4, 5);
463
464 store
465 .save_report("http://localhost:3000", &report)
466 .expect("first target report should save");
467 store
468 .save_report("http://localhost:4000", &report)
469 .expect("second target report should save");
470 store
471 .save_report("http://localhost:3000", &report)
472 .expect("third target report should save");
473
474 let runs = store
475 .recent_results_for_target("http://localhost:3000", 10)
476 .expect("filtered results should load");
477
478 assert_eq!(runs.len(), 2);
479 assert!(runs.iter().all(|run| run.target == "http://localhost:3000"));
480 }
481
482 #[test]
483 fn delete_results_for_target_removes_matching_rows() {
484 let store = SqliteStore::in_memory().expect("store should initialize");
485 let report = build_report(1, 2, 3, 4, 5);
486
487 store
488 .save_report("http://localhost:3000", &report)
489 .expect("first target report should save");
490 store
491 .save_report("http://localhost:4000", &report)
492 .expect("second target report should save");
493 store
494 .save_report("http://localhost:3000", &report)
495 .expect("third target report should save");
496
497 let deleted = store
498 .delete_results_for_target("http://localhost:3000")
499 .expect("deletion should succeed");
500
501 assert_eq!(deleted, 2);
502 assert_eq!(
503 store
504 .recent_results_for_target("http://localhost:3000", 10)
505 .expect("filtered results should load")
506 .len(),
507 0
508 );
509 assert_eq!(
510 store
511 .recent_results_for_target("http://localhost:4000", 10)
512 .expect("remaining results should load")
513 .len(),
514 1
515 );
516 assert_eq!(
517 store
518 .count_results_for_target("http://localhost:3000")
519 .expect("count should load"),
520 0
521 );
522 assert_eq!(
523 store
524 .count_results_for_target("http://localhost:4000")
525 .expect("remaining count should load"),
526 1
527 );
528 }
529
530 #[test]
531 fn save_result_persists_failure_details() {
532 let store = SqliteStore::in_memory().expect("store should initialize");
533 let result = NewProbeResult {
534 target: "http://localhost:3000".to_string(),
535 request_path: "/".to_string(),
536 report: build_report(0, 0, 0, 0, 0),
537 status: ProbeStatus::Failed,
538 response_status_code: None,
539 response_status_text: None,
540 response_bytes: None,
541 error_message: Some("request timed out".to_string()),
542 tls: Some(TlsReport {
543 enabled: true,
544 status: TlsStatus::HandshakeFailed,
545 version: None,
546 cipher_suite: None,
547 certificate: None,
548 error_message: Some("request timed out".to_string()),
549 }),
550 };
551
552 let saved = store
553 .save_result(&result)
554 .expect("result should save");
555
556 let loaded = store
557 .latest_result_for_target("http://localhost:3000")
558 .expect("latest result should load")
559 .expect("target should have a result");
560
561 assert_eq!(loaded.id, saved.id);
562 assert_eq!(loaded.status, ProbeStatus::Failed);
563 assert_eq!(loaded.request_path, "/");
564 assert_eq!(loaded.error_message.as_deref(), Some("request timed out"));
565 assert_eq!(
566 loaded.tls.expect("tls should be stored").status,
567 TlsStatus::HandshakeFailed
568 );
569 }
570}