Coverage for src/kwai/api/app.py: 67%

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 loguru import logger 

13 

14from kwai.api.v1.auth.api import api_router as auth_api_router 

15from kwai.api.v1.club.api import api_router as club_api_router 

16from kwai.api.v1.news.api import api_router as news_api_router 

17from kwai.api.v1.pages.api import api_router as pages_api_router 

18from kwai.api.v1.portal.api import api_router as portal_api_router 

19from kwai.api.v1.teams.api import router as teams_api_router 

20from kwai.api.v1.trainings.api import api_router as training_api_router 

21from kwai.core.settings import LoggerSettings, Settings, get_settings 

22 

23 

24APP_NAME = "kwai API" 

25 

26 

27@asynccontextmanager 

28async def lifespan(app: FastAPI): 

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

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

31 yield 

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

33 

34 

35def configure_logger(settings: LoggerSettings): 

36 """Configure the logger.""" 

37 try: 

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

39 except ValueError: 

40 pass 

41 

42 def log_format(record): 

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

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

45 new_format = ( 

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

47 ) 

48 else: 

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

50 if record["exception"]: 

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

52 return new_format 

53 

54 logger.add( 

55 settings.file or sys.stderr, 

56 format=log_format, 

57 level=settings.level, 

58 colorize=True, 

59 retention=settings.retention, 

60 rotation=settings.rotation, 

61 backtrace=False, 

62 diagnose=False, 

63 ) 

64 

65 

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

67 """Create the FastAPI application. 

68 

69 Args: 

70 settings: Settings to use in this application. 

71 """ 

72 app = FastAPI(lifespan=lifespan, separate_input_output_schemas=False) 

73 

74 if settings is None: 

75 settings = get_settings() 

76 

77 @app.middleware("http") 

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

79 """Middleware for logging the requests.""" 

80 request_id = str(uuid.uuid4()) 

81 with logger.contextualize(request_id=request_id): 

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

83 

84 response = None # Make pylint happy... 

85 try: 

86 response = await call_next(request) 

87 except Exception as ex: 

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

89 logger.exception(ex) 

90 response = JSONResponse( 

91 content={ 

92 "detail": str(ex), 

93 }, 

94 status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, 

95 ) 

96 finally: 

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

98 logger.info( 

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

100 f"{response.status_code}" 

101 ) 

102 

103 return response 

104 

105 # Setup CORS 

106 if settings.cors: 

107 app.add_middleware( 

108 CORSMiddleware, 

109 allow_origins=settings.cors.origins, 

110 allow_credentials=True, 

111 allow_methods=settings.cors.methods, 

112 allow_headers=settings.cors.headers, 

113 ) 

114 

115 # Setup the logger. 

116 if settings.logger: 

117 configure_logger(settings.logger) 

118 

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

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

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

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

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

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

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

126 

127 return app