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

1#!/usr/bin/env python3

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

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

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

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

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

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

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

9

10 Exit 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

17 Features

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

27 Environment variables

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

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

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

31 command-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

43 Usage examples

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

45 Shell ::

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

51 Python ::

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

60 from __future__ import annotations

61

62# Standard

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

64# Standard library imports

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

66 import argparse

67 import asyncio

68 import logging

69 import os

70 import re

71 import sys

72 import time

73 from typing import Any, Dict, Final, Optional

74

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

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

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

78 try:

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

84 except 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# ---------------------------------------------------------------------------

91 try:

92 # First-Party

93 from mcpgateway.config import settings

94 except 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# ---------------------------------------------------------------------------

107 ENV_DB_URL: Final[str] = "DATABASE_URL"

108 ENV_MAX_TRIES: Final[str] = "DB_WAIT_MAX_TRIES"

109 ENV_INTERVAL: Final[str] = "DB_WAIT_INTERVAL"

110 ENV_TIMEOUT: Final[str] = "DB_CONNECT_TIMEOUT"

111

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

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

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

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

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

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

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

119 DEFAULT_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

128 def _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

143 def _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

167 def 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

265 def _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

288 def 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

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

324 main()