diff --git a/dbschema/default.gel b/dbschema/default.gel index 0606474..078df8b 100644 --- a/dbschema/default.gel +++ b/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; diff --git a/my_modules/EdgeDB.py b/my_modules/EdgeDB.py index 0a6aaac..a6b6a66 100644 --- a/my_modules/EdgeDB.py +++ b/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 := $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 - }; - """, - 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 := $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 + }; + """, + 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( diff --git a/my_modules/file_helper_functions.py b/my_modules/file_helper_functions.py index a8fb5d2..e10e909 100644 --- a/my_modules/file_helper_functions.py +++ b/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 diff --git a/routes/side/upload.py b/routes/side/upload.py index 0ac4df5..d74429e 100644 --- a/routes/side/upload.py +++ b/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,