Fix background task not executed (#86)

* Fix login auth failure exception

* Add background task to custom exception class

* Fix branch conflicts

* Update login log status to enum.

Author: wu-clan

backend/app/common/enums.py

@@ -48,3 +48,10 @@ class MethodType(StrEnum):
     PUT = 'PUT'
     DELETE = 'DELETE'
     PATCH = 'PATCH'
+
+
+class LoginLogStatus(IntEnum):
+    """登陆日志状态"""
+
+    fail = 0
+    success = 1

backend/app/common/exception/errors.py

@@ -3,16 +3,19 @@
 from typing import Any
 
 from fastapi import HTTPException
+from starlette.background import BackgroundTask
 
 from backend.app.common.response.response_code import CustomCode
 
 
 class BaseExceptionMixin(Exception):
     code: int
 
-    def __init__(self, *, msg: str = None, data: Any = None):
+    def __init__(self, *, msg: str = None, data: Any = None, background: BackgroundTask | None = None):
         self.msg = msg
         self.data = data
+        # The original background task: https://www.starlette.io/background/
+        self.background = background
 
 
 class HTTPError(HTTPException):
@@ -21,51 +24,53 @@ def __init__(self, *, code: int, msg: Any = None, headers: dict[str, Any] | None
 
 
 class CustomError(BaseExceptionMixin):
-    def __init__(self, *, error: CustomCode, data: Any = None):
+    def __init__(self, *, error: CustomCode, data: Any = None, background: BackgroundTask | None = None):
         self.code = error.code
-        super().__init__(msg=error.msg, data=data)
+        super().__init__(msg=error.msg, data=data, background=background)
 
 
 class RequestError(BaseExceptionMixin):
     code = 400
 
-    def __init__(self, *, msg: str = 'Bad Request', data: Any = None):
-        super().__init__(msg=msg, data=data)
+    def __init__(self, *, msg: str = 'Bad Request', data: Any = None, background: BackgroundTask | None = None):
+        super().__init__(msg=msg, data=data, background=background)
 
 
 class ForbiddenError(BaseExceptionMixin):
     code = 403
 
-    def __init__(self, *, msg: str = 'Forbidden', data: Any = None):
-        super().__init__(msg=msg, data=data)
+    def __init__(self, *, msg: str = 'Forbidden', data: Any = None, background: BackgroundTask | None = None):
+        super().__init__(msg=msg, data=data, background=background)
 
 
 class NotFoundError(BaseExceptionMixin):
     code = 404
 
-    def __init__(self, *, msg: str = 'Not Found', data: Any = None):
-        super().__init__(msg=msg, data=data)
+    def __init__(self, *, msg: str = 'Not Found', data: Any = None, background: BackgroundTask | None = None):
+        super().__init__(msg=msg, data=data, background=background)
 
 
 class ServerError(BaseExceptionMixin):
     code = 500
 
-    def __init__(self, *, msg: str = 'Internal Server Error', data: Any = None):
-        super().__init__(msg=msg, data=data)
+    def __init__(
+        self, *, msg: str = 'Internal Server Error', data: Any = None, background: BackgroundTask | None = None
+    ):
+        super().__init__(msg=msg, data=data, background=background)
 
 
 class GatewayError(BaseExceptionMixin):
     code = 502
 
-    def __init__(self, *, msg: str = 'Bad Gateway', data: Any = None):
-        super().__init__(msg=msg, data=data)
+    def __init__(self, *, msg: str = 'Bad Gateway', data: Any = None, background: BackgroundTask | None = None):
+        super().__init__(msg=msg, data=data, background=background)
 
 
 class AuthorizationError(BaseExceptionMixin):
     code = 401
 
-    def __init__(self, *, msg: str = 'Permission denied', data: Any = None):
-        super().__init__(msg=msg, data=data)
+    def __init__(self, *, msg: str = 'Permission denied', data: Any = None, background: BackgroundTask | None = None):
+        super().__init__(msg=msg, data=data, background=background)
 
 
 class TokenError(HTTPError):

backend/app/common/exception/exception_handler.py

@@ -104,15 +104,14 @@ def all_exception_handler(request: Request, exc: Exception):
             return JSONResponse(
                 status_code=_get_exception_code(exc.code),
                 content=response_base.fail(code=exc.code, msg=str(exc.msg), data=exc.data if exc.data else None),
+                background=exc.background,
             )
 
         else:
-            import traceback
-
-            log.exception(traceback.format_exc())
+            log.error(exc)
             return JSONResponse(
                 status_code=500,
-                content=response_base.fail(code=500, msg=traceback.format_exc())
+                content=response_base.fail(code=500, msg=str(exc))
                 if settings.ENVIRONMENT == 'dev'
                 else response_base.fail(code=500, msg='Internal Server Error'),
             )

backend/app/services/auth_service.py

@@ -6,9 +6,10 @@
 from fastapi import Request
 from fastapi.security import OAuth2PasswordRequestForm
 from pydantic.datetime_parse import parse_datetime
-from starlette.background import BackgroundTasks
+from starlette.background import BackgroundTasks, BackgroundTask
 
 from backend.app.common import jwt
+from backend.app.common.enums import LoginLogStatus
 from backend.app.common.exception import errors
 from backend.app.common.jwt import get_token
 from backend.app.common.redis import redis_client
@@ -62,19 +63,31 @@ async def login(self, *, request: Request, obj: Auth, background_tasks: Backgrou
                 refresh_token, refresh_token_expire_time = await jwt.create_refresh_token(
                     str(user.id), access_token_expire_time, role_ids=user_role_ids, multi_login=user.is_multi_login
                 )
-                login_logs_params = dict(
-                    db=db, request=request, user=user, login_time=self.login_time, status=1, msg='登录成功'
-                )
             except errors.NotFoundError as e:
                 raise errors.NotFoundError(msg=e.msg)
             except errors.AuthorizationError as e:
-                login_logs_params.update({'status': 0, 'msg': e.msg})
-                background_tasks.add_task(LoginLogService.create, **login_logs_params)
-                raise errors.AuthorizationError(msg=e.msg)
+                err_log_info = dict(
+                    db=db,
+                    request=request,
+                    user=current_user,
+                    login_time=self.login_time,
+                    status=LoginLogStatus.fail,
+                    msg=e.msg,
+                )
+                task = BackgroundTask(LoginLogService.create, **err_log_info)
+                raise errors.AuthorizationError(msg=e.msg, background=task)
             except Exception as e:
                 raise e
             else:
-                background_tasks.add_task(LoginLogService.create, **login_logs_params)
+                log_info = dict(
+                    db=db,
+                    request=request,
+                    user=user,
+                    login_time=self.login_time,
+                    status=LoginLogStatus.success,
+                    msg='登录成功',
+                )
+                background_tasks.add_task(LoginLogService.create, **log_info)
                 return access_token, refresh_token, access_token_expire_time, refresh_token_expire_time, user
 
     @staticmethod

backend/app/services/login_log_service.py

@@ -49,7 +49,7 @@ async def create(
             )
             await LoginLogDao.create(db, obj_in)
         except Exception as e:
-            log.error(f'登录日志创建失败: {e}')
+            log.exception(f'登录日志创建失败: {e}')
 
     @staticmethod
     async def delete(pk: list[int]) -> int: