From 48dc6bd8ae54089bbd5ca9ec62e67a478caec20f Mon Sep 17 00:00:00 2001
From: Mawoka <git@mawoka.eu>
Date: Mon, 29 May 2023 19:21:15 +0200
Subject: [PATCH] :sparkles: Added storage_used to user model

---
 classquiz/db/models.py                        |  1 +
 classquiz/worker/storage.py                   |  9 +++--
 ...01d7a503e360_added_storage_used_to_user.py | 33 +++++++++++++++++++
 3 files changed, 41 insertions(+), 2 deletions(-)
 create mode 100644 migrations/versions/01d7a503e360_added_storage_used_to_user.py

diff --git a/classquiz/db/models.py b/classquiz/db/models.py
index 7f1d523..ace21a3 100644
--- a/classquiz/db/models.py
+++ b/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"
diff --git a/classquiz/worker/storage.py b/classquiz/worker/storage.py
index 2c2b003..2ada401 100644
--- a/classquiz/worker/storage.py
+++ b/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):
diff --git a/migrations/versions/01d7a503e360_added_storage_used_to_user.py b/migrations/versions/01d7a503e360_added_storage_used_to_user.py
new file mode 100644
index 0000000..219de26
--- /dev/null
+++ b/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 ###