Manage cli add fsck command

app/core/views.py

@@ -1,11 +1,9 @@
-from typing import Optional
-
 from flask import g
 from flask_apikit.views import APIView
 from app.models.user import User
 
 
 class MoeAPIView(APIView):
     @property
-    def current_user(self) -> Optional[User]:
+    def current_user(self) -> User:
         return g.get("current_user")

app/models/user.py

@@ -36,7 +36,7 @@
 from app.models.message import Message
 from app.models.project import Project, ProjectRole, ProjectUserRelation
 from app.models.site_setting import SiteSetting
-from app.models.team import Team, TeamPermission, TeamUserRelation
+from app.models.team import Team, TeamPermission, TeamRole, TeamUserRelation
 from app.regexs import EMAIL_REGEX, USER_NAME_REGEX
 from app.constants.locale import Locale
 from app.utils.hash import md5
@@ -339,7 +339,7 @@ def projects(
         projects = mongo_slice(projects, skip, limit)
         return projects
 
-    def get_project_relation(self, project):
+    def get_project_relation(self, project) -> ProjectUserRelation | None:
         """获取与某个项目的关系"""
         relation = ProjectUserRelation.objects(user=self, group=project).first()
         if relation:
@@ -354,14 +354,18 @@ def join_project(self, project, role=None):
             return ProjectUserRelation(user=self, group=project, role=role).save()
 
     # =====加入流程=====
-    def invitations(self, group=None, status=None, skip=None, limit=None):
+    def invitations(
+        self, group=None, status=None, skip=None, limit=None
+    ) -> list[Invitation]:
         """获取对于个人的所有的邀请"""
         invitations = Invitation.get(
             user=self, group=group, status=status, skip=skip, limit=limit
         )
         return invitations
 
-    def applications(self, group=None, status=None, skip=None, limit=None):
+    def applications(
+        self, group=None, status=None, skip=None, limit=None
+    ) -> list[Application]:
         """获取自己发出的所有的申请"""
         applications = Application.get(
             user=self,
@@ -372,7 +376,13 @@ def applications(self, group=None, status=None, skip=None, limit=None):
         )
         return applications
 
-    def invite(self, user, group, role, message=""):
+    def invite(
+        self,
+        user: "User",
+        group: Team | Project,
+        role: ProjectRole | TeamRole,
+        message="",
+    ):
         """邀请某个用户加入某个团体"""
         # 判断团体是否已满员
         if group.is_full():
@@ -478,7 +488,7 @@ def apply(self, group, message=""):
         return {"message": gettext("申请成功，请等待管理员审核")}
 
     # =====自动鉴别部分=====
-    def get_relation(self, group):
+    def get_relation(self, group: Team | Project):
         """返回与目标的关系，返回关系对象或None，以此判断是否加入"""
         if isinstance(group, Team):
             return self.get_team_relation(group)

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,
+    )

app/utils/mongo.py

@@ -3,7 +3,11 @@
 T = TypeVar("T")
 
 
-def mongo_order(objects: List[T], order_by, default_order_by) -> List[T]:
+def mongo_order(
+    objects: List[T],
+    order_by: None | list[str] | str,
+    default_order_by: str | list[str],
+) -> List[T]:
     """处理排序"""
     # 设置排序默认值
     if order_by is None or order_by == []:

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()