veltro-cli 0.15.7__py3-none-any.whl

veltro_cli/lib/compose.py

@@ -0,0 +1,276 @@
+"""Compose argv derivation shared by the CLI and Makefile."""
+
+from __future__ import annotations
+
+import os
+import shutil
+import subprocess
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Callable, Mapping, Sequence
+
+from veltro_cli.lib.errors import VeltroCliError
+from veltro_cli.lib.paths import enrollment_token_dir
+from veltro_cli.lib.products import PRODUCTS
+from veltro_cli.lib.release import is_release_tag
+from veltro_cli.lib.toml_config import TLS_MODES, VELTRO_TOML, read_config
+
+Runner = Callable[[Sequence[str], Mapping[str, str] | None, Path | None], "RunResult"]
+
+
+@dataclass(frozen=True)
+class RunResult:
+    returncode: int
+    stdout: str = ""
+    stderr: str = ""
+
+
+@dataclass(frozen=True)
+class ComposeInvocation:
+    """Exact process environment + argv that should be executed."""
+
+    argv: tuple[str, ...]
+    env: dict[str, str] = field(default_factory=dict)
+    cwd: str | None = None
+
+    def as_shell(self) -> str:
+        parts: list[str] = []
+        for key in sorted(self.env):
+            parts.append(f"{key}={_shell_quote(self.env[key])}")
+        parts.extend(_shell_quote(a) for a in self.argv)
+        return " ".join(parts)
+
+
+def _shell_quote(value: str) -> str:
+    if not value or re_needs_quote(value):
+        return "'" + value.replace("'", "'\"'\"'") + "'"
+    return value
+
+
+def re_needs_quote(value: str) -> bool:
+    return any(ch in value for ch in " \t\n\"'$&|<>;`\\*?[]{}()#!")
+
+
+def profile_flags(
+    products: tuple[str, ...],
+    *, inspect: bool = False, tls_mode: str = "off", monitoring: bool = False,
+) -> list[str]:
+    unknown = set(products) - set(PRODUCTS)
+    if unknown:
+        raise VeltroCliError(f"unknown products: {', '.join(sorted(unknown))}")
+    flags = ["--profile", "core"]
+    for product in PRODUCTS:
+        if product in products:
+            flags.extend(["--profile", product])
+    if "chad" in products:
+        flags.extend(["--profile", "opensearch"])
+    if inspect and "vectorflow" in products:
+        flags.extend(["--profile", "agent"])
+    if tls_mode and tls_mode != "off":
+        if tls_mode not in TLS_MODES:
+            raise VeltroCliError(
+                f"unknown tls mode {tls_mode!r}",
+                remedy=f"choose one of: {', '.join(TLS_MODES)}",
+            )
+        flags.extend(["--profile", "tls"])
+    if monitoring:
+        flags.extend(["--profile", "monitoring"])
+    return flags
+
+
+def compose_env(
+    products: tuple[str, ...],
+    *, environ: Mapping[str, str] | None = None,
+) -> dict[str, str]:
+    """Environment variables the Makefile prefixes onto the compose process."""
+    base = dict(environ if environ is not None else os.environ)
+    env: dict[str, str] = {"VELTRO_PRODUCTS": ",".join(products)}
+    if "chad" in products:
+        env["OPENSEARCH_SEED_URL"] = base.get("OPENSEARCH_SEED_URL") or "https://opensearch:9200"
+    return env
+
+
+def compose_base_argv(
+    root: Path,
+    *,
+    versions_file: str = "versions.env",
+    env_file: str = ".env",
+) -> list[str]:
+    return [
+        "docker",
+        "compose",
+        "--env-file",
+        str(root / versions_file) if not Path(versions_file).is_absolute() else versions_file,
+        "--env-file",
+        str(root / env_file) if not Path(env_file).is_absolute() else env_file,
+    ]
+
+
+def build_compose_invocation(
+    root: Path,
+    products: tuple[str, ...],
+    compose_args: Sequence[str],
+    *,
+    inspect: bool = False,
+    tls_mode: str = "off",
+    monitoring: bool = False,
+    environ: Mapping[str, str] | None = None,
+    extra_env: Mapping[str, str] | None = None,
+) -> ComposeInvocation:
+    """Build the exact argv + env matching `make <target> PRODUCTS=…`."""
+    argv = compose_base_argv(root)
+    argv.extend(profile_flags(products, inspect=inspect, tls_mode=tls_mode, monitoring=monitoring))
+    argv.extend(compose_args)
+    env = compose_env(products, environ=environ)
+    if extra_env:
+        env.update(dict(extra_env))
+    return ComposeInvocation(argv=tuple(argv), env=env, cwd=str(root))
+
+
+def default_runner(
+    argv: Sequence[str],
+    env: Mapping[str, str] | None = None,
+    cwd: Path | None = None,
+) -> RunResult:
+    merged = dict(os.environ)
+    if env:
+        merged.update(env)
+    try:
+        completed = subprocess.run(
+            list(argv),
+            cwd=str(cwd) if cwd else None,
+            env=merged,
+            capture_output=True,
+            text=True,
+            check=False,
+        )
+    except FileNotFoundError as exc:
+        raise VeltroCliError(
+            f"executable not found: {argv[0]}",
+            remedy="install Docker Desktop / Engine and ensure `docker` is on PATH",
+        ) from exc
+    return RunResult(completed.returncode, completed.stdout or "", completed.stderr or "")
+
+def streaming_runner(
+    argv: Sequence[str],
+    env: Mapping[str, str] | None = None,
+    cwd: Path | None = None,
+) -> RunResult:
+    """Run a long-lived command with the caller's terminal attached."""
+    merged = dict(os.environ)
+    if env:
+        merged.update(env)
+    try:
+        completed = subprocess.run(
+            list(argv),
+            cwd=str(cwd) if cwd else None,
+            env=merged,
+            check=False,
+        )
+    except FileNotFoundError as exc:
+        raise VeltroCliError(
+            f"executable not found: {argv[0]}",
+            remedy="install Docker Desktop / Engine and ensure `docker` is on PATH",
+        ) from exc
+    return RunResult(completed.returncode)
+
+
+class ComposeClient:
+    """Thin wrapper that records or runs compose invocations."""
+
+    def __init__(
+        self,
+        root: Path,
+        *,
+        products: tuple[str, ...] = PRODUCTS,
+        tls_mode: str = "off",
+        monitoring: bool = False,
+        runner: Runner | None = None,
+        stream_runner: Runner | None = None,
+    ) -> None:
+        self.root = root
+        self.products = products
+        self.tls_mode = tls_mode
+        self.monitoring = monitoring
+        self.runner = runner or default_runner
+        self.stream_runner = stream_runner or streaming_runner
+        self.calls: list[ComposeInvocation] = []
+
+    def invoke(
+        self,
+        compose_args: Sequence[str],
+        *,
+        inspect: bool = False,
+        extra_env: Mapping[str, str] | None = None,
+        dry_run: bool = False,
+        inherit_stdio: bool = False,
+    ) -> RunResult:
+        runtime_env = dict(extra_env or {})
+        config = read_config(self.root / VELTRO_TOML)
+        if config and is_release_tag(config.installed_suite_version):
+            runtime_env.setdefault("VELTRO_SUITE_VERSION", config.installed_suite_version)
+        inv = build_compose_invocation(
+            self.root,
+            self.products,
+            compose_args,
+            inspect=inspect,
+            tls_mode=self.tls_mode,
+            monitoring=self.monitoring,
+            extra_env=runtime_env,
+        )
+        self.calls.append(inv)
+        if dry_run:
+            return RunResult(0, inv.as_shell() + "\n", "")
+        merged_env = dict(os.environ)
+        merged_env.update(inv.env)
+        runner = self.stream_runner if inherit_stdio else self.runner
+        return runner(inv.argv, merged_env, Path(inv.cwd) if inv.cwd else None)
+
+    def up(self, *extra: str, wait: bool = True, dry_run: bool = False) -> RunResult:
+        args = ["up", "-d"]
+        if wait:
+            args.append("--wait")
+        args.extend(extra)
+        return self.invoke(args, dry_run=dry_run)
+
+    def down(self, *extra: str, dry_run: bool = False) -> RunResult:
+        return self.invoke(("down", *extra), inspect=True, dry_run=dry_run)
+
+    def ps(self, *extra: str, dry_run: bool = False) -> RunResult:
+        return self.invoke(("ps", *extra), inspect=True, dry_run=dry_run)
+
+    def logs(self, *extra: str, dry_run: bool = False) -> RunResult:
+        args = ["logs", "-f", "--tail=200", *extra]
+        return self.invoke(args, inspect=True, dry_run=dry_run, inherit_stdio=not dry_run)
+
+    def config(self, *extra: str, dry_run: bool = False) -> RunResult:
+        return self.invoke(("config", *extra), dry_run=dry_run)
+
+    def pull(self, *extra: str, dry_run: bool = False) -> RunResult:
+        return self.invoke(("pull", *extra), inspect=True, dry_run=dry_run)
+
+    def restart(self, *extra: str, dry_run: bool = False) -> RunResult:
+        return self.invoke(("restart", *extra), inspect=True, dry_run=dry_run)
+
+    def agent_up(self, *extra: str, dry_run: bool = False) -> RunResult:
+        token_dir = enrollment_token_dir(self.root)
+        if not dry_run:
+            token_dir.mkdir(parents=True, exist_ok=True)
+            os.chmod(token_dir, 0o700)
+        # agent profile only — matches Makefile agent-up (no --wait).
+        argv = compose_base_argv(self.root) + ["--profile", "agent", "up", "-d", *extra]
+        env = compose_env(self.products)
+        env["VF_ENROLLMENT_TOKEN_DIR"] = str(token_dir)
+        inv = ComposeInvocation(argv=tuple(argv), env=env, cwd=str(self.root))
+        self.calls.append(inv)
+        if dry_run:
+            return RunResult(0, inv.as_shell() + "\n", "")
+        merged = dict(os.environ)
+        merged.update(inv.env)
+        return self.runner(inv.argv, merged, self.root)
+
+
+
+
+def which_docker() -> str | None:
+    return shutil.which("docker")