diff --git a/api/crud.py b/api/crud.py index d0ebfae..f6f749d 100644 --- a/api/crud.py +++ b/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] diff --git a/api/keygen.py b/api/keygen.py index b74991b..a27fdef 100644 --- a/api/keygen.py +++ b/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.") \ No newline at end of file diff --git a/api/main.py b/api/main.py index 3b5be52..b16c161 100644 --- a/api/main.py +++ b/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