stackpilot/backend/services/volume_service.py

"""Docker volume management + Compose volume YAML generation."""
from __future__ import annotations

from typing import Literal, Optional

import yaml

from docker_client import get_client, safe_call

COMPOSE_PROJECT_LABEL = "com.docker.compose.project"


# --------------------------------------------------------------------------- #
# Listing / pruning
# --------------------------------------------------------------------------- #


def list_volumes() -> list[dict]:
    client = get_client()
    volumes = safe_call(client.volumes.list)
    # Map volume name -> set of containers using it.
    containers = safe_call(client.containers.list, all=True)
    usage: dict[str, list[str]] = {}
    for c in containers:
        for mount in c.attrs.get("Mounts", []) or []:
            if mount.get("Type") == "volume" and mount.get("Name"):
                usage.setdefault(mount["Name"], []).append(c.name)

    result = []
    for v in volumes:
        attrs = v.attrs
        name = v.name
        result.append(
            {
                "name": name,
                "driver": attrs.get("Driver"),
                "mountpoint": attrs.get("Mountpoint"),
                "created_at": attrs.get("CreatedAt"),
                "labels": attrs.get("Labels") or {},
                "scope": attrs.get("Scope"),
                "stack": (attrs.get("Labels") or {}).get(COMPOSE_PROJECT_LABEL),
                "used_by": usage.get(name, []),
                "in_use": bool(usage.get(name)),
            }
        )
    return result


def orphaned_volumes() -> list[dict]:
    return [v for v in list_volumes() if not v["in_use"]]


def remove_volume(name: str, force: bool = False) -> None:
    client = get_client()
    volume = safe_call(client.volumes.get, name)
    safe_call(volume.remove, force=force)


def prune_volumes() -> dict:
    client = get_client()
    return safe_call(client.volumes.prune)


# --------------------------------------------------------------------------- #
# YAML generation
# --------------------------------------------------------------------------- #

VolumeType = Literal["bind", "named", "nfs", "smb", "tmpfs"]


def _nfs_options(server: str, opts: dict) -> str:
    """Build the `o:` option string for an NFS driver_opts block."""
    parts = [f"addr={server}"]
    # rw / ro
    if opts.get("rw", True):
        parts.append("rw")
    else:
        parts.append("ro")
    for flag in ("soft", "hard", "nolock", "noatime", "noacl", "nocto", "bg"):
        if opts.get(flag):
            parts.append(flag)
    version = opts.get("nfsvers") or opts.get("version")
    if version:
        parts.append(f"nfsvers={version}")
    if opts.get("timeo") is not None:
        parts.append(f"timeo={opts['timeo']}")
    if opts.get("retrans") is not None:
        parts.append(f"retrans={opts['retrans']}")
    extra = (opts.get("extra") or "").strip()
    if extra:
        parts.append(extra.lstrip(","))
    return ",".join(parts)


def _smb_options(opts: dict) -> str:
    parts = []
    if opts.get("username"):
        parts.append(f"username={opts['username']}")
    if opts.get("password") is not None:
        parts.append(f"password={opts['password']}")
    if opts.get("uid") is not None:
        parts.append(f"uid={opts['uid']}")
    if opts.get("gid") is not None:
        parts.append(f"gid={opts['gid']}")
    version = opts.get("vers")
    if version:
        parts.append(f"vers={version}")
    if opts.get("noperm"):
        parts.append("noperm")
    if opts.get("file_mode"):
        parts.append(f"file_mode={opts['file_mode']}")
    if opts.get("dir_mode"):
        parts.append(f"dir_mode={opts['dir_mode']}")
    extra = (opts.get("extra") or "").strip()
    if extra:
        parts.append(extra.lstrip(","))
    return ",".join(parts)


def build_volume_definition(spec: dict) -> dict:
    """Return the top-level `volumes:` declaration for a named/nfs/smb volume.

    Returns {} for bind/tmpfs (those live entirely on the service).
    """
    vtype: VolumeType = spec["type"]
    name = spec.get("volume_name") or spec.get("name")

    if vtype == "named":
        decl: dict = {"driver": spec.get("driver", "local")}
        if spec.get("driver_opts"):
            decl["driver_opts"] = spec["driver_opts"]
        if not decl.get("driver_opts") and decl["driver"] == "local":
            # A plain named volume can be declared as null.
            return {name: None}
        return {name: decl}

    if vtype == "nfs":
        o = _nfs_options(spec["nfs_server"], spec.get("options", {}))
        device = spec["nfs_path"]
        if not device.startswith(":"):
            device = ":" + device
        return {
            name: {
                "driver": "local",
                "driver_opts": {"type": "nfs", "o": o, "device": device},
            }
        }

    if vtype == "smb":
        o = _smb_options(spec.get("options", {}))
        device = spec["smb_share"]  # //server/share
        return {
            name: {
                "driver": "local",
                "driver_opts": {"type": "cifs", "o": o, "device": device},
            }
        }

    return {}


def build_service_mount(spec: dict) -> str | dict:
    """Return the entry to add to a service's `volumes:` list."""
    vtype: VolumeType = spec["type"]
    if vtype == "bind":
        mode = "ro" if not spec.get("options", {}).get("rw", True) else "rw"
        return f"{spec['host_path']}:{spec['container_path']}:{mode}"
    if vtype == "tmpfs":
        return {
            "type": "tmpfs",
            "target": spec["container_path"],
            "tmpfs": {
                k: v
                for k, v in {
                    "size": spec.get("options", {}).get("size"),
                    "mode": spec.get("options", {}).get("mode"),
                }.items()
                if v is not None
            },
        }
    # named / nfs / smb
    name = spec.get("volume_name") or spec.get("name")
    return f"{name}:{spec['container_path']}"


def generate_yaml(spec: dict) -> str:
    """Produce a ready-to-insert YAML fragment for the given volume spec.

    Includes the top-level `volumes:` block (if any) and a `services:` example
    showing the mount, mirroring the wizard preview.
    """
    doc: dict = {}
    definition = build_volume_definition(spec)
    if definition:
        doc["volumes"] = definition
    service_name = spec.get("service", "service")
    doc["services"] = {service_name: {"volumes": [build_service_mount(spec)]}}
    return yaml.safe_dump(doc, sort_keys=False, default_flow_style=False)