diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 3bf4dfd..57ea60a 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -1,4 +1,4 @@ -default_stages: [commit] +default_stages: [pre-commit] repos: - repo: https://github.com/asottile/pyupgrade rev: v3.3.1 diff --git a/apps/account/views.py b/apps/account/views.py index f56048d..fade9f2 100644 --- a/apps/account/views.py +++ b/apps/account/views.py @@ -40,6 +40,7 @@ WeChatLoginReqSerializer, ) from core.auth import ApplicationAuthenticate +from core.threadpool import db_executor from core.utils import is_wechat USER_MODEL: User = get_user_model() @@ -81,7 +82,7 @@ async def sign_in(self, request, *args, **kwargs): request_data = request_serializer.validated_data # login - user: User = await database_sync_to_async(auth.authenticate)(request, **request_data) + user: User = await database_sync_to_async(auth.authenticate, executor=db_executor)(request, **request_data) if not user: raise WrongSignInParam() @@ -90,7 +91,7 @@ async def sign_in(self, request, *args, **kwargs): await self.update_user_by_wechat(user, request_data["wechat_code"]) # auth session - await database_sync_to_async(auth.login)(request, user) + await database_sync_to_async(auth.login, executor=db_executor)(request, user) # oauth if request_data["is_oauth"]: @@ -104,7 +105,7 @@ async def sign_out(self, request, *args, **kwargs): Sign out """ - await database_sync_to_async(auth.logout)(request) + await database_sync_to_async(auth.logout, executor=db_executor)(request) return Response() @action(methods=["POST"], detail=False, authentication_classes=[SessionAuthenticate]) @@ -119,7 +120,7 @@ async def sign_up(self, request, *args, **kwargs): request_data = request_serializer.validated_data # save - user = await database_sync_to_async(USER_MODEL.objects.create_user)( + user = await database_sync_to_async(USER_MODEL.objects.create_user, executor=db_executor)( last_login=datetime.datetime.now(), username=request_data["username"], password=request_data["password"], @@ -132,7 +133,7 @@ async def sign_up(self, request, *args, **kwargs): await self.update_user_by_wechat(user, request_data["wechat_code"]) # login session - await database_sync_to_async(auth.login)(request, user) + await database_sync_to_async(auth.login, executor=db_executor)(request, user) # oauth if request_data["is_oauth"]: @@ -183,7 +184,9 @@ async def verify_code(self, request, *args, **kwargs): request_data = request_serializer.validated_data # load user - is_success, user = await database_sync_to_async(USER_MODEL.check_oauth_code)(request_data["code"]) + is_success, user = await database_sync_to_async(USER_MODEL.check_oauth_code, executor=db_executor)( + request_data["code"] + ) if is_success: return Response(await UserInfoSerializer(instance=user).adata) raise WrongToken() @@ -274,10 +277,12 @@ async def wechat_login(self, request, *args, **kwargs): cache.set(cache_key, json.dumps(user_info, ensure_ascii=False), timeout=settings.WECHAT_SCOPE_TIMEOUT) # load user - user: User = await database_sync_to_async(USER_MODEL.load_user_by_union_id)(union_id=user_info["unionid"]) + user: User = await database_sync_to_async(USER_MODEL.load_user_by_union_id, executor=db_executor)( + union_id=user_info["unionid"] + ) if user: await self.update_user_by_wechat(user, code) - await database_sync_to_async(auth.login)(request, user) + await database_sync_to_async(auth.login, executor=db_executor)(request, user) return Response({"code": user.generate_oauth_code() if request_data["is_oauth"] else ""}) # need registry @@ -299,7 +304,9 @@ async def update_user_by_wechat(self, user: User, wechat_code: str) -> None: user.wechat_union_id = user_info["unionid"] user.wechat_open_id = user_info["openid"] user.avatar = user_info["headimgurl"] - await database_sync_to_async(user.save)(update_fields=["wechat_union_id", "wechat_open_id", "avatar"]) + await database_sync_to_async(user.save, executor=db_executor)( + update_fields=["wechat_union_id", "wechat_open_id", "avatar"] + ) @action(methods=["POST"], detail=False, authentication_classes=[SessionAuthenticate]) async def reset_password(self, request, *args, **kwargs) -> Response: @@ -314,14 +321,14 @@ async def reset_password(self, request, *args, **kwargs) -> Response: # load user try: - user: User = await database_sync_to_async(USER_MODEL.objects.get)( + user: User = await database_sync_to_async(USER_MODEL.objects.get, executor=db_executor)( username=request_data["username"], phone_number=request_data["phone_number"] ) except USER_MODEL.DoesNotExist as err: raise UserNotExist() from err # set new password - await database_sync_to_async(user.reset_password)(request_data["password"]) + await database_sync_to_async(user.reset_password, executor=db_executor)(request_data["password"]) return Response() diff --git a/apps/notice/utils/base.py b/apps/notice/utils/base.py index 59b31fc..4b9ad69 100644 --- a/apps/notice/utils/base.py +++ b/apps/notice/utils/base.py @@ -13,6 +13,7 @@ from apps.cel.tasks import send_notice from apps.notice.constants import NoticeWayChoices from apps.notice.models import NoticeLog +from core.threadpool import db_executor USER_MODEL: User = get_user_model() @@ -61,7 +62,7 @@ async def send(self) -> None: msg = traceback.format_exc() logger.error("[%s SendNoticeFailed] Err => %s; Detail => %s", self.__class__.__name__, err, msg) result = {"err": str(err)} - await database_sync_to_async(NoticeLog.objects.create)( + await database_sync_to_async(NoticeLog.objects.create, executor=db_executor)( receivers=self.receivers, content=self.content, extra_params=self.kwargs, result=str(result) ) return result diff --git a/apps/notice/views.py b/apps/notice/views.py index d655328..c55fd2d 100644 --- a/apps/notice/views.py +++ b/apps/notice/views.py @@ -15,6 +15,7 @@ ) from apps.notice.utils import NoticeBase from core.auth import ApplicationAuthenticate +from core.threadpool import db_executor class NoticeViewSet(MainViewSet): @@ -80,7 +81,9 @@ async def registry_robot(self, request, *args, **kwargs): # get instance instance = None if request.data.get("id"): - instance = await database_sync_to_async(get_object_or_404)(Robot, pk=request.data.pop("id")) + instance = await database_sync_to_async(get_object_or_404, executor=db_executor)( + Robot, pk=request.data.pop("id") + ) # validate request request_serializer = RegistryRobotSerializer(instance=instance, data=request.data, partial=bool(instance)) diff --git a/apps/tcloud/views.py b/apps/tcloud/views.py index a10fddd..e71057d 100644 --- a/apps/tcloud/views.py +++ b/apps/tcloud/views.py @@ -10,6 +10,7 @@ from apps.notice.constants import NoticeWayChoices from apps.tcloud.models import AuditCallback from apps.tcloud.serializers import TCICallbackSerializer +from core.threadpool import db_executor class AuditCallbackViewSet(MainViewSet): @@ -27,7 +28,9 @@ async def create(self, request: Request, *args, **kwargs) -> Response: req_slz = TCICallbackSerializer(data=request.data) req_slz.is_valid(raise_exception=True) # save - callback = await database_sync_to_async(AuditCallback.add_callback)(req_slz.validated_data) + callback = await database_sync_to_async(AuditCallback.add_callback, executor=db_executor)( + req_slz.validated_data + ) if callback.is_sensitive: send_notice.delay( notice_type=NoticeWayChoices.ROBOT, diff --git a/core/auth.py b/core/auth.py index 91a8ab2..f619047 100644 --- a/core/auth.py +++ b/core/auth.py @@ -9,6 +9,7 @@ from apps.application.models import Application from core.constants import APP_AUTH_HEADER_KEY from core.exceptions import AppAuthFailed +from core.threadpool import db_executor USER_MODEL = get_user_model() @@ -34,7 +35,7 @@ async def authenticate(self, request) -> (Application, None): raise AppAuthFailed(gettext("App Auth Params Not Exist")) # varify app try: - app = await database_sync_to_async(Application.objects.get)(pk=app_code) + app = await database_sync_to_async(Application.objects.get, executor=db_executor)(pk=app_code) except Application.DoesNotExist as err: # pylint: disable=E1101 raise AppAuthFailed(gettext("App Not Exist")) from err # verify secret diff --git a/core/threadpool.py b/core/threadpool.py new file mode 100644 index 0000000..394d645 --- /dev/null +++ b/core/threadpool.py @@ -0,0 +1,5 @@ +from concurrent.futures.thread import ThreadPoolExecutor + +from django.conf import settings + +db_executor = ThreadPoolExecutor(max_workers=settings.DB_EXECUTOR_SIZE) diff --git a/entry/settings.py b/entry/settings.py index f632bdb..3e39f52 100644 --- a/entry/settings.py +++ b/entry/settings.py @@ -255,3 +255,6 @@ CAPTCHA_APP_ID = int(os.getenv("CAPTCHA_APP_ID", "0")) CAPTCHA_APP_SECRET = os.getenv("CAPTCHA_APP_SECRET", "") CAPTCHA_APP_INFO_TIMEOUT = int(os.getenv("CAPTCHA_APP_INFO_TIMEOUT", str(60 * 10))) + +# Thread +DB_EXECUTOR_SIZE = int(os.getenv("DB_EXECUTOR_SIZE", str(10)))