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
« 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.
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``).
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*).
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.
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+------------------------+----------------------------------------------+-----------+
43Usage examples
44--------------
45Shell ::
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
51Python ::
53 from db_isready import wait_for_db_ready
55 await wait_for_db_ready() # asynchronous
56 wait_for_db_ready(sync=True) # synchronous / blocking
57"""
59# Future
60from __future__ import annotations
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
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)
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
96 class _Settings:
97 """Fallback dummy settings when *mcpgateway* is not import-able."""
99 database_url: str = "sqlite:///./mcp.db"
100 log_level: str = "INFO"
102 settings = _Settings() # type: ignore
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"
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()
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]+)")
128def _sanitize(txt: str) -> str:
129 """Hide credentials contained in connection strings or driver errors.
131 Args:
132 txt: Arbitrary text that may contain a DB DSN or ``password=...``
133 parameter.
135 Returns:
136 Same *txt* but with credentials replaced by ``***``.
137 """
139 redacted = _CRED_RE.sub(r"://\\1:***@", txt)
140 return _PWD_RE.sub(r"\\1=***", redacted)
143def _format_target(url: URL) -> str:
144 """Return a concise *host[:port]/db* representation for logging.
146 Args:
147 url: A parsed :class:`sqlalchemy.engine.url.URL` instance.
149 Returns:
150 Human-readable connection target string suitable for log messages.
151 """
153 if url.get_backend_name() == "sqlite":
154 return url.database or "<memory>"
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}"
162# ---------------------------------------------------------------------------
163# Public API - *wait_for_db_ready*
164# ---------------------------------------------------------------------------
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``.
178 The helper can be awaited **asynchronously** *or* called in *blocking*
179 mode by passing ``sync=True``.
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!
194 Raises:
195 RuntimeError: If *invalid* parameters are supplied or the database is
196 still unavailable after the configured number of attempts.
197 """
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 )
207 if max_tries < 1 or interval <= 0 or timeout <= 0:
208 raise RuntimeError("Invalid max_tries / interval / timeout values")
210 url_obj: URL = make_url(database_url)
211 backend: str = url_obj.get_backend_name()
212 target: str = _format_target(url_obj)
214 log.info(f"Probing {backend} at {target} (timeout={timeout}s, interval={interval}s, max_tries={max_tries})")
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
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 )
229 def _probe() -> None: # noqa: D401 - internal helper
230 """Inner synchronous probe running in either the current or a thread.
232 Returns:
233 None - the function exits successfully once the DB answers.
235 Raises:
236 RuntimeError: Forwarded after exhausting ``max_tries`` attempts.
237 """
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")
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))
260# ---------------------------------------------------------------------------
261# CLI helpers
262# ---------------------------------------------------------------------------
265def _parse_cli() -> argparse.Namespace:
266 """Parse command-line arguments for the *db_isready* CLI wrapper.
268 Returns:
269 Parsed :class:`argparse.Namespace` holding all CLI options.
270 """
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()
288def main() -> None: # pragma: no cover
289 """CLI entry-point.
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:
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()
304 log = logging.getLogger("db_isready")
305 log.setLevel(cli_args.log_level.upper())
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)
320 sys.exit(0)
323if __name__ == "__main__": # pragma: no cover
324 main()