Coverage for fss\starter\server.py: 70%

71 statements  

« prev     ^ index     » next       coverage.py v7.4.4, created at 2024-04-11 19:09 +0800

1"""Server init""" 

2 

3import uvicorn 

4from fastapi import Request 

5from fastapi.exception_handlers import ( 

6 http_exception_handler, 

7 request_validation_exception_handler, 

8) 

9from fastapi.exceptions import RequestValidationError 

10from fastapi.utils import is_body_allowed_for_status_code 

11from fastapi_offline import FastAPIOffline 

12from jose import JWTError 

13from loguru import logger 

14from sqlalchemy import NullPool 

15from starlette.exceptions import HTTPException as StarletteHTTPException 

16from starlette.middleware.cors import CORSMiddleware 

17from starlette.responses import JSONResponse, Response 

18 

19from fss.common import config 

20from fss.common.config import configs 

21from fss.common.exception.exception import ServiceException 

22from fss.common.util.security import is_valid_token 

23from fss.middleware.db_session_middleware import SQLAlchemyMiddleware 

24from fss.starter.system.router.system import system_router 

25 

26# FastAPIOffline setting 

27app = FastAPIOffline( 

28 title=configs.app_name, 

29 openapi_url=f"{configs.api_version}/openapi.json", 

30 description=configs.app_desc, 

31 default_response_model_exclude_unset=True, 

32) 

33 

34# Add project routing 

35app.include_router(system_router, prefix=configs.api_version) 

36 

37# Add SQLAlchemyMiddleware 

38app.add_middleware( 

39 SQLAlchemyMiddleware, 

40 db_url=str(configs.sqlalchemy_database_url), 

41 engine_args={"echo": True, "poolclass": NullPool}, 

42) 

43 

44 

45# Global exception handling 

46@app.exception_handler(ServiceException) 

47async def service_exception_handler(request: Request, exc: ServiceException): 

48 """ 

49 Asynchronous serviceException handler 

50 :param request: The request instance containing all request details 

51 :param exc: ServiceException instance 

52 :return: A Starlette Response object that could be a basic Response or a 

53 JSONResponse, depending on whether a response body is allowed for 

54 the given status code. 

55 """ 

56 logger.error( 

57 f"Exception occurred: {exc} \n" 

58 f"Request path: {request.url.path} \n" 

59 f"Request method: {request.method} \n" 

60 f"Request headers: {request.headers} \n" 

61 f"Request body: {await request.body()}" 

62 ) 

63 headers = getattr(exc, "headers", None) 

64 if not is_body_allowed_for_status_code(exc.status_code): 

65 return Response(status_code=exc.status_code, headers=headers) 

66 return JSONResponse( 

67 {"code": exc.code, "detail": exc.detail}, 

68 status_code=exc.status_code, 

69 headers=headers, 

70 ) 

71 

72 

73@app.exception_handler(StarletteHTTPException) 

74async def custom_http_exception_handler(request: Request, exc: StarletteHTTPException): 

75 """ 

76 Asynchronous handler for StarletteHTTPException 

77 :param request: The request instance containing all request details 

78 :param exc: StarletteHTTPException instance 

79 :return: A Starlette Response object that could be a basic Response or a 

80 JSONResponse, depending on whether a response body is allowed for 

81 the given status code. 

82 """ 

83 logger.error( 

84 f"Exception occurred: {exc} \n" 

85 f"Request path: {request.url.path} \n" 

86 f"Request method: {request.method} \n" 

87 f"Request headers: {request.headers} \n" 

88 f"Request body: {await request.body()}" 

89 ) 

90 return await http_exception_handler(request, exc) 

91 

92 

93@app.exception_handler(RequestValidationError) 

94async def validation_exception_handler(request: Request, exc: RequestValidationError): 

95 """ 

96 Asynchronous handler for RequestValidationError 

97 :param request: The request instance containing all request details 

98 :param exc: StarletteHTTPException instance 

99 :return: A Starlette Response object. 

100 """ 

101 logger.error( 

102 f"Exception occurred: {exc} \n" 

103 f"Request path: {request.url.path} \n" 

104 f"Request method: {request.method} \n" 

105 f"Request headers: {request.headers} \n" 

106 f"Request body: {await request.body()}" 

107 ) 

108 return await request_validation_exception_handler(request, exc) 

109 

110 

111WHITE_LIST_ROUTES = set() 

112for router in configs.white_list_routes.split(","): 

113 WHITE_LIST_ROUTES.add(router.strip()) 

114 

115 

116# Jwt middleware 

117@app.middleware("http") 

118async def jwt_middleware(request: Request, call_next): 

119 raw_url_path = request.url.path 

120 if not raw_url_path.__contains__(configs.api_version) or raw_url_path.__contains__( 

121 ".json" 

122 ): 

123 if configs.enable_swagger: 

124 return await call_next(request) 

125 else: 

126 return JSONResponse( 

127 status_code=403, 

128 content={"detail": "The documentation isn't ready yet."}, 

129 ) 

130 request_url_path = configs.api_version + raw_url_path.split(configs.api_version)[1] 

131 if request_url_path in WHITE_LIST_ROUTES: 

132 return await call_next(request) 

133 

134 auth_header = request.headers.get("Authorization") 

135 if auth_header: 

136 try: 

137 token = auth_header.split(" ")[-1] 

138 await is_valid_token(token) 

139 except JWTError as e: 

140 logger.error(f"{e}") 

141 return JSONResponse( 

142 status_code=401, 

143 content={"detail": "Invalid token or expired token"}, 

144 ) 

145 

146 else: 

147 return JSONResponse( 

148 status_code=401, 

149 content={"detail": "Missing Authentication token"}, 

150 ) 

151 

152 return await call_next(request) 

153 

154 

155# Cors processing 

156if configs.backend_cors_origins: 

157 origins = [] 

158 for origin in configs.backend_cors_origins.split(","): 

159 origins.append(origin.strip()) 

160 app.add_middleware( 

161 CORSMiddleware, 

162 allow_origins=origins, 

163 allow_credentials=True, 

164 allow_methods=["*"], 

165 allow_headers=["*"], 

166 ) 

167 

168 

169# Prepare run 

170def prepare_run(): 

171 config.init_log() 

172 config.init_tz() 

173 return config.complete() 

174 

175 

176# Project run 

177def run() -> None: 

178 port, workers = prepare_run() 

179 uvicorn.run( 

180 app="fss.starter.server:app", host="0.0.0.0", port=port, workers=workers 

181 )