Ensure code uniqueness

interclip · Mar 29, 2024 · e90dcde · e90dcde
1 parent 5f2fcde
commit e90dcde
Show file tree

Hide file tree

Showing 3 changed files with 59 additions and 19 deletions.
diff --git a/migrations/2024-01-07-130330_create_clips/up.sql b/migrations/2024-01-07-130330_create_clips/up.sql
@@ -4,4 +4,6 @@ CREATE TABLE clips (
     code TEXT NOT NULL,
     created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
     expires_at TIMESTAMP
-);
+);
+
+ALTER TABLE clips ADD CONSTRAINT code_unique UNIQUE (code);
diff --git a/src/main.rs b/src/main.rs
@@ -108,10 +108,10 @@ fn status(_rate_limiter: RateLimiter) -> Result<Json<APIResponse>, Custom<Json<A
         }
     };
 
-    let code = "test".to_string();
     let url = "https://github.com".to_string();
 
-    if let Err(e) = db::insert_clip(&mut db_connection, url, code.clone()) {
+    let insert_result = db::insert_clip(&mut db_connection, url);
+    if let Err(e) = insert_result {
         error!("{}", e);
         let response = APIResponse {
             status: APIStatus::Error,
@@ -120,7 +120,7 @@ fn status(_rate_limiter: RateLimiter) -> Result<Json<APIResponse>, Custom<Json<A
         return Err(Custom(Status::InternalServerError, Json(response)));
     }
 
-    let result = db::get_clip(&mut db_connection, code);
+    let result = db::get_clip(&mut db_connection, insert_result.unwrap().code);
 
     match result {
         Ok(_) => {
@@ -240,13 +240,12 @@ fn set_clip(
         return Err(Custom(Status::InternalServerError, Json(response)));
     }
 
-    let code = gen_id(5);
-    let result = db::insert_clip(&mut db_connection, url.to_string(), code.clone());
+    let result = db::insert_clip(&mut db_connection, url.to_string());
     match result {
         Ok(_) => {
             let response = APIResponse {
                 status: APIStatus::Success,
-                result: code,
+                result: result.unwrap().code,
             };
             Ok(Json(response))
         }

diff --git a/src/utils/db.rs b/src/utils/db.rs
@@ -3,8 +3,13 @@ use crate::schema::*;
 
 use diesel::pg::PgConnection;
 use diesel::prelude::*;
+use diesel::result::DatabaseErrorKind;
+use diesel::result::Error;
 
 use std::env;
+use std::fmt;
+
+use super::id::gen_id;
 
 /// Tries to connect to the database and if it doesn't exist, creates it from the current schema
 /// Returns the connection
@@ -50,26 +55,60 @@ pub fn get_clip_by_url(
         .optional()
 }
 
+#[derive(Debug)]
+pub enum InsertClipError {
+    DieselError(diesel::result::Error),
+    MaxAttemptsExceeded,
+}
+impl From<diesel::result::Error> for InsertClipError {
+    fn from(error: diesel::result::Error) -> Self {
+        InsertClipError::DieselError(error)
+    }
+}
+impl fmt::Display for InsertClipError {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match self {
+            InsertClipError::DieselError(e) => write!(f, "Database error: {}", e),
+            InsertClipError::MaxAttemptsExceeded => write!(f, "Exceeded maximum attempts to generate a unique code."),
+        }
+    }
+}
+
 /// Inserts a clip into the database
 /// Returns the inserted clip
 pub fn insert_clip(
     connection: &mut PgConnection,
     url: String,
-    code: String,
-) -> Result<Clip, diesel::result::Error> {
+) -> Result<Clip, InsertClipError> {
     let expiry_date = chrono::Local::now().naive_local() + chrono::Duration::days(7);
-    let new_clip = NewClip {
-        url,
-        code,
-        created_at: expiry_date,
-        expires_at: None,
-    };
-
-    diesel::insert_into(clips::table)
-        .values(&new_clip)
-        .get_result::<Clip>(connection)
+    let mut attempts = 0;
+    const MAX_ATTEMPTS: usize = 10; // Maximum attempts to generate a unique code
+
+    while attempts < MAX_ATTEMPTS {
+        let code = gen_id(5);
+        let new_clip = NewClip {
+            url: url.clone(),
+            code: code.clone(),
+            created_at: chrono::Local::now().naive_local(),
+            expires_at: Some(expiry_date),
+        };
+
+        match diesel::insert_into(clips::table)
+            .values(&new_clip)
+            .get_result::<Clip>(connection).map_err(InsertClipError::from) {
+                Ok(clip) => return Ok(clip),
+                Err(InsertClipError::DieselError(Error::DatabaseError(DatabaseErrorKind::UniqueViolation, _))) => {
+                    attempts += 1;
+                    continue;
+                },
+                Err(e) => return Err(e), // For any other diesel error, return immediately
+        }
+    }
+
+    Err(InsertClipError::MaxAttemptsExceeded)
 }
 
+
 /// Deletes expired clips from the database
 pub fn collect_garbage(connection: &mut PgConnection) -> Result<usize, diesel::result::Error> {
     use crate::schema::clips::dsl::*;