added max attempts to create unique random key

api/crud.py

@@ -6,19 +6,12 @@
 from . import keygen, models, schemas
 
 
-def create_unique_random_key(db: Session) -> str:
-    key = keygen.create_random_key()
-    while get_db_paste_by_key(db, key):
-        key = keygen.create_random_key()
-    return key
-
-
 def get_db_paste_by_key(db: Session, key: str) -> Optional[schemas.PasteInfo]:
     return db.query(models.Paste).filter(models.Paste.key==key).first()
 
 
 def create_db_paste(db: Session, paste: schemas.Paste) -> schemas.PasteInfo:
-    key = create_unique_random_key(db)
+    key = keygen.create_unique_random_key(db)
 
     db_paste = models.Paste(key=key, expiry=paste.expiry)
     db_files = [models.File(text=f.text, kind=f.kind, name=f.name) for f in paste.files]

api/keygen.py

@@ -1,7 +1,20 @@
 import string
 import secrets
 
+from sqlalchemy.orm import Session
+
+from . import crud
+
 
 def create_random_key(length: int = 4) -> str:
     chars = string.ascii_lowercase + string.digits
     return "".join(secrets.choice(chars) for _ in range(length))
+
+
+def create_unique_random_key(db: Session) -> str:
+    ATTEMPTS = 10
+    for _ in range(ATTEMPTS):
+        key = create_random_key()
+        if crud.get_db_paste_by_key(db, key) is None:
+            return key
+    raise RuntimeError(f"Failed to generate unique random key in {ATTEMPTS}attempts.")

api/main.py

@@ -30,7 +30,11 @@ def root():
 
 @app.post("/")
 def create_paste(paste: schemas.Paste, db: Session = Depends(get_db)):
-    db_paste = crud.create_db_paste(db, paste)
+    try:
+        db_paste = crud.create_db_paste(db, paste)
+    except RuntimeError as error:
+        raise HTTPException(status_code=500, detail=str(error))
+
     url =  get_settings().base_url + f"/{db_paste.key}"
     return url