dbhub_reload.py

#!/usr/bin/env python3
"""
DBHub Hot-Reload Wrapper (Python Version)
Automatically restarts DBHub when dbhub.toml changes

Usage: python dbhub_reload.py
"""

import os
import sys
import time
import signal
import subprocess
from pathlib import Path
from watchdog.observers import Observer
from watchdog.events import FileSystemEventHandler
from datetime import datetime

# ─────────────────────────────────────────────────────────────
# Configuration — update this path if your project folder is elsewhere
# ─────────────────────────────────────────────────────────────
CONFIG_FILE = Path("C:/Users/LENOVO/Downloads/dbhub_project/dbhub.toml")

# Full path to npx.cmd — needed because Claude Desktop may have a restricted PATH
NPX_CMD = r"C:\Program Files\nodejs\npx.cmd"

# Only used when TRANSPORT_MODE is "http"
HTTP_PORT = 8080

# Auto-detect transport mode:
#   - If stdin is NOT a TTY (e.g. called by Claude Desktop) → use stdio
#   - If stdin IS a TTY (run manually in terminal)          → use http
TRANSPORT_MODE = "stdio" if not sys.stdin.isatty() else "http"

# How many consecutive crashes before we stop auto-restarting
MAX_CRASH_RETRIES = 5

# Minimum seconds a process must run before we reset the crash counter
MIN_UPTIME_SECONDS = 5

# ─────────────────────────────────────────────────────────────

# Global state
dbhub_process = None
is_restarting = False
crash_count = 0
process_start_time = None
shutdown_requested = False   # Set True on Ctrl+C so the loop doesn't restart


def build_command():
    """Build the DBHub command based on transport mode."""
    # Use cmd.exe /C so Windows can resolve the .cmd extension correctly
    # Use full path to npx.cmd to avoid PATH issues when launched by Claude Desktop
    cmd = ["cmd.exe", "/C", NPX_CMD, "-y", "@bytebase/dbhub@latest",
           "--transport", TRANSPORT_MODE,
           "--config", str(CONFIG_FILE)]
    if TRANSPORT_MODE == "http":
        cmd += ["--port", str(HTTP_PORT)]
    return cmd


def log(message, level="INFO"):
    """Log messages with timestamp to stderr."""
    timestamp = datetime.now().strftime("%Y-%m-%dT%H:%M:%S")
    print(f"[{timestamp}] [{level}] {message}", file=sys.stderr, flush=True)


# ─────────────────────────────────────────────────────────────
# Process management
# ─────────────────────────────────────────────────────────────

def start_dbhub():
    """Start the DBHub process."""
    global dbhub_process, process_start_time

    if dbhub_process and dbhub_process.poll() is None:
        log("DBHub is already running", "WARN")
        return

    if not CONFIG_FILE.exists():
        log(f"Config file not found: {CONFIG_FILE}", "ERROR")
        return

    log(f"Starting DBHub (transport={TRANSPORT_MODE})...")
    cmd = build_command()

    try:
        if TRANSPORT_MODE == "stdio":
            # In stdio mode forward stdin/stdout for the MCP client
            dbhub_process = subprocess.Popen(
                cmd,
                stdin=sys.stdin,
                stdout=sys.stdout,
                stderr=sys.stderr,
                bufsize=0
            )
        else:
            # In SSE / HTTP mode keep stdin detached so the process stays alive
            dbhub_process = subprocess.Popen(
                cmd,
                stdin=subprocess.DEVNULL,
                stdout=sys.stdout,
                stderr=sys.stderr,
                bufsize=0
            )

        process_start_time = time.time()
        log(f"DBHub started (PID: {dbhub_process.pid})")

        if TRANSPORT_MODE == "http":
            log(f"DBHub HTTP server listening on http://localhost:{HTTP_PORT}")

    except FileNotFoundError:
        log("'npx' not found. Make sure Node.js / npm is installed and on PATH.", "ERROR")
        dbhub_process = None
    except Exception as e:
        log(f"Failed to start DBHub: {e}", "ERROR")
        dbhub_process = None


def stop_dbhub():
    """Stop the DBHub process gracefully."""
    global dbhub_process

    if not dbhub_process or dbhub_process.poll() is not None:
        return

    log("Stopping DBHub...")
    try:
        dbhub_process.terminate()
        try:
            dbhub_process.wait(timeout=5)
            log("DBHub stopped gracefully")
        except subprocess.TimeoutExpired:
            log("DBHub did not stop in time, force-killing...", "WARN")
            dbhub_process.kill()
            dbhub_process.wait()
            log("DBHub force killed")
    except Exception as e:
        log(f"Error stopping DBHub: {e}", "ERROR")
    finally:
        dbhub_process = None


def restart_dbhub(reason="config change"):
    """Restart the DBHub process."""
    global is_restarting, crash_count

    if is_restarting:
        log("Restart already in progress, skipping", "WARN")
        return

    is_restarting = True
    log(f"Restarting DBHub ({reason})...")

    try:
        stop_dbhub()
        time.sleep(0.5)
        start_dbhub()
        crash_count = 0          # reset crash counter on intentional restart
        log("DBHub restarted successfully")
    except Exception as e:
        log(f"Restart failed: {e}", "ERROR")
    finally:
        is_restarting = False


# ─────────────────────────────────────────────────────────────
# File watcher
# ─────────────────────────────────────────────────────────────

class ConfigFileHandler(FileSystemEventHandler):
    """Watch dbhub.toml for changes and trigger a restart."""

    def __init__(self):
        self._last_trigger = 0
        self._debounce_sec = 1.0   # ignore duplicate events within 1 s

    def on_modified(self, event):
        target = str(CONFIG_FILE.resolve())
        src    = str(Path(event.src_path).resolve())

        if src != target:
            return                 # not our file — ignore

        now = time.time()
        if now - self._last_trigger < self._debounce_sec:
            return                 # debounce

        self._last_trigger = now
        log("dbhub.toml changed — scheduling reload...")
        time.sleep(0.3)           # let the file finish writing
        restart_dbhub(reason="config change")


def watch_config_file():
    """Set up and start the config file watcher."""
    if not CONFIG_FILE.exists():
        log(f"Config file not found: {CONFIG_FILE}", "ERROR")
        log("Please make sure dbhub.toml exists at the path above", "ERROR")
        sys.exit(1)

    log(f"Watching: {CONFIG_FILE}")
    handler  = ConfigFileHandler()
    observer = Observer()
    observer.schedule(handler, str(CONFIG_FILE.parent), recursive=False)
    observer.start()
    return observer


# ─────────────────────────────────────────────────────────────
# Signal handling
# ─────────────────────────────────────────────────────────────

def handle_shutdown(signum, frame):
    """Handle Ctrl+C / SIGTERM gracefully."""
    global shutdown_requested
    shutdown_requested = True
    log("Shutdown requested, stopping...")
    stop_dbhub()
    sys.exit(0)


# ─────────────────────────────────────────────────────────────
# Main
# ─────────────────────────────────────────────────────────────

def main():
    global crash_count, process_start_time

    log("DBHub Hot-Reload Wrapper starting...")
    log(f"Python: {sys.version.split()[0]}")
    log(f"Config: {CONFIG_FILE}")
    log(f"Transport: {TRANSPORT_MODE}")
    log("Databases configured:")
    log("  • my-local-db  → postgres://localhost:5432/temp")
    log("  • practice-db  → postgres://localhost:5432/Practice_db")
    log("  • analytics    → postgres://localhost:5432/analytics")

    # Validate watchdog
    try:
        import watchdog  # noqa: F401
        log("watchdog package OK")
    except ImportError:
        log("ERROR: 'watchdog' not installed. Run: pip install watchdog", "ERROR")
        sys.exit(1)

    # Register signal handlers
    signal.signal(signal.SIGINT,  handle_shutdown)
    signal.signal(signal.SIGTERM, handle_shutdown)

    # Start file watcher
    observer = watch_config_file()

    # Initial DBHub start
    start_dbhub()

    if TRANSPORT_MODE == "stdio":
        log("stdio mode: called by Claude Desktop — proxying MCP over stdin/stdout.")
        log("Wrapper will exit when DBHub exits.")
    else:
        log(f"HTTP mode: DBHub running on http://localhost:{HTTP_PORT}")
        log("Press Ctrl+C to stop. Edit dbhub.toml to trigger a hot-reload.")

    # ── Main monitoring loop ──────────────────────────────────
    try:
        while not shutdown_requested:
            time.sleep(1)

            # In stdio mode: just wait for DBHub to finish — Claude Desktop
            # manages the lifecycle. Do NOT restart.
            if TRANSPORT_MODE == "stdio":
                if dbhub_process and dbhub_process.poll() is not None:
                    exit_code = dbhub_process.returncode
                    log(f"DBHub (stdio) exited with code {exit_code} — wrapper stopping.")
                    break
                continue   # nothing more to do in stdio mode

            # HTTP mode: unexpected crash → limited auto-restart
            if (
                not is_restarting
                and dbhub_process
                and dbhub_process.poll() is not None
            ):
                uptime = time.time() - (process_start_time or 0)
                if uptime > MIN_UPTIME_SECONDS:
                    crash_count = 0   # long-lived run resets counter

                crash_count += 1
                if crash_count > MAX_CRASH_RETRIES:
                    log(
                        f"DBHub crashed {crash_count} times in a row. "
                        "Giving up — fix the issue and restart the wrapper.",
                        "ERROR"
                    )
                    break

                log(
                    f"DBHub exited unexpectedly (attempt {crash_count}/{MAX_CRASH_RETRIES}). "
                    f"Uptime was {uptime:.1f}s. Restarting in 3 s...",
                    "WARN"
                )
                time.sleep(3)
                start_dbhub()

    except KeyboardInterrupt:
        pass
    finally:
        observer.stop()
        observer.join()
        stop_dbhub()
        log("Wrapper exited cleanly.")


if __name__ == "__main__":
    main()