Coverage for app.py: 95%

1import logging 

2import time 

3from contextlib import asynccontextmanager 

4 

5import psycopg_pool 

6from apscheduler.schedulers.asyncio import AsyncIOScheduler 

7from decouple import config 

8from fastapi import FastAPI, Request, Response 

9from fastapi.middleware.cors import CORSMiddleware 

10from fastapi.middleware.gzip import GZipMiddleware 

11from psycopg import AsyncConnection 

12from psycopg.rows import dict_row 

13from starlette.responses import StreamingResponse 

14 

15from chalicelib.utils import helper 

16from chalicelib.utils import pg_client, ch_client 

17from chalicelib.utils.log import sanitize 

18from crons import core_crons, core_dynamic_crons 

19from routers import core, core_dynamic 

20from routers.subs import health, spot, mcp 

21 

22loglevel = config("LOGLEVEL", default=logging.WARNING) 

23print(f">Loglevel set to: {loglevel}") 

24logging.basicConfig(level=loglevel) 

25logger = logging.getLogger(__name__) 

26 

27 

28class ORPYAsyncConnection(AsyncConnection): 

29 def __init__(self, *args, **kwargs): 

30 super().__init__(*args, row_factory=dict_row, **kwargs) 

31 

32 

33@asynccontextmanager 

34async def lifespan(app: FastAPI): 

35 # Startup 

36 logging.info(">>>>> starting up <<<<<") 

37 ap_logger = logging.getLogger("apscheduler") 

38 ap_logger.setLevel(loglevel) 

39 

40 app.schedule = AsyncIOScheduler() 

41 await pg_client.init() 

42 await ch_client.init() 

43 app.schedule.start() 

44 

45 for job in core_crons.cron_jobs + core_dynamic_crons.cron_jobs: 

46 app.schedule.add_job(id=job["func"].__name__, **job) 

47 

48 ap_logger.info(">Scheduled jobs:") 

49 for job in app.schedule.get_jobs(): 

50 ap_logger.info( 

51 { 

52 "Name": str(job.id), 

53 "Run Frequency": str(job.trigger), 

54 "Next Run": str(job.next_run_time), 

55 } 

56 ) 

57 

58 database = { 

59 "host": config("pg_host", default="localhost"), 

60 "dbname": config("pg_dbname", default="orpy"), 

61 "user": config("pg_user", default="orpy"), 

62 "password": config("pg_password", default="orpy"), 

63 "port": config("pg_port", cast=int, default=5432), 

64 "application_name": "AIO" + config("APP_NAME", default="PY"), 

65 } 

66 

67 database = psycopg_pool.AsyncConnectionPool( 

68 kwargs=database, 

69 connection_class=ORPYAsyncConnection, 

70 min_size=config("PG_AIO_MINCONN", cast=int, default=1), 

71 max_size=config("PG_AIO_MAXCONN", cast=int, default=5), 

72 ) 

73 await database.open() 

74 app.state.postgresql = database 

75 

76 # App listening 

77 yield 

78 

79 # Shutdown 

80 await database.close() 

81 logging.info(">>>>> shutting down <<<<<") 

82 app.schedule.shutdown(wait=True) 

83 await pg_client.terminate() 

84 

85 

86app = FastAPI( 

87 root_path=config("root_path", default="/api"), 

88 docs_url=config("docs_url", default=""), 

89 redoc_url=config("redoc_url", default=""), 

90 lifespan=lifespan, 

91) 

92app.add_middleware(GZipMiddleware, minimum_size=1000) 

93 

94IGNORE_ENDPOINT_LOG = ["/"] 

95 

96 

97@app.middleware("http") 

98async def log_all_requests(request: Request, call_next): 

99 method = sanitize(request.method, max_length=16) 

100 endpoint = sanitize(request.url.path) 

101 response: Response = await call_next(request) 

102 # Log all endpoints except health check 

103 if request.url.path not in IGNORE_ENDPOINT_LOG or response.status_code != 200: 

104 logger.info(f"{method}:{endpoint} {response.status_code}") 

105 return response 

106 

107 

108@app.middleware("http") 

109async def or_middleware(request: Request, call_next): 

110 if helper.TRACK_TIME: 110 ↛ 112line 110 didn't jump to line 112 because the condition on line 110 was always true

111 now = time.time() 

112 try: 

113 response: StreamingResponse = await call_next(request) 

114 except: 

115 logging.error(f"{sanitize(request.method, max_length=16)}: {sanitize(request.url.path)} FAILED!") 

116 raise 

117 if response.status_code // 100 != 2: 

118 logging.warning( 

119 f"{sanitize(request.method, max_length=16)}:{sanitize(request.url.path)} {response.status_code}!") 

120 if helper.TRACK_TIME: 120 ↛ 126line 120 didn't jump to line 126 because the condition on line 120 was always true

121 now = time.time() - now 

122 if now > 2: 

123 now = round(now, 2) 

124 logging.warning( 

125 f"Execution time: {now} s for {sanitize(request.method, max_length=16)}: {sanitize(request.url.path)}") 

126 response.headers["x-robots-tag"] = 'noindex, nofollow' 

127 return response 

128 

129 

130origins = [ 

131 "*", 

132] 

133 

134app.add_middleware( 

135 CORSMiddleware, 

136 allow_origins=origins, 

137 allow_credentials=True, 

138 allow_methods=["*"], 

139 allow_headers=["*"], 

140) 

141app.include_router(core.public_app) 

142app.include_router(core.app) 

143app.include_router(core.app_apikey) 

144app.include_router(core_dynamic.public_app) 

145app.include_router(core_dynamic.app) 

146app.include_router(core_dynamic.app_apikey) 

147app.include_router(health.public_app) 

148app.include_router(health.app) 

149app.include_router(health.app_apikey) 

150 

151app.include_router(spot.public_app) 

152app.include_router(spot.app) 

153app.include_router(spot.app_apikey) 

154 

155app.include_router(mcp.app) 

156app.include_router(mcp.public_app)