do sql inline

1 files changed, 15 insertions, 31 deletions

diff --git a/src/main.rs b/src/main.rs
index 5e0fd4a..ca69b7f 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -9,25 +9,6 @@ use rand::{thread_rng, Rng, distributions::Alphanumeric};
 
 
 
-
-// Get URL from database
-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 {
-        db.execute("UPDATE shorts SET count = count + 1, last_visited = now() WHERE short = $1;", &[&short])?;
-    }
-
-    let url: String = row.get("url");
-
-    Ok(url)
-}
-
-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)
-}
-
 fn status<'a>(code:u16) -> &'a str {
     // I've only implemented statuscodes I *might* use
     return match code {
@@ -137,8 +118,9 @@ fn main() -> std::io::Result<()> {
     // Connect to db
     let dburl:String = dotenv::var("DATABASE_URL").unwrap();
     let mut db = Client::connect(&dburl, NoTls).unwrap();
+    // TODO: Close connection when done.
 
-
+    //
     let shortprefix = dotenv::var("SHORTPREFIX").unwrap();
 
     // Gather all request data from the environment and stdin.
@@ -169,10 +151,14 @@ fn main() -> std::io::Result<()> {
         }
 
         // Fetch URL from postgres, and redirect or return 404
-        match get_url(&mut db, docuri, true) {
-            Ok(url) => respond(301, &url),
-            Err(_e) => respond(404, ""),
-        };
+
+        match db.query_opt("SELECT url FROM shorts WHERE short = $1 LIMIT 1", &[&docuri]).unwrap() {
+            None => respond(404, ""),
+            Some(row) => {
+                db.execute("UPDATE shorts SET count = count + 1, last_visited = now() WHERE short = $1;", &[&docuri]).unwrap();
+                respond(301, row.get("url"));
+            },
+        }
     }
 
     else if action == "create" {
@@ -191,20 +177,18 @@ fn main() -> std::io::Result<()> {
                 let mut short:String = gen_short();
 
                 for i in 1..5 {
-                    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
-                                          // returns an error
+                    match db.query_opt("SELECT url FROM shorts WHERE short = $1 LIMIT 1", &[&short]).unwrap() {
+                        Some(_row) => short = gen_short(), // If a row was returned, the short was not unique. Continue loop
+                        None => break,   // If nothing was returned, the short IS unique. Break out of loop
                     }
 
                     // Throw error if we couldn't create a unique short in fire tries
                     if i == 5 { respond(500, "Could not find unique short"); }
                 }
 
-                match insert_short(&mut db, url, &short, user) {
+                match db.execute("INSERT INTO shorts (url, short, created_by) VALUES ($1, $2, $3);", &[&url, &short, &user]) {
                     Ok(_v) => respond(200, &format!("{}{}", shortprefix, short)),
-                    Err(_e) => respond(400, "looool"),
+                    Err(_e) => respond(500, "Could not save shortened url to database"),
                 };
                 exit(0);