Coverage for apps/kwai-api/src/kwai_api/app.py: 81%

1"""Module that implements a factory method for a FastAPI application.""" 

2 

3import os 

4import sys 

5import uuid 

6 

7from contextlib import asynccontextmanager 

8 

9from fastapi import FastAPI, Request, status 

10from fastapi.middleware.cors import CORSMiddleware 

11from fastapi.responses import JSONResponse 

12from kwai_core.settings import LoggerSettings, Settings, get_settings 

13from loguru import logger 

14 

15from kwai_api.frontend.app import create_frontend 

16from kwai_api.v1.auth.api import api_router as auth_api_router 

17from kwai_api.v1.club.api import api_router as club_api_router 

18from kwai_api.v1.news.api import api_router as news_api_router 

19from kwai_api.v1.pages.api import api_router as pages_api_router 

20from kwai_api.v1.portal.api import api_router as portal_api_router 

21from kwai_api.v1.teams.api import router as teams_api_router 

22from kwai_api.v1.trainings.api import api_router as training_api_router 

23 

24 

25APP_NAME = "kwai API" 

26 

27 

28@asynccontextmanager 

29async def lifespan(app: FastAPI): 

30 """Log the start/stop of the application.""" 

31 logger.info(f"{APP_NAME} is starting") 

32 yield 

33 logger.warning(f"{APP_NAME} has ended!") 

34 

35 

36def configure_logger(settings: LoggerSettings): 

37 """Configure the logger.""" 

38 try: 

39 logger.remove(0) # Remove the default logger 

40 except ValueError: 

41 pass 

42 

43 def log_format(record): 

44 """Change the format when a request_id is set in extra.""" 

45 if "request_id" in record["extra"]: 

46 new_format = ( 

47 "{time} - {level} - ({extra[request_id]}) - {message}" + os.linesep 

48 ) 

49 else: 

50 new_format = "{time} - {level} - {message}" + os.linesep 

51 if record["exception"]: 

52 new_format += "{exception}" + os.linesep 

53 return new_format 

54 

55 logger.add( 

56 settings.file or sys.stderr, 

57 format=log_format, 

58 level=settings.level, 

59 colorize=True, 

60 retention=settings.retention, 

61 rotation=settings.rotation, 

62 backtrace=False, 

63 diagnose=False, 

64 ) 

65 

66 

67def create_api(settings: Settings | None = None) -> FastAPI: 

68 """Create the FastAPI application. 

69 

70 Args: 

71 settings: Settings to use in this application. 

72 """ 

73 settings = settings or get_settings() 

74 

75 app = FastAPI( 

76 lifespan=lifespan, 

77 separate_input_output_schemas=False, 

78 openapi_url=settings.openapi_url, 

79 ) 

80 

81 @app.middleware("http") 

82 async def log(request: Request, call_next): 

83 """Middleware for logging the requests.""" 

84 request_id = str(uuid.uuid4()) 

85 with logger.contextualize(request_id=request_id): 

86 logger.info(f"{request.url} - {request.method} - Request started") 

87 

88 response = None # Make pylint happy... 

89 try: 

90 response = await call_next(request) 

91 except Exception as ex: 

92 logger.error(f"{request.url} - Request failed: {ex}") 

93 logger.exception(ex) 

94 response = JSONResponse( 

95 content={ 

96 "detail": str(ex), 

97 }, 

98 status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, 

99 ) 

100 finally: 

101 if response: 

102 response.headers["X-Request-ID"] = request_id 

103 logger.info( 

104 f"{request.url} - {request.method} - Request ended: " 

105 f"{response.status_code}" 

106 ) 

107 

108 return response 

109 

110 # Setup CORS 

111 if settings.cors: 

112 app.add_middleware( 

113 CORSMiddleware, 

114 allow_origins=settings.cors.origins, 

115 allow_credentials=True, 

116 allow_methods=settings.cors.methods, 

117 allow_headers=settings.cors.headers, 

118 ) 

119 

120 # Setup the logger. 

121 if settings.logger: 

122 configure_logger(settings.logger) 

123 

124 app.include_router(auth_api_router, prefix="/v1") 

125 app.include_router(portal_api_router, prefix="/v1") 

126 app.include_router(pages_api_router, prefix="/v1") 

127 app.include_router(news_api_router, prefix="/v1") 

128 app.include_router(teams_api_router, prefix="/v1") 

129 app.include_router(training_api_router, prefix="/v1") 

130 app.include_router(club_api_router, prefix="/v1") 

131 

132 return app 

133 

134 

135def create_app() -> FastAPI: 

136 """Create the FastAPI application for API and frontend.""" 

137 main_app = FastAPI(title=APP_NAME, lifespan=lifespan) 

138 

139 api_app = create_api() 

140 main_app.mount("/api", api_app) 

141 

142 frontend_app = create_frontend() 

143 main_app.mount("/", frontend_app) 

144 

145 return main_app