From fffceac9cee8c5f2c08a6ef4565952ffc9989378 Mon Sep 17 00:00:00 2001
From: Wang Guan <momocraft@gmail.com>
Date: Fri, 29 Nov 2024 23:42:19 +0900
Subject: [PATCH] new cli command: fsck

---
 app/scripts/__init__.py |  0
 app/scripts/fsck.py     | 68 +++++++++++++++++++++++++++++++++++++++++
 manage.py               |  3 ++
 3 files changed, 71 insertions(+)
 create mode 100644 app/scripts/__init__.py
 create mode 100644 app/scripts/fsck.py

diff --git a/app/scripts/__init__.py b/app/scripts/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/app/scripts/fsck.py b/app/scripts/fsck.py
new file mode 100644
index 0000000..58f94aa
--- /dev/null
+++ b/app/scripts/fsck.py
@@ -0,0 +1,68 @@
+from dataclasses import dataclass
+from typing import Literal
+import io
+import click
+from app.models.file import File
+from app.models.team import Team
+from app.models.project import ProjectSet, Project
+import csv
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(frozen=True, kw_only=True)
+class SizeEntry:
+    team_id: str
+    team_name: str
+    project_set_id: str
+    project_set_name: str
+    project_id: str
+    project_name: str
+    num_files: int
+    total_size_kb: int
+
+
+@click.command("fsck")
+@click.option("--export-csv", type=click.File("w"), help="export statistics to csv")
+@click.option(
+    "--sum-by",
+    type=click.Choice(["team", "project_set", "project"]),
+    default="team",
+    help="export statistics to csv",
+)
+def fsck(
+    *, export_csv: io.FileIO | None, sum_by: Literal["team", "project_set", "project"]
+):
+    logger.info("Running fsck")
+    size_entries: list[SizeEntry] = [
+        _sum_by_project(t, ps, p)
+        for t in Team.objects()
+        for ps in ProjectSet.objects(team=t)
+        for p in Project.objects(project_set=ps)
+    ]
+    if export_csv:
+        writer = csv.DictWriter(export_csv, fieldnames=SizeEntry.__dataclass_fields__)
+        writer.writeheader()
+        for entry in size_entries:
+            writer.writerow(entry.__dict__)
+        export_csv.close()
+    else:
+        for entry in size_entries:
+            logger.info(entry)
+
+
+def _sum_by_project(t: Team, ps: ProjectSet, p: Project) -> SizeEntry:
+    files = File.objects(project=p)
+    num_files = len(files)
+    size_kb = sum(f.file_size for f in files)
+    return SizeEntry(
+        team_id=t.id,
+        team_name=t.name,
+        project_set_id=ps.id,
+        project_set_name=ps.name,
+        project_id=p.id,
+        project_name=p.name,
+        total_size_kb=size_kb,
+        num_files=num_files,
+    )
diff --git a/manage.py b/manage.py
index 3ff8a6a..7f3ce85 100644
--- a/manage.py
+++ b/manage.py
@@ -4,6 +4,7 @@
 import logging
 from app import flask_app
 from app.factory import init_db
+from app.scripts.fsck import fsck
 
 logging.basicConfig(
     level=logging.INFO,
@@ -81,6 +82,8 @@ def mit_preprocess_dir(dir: str):
 main.add_command(list_translations)
 main.add_command(mit_preprocess_file)
 main.add_command(mit_preprocess_dir)
+main.add_command(fsck)
+
 
 if __name__ == "__main__":
     main()