zxdb/zxdb/setup/docker_manager.py

import logging
import subprocess
import time
from datetime import datetime, timezone
from pathlib import Path

logger = logging.getLogger(__name__)

CONTAINER_NAME = "mysql"


def _run(args: list[str], **kwargs) -> subprocess.CompletedProcess:
    return subprocess.run(args, capture_output=True, text=True, **kwargs)


def ensure_container() -> str:
    """Ensure the MySQL Docker container exists and is running. Returns container name."""
    import sys
    sys.path.insert(0, str(Path(__file__).parent.parent.parent))
    import config

    result = _run(["docker", "ps", "-a", "--filter", f"name=^{CONTAINER_NAME}$", "--format", "{{.Names}}"])
    existing = result.stdout.strip()

    if existing == CONTAINER_NAME:
        # Check if it's running
        running = _run(["docker", "ps", "--filter", f"name=^{CONTAINER_NAME}$", "--format", "{{.Names}}"])
        if running.stdout.strip() != CONTAINER_NAME:
            logger.info("Starting existing container '%s'...", CONTAINER_NAME)
            _run(["docker", "start", CONTAINER_NAME], check=True)
        else:
            logger.info("Container '%s' is already running.", CONTAINER_NAME)
    else:
        logger.info("Creating new MySQL container '%s'...", CONTAINER_NAME)
        _run([
            "docker", "run", "-d",
            "--name", CONTAINER_NAME,
            "-p", "3306:3306",
            "-e", f"MYSQL_ROOT_PASSWORD={config.DB_PASSWORD}",
            "mysql:latest",
        ], check=True)

    return CONTAINER_NAME


def wait_for_mysql(container: str, timeout: int = 60) -> bool:
    """Poll until MySQL is ready or timeout (seconds) is reached. Returns True if ready."""
    import sys
    sys.path.insert(0, str(Path(__file__).parent.parent.parent))
    import config

    logger.info("Waiting for MySQL to be ready (timeout=%ds)...", timeout)
    deadline = time.time() + timeout
    while time.time() < deadline:
        result = _run([
            "docker", "exec", container,
            "mysqladmin", "ping", "-h", "127.0.0.1",
            f"-p{config.DB_PASSWORD}", "--silent",
        ])
        if result.returncode == 0:
            logger.info("MySQL is ready.")
            return True
        time.sleep(2)

    logger.error("Timed out waiting for MySQL.")
    return False


def db_exists(container: str) -> bool:
    """Return True if the 'zxdb' database exists in the container."""
    import sys
    sys.path.insert(0, str(Path(__file__).parent.parent.parent))
    import config

    result = _run([
        "docker", "exec", container,
        "mysql", "-u", "root", f"-p{config.DB_PASSWORD}",
        "-e", "SHOW DATABASES LIKE 'zxdb';",
    ])
    return "zxdb" in result.stdout


def get_db_creation_date(container: str) -> datetime | None:
    """Return the approximate import datetime of the zxdb database, or None."""
    import sys
    sys.path.insert(0, str(Path(__file__).parent.parent.parent))
    import config

    result = _run([
        "docker", "exec", container,
        "mysql", "-u", "root", f"-p{config.DB_PASSWORD}", "--skip-column-names", "-e",
        "SELECT MIN(CREATE_TIME) FROM information_schema.TABLES WHERE TABLE_SCHEMA = 'zxdb' AND CREATE_TIME IS NOT NULL;",
    ])
    value = result.stdout.strip()
    if not value or value == "NULL":
        return None
    # MySQL format: "2024-11-12 17:07:23"
    return datetime.strptime(value, "%Y-%m-%d %H:%M:%S").replace(tzinfo=timezone.utc)


def stop_container(container: str) -> None:
    """Stop the MySQL Docker container."""
    logger.info("Stopping container '%s'...", container)
    _run(["docker", "stop", container])


def import_sql(container: str, sql_path: Path) -> bool:
    """Create the zxdb database if needed and import the SQL file. Returns True on success."""
    import sys
    sys.path.insert(0, str(Path(__file__).parent.parent.parent))
    import config

    # Drop and recreate database to ensure a clean import
    create_result = _run([
        "docker", "exec", container,
        "mysql", "-u", "root", f"-p{config.DB_PASSWORD}",
        "-e", "DROP DATABASE IF EXISTS zxdb; CREATE DATABASE zxdb CHARACTER SET utf8mb4;",
    ])
    if create_result.returncode != 0:
        logger.error("Failed to recreate database: %s", create_result.stderr)
        return False

    logger.info("Importing %s into zxdb (this may take several minutes)...", sql_path)
    start = time.time()
    try:
        with open(sql_path, "rb") as f:
            result = subprocess.run(
                ["docker", "exec", "-i", container,
                 "mysql", "-u", "root", f"-p{config.DB_PASSWORD}", "zxdb"],
                stdin=f,
                check=True,
            )
        elapsed = time.time() - start
        logger.info("Import completed in %.1f seconds.", elapsed)
        return True
    except subprocess.CalledProcessError as e:
        logger.error("Import failed: %s", e)
        return False