1 files changed, 15 insertions, 18 deletions

diff --git a/src/main.rs b/src/main.rs
index da58431..5e0fd4a 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -11,29 +11,20 @@ use rand::{thread_rng, Rng, distributions::Alphanumeric};
 
 
 // Get URL from database
-fn get_url(short:&str, update:bool) -> Result<String, postgres::Error> {
-    // Connect to db
-    let dburl:String = dotenv::var("DATABASE_URL").unwrap();
-    let mut client = Client::connect(&dburl, NoTls)?;
-
-    let row = client.query_one("SELECT url FROM shorts WHERE short = $1 LIMIT 1", &[&short])?;
+fn get_url(db:&mut Client, short:&str, update:bool) -> Result<String, postgres::Error> {
+    let row = db.query_one("SELECT url FROM shorts WHERE short = $1 LIMIT 1", &[&short])?;
 
     if row.len() == 1 && update {
-        client.execute("UPDATE shorts SET count = count + 1, last_visited = now() WHERE short = $1;", &[&short])?;
+        db.execute("UPDATE shorts SET count = count + 1, last_visited = now() WHERE short = $1;", &[&short])?;
     }
-    client.close()?;
 
     let url: String = row.get("url");
 
     Ok(url)
 }
 
-fn insert_short(url: &str, short: &str, user: &str) -> Result<u64, postgres::Error> {
-    let dburl:String = dotenv::var("DATABASE_URL").unwrap();
-    let mut client = Client::connect(&dburl, NoTls)?;
-
-    let n = client.execute("INSERT INTO shorts (url, short, created_by) VALUES ($1, $2, $3);", &[&url, &short, &user])?;
-    client.close()?;
+fn insert_short(db:&mut Client, url: &str, short: &str, user: &str) -> Result<u64, postgres::Error> {
+    let n = db.execute("INSERT INTO shorts (url, short, created_by) VALUES ($1, $2, $3);", &[&url, &short, &user])?;
     Ok(n)
 }
 
@@ -140,7 +131,12 @@ fn main() -> std::io::Result<()> {
     // shorturl = full shortened url
     // shortprefix = part of url that comes before short id
     // docuri = DOCUMENT_URI from http request. This is everything after the domain.
-        // Example: "/hey" in "https://example.com/hey"
+    //          Example: "/hey" in "https://example.com/hey"
+
+
+    // Connect to db
+    let dburl:String = dotenv::var("DATABASE_URL").unwrap();
+    let mut db = Client::connect(&dburl, NoTls).unwrap();
 
 
     let shortprefix = dotenv::var("SHORTPREFIX").unwrap();
@@ -156,6 +152,7 @@ fn main() -> std::io::Result<()> {
     }
 
     else if action == "redirect" {
+
         // get DOCUMENT_URI
         let docuri:&str = req.var("DOCUMENT_URI").unwrap_or("");
         if docuri == "" {
@@ -172,7 +169,7 @@ fn main() -> std::io::Result<()> {
         }
 
         // Fetch URL from postgres, and redirect or return 404
-        match get_url(docuri, true) {
+        match get_url(&mut db, docuri, true) {
             Ok(url) => respond(301, &url),
             Err(_e) => respond(404, ""),
         };
@@ -194,7 +191,7 @@ fn main() -> std::io::Result<()> {
                 let mut short:String = gen_short();
 
                 for i in 1..5 {
-                    match get_url(&short, false) {
+                    match get_url(&mut db, &short, false) {
                         Ok(_url) => short = gen_short(), // If Ok, then short is already in
                                                          // use. Try a new one.
                         Err(_e) => break, // we assume that we found a unique short if get_url
@@ -205,7 +202,7 @@ fn main() -> std::io::Result<()> {
                     if i == 5 { respond(500, "Could not find unique short"); }
                 }
 
-                match insert_short(url, &short, user) {
+                match insert_short(&mut db, url, &short, user) {
                     Ok(_v) => respond(200, &format!("{}{}", shortprefix, short)),
                     Err(_e) => respond(400, "looool"),
                 };