From 460e6482ca11ea9c3b8ea72721221dea64cb3e0a Mon Sep 17 00:00:00 2001
From: Brandon Simpson <simpsob6@wwu.edu>
Date: Mon, 10 Jul 2023 09:38:27 -0700
Subject: [PATCH] save request results into sqlite database

---
 Cargo.toml  |  1 +
 src/db.rs   | 71 +++++++++++++++++++++++++++++++++++++++++++++++++++++
 src/main.rs |  7 ++++++
 3 files changed, 79 insertions(+)
 create mode 100644 src/db.rs

diff --git a/Cargo.toml b/Cargo.toml
index d9cfdf8f..36c029d5 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -53,6 +53,7 @@ trust-dns-resolver = "0.22.0"
 rand_regex = "0.15.1"
 regex-syntax = "0.6.22"
 url = "2.4.0"
+rusqlite = "0.29.0"
 
 [target.'cfg(unix)'.dependencies]
 rlimit = "0.9.0"
diff --git a/src/db.rs b/src/db.rs
new file mode 100644
index 00000000..81246be4
--- /dev/null
+++ b/src/db.rs
@@ -0,0 +1,71 @@
+use rusqlite::Connection;
+
+use crate::client::{ClientError, RequestResult};
+
+fn create_db(conn: &Connection) -> Result<usize, rusqlite::Error> {
+    conn.execute(
+        "CREATE TABLE loadtest (
+            url TEXT NOT NULL,
+            duration REAL,
+            status INTEGER,
+            len_bytes INTEGER
+        )",
+        (),
+    )
+}
+
+pub fn store(
+    db_url: &str,
+    req_url: String,
+    request_records: Vec<Result<RequestResult, ClientError>>,
+) -> Result<usize, rusqlite::Error> {
+    let conn = Connection::open(db_url)?;
+    _ = create_db(&conn);
+
+    let request_url = req_url
+        .replace("https", "")
+        .replace("http", "")
+        .replace("://", "");
+
+    let affected_rows = request_records
+        .into_iter()
+        .map(|req| {
+            match req {
+      Ok(req) => {
+        conn.execute(
+          "INSERT INTO loadtest (url, duration, status, len_bytes) VALUES (?1, ?2, ?3, ?4)",
+        (&request_url, req.duration().as_secs_f32(), req.status.as_u16() as u32, req.len_bytes),
+        ).unwrap_or(0)
+      },
+      Err(_) => {
+        0
+      }
+    }
+        })
+        .sum();
+
+    Ok(affected_rows)
+}
+
+#[cfg(test)]
+mod test_db {
+    use super::*;
+
+    #[test]
+    fn test_store() {
+        let conn = Connection::open_in_memory().unwrap();
+        let _ = create_db(&conn);
+        let test_val = RequestResult {
+            status: hyper::StatusCode::OK,
+            len_bytes: 100,
+            start_latency_correction: None,
+            start: std::time::Instant::now(),
+            connection_time: None,
+            end: std::time::Instant::now(),
+        };
+        let test_vec = vec![Ok(test_val.clone()), Ok(test_val.clone())];
+        let result = store("test.db", "test.com".to_owned(), test_vec);
+        assert_eq!(result.unwrap(), 2);
+        std::fs::remove_file("test.db").unwrap();
+    }
+}
\ No newline at end of file
diff --git a/src/main.rs b/src/main.rs
index d087e05f..8067a3e3 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -16,6 +16,7 @@ mod monitor;
 mod printer;
 mod timescale;
 mod url_generator;
+mod db;
 
 #[cfg(unix)]
 #[global_allocator]
@@ -136,6 +137,8 @@ Examples: -z 10s -z 3m.",
         long = "unix-socket"
     )]
     unix_socket: Option<std::path::PathBuf>,
+    #[clap(help = "sqlite datebase url E.G test.db", long = "db-url")]
+    db_url: Option<String>
 }
 
 /// An entry specified by `connect-to` to override DNS resolution and default
@@ -473,5 +476,9 @@ async fn main() -> anyhow::Result<()> {
 
     printer::print_result(&mut std::io::stdout(), print_mode, start, &res, duration)?;
 
+    if let Some(db_url) = opts.db_url{
+        let _ = db::store(&db_url, opts.url, res);
+    }
+
     Ok(())
 }