generate longer file_ids and retry if they is a collison in the database

dbschema/default.gel

@@ -14,7 +14,9 @@ module default {
   }
 
   type files {
-    required file_id: str;
+    required file_id: str {
+      constraint exclusive;
+    };
     required file_name: str;
     required file_size: str;
     required note: str;

my_modules/EdgeDB.py

@@ -1,3 +1,4 @@
+from my_modules.file_helper_functions import generate_short_id
 from my_modules.app.logger import logger
 
 import asyncio, gel
@@ -108,30 +109,36 @@ class EdgeDB:
       "expires_at": i.expires_at if i.expires_at else '',
     } for i in data]
 
-  async def add_file(self, file_id, file_name, file_size, note, content_type, uploaded_at, expires_at, user_id:str):
-    return await self.run_query_with_reconnection(
-      self.client.query,
-      """
-      insert files {
-        file_id := <str>$file_id,
-        file_name := <str>$file_name,
-        file_size := <str>$file_size,
-        note := <str>$note,
-        content_type := <str>$content_type,
-        uploaded_at := <datetime>$uploaded_at,
-        expires_at := <optional datetime>$expires_at,
-        user_id := <str>$user_id
-      };
-      """,
-      file_id=file_id,
-      file_name=file_name,
-      file_size=file_size,
-      note=note,
-      content_type=content_type,
-      uploaded_at=uploaded_at,
-      expires_at=expires_at,
-      user_id=user_id,
-    )
+  async def add_file(self, file_name:str, file_size:str, note:str, content_type:str, uploaded_at, expires_at, user_id:str):
+    for attempt in range(10):
+      try:
+        return await self.run_query_with_reconnection(
+          self.client.query_single,
+          """
+          insert files {
+            file_id := <str>$file_id,
+            file_name := <str>$file_name,
+            file_size := <str>$file_size,
+            note := <str>$note,
+            content_type := <str>$content_type,
+            uploaded_at := <datetime>$uploaded_at,
+            expires_at := <optional datetime>$expires_at,
+            user_id := <str>$user_id
+          };
+          """,
+          file_id=generate_short_id(),
+          file_name=file_name,
+          file_size=file_size,
+          note=note,
+          content_type=content_type,
+          uploaded_at=uploaded_at,
+          expires_at=expires_at,
+          user_id=user_id,
+        )
+      except gel.errors.ConstraintViolationError as e:
+        await logger.warning(f'file_id collision on attempt {attempt+1}, regenerating…')
+        continue
+    raise RuntimeError("Could not allocate unique file_id after multiple retries")
 
   async def update_file(self, file_id:str, file_name:str, note:str, expires_at, user_id:str):
     return await self.run_query_with_reconnection(

my_modules/file_helper_functions.py

@@ -1,6 +1,6 @@
 from my_modules.app.constens import SECRET_KEY
 
-import hmac, hashlib, base64, secrets, time
+import hmac, hashlib, base64, secrets, string, time
 from datetime import datetime, timezone
 
 def base64url_encode(data: bytes) -> str:
@@ -10,9 +10,12 @@ def base64url_decode(data: str) -> bytes:
   padding = '=' * (-len(data) % 4)
   return base64.urlsafe_b64decode(data + padding)
 
-def generate_short_id(length=8):
-  token = base64.urlsafe_b64encode(secrets.token_bytes(length)).decode('utf-8')
-  return token.replace('=', '').replace('-', '').replace('_', '')[:length]
+def generate_short_id(length=11):
+  alphabet = string.ascii_letters + string.digits + '_-'
+  while True:
+    token = ''.join(secrets.choice(alphabet) for _ in range(length))
+    if not token.startswith('-'):
+      return token
 
 def generate_signed_url(file_id: str) -> str:
   # signature based only on the file_id

routes/side/upload.py

@@ -1,4 +1,3 @@
-from my_modules.file_helper_functions import generate_short_id
 from my_modules.decoratory.header import login_required
 
 from quart import Blueprint, request, jsonify, current_app
@@ -134,7 +133,6 @@ async def api_upload(user):
     file_size_pretty = format_size(size_bytes)
 
     await current_app.edgedb.add_file(
-      file_id=generate_short_id(),
       file_name=fname,
       file_size=file_size_pretty,
       note=note,
@@ -157,7 +155,6 @@ async def api_upload(user):
     file_size_pretty = format_size(size_bytes)
 
     await current_app.edgedb.add_file(
-      file_id=generate_short_id(),
       file_name=fname,
       file_size=file_size_pretty,
       note=note,