Only generate new IDs on failure

Author: EvieePy

core/database.py

@@ -129,12 +129,24 @@ async def create_paste(self, *, data: dict[str, Any]) -> PasteModel:
         password: str | None = data["password"]
 
         async with self.pool.acquire() as connection:
-            identifier: str = await utils.generate_id(self)
-            safety: str = utils.generate_safety_token()
+            while True:
+                identifier: str = await utils.generate_id()
+                safety: str = utils.generate_safety_token()
+
+                try:
+                    paster: asyncpg.Record | None = await connection.fetchrow(
+                        paste_query,
+                        identifier,
+                        expiry,
+                        password,
+                        safety,
+                    )
+                except asyncpg.exceptions.UniqueViolationError:
+                    continue
+                else:
+                    break
 
-            paster: asyncpg.Record | None = await connection.fetchrow(paste_query, identifier, expiry, password, safety)
-            if not paster:
-                raise RuntimeError("Unable to create new paste.")
+            assert paster
 
             paste: PasteModel = PasteModel(paster)
             async with connection.transaction():

core/utils.py

@@ -23,26 +23,18 @@
 import json
 import re
 import secrets
-from typing import TYPE_CHECKING, Any
+from typing import Any
 
 import starlette_plus
 
 from core import CONFIG
 
 
-if TYPE_CHECKING:
-    from .database import Database
-
-
 TOKEN_REGEX = re.compile(r"[a-zA-Z0-9_-]{23,28}\.[a-zA-Z0-9_-]{6,7}\.[a-zA-Z0-9_-]{27,}")
 
 
-async def generate_id(database: Database) -> str:
-    while True:
-        identifier: str = secrets.token_hex(8)
-
-        if not await database.fetch_paste(identifier, password=None):
-            return identifier
+async def generate_id() -> str:
+    return secrets.token_hex(8)
 
 
 def generate_safety_token() -> str: