Coverage for kwai/api/app.py: 58%

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

2import os 

3import sys 

4import uuid 

5from contextlib import asynccontextmanager 

6 

7from fastapi import FastAPI, Request, status 

8from fastapi.middleware.cors import CORSMiddleware 

9from fastapi.responses import JSONResponse 

10from loguru import logger 

11 

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

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

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

15from kwai.core.dependencies import container 

16from kwai.core.settings import LoggerSettings, Settings, SettingsException 

17 

18 

19@asynccontextmanager 

20async def lifespan(app: FastAPI): 

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

22 logger.info("kwai is starting") 

23 yield 

24 logger.warning("kwai has ended!") 

25 

26 

27def configure_logger(settings: LoggerSettings): 

28 """Configure the logger.""" 

29 try: 

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

31 except ValueError: 

32 pass 

33 

34 def log_format(record): 

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

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

37 new_format = ( 

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

39 ) 

40 else: 

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

42 if record["exception"]: 

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

44 return new_format 

45 

46 logger.add( 

47 settings.file or sys.stderr, 

48 format=log_format, 

49 level=settings.level, 

50 colorize=True, 

51 retention=settings.retention, 

52 rotation=settings.rotation, 

53 backtrace=False, 

54 diagnose=False, 

55 ) 

56 

57 

58def create_app(settings: Settings | None = None) -> FastAPI: 

59 """Create the FastAPI application. 

60 

61 Args: 

62 settings: Settings to use in this application. 

63 

64 When settings is None (the default), the dependency container will 

65 load the settings. 

66 """ 

67 if settings is None: 

68 try: 

69 settings = container[Settings] 

70 except SettingsException as ex: 

71 logger.error(f"Could not load settings: {ex}") 

72 sys.exit(0) 

73 

74 app = FastAPI(title="kwai", lifespan=lifespan) 

75 

76 # Setup CORS 

77 if settings.cors: 

78 app.add_middleware( 

79 CORSMiddleware, 

80 allow_origins=settings.cors.origins, 

81 allow_credentials=True, 

82 allow_methods=settings.cors.methods, 

83 allow_headers=settings.cors.headers, 

84 ) 

85 

86 # Setup the logger. 

87 if settings.logger: 

88 configure_logger(settings.logger) 

89 

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

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

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

93 

94 @app.middleware("http") 

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

96 """Middleware for logging the requests.""" 

97 request_id = str(uuid.uuid4()) 

98 with logger.contextualize(request_id=request_id): 

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

100 

101 response = None # Make pylint happy... 

102 try: 

103 response = await call_next(request) 

104 except Exception as ex: 

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

106 logger.exception(ex) 

107 response = JSONResponse( 

108 content={ 

109 "detail": str(ex), 

110 }, 

111 status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, 

112 ) 

113 finally: 

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

115 logger.info( 

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

117 f"{response.status_code}" 

118 ) 

119 

120 return response 

121 

122 return app