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,9 +109,11 @@ 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):
+  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,
+          self.client.query_single,
           """
           insert files {
             file_id := <str>$file_id,
@@ -123,7 +126,7 @@ class EdgeDB:
             user_id := <str>$user_id
           };
           """,
-      file_id=file_id,
+          file_id=generate_short_id(),
           file_name=file_name,
           file_size=file_size,
           note=note,
@@ -132,6 +135,10 @@ class EdgeDB:
           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,