Merge pull request #204 from IBM/update-helm-db-migration

Add init container variant

Author: crivetimihai

docker-compose.yml

@@ -122,6 +122,8 @@ services:
   #   networks: [mcpnet]
 
   migration:
+    #image: ghcr.io/ibm/mcp-context-forge:0.2.0 # Use the release MCP Context Forge image
+    image: mcpgateway/mcpgateway:latest # Use the local latest image. Run `make docker-prod` to build it.
     build:
       context: .
       dockerfile: Containerfile

mcpgateway/utils/db_isready.py

@@ -0,0 +1,323 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""db_isready - Wait until the configured database is ready
+==========================================================
+This helper blocks until the given database (defined by an **SQLAlchemy** URL)
+successfully answers a trivial round-trip - ``SELECT 1`` - and then returns.
+It is useful as a container **readiness/health probe** or imported from Python
+code to delay start-up of services that depend on the DB.
+
+Exit codes when executed as a script
+-----------------------------------
+* ``0`` - database ready.
+* ``1`` - all attempts exhausted / timed-out.
+* ``2`` - :pypi:`SQLAlchemy` is **not** installed.
+* ``3`` - invalid parameter combination (``max_tries``/``interval``/``timeout``).
+
+Features
+--------
+* Accepts **any** SQLAlchemy URL supported by the installed version.
+* Timing knobs (tries, interval, connect-timeout) configurable through
+  *environment variables* **or** *CLI flags* - see below.
+* Works **synchronously** (blocking) or **asynchronously** - simply
+  ``await wait_for_db_ready()``.
+* Credentials appearing in log lines are automatically **redacted**.
+* Depends only on ``sqlalchemy`` (already required by *mcpgateway*).
+
+Environment variables
+---------------------
+The script falls back to :pydata:`mcpgateway.config.settings`, but the values
+below can be overridden via environment variables *or* the corresponding
+command-line options.
+
++------------------------+----------------------------------------------+-----------+
+| Name                   | Description                                  | Default   |
++========================+==============================================+===========+
+| ``DATABASE_URL``       | SQLAlchemy connection URL                    | ``sqlite:///./mcp.db`` |
+| ``DB_WAIT_MAX_TRIES``  | Maximum attempts before giving up            | ``30``    |
+| ``DB_WAIT_INTERVAL``   | Delay between attempts *(seconds)*           | ``2``     |
+| ``DB_CONNECT_TIMEOUT`` | Per-attempt connect timeout *(seconds)*      | ``2``     |
+| ``LOG_LEVEL``          | Log verbosity when not set via ``--log-level`` | ``INFO`` |
++------------------------+----------------------------------------------+-----------+
+
+Usage examples
+--------------
+Shell ::
+
+    python db_isready.py
+    python db_isready.py --database-url "postgresql://user:pw@db:5432/mcp" \
+                         --max-tries 20 --interval 1 --timeout 1
+
+Python ::
+
+    from db_isready import wait_for_db_ready
+
+    await wait_for_db_ready()          # asynchronous
+    wait_for_db_ready(sync=True)       # synchronous / blocking
+"""
+# Future
+from __future__ import annotations
+
+# Standard
+# ---------------------------------------------------------------------------
+# Standard library imports
+# ---------------------------------------------------------------------------
+import argparse
+import asyncio
+import logging
+import os
+import re
+import sys
+import time
+from typing import Any, Dict, Final, Optional
+
+# ---------------------------------------------------------------------------
+# Third-party imports - abort early if SQLAlchemy is missing
+# ---------------------------------------------------------------------------
+try:
+    # Third-Party
+    from sqlalchemy import create_engine, text
+    from sqlalchemy.engine import Engine, URL
+    from sqlalchemy.engine.url import make_url
+    from sqlalchemy.exc import OperationalError
+except ImportError:  # pragma: no cover - handled at runtime for the CLI
+    sys.stderr.write("SQLAlchemy not installed - aborting (pip install sqlalchemy)\n")
+    sys.exit(2)
+
+# ---------------------------------------------------------------------------
+# Optional project settings (silently ignored if mcpgateway package is absent)
+# ---------------------------------------------------------------------------
+try:
+    # First-Party
+    from mcpgateway.config import settings
+except Exception:  # pragma: no cover - fallback minimal settings
+
+    class _Settings:
+        """Fallback dummy settings when *mcpgateway* is not import-able."""
+
+        database_url: str = "sqlite:///./mcp.db"
+        log_level: str = "INFO"
+
+    settings = _Settings()  # type: ignore
+
+# ---------------------------------------------------------------------------
+# Environment variable names
+# ---------------------------------------------------------------------------
+ENV_DB_URL: Final[str] = "DATABASE_URL"
+ENV_MAX_TRIES: Final[str] = "DB_WAIT_MAX_TRIES"
+ENV_INTERVAL: Final[str] = "DB_WAIT_INTERVAL"
+ENV_TIMEOUT: Final[str] = "DB_CONNECT_TIMEOUT"
+
+# ---------------------------------------------------------------------------
+# Defaults - overridable via env-vars or CLI flags
+# ---------------------------------------------------------------------------
+DEFAULT_DB_URL: Final[str] = os.getenv(ENV_DB_URL, settings.database_url)
+DEFAULT_MAX_TRIES: Final[int] = int(os.getenv(ENV_MAX_TRIES, "30"))
+DEFAULT_INTERVAL: Final[float] = float(os.getenv(ENV_INTERVAL, "2"))
+DEFAULT_TIMEOUT: Final[int] = int(os.getenv(ENV_TIMEOUT, "2"))
+DEFAULT_LOG_LEVEL: Final[str] = os.getenv("LOG_LEVEL", settings.log_level).upper()
+
+# ---------------------------------------------------------------------------
+# Helpers - sanitising / formatting util functions
+# ---------------------------------------------------------------------------
+_CRED_RE: Final[re.Pattern[str]] = re.compile(r"://([^:/?#]+):([^@]+)@")
+_PWD_RE: Final[re.Pattern[str]] = re.compile(r"(?i)(password|pwd)=([^\s]+)")
+
+
+def _sanitize(txt: str) -> str:
+    """Hide credentials contained in connection strings or driver errors.
+
+    Args:
+        txt: Arbitrary text that may contain a DB DSN or ``password=…``
+            parameter.
+
+    Returns:
+        Same *txt* but with credentials replaced by ``***``.
+    """
+
+    redacted = _CRED_RE.sub(r"://\\1:***@", txt)
+    return _PWD_RE.sub(r"\\1=***", redacted)
+
+
+def _format_target(url: URL) -> str:
+    """Return a concise *host[:port]/db* representation for logging.
+
+    Args:
+        url: A parsed :class:`sqlalchemy.engine.url.URL` instance.
+
+    Returns:
+        Human-readable connection target string suitable for log messages.
+    """
+
+    if url.get_backend_name() == "sqlite":
+        return url.database or "<memory>"
+
+    host: str = url.host or "localhost"
+    port: str = f":{url.port}" if url.port else ""
+    db: str = f"/{url.database}" if url.database else ""
+    return f"{host}{port}{db}"
+
+
+# ---------------------------------------------------------------------------
+# Public API - *wait_for_db_ready*
+# ---------------------------------------------------------------------------
+
+
+def wait_for_db_ready(
+    *,
+    database_url: str = DEFAULT_DB_URL,
+    max_tries: int = DEFAULT_MAX_TRIES,
+    interval: float = DEFAULT_INTERVAL,
+    timeout: int = DEFAULT_TIMEOUT,
+    logger: Optional[logging.Logger] = None,
+    sync: bool = False,
+) -> None:
+    """Block until the database replies to ``SELECT 1``.
+
+    The helper can be awaited **asynchronously** *or* called in *blocking*
+    mode by passing ``sync=True``.
+
+    Args:
+        database_url: SQLAlchemy URL to probe. Falls back to ``$DATABASE_URL``
+            or the project default (usually an on-disk SQLite file).
+        max_tries: Total number of connection attempts before giving up.
+        interval: Delay *in seconds* between attempts.
+        timeout: Per-attempt connection timeout in seconds (passed to the DB
+            driver when supported).
+        logger: Optional custom :class:`logging.Logger`. If omitted, a default
+            one named ``"db_isready"`` is lazily configured.
+        sync: When *True*, run in the **current** thread instead of scheduling
+            the probe inside an executor. Setting this flag from inside a
+            running event-loop will block that loop!
+
+    Raises:
+        RuntimeError: If *invalid* parameters are supplied or the database is
+            still unavailable after the configured number of attempts.
+    """
+
+    log = logger or logging.getLogger("db_isready")
+    if not log.handlers:  # basicConfig **once** - respects *log.setLevel* later
+        logging.basicConfig(
+            level=getattr(logging, DEFAULT_LOG_LEVEL, logging.INFO),
+            format="%(asctime)s [%(levelname)s] %(message)s",
+            datefmt="%Y-%m-%dT%H:%M:%S",
+        )
+
+    if max_tries < 1 or interval <= 0 or timeout <= 0:
+        raise RuntimeError("Invalid max_tries / interval / timeout values")
+
+    url_obj: URL = make_url(database_url)
+    backend: str = url_obj.get_backend_name()
+    target: str = _format_target(url_obj)
+
+    log.info(f"Probing {backend} at {target} (timeout={timeout}s, interval={interval}s, max_tries={max_tries})")
+
+    connect_args: Dict[str, Any] = {}
+    if backend.startswith(("postgresql", "mysql")):
+        # Most drivers honour this parameter - harmless for others.
+        connect_args["connect_timeout"] = timeout
+
+    engine: Engine = create_engine(
+        database_url,
+        pool_pre_ping=True,
+        pool_size=1,
+        max_overflow=0,
+        connect_args=connect_args,
+    )
+
+    def _probe() -> None:  # noqa: D401 - internal helper
+        """Inner synchronous probe running in either the current or a thread.
+
+        Returns:
+            None - the function exits successfully once the DB answers.
+
+        Raises:
+            RuntimeError: Forwarded after exhausting ``max_tries`` attempts.
+        """
+
+        start = time.perf_counter()
+        for attempt in range(1, max_tries + 1):
+            try:
+                with engine.connect() as conn:
+                    conn.execute(text("SELECT 1"))
+                elapsed = time.perf_counter() - start
+                log.info(f"Database ready after {elapsed:.2f}s (attempt {attempt})")
+                return
+            except OperationalError as exc:
+                log.debug(f"Attempt {attempt}/{max_tries} failed ({_sanitize(str(exc))}) - retrying in {interval:.1f}s")
+            time.sleep(interval)
+        raise RuntimeError(f"Database not ready after {max_tries} attempts")
+
+    if sync:
+        _probe()
+    else:
+        loop = asyncio.get_event_loop()
+        # Off-load to default executor to avoid blocking the event-loop.
+        loop.run_until_complete(loop.run_in_executor(None, _probe))
+
+
+# ---------------------------------------------------------------------------
+# CLI helpers
+# ---------------------------------------------------------------------------
+
+
+def _parse_cli() -> argparse.Namespace:
+    """Parse command-line arguments for the *db_isready* CLI wrapper.
+
+    Returns:
+        Parsed :class:`argparse.Namespace` holding all CLI options.
+    """
+
+    parser = argparse.ArgumentParser(
+        description="Wait until the configured database is ready.",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+    parser.add_argument(
+        "--database-url",
+        default=DEFAULT_DB_URL,
+        help="SQLAlchemy URL (env DATABASE_URL)",
+    )
+    parser.add_argument("--max-tries", type=int, default=DEFAULT_MAX_TRIES, help="Maximum connection attempts")
+    parser.add_argument("--interval", type=float, default=DEFAULT_INTERVAL, help="Delay between attempts in seconds")
+    parser.add_argument("--timeout", type=int, default=DEFAULT_TIMEOUT, help="Per-attempt connect timeout in seconds")
+    parser.add_argument("--log-level", default=DEFAULT_LOG_LEVEL, help="Logging level (DEBUG, INFO, …)")
+    return parser.parse_args()
+
+
+def main() -> None:  # pragma: no cover
+    """CLI entry-point.
+
+    * Parses command-line options.
+    * Applies ``--log-level`` to the *db_isready* logger **before** the first
+      message is emitted.
+    * Delegates the actual probing to :func:`wait_for_db_ready`.
+    * Exits with:
+
+        * ``0`` - database became ready.
+        * ``1`` - connection attempts exhausted.
+        * ``2`` - SQLAlchemy missing (handled on import).
+        * ``3`` - invalid parameter combination.
+    """
+    cli_args = _parse_cli()
+
+    log = logging.getLogger("db_isready")
+    log.setLevel(cli_args.log_level.upper())
+
+    try:
+        wait_for_db_ready(
+            database_url=cli_args.database_url,
+            max_tries=cli_args.max_tries,
+            interval=cli_args.interval,
+            timeout=cli_args.timeout,
+            sync=True,
+            logger=log,
+        )
+    except RuntimeError as exc:
+        log.error(f"Database unavailable: {exc}")
+        sys.exit(1)
+
+    sys.exit(0)
+
+
+if __name__ == "__main__":  # pragma: no cover
+    main()