Fix bug

fastapi_admin/cli.py

@@ -6,8 +6,8 @@ from prompt_toolkit import PromptSession
 from tortoise import Tortoise, run_async
 
 from fastapi_admin import enums, version
-from fastapi_admin.common import import_obj, pwd_context
-from fastapi_admin.models import Permission
+from fastapi_admin.common import get_all_models, import_obj, pwd_context
+from fastapi_admin.models import AbstractPermission
 
 init(autoreset=True)
 
@@ -33,17 +33,23 @@ async def init_tortoise(args):
 async def register_permissions(args):
     await init_tortoise(args)
     await Tortoise.generate_schemas()
-
+    permission_model = None
+    for model_name, model in get_all_models():
+        if issubclass(model, AbstractPermission):
+            permission_model = model
+            break
+    if not permission_model:
+        Logger.error("No Permission Model Founded.")
+        return
     if args.clean:
-        await Permission.all().delete()
+        await permission_model.all().delete()
         Logger.waring("Cleaned all permissions success.")
-    models = Tortoise.apps.get("models").keys()
-    models = list(models)
-    for model in models:
+
+    for model, _ in get_all_models():
         for action in enums.PermissionAction:
             label = f"{enums.PermissionAction.choices().get(action)} {model}"
             defaults = dict(label=label, model=model, action=action,)
-            _, created = await Permission.get_or_create(**defaults,)
+            _, created = await permission_model.get_or_create(**defaults,)
             if created:
                 Logger.success(f"Create permission {label} success.")
 

fastapi_admin/common.py

@@ -2,6 +2,7 @@ import importlib
 from copy import deepcopy
 
 from passlib.context import CryptContext
+from tortoise import Tortoise
 
 pwd_context = CryptContext(schemes=["bcrypt"], deprecated="auto")
 
@@ -57,3 +58,13 @@ def import_obj(path: str):
     module = ".".join(splits[:-1])
     class_name = splits[-1]
     return getattr(importlib.import_module(module), class_name)
+
+
+def get_all_models():
+    """
+    get all tortoise models
+    :return:
+    """
+    for tortoise_app, models in Tortoise.apps.items():
+        for model_item in models.items():
+            yield model_item

fastapi_admin/factory.py

@@ -3,9 +3,9 @@ from typing import Dict, List, Optional, Type
 import jwt
 from fastapi import FastAPI, HTTPException
 from starlette.status import HTTP_403_FORBIDDEN
-from tortoise import Model, Tortoise
+from tortoise import Model
 
-from .common import import_obj, pwd_context
+from .common import get_all_models, import_obj, pwd_context
 from .exceptions import exception_handler
 from .models import AbstractPermission, AbstractRole, AbstractUser
 from .schemas import LoginIn
@@ -35,6 +35,8 @@ class AdminApp(FastAPI):
     models: Dict[str, Type[Model]] = {}
     admin_secret: str
     user_model: Type[Model]
+    permission_model: Type[Model]
+    role_model: Type[Model]
     site: Site
     permission: bool
     _inited: bool = False
@@ -76,7 +78,7 @@ class AdminApp(FastAPI):
 
     def _build_content_menus(self) -> List[Menu]:
         menus = []
-        for model_name, model in self._get_all_models():
+        for model_name, model in get_all_models():
             if issubclass(model, (AbstractUser, AbstractPermission, AbstractRole)):
                 continue
             menu = Menu(
@@ -145,7 +147,7 @@ class AdminApp(FastAPI):
         self.site = site
         self.permission = permission
         self.admin_secret = admin_secret
-        for model_name, model in self._get_all_models():
+        for model_name, model in get_all_models():
             if issubclass(model, AbstractUser):
                 self.user_model = model
             self.models[model_name] = model
@@ -158,11 +160,6 @@ class AdminApp(FastAPI):
         else:
             self.add_api_route("/login", login, methods=["POST"])
 
-    def _get_all_models(self):
-        for tortoise_app, models in Tortoise.apps.items():
-            for model_item in models.items():
-                yield model_item
-
     def _exclude_field(self, resource: str, field: str):
         """
         exclude field by menu include and exclude