Coverage for mcpgateway/utils/db_isready.py: 99%

76 statements  

« prev     ^ index     » next       coverage.py v7.9.2, created at 2025-07-09 11:03 +0100

1#!/usr/bin/env python3 

2# -*- coding: utf-8 -*- 

3"""db_isready - Wait until the configured database is ready 

4========================================================== 

5This helper blocks until the given database (defined by an **SQLAlchemy** URL) 

6successfully answers a trivial round-trip - ``SELECT 1`` - and then returns. 

7It is useful as a container **readiness/health probe** or imported from Python 

8code to delay start-up of services that depend on the DB. 

9 

10Exit codes when executed as a script 

11----------------------------------- 

12* ``0`` - database ready. 

13* ``1`` - all attempts exhausted / timed-out. 

14* ``2`` - :pypi:`SQLAlchemy` is **not** installed. 

15* ``3`` - invalid parameter combination (``max_tries``/``interval``/``timeout``). 

16 

17Features 

18-------- 

19* Accepts **any** SQLAlchemy URL supported by the installed version. 

20* Timing knobs (tries, interval, connect-timeout) configurable through 

21 *environment variables* **or** *CLI flags* - see below. 

22* Works **synchronously** (blocking) or **asynchronously** - simply 

23 ``await wait_for_db_ready()``. 

24* Credentials appearing in log lines are automatically **redacted**. 

25* Depends only on ``sqlalchemy`` (already required by *mcpgateway*). 

26 

27Environment variables 

28--------------------- 

29The script falls back to :pydata:`mcpgateway.config.settings`, but the values 

30below can be overridden via environment variables *or* the corresponding 

31command-line options. 

32 

33+------------------------+----------------------------------------------+-----------+ 

34| Name | Description | Default | 

35+========================+==============================================+===========+ 

36| ``DATABASE_URL`` | SQLAlchemy connection URL | ``sqlite:///./mcp.db`` | 

37| ``DB_WAIT_MAX_TRIES`` | Maximum attempts before giving up | ``30`` | 

38| ``DB_WAIT_INTERVAL`` | Delay between attempts *(seconds)* | ``2`` | 

39| ``DB_CONNECT_TIMEOUT`` | Per-attempt connect timeout *(seconds)* | ``2`` | 

40| ``LOG_LEVEL`` | Log verbosity when not set via ``--log-level`` | ``INFO`` | 

41+------------------------+----------------------------------------------+-----------+ 

42 

43Usage examples 

44-------------- 

45Shell :: 

46 

47 python db_isready.py 

48 python db_isready.py --database-url "postgresql://user:pw@db:5432/mcp" \ 

49 --max-tries 20 --interval 1 --timeout 1 

50 

51Python :: 

52 

53 from db_isready import wait_for_db_ready 

54 

55 await wait_for_db_ready() # asynchronous 

56 wait_for_db_ready(sync=True) # synchronous / blocking 

57""" 

58 

59# Future 

60from __future__ import annotations 

61 

62# Standard 

63# --------------------------------------------------------------------------- 

64# Standard library imports 

65# --------------------------------------------------------------------------- 

66import argparse 

67import asyncio 

68import logging 

69import os 

70import re 

71import sys 

72import time 

73from typing import Any, Dict, Final, Optional 

74 

75# --------------------------------------------------------------------------- 

76# Third-party imports - abort early if SQLAlchemy is missing 

77# --------------------------------------------------------------------------- 

78try: 

79 # Third-Party 

80 from sqlalchemy import create_engine, text 

81 from sqlalchemy.engine import Engine, URL 

82 from sqlalchemy.engine.url import make_url 

83 from sqlalchemy.exc import OperationalError 

84except ImportError: # pragma: no cover - handled at runtime for the CLI 

85 sys.stderr.write("SQLAlchemy not installed - aborting (pip install sqlalchemy)\n") 

86 sys.exit(2) 

87 

88# --------------------------------------------------------------------------- 

89# Optional project settings (silently ignored if mcpgateway package is absent) 

90# --------------------------------------------------------------------------- 

91try: 

92 # First-Party 

93 from mcpgateway.config import settings 

94except Exception: # pragma: no cover - fallback minimal settings 

95 

96 class _Settings: 

97 """Fallback dummy settings when *mcpgateway* is not import-able.""" 

98 

99 database_url: str = "sqlite:///./mcp.db" 

100 log_level: str = "INFO" 

101 

102 settings = _Settings() # type: ignore 

103 

104# --------------------------------------------------------------------------- 

105# Environment variable names 

106# --------------------------------------------------------------------------- 

107ENV_DB_URL: Final[str] = "DATABASE_URL" 

108ENV_MAX_TRIES: Final[str] = "DB_WAIT_MAX_TRIES" 

109ENV_INTERVAL: Final[str] = "DB_WAIT_INTERVAL" 

110ENV_TIMEOUT: Final[str] = "DB_CONNECT_TIMEOUT" 

111 

112# --------------------------------------------------------------------------- 

113# Defaults - overridable via env-vars or CLI flags 

114# --------------------------------------------------------------------------- 

115DEFAULT_DB_URL: Final[str] = os.getenv(ENV_DB_URL, settings.database_url) 

116DEFAULT_MAX_TRIES: Final[int] = int(os.getenv(ENV_MAX_TRIES, "30")) 

117DEFAULT_INTERVAL: Final[float] = float(os.getenv(ENV_INTERVAL, "2")) 

118DEFAULT_TIMEOUT: Final[int] = int(os.getenv(ENV_TIMEOUT, "2")) 

119DEFAULT_LOG_LEVEL: Final[str] = os.getenv("LOG_LEVEL", settings.log_level).upper() 

120 

121# --------------------------------------------------------------------------- 

122# Helpers - sanitising / formatting util functions 

123# --------------------------------------------------------------------------- 

124_CRED_RE: Final[re.Pattern[str]] = re.compile(r"://([^:/?#]+):([^@]+)@") 

125_PWD_RE: Final[re.Pattern[str]] = re.compile(r"(?i)(password|pwd)=([^\s]+)") 

126 

127 

128def _sanitize(txt: str) -> str: 

129 """Hide credentials contained in connection strings or driver errors. 

130 

131 Args: 

132 txt: Arbitrary text that may contain a DB DSN or ``password=...`` 

133 parameter. 

134 

135 Returns: 

136 Same *txt* but with credentials replaced by ``***``. 

137 """ 

138 

139 redacted = _CRED_RE.sub(r"://\\1:***@", txt) 

140 return _PWD_RE.sub(r"\\1=***", redacted) 

141 

142 

143def _format_target(url: URL) -> str: 

144 """Return a concise *host[:port]/db* representation for logging. 

145 

146 Args: 

147 url: A parsed :class:`sqlalchemy.engine.url.URL` instance. 

148 

149 Returns: 

150 Human-readable connection target string suitable for log messages. 

151 """ 

152 

153 if url.get_backend_name() == "sqlite": 

154 return url.database or "<memory>" 

155 

156 host: str = url.host or "localhost" 

157 port: str = f":{url.port}" if url.port else "" 

158 db: str = f"/{url.database}" if url.database else "" 

159 return f"{host}{port}{db}" 

160 

161 

162# --------------------------------------------------------------------------- 

163# Public API - *wait_for_db_ready* 

164# --------------------------------------------------------------------------- 

165 

166 

167def wait_for_db_ready( 

168 *, 

169 database_url: str = DEFAULT_DB_URL, 

170 max_tries: int = DEFAULT_MAX_TRIES, 

171 interval: float = DEFAULT_INTERVAL, 

172 timeout: int = DEFAULT_TIMEOUT, 

173 logger: Optional[logging.Logger] = None, 

174 sync: bool = False, 

175) -> None: 

176 """Block until the database replies to ``SELECT 1``. 

177 

178 The helper can be awaited **asynchronously** *or* called in *blocking* 

179 mode by passing ``sync=True``. 

180 

181 Args: 

182 database_url: SQLAlchemy URL to probe. Falls back to ``$DATABASE_URL`` 

183 or the project default (usually an on-disk SQLite file). 

184 max_tries: Total number of connection attempts before giving up. 

185 interval: Delay *in seconds* between attempts. 

186 timeout: Per-attempt connection timeout in seconds (passed to the DB 

187 driver when supported). 

188 logger: Optional custom :class:`logging.Logger`. If omitted, a default 

189 one named ``"db_isready"`` is lazily configured. 

190 sync: When *True*, run in the **current** thread instead of scheduling 

191 the probe inside an executor. Setting this flag from inside a 

192 running event-loop will block that loop! 

193 

194 Raises: 

195 RuntimeError: If *invalid* parameters are supplied or the database is 

196 still unavailable after the configured number of attempts. 

197 """ 

198 

199 log = logger or logging.getLogger("db_isready") 

200 if not log.handlers: # basicConfig **once** - respects *log.setLevel* later 200 ↛ 207line 200 didn't jump to line 207 because the condition on line 200 was always true

201 logging.basicConfig( 

202 level=getattr(logging, DEFAULT_LOG_LEVEL, logging.INFO), 

203 format="%(asctime)s [%(levelname)s] %(message)s", 

204 datefmt="%Y-%m-%dT%H:%M:%S", 

205 ) 

206 

207 if max_tries < 1 or interval <= 0 or timeout <= 0: 

208 raise RuntimeError("Invalid max_tries / interval / timeout values") 

209 

210 url_obj: URL = make_url(database_url) 

211 backend: str = url_obj.get_backend_name() 

212 target: str = _format_target(url_obj) 

213 

214 log.info(f"Probing {backend} at {target} (timeout={timeout}s, interval={interval}s, max_tries={max_tries})") 

215 

216 connect_args: Dict[str, Any] = {} 

217 if backend.startswith(("postgresql", "mysql")): 

218 # Most drivers honour this parameter - harmless for others. 

219 connect_args["connect_timeout"] = timeout 

220 

221 engine: Engine = create_engine( 

222 database_url, 

223 pool_pre_ping=True, 

224 pool_size=1, 

225 max_overflow=0, 

226 connect_args=connect_args, 

227 ) 

228 

229 def _probe() -> None: # noqa: D401 - internal helper 

230 """Inner synchronous probe running in either the current or a thread. 

231 

232 Returns: 

233 None - the function exits successfully once the DB answers. 

234 

235 Raises: 

236 RuntimeError: Forwarded after exhausting ``max_tries`` attempts. 

237 """ 

238 

239 start = time.perf_counter() 

240 for attempt in range(1, max_tries + 1): 

241 try: 

242 with engine.connect() as conn: 

243 conn.execute(text("SELECT 1")) 

244 elapsed = time.perf_counter() - start 

245 log.info(f"Database ready after {elapsed:.2f}s (attempt {attempt})") 

246 return 

247 except OperationalError as exc: 

248 log.debug(f"Attempt {attempt}/{max_tries} failed ({_sanitize(str(exc))}) - retrying in {interval:.1f}s") 

249 time.sleep(interval) 

250 raise RuntimeError(f"Database not ready after {max_tries} attempts") 

251 

252 if sync: 

253 _probe() 

254 else: 

255 loop = asyncio.get_event_loop() 

256 # Off-load to default executor to avoid blocking the event-loop. 

257 loop.run_until_complete(loop.run_in_executor(None, _probe)) 

258 

259 

260# --------------------------------------------------------------------------- 

261# CLI helpers 

262# --------------------------------------------------------------------------- 

263 

264 

265def _parse_cli() -> argparse.Namespace: 

266 """Parse command-line arguments for the *db_isready* CLI wrapper. 

267 

268 Returns: 

269 Parsed :class:`argparse.Namespace` holding all CLI options. 

270 """ 

271 

272 parser = argparse.ArgumentParser( 

273 description="Wait until the configured database is ready.", 

274 formatter_class=argparse.ArgumentDefaultsHelpFormatter, 

275 ) 

276 parser.add_argument( 

277 "--database-url", 

278 default=DEFAULT_DB_URL, 

279 help="SQLAlchemy URL (env DATABASE_URL)", 

280 ) 

281 parser.add_argument("--max-tries", type=int, default=DEFAULT_MAX_TRIES, help="Maximum connection attempts") 

282 parser.add_argument("--interval", type=float, default=DEFAULT_INTERVAL, help="Delay between attempts in seconds") 

283 parser.add_argument("--timeout", type=int, default=DEFAULT_TIMEOUT, help="Per-attempt connect timeout in seconds") 

284 parser.add_argument("--log-level", default=DEFAULT_LOG_LEVEL, help="Logging level (DEBUG, INFO, ...)") 

285 return parser.parse_args() 

286 

287 

288def main() -> None: # pragma: no cover 

289 """CLI entry-point. 

290 

291 * Parses command-line options. 

292 * Applies ``--log-level`` to the *db_isready* logger **before** the first 

293 message is emitted. 

294 * Delegates the actual probing to :func:`wait_for_db_ready`. 

295 * Exits with: 

296 

297 * ``0`` - database became ready. 

298 * ``1`` - connection attempts exhausted. 

299 * ``2`` - SQLAlchemy missing (handled on import). 

300 * ``3`` - invalid parameter combination. 

301 """ 

302 cli_args = _parse_cli() 

303 

304 log = logging.getLogger("db_isready") 

305 log.setLevel(cli_args.log_level.upper()) 

306 

307 try: 

308 wait_for_db_ready( 

309 database_url=cli_args.database_url, 

310 max_tries=cli_args.max_tries, 

311 interval=cli_args.interval, 

312 timeout=cli_args.timeout, 

313 sync=True, 

314 logger=log, 

315 ) 

316 except RuntimeError as exc: 

317 log.error(f"Database unavailable: {exc}") 

318 sys.exit(1) 

319 

320 sys.exit(0) 

321 

322 

323if __name__ == "__main__": # pragma: no cover 

324 main()