✨ Added storage_used to user model

classquiz/db/models.py

@@ -42,6 +42,7 @@ class User(ormar.Model):
     require_password: bool = ormar.Boolean(default=True, nullable=False)
     backup_code: str = ormar.String(max_length=64, min_length=64, nullable=False, default=os.urandom(32).hex())
     totp_secret: str = ormar.String(max_length=32, min_length=32, nullable=True, default=None)
+    storage_used: int = ormar.BigInteger(nullable=False, default=0, minimum=0)
 
     class Meta:
         tablename = "users"

classquiz/worker/storage.py

@@ -10,7 +10,7 @@ import xxhash
 from classquiz.config import redis, storage
 from tempfile import SpooledTemporaryFile
 
-from classquiz.db.models import StorageItem, Quiz
+from classquiz.db.models import StorageItem, Quiz, User
 from classquiz.helpers import extract_image_ids_from_quiz
 from classquiz.storage.errors import DeletionFailedError
 from thumbhash import image_to_thumbhash
@@ -34,7 +34,7 @@ async def clean_editor_images_up(ctx):
 
 async def calculate_hash(ctx, file_id_as_str: str):
     file_id = uuid.UUID(file_id_as_str)
-    file_data: StorageItem = await StorageItem.objects.get(id=file_id)
+    file_data: StorageItem = await StorageItem.objects.select_related(StorageItem.user).get(id=file_id)
     file_path = file_id.hex
     if file_data.storage_path is not None:
         file_path = file_data.storage_path
@@ -61,6 +61,11 @@ async def calculate_hash(ctx, file_id_as_str: str):
     print("Got hash!")
     await file_data.update()
     file.close()
+    user: User | None = await User.objects.get_or_none(id=file_data.user.id)
+    if user is None:
+        return
+    user.storage_used += file_data.size
+    await user.update()
 
 
 async def quiz_update(ctx, old_quiz: Quiz, quiz_id: uuid.UUID):

migrations/versions/01d7a503e360_added_storage_used_to_user.py

@@ -0,0 +1,33 @@
+"""Added storage_used to User
+
+Revision ID: 01d7a503e360
+Revises: 1310317108bd
+Create Date: 2023-05-29 19:05:04.142548
+
+"""
+from alembic import op
+import sqlalchemy as sa
+import ormar
+from sqlalchemy.orm import Session
+
+# revision identifiers, used by Alembic.
+revision = "01d7a503e360"
+down_revision = "1310317108bd"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.add_column("users", sa.Column("storage_used", sa.BigInteger(), nullable=True))
+    conn = op.get_bind()
+    session = Session(bind=conn)
+    session.execute("UPDATE users SET storage_used=0 WHERE storage_used is NULL")
+    op.alter_column("users", "storage_used", nullable=False)
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_column("users", "storage_used")
+    # ### end Alembic commands ###