docker_helper.py

import os
import shutil
import subprocess
from dataclasses import dataclass
from urllib.parse import quote


@dataclass
class ExecResult:
    stdout: str
    stderr: str
    returncode: int

    @property
    def ok(self) -> bool:
        return self.returncode == 0


class DockerHelper:
    def __init__(self, cli: str | None = None):
        if cli is None:
            # Treat an empty CONTAINER_CLI as unset so it falls back to auto-detection.
            cli = os.environ.get("CONTAINER_CLI") or None
        if cli is not None:
            # An explicitly chosen CLI (argument or CONTAINER_CLI) must actually exist on
            # PATH, otherwise every command later fails with an opaque FileNotFoundError.
            if not shutil.which(cli):
                raise RuntimeError(
                    f"Container CLI {cli!r} (from the cli argument or CONTAINER_CLI) was "
                    "not found on PATH. Install it or set CONTAINER_CLI correctly."
                )
        else:
            for candidate in ("docker", "podman"):
                if shutil.which(candidate):
                    cli = candidate
                    break
        if not cli:
            raise RuntimeError(
                "No container CLI found. Install docker or podman, or set CONTAINER_CLI."
            )
        self.cli = cli

    def _run(
        self,
        args: list[str],
        check: bool = True,
        input_text: str | None = None,
        timeout: float | None = None,
    ) -> ExecResult:
        """Run a container CLI command and return its result, raising on failure when check is set.

        A timeout (seconds) bounds how long the call may block; on timeout the process is killed and
        a non-ok result is returned (or raised when check is set), so a hung connection can't stall forever.
        """
        try:
            proc = subprocess.run(
                [self.cli, *args],
                capture_output=True,
                text=True,
                input=input_text,
                timeout=timeout,
            )
        except subprocess.TimeoutExpired as exc:
            result = ExecResult(
                stdout=exc.stdout or "",
                stderr=(exc.stderr or "").strip() or f"timed out after {timeout}s",
                returncode=124,
            )
        else:
            result = ExecResult(stdout=proc.stdout, stderr=proc.stderr, returncode=proc.returncode)
        if check and not result.ok:
            raise RuntimeError(
                f"{self.cli} {' '.join(args)} failed (exit {result.returncode})\n"
                f"stdout: {result.stdout}\nstderr: {result.stderr}"
            )
        return result

    def create(
        self,
        image: str,
        name: str,
        hostname: str | None = None,
        environment: dict[str, str] | None = None,
        volumes: list[str] | None = None,
        network: str | None = None,
        ports: list[str] | None = None,
        entrypoint: str | None = None,
        command: list[str] | None = None,
        detach: bool = True,
        restart: str | None = None,
        platform: str | None = None,
        cap_add: list[str] | None = None,
    ) -> ExecResult:
        """Create and start a long-lived (detached) container with the given config."""
        # These containers are long-lived (no --rm), so a run that crashed before teardown
        # leaves one with this name behind, making `run --name` fail with a conflict.
        # Remove any leftover first so reruns are idempotent (named data volumes survive
        # `rm -f -v`, so this only drops the stale container, not its data).
        self.destroy(name)
        args = ["run"]
        if detach:
            args.append("-d")
        args.extend(["--name", name])
        if platform:
            # Needed for multi-arch manifest-list images with no native variant
            # (e.g. an amd64-only image on Apple Silicon).
            args.extend(["--platform", platform])
        if hostname:
            args.extend(["--hostname", hostname])
        if entrypoint:
            args.extend(["--entrypoint", entrypoint])
        if restart:
            args.extend(["--restart", restart])
        for cap in cap_add or []:
            args.extend(["--cap-add", cap])
        for k, v in (environment or {}).items():
            args.extend(["-e", f"{k}={v}"])
        for vol in volumes or []:
            args.extend(["-v", vol])
        if network:
            args.extend(["--network", network])
        for port in ports or []:
            args.extend(["-p", port])
        args.append(image)
        if command:
            args.extend(command)
        return self._run(args)

    def run(
        self,
        image: str,
        name: str | None = None,
        network: str | None = None,
        entrypoint: str | None = None,
        command: list[str] | None = None,
        volumes: list[str] | None = None,
        user: str | None = None,
        platform: str | None = None,
        remove: bool = True,
        check: bool = True,
    ) -> ExecResult:
        """Run a one-off (by default --rm) container, e.g. an ephemeral helper task."""
        # --rm only cleans up on a clean exit; an interrupted prior run can leave a
        # container with this name behind, making `run --name` fail with a conflict.
        # Remove any leftover first so reruns are idempotent.
        if name:
            self.destroy(name)
        args = ["run"]
        if remove:
            args.append("--rm")
        if name:
            args.extend(["--name", name])
        if platform:
            args.extend(["--platform", platform])
        if user:
            args.extend(["--user", user])
        if entrypoint:
            args.extend(["--entrypoint", entrypoint])
        for vol in volumes or []:
            args.extend(["-v", vol])
        if network:
            args.extend(["--network", network])
        args.append(image)
        if command:
            args.extend(command)
        return self._run(args, check=check)

    def destroy(self, name: str) -> ExecResult:
        """Force-remove a container and its anonymous volumes, ignoring errors if absent.

        The -v flag drops anonymous volumes the image declares (e.g. /var/log/mysql on
        the server image, /etc/haproxy/pxc on the HAProxy image) that we never bind to a
        named volume — otherwise they leak on every run. Named volumes (mounted by name)
        are not affected and are removed explicitly via volume_remove.
        """
        return self._run(["rm", "-f", "-v", name], check=False)

    def start(self, name: str) -> ExecResult:
        """Start an existing stopped container."""
        return self._run(["start", name])

    def stop(self, name: str) -> ExecResult:
        """Stop a running container."""
        return self._run(["stop", name])

    def exec_command(self, name: str, command: str, check: bool = False) -> ExecResult:
        """Run a shell command inside a running container."""
        return self._run(["exec", name, "sh", "-c", command], check=check)

    def exec_mysql(
        self,
        name: str,
        sql: str,
        user: str = "root",
        password: str = "rootpass",
        database: str | None = None,
        host: str | None = None,
        port: int | None = None,
        check: bool = True,
        timeout: float | None = None,
    ) -> ExecResult:
        """Run a SQL statement inside a container using the mysql client.

        With host/port omitted the client uses the container's local socket. Pass
        host/port to connect over TCP instead (e.g. through a proxy). A timeout (seconds)
        bounds the call so a connection that hangs (e.g. a proxy with no live backend) can't block.
        """
        args = ["exec", name, "mysql", f"-u{user}", f"-p{password}", "-N", "-B"]
        if host:
            args.append(f"-h{host}")
        if port:
            # Force TCP so a host such as "localhost" is not silently swapped for the socket.
            args.extend([f"-P{port}", "--protocol=TCP"])
        if database:
            args.extend(["-D", database])
        args.extend(["-e", sql])
        return self._run(args, check=check, timeout=timeout)

    def exec_mysqlsh(
        self,
        name: str,
        script: str,
        user: str = "root",
        password: str = "rootpass",
        host: str = "localhost",
        port: int = 3306,
        language: str = "js",
        check: bool = True,
        timeout: float | None = None,
    ) -> ExecResult:
        """Run a MySQL Shell (mysqlsh) script inside a container against the given URI.

        A timeout (seconds) bounds the call so a stuck mysqlsh (e.g. waiting on a
        connection) can't hang the run until an outer pytest timeout.
        """
        # Percent-encode the credentials so a user/password containing URI-reserved
        # characters (@ : / # ?) doesn't make mysqlsh misparse the connection string.
        uri = f"{quote(user, safe='')}:{quote(password, safe='')}@{host}:{port}"
        args = [
            "exec",
            "-i",
            name,
            "mysqlsh",
            "--no-wizard",
            "--uri",
            uri,
            f"--{language}",
            "-e",
            script,
        ]
        return self._run(args, check=check, timeout=timeout)

    def network_create(self, name: str) -> ExecResult:
        """Create a container network, reusing it if one with the same name already exists."""
        existing = self._run(
            ["network", "ls", "--filter", f"name=^{name}$", "--format", "{{.Name}}"],
            check=False,
        )
        if existing.ok and existing.stdout.strip() == name:
            return existing
        return self._run(["network", "create", name])

    def network_remove(self, name: str) -> ExecResult:
        """Remove a container network, ignoring errors if it does not exist."""
        return self._run(["network", "rm", name], check=False)

    def container_networks(self, name: str) -> list[str]:
        """Return the names of the networks a container is currently attached to.

        Empty means attached to nothing — the normal state of a node that
        network_disconnect() has isolated, not an error. An inspect that fails (no such
        container, no daemon) raises rather than returning [], so callers can act on the
        answer instead of guessing: reporting a failure as "attached to nothing" would make
        network_disconnect() skip the disconnect and report success, leaving a partition
        test reasoning about a partition that never happened.
        """
        result = self._run(
            [
                "inspect",
                "-f",
                '{{range $net, $_ := .NetworkSettings.Networks}}{{$net}}{{"\\n"}}{{end}}',
                name,
            ],
        )
        return [line.strip() for line in result.stdout.splitlines() if line.strip()]

    def network_connect(self, network: str, name: str) -> ExecResult | None:
        """Attach a running container to a network, doing nothing if it is already attached.

        The idempotence matters for reruns and for healing a partition that was only
        partially applied: connecting twice otherwise fails with "already exists in network".
        Returns None when the container was already attached.
        """
        if network in self.container_networks(name):
            return None
        return self._run(["network", "connect", network, name])

    def network_disconnect(self, network: str, name: str, force: bool = False) -> ExecResult | None:
        """Detach a running container from a network, doing nothing if it is not attached.

        Unlike stop(), the process inside the container is untouched — it simply loses
        connectivity — which is what makes this usable for network-partition tests.
        Returns None when the container was not attached in the first place.

        Note: reconnecting later does not restore the container's published host port
        mappings (the -p flags given at create time). Nothing in this suite reaches nodes
        from the host, but a healed node is no longer reachable on its host port.
        """
        if network not in self.container_networks(name):
            return None
        args = ["network", "disconnect"]
        if force:
            args.append("--force")
        args.extend([network, name])
        return self._run(args)

    def volume_remove(self, name: str) -> ExecResult:
        """Remove a container volume, ignoring errors if it does not exist."""
        return self._run(["volume", "rm", name], check=False)

    def container_ip(self, name: str, network: str) -> str:
        """Return a container's IPv4 address on the given network, or "" if it is not attached."""
        template = f'{{{{(index .NetworkSettings.Networks "{network}").IPAddress}}}}'
        result = self._run(["inspect", "-f", template, name], check=False)
        return result.stdout.strip() if result.ok else ""

    def container_state(self, name: str) -> str:
        """Return a container's status and restart count ("running restarts=0"), or "" if unknown.

        Handy in a timeout message: it distinguishes a container that is up but not yet
        serving from one that has died or is stuck in a restart loop.
        """
        result = self._run(
            ["inspect", "-f", "{{.State.Status}} restarts={{.RestartCount}}", name],
            check=False,
        )
        return result.stdout.strip() if result.ok else ""

    def container_exists(self, name: str) -> bool:
        """Return True if a container with the exact given name exists (running or stopped)."""
        result = self._run(
            ["ps", "-a", "--filter", f"name=^{name}$", "--format", "{{.Names}}"],
            check=False,
        )
        return result.ok and result.stdout.strip() == name