stackpilot/backend/services/image_service.py

"""Image listing and pruning for the images router."""
from __future__ import annotations

import logging
from typing import Callable, Optional

from docker_client import DockerError, get_client, get_long_client, safe_call
from services import update_service
from services.prune_job import PruneJob, PrunePlan

logger = logging.getLogger(__name__)

COMPOSE_PROJECT_LABEL = "com.docker.compose.project"


def list_images() -> list[dict]:
    """Return one row per image tag, annotated with using-stacks + update status."""
    try:
        client = get_client()
        images = safe_call(client.images.list)
        containers = safe_call(client.containers.list, all=True)
    except DockerError:
        return []

    # Map image ref -> stacks using it.
    usage: dict[str, set[str]] = {}
    for c in containers:
        ref = c.attrs.get("Config", {}).get("Image")
        stack = c.labels.get(COMPOSE_PROJECT_LABEL)
        if ref:
            usage.setdefault(ref, set())
            if stack:
                usage[ref].add(stack)

    cache = update_service.get_cache()
    result = []
    for img in images:
        tags = img.tags or []
        if not tags:
            continue
        for tag in tags:
            result.append(
                {
                    "id": img.short_id,
                    "tag": tag,
                    "size": img.attrs.get("Size", 0),
                    "created": img.attrs.get("Created"),
                    "stacks": sorted(usage.get(tag, set())),
                    "update": cache.get(tag),
                }
            )
    result.sort(key=lambda r: r["tag"])
    return result


# --------------------------------------------------------------------------- #
# Prune
# --------------------------------------------------------------------------- #

_prune = PruneJob("image")


def prune_status() -> dict:
    return _prune.status()


def _layers_size(client) -> Optional[int]:
    """Total size of all image layers on the host, or None if unavailable.

    Taken before and after the prune, the difference is the space actually
    freed. Summing the sizes of the deleted images would count base layers
    they share with the images that stay. ``type=image`` keeps the daemon from
    also walking every volume (API 1.42+; older daemons ignore it).
    """
    try:
        api = client.api
        response = api._get(api._url("/system/df"), params={"type": "image"})
        return api._result(response, True).get("LayersSize")
    except Exception as exc:  # noqa: BLE001 - the size is a nicety, not the job
        logger.warning("Could not read image disk usage: %s", exc)
        return None


def _candidates(client, all_unused: bool) -> list[dict]:
    """The images ``docker image prune [-a]`` would remove, largest first."""
    in_use = {c.get("ImageID") for c in safe_call(client.api.containers, all=True)}
    filters = None if all_unused else {"dangling": True}
    images = safe_call(client.api.images, filters=filters)
    unused = [img for img in images if img["Id"] not in in_use]
    return sorted(unused, key=lambda img: img.get("Size", 0), reverse=True)


def _label(image: dict) -> str:
    tags = [t for t in image.get("RepoTags") or [] if t != "<none>:<none>"]
    if tags:
        return tags[0] if len(tags) == 1 else f"{tags[0]} (+{len(tags) - 1})"
    return image["Id"].removeprefix("sha256:")[:12]


def _remove(client, image: dict) -> list[dict]:
    """Delete one image the way the prune would, returning Docker's entries.

    Tags are removed one by one rather than with force=True: the last one takes
    the image with it, and a container created since the candidates were listed
    still makes the daemon refuse instead of losing its image.
    """
    tags = [t for t in image.get("RepoTags") or [] if t != "<none>:<none>"]
    entries: list[dict] = []
    for ref in tags or [image["Id"]]:
        entries.extend(safe_call(client.api.remove_image, ref) or [])
    return entries


def _plan(all_unused: bool) -> PrunePlan:
    client = get_long_client()
    measured: dict = {}

    def prepare() -> None:
        measured["before"] = _layers_size(client)

    def remove(image: dict) -> tuple[list, bool]:
        entries = _remove(client, image)
        return entries, any("Deleted" in e for e in entries)

    def finish(entries: list) -> tuple[dict, Optional[int]]:
        # Let the daemon sweep up whatever the per-image pass left, such as
        # untagged images referenced from several repositories.
        filters = {"dangling": False} if all_unused else {"dangling": True}
        sweep = safe_call(client.images.prune, filters=filters)
        entries = entries + (sweep.get("ImagesDeleted") or [])
        before, after = measured.get("before"), _layers_size(client)
        if before is not None and after is not None:
            reclaimed = max(before - after, 0)
        else:
            reclaimed = sweep.get("SpaceReclaimed", 0)
        return {"ImagesDeleted": entries, "SpaceReclaimed": reclaimed}, reclaimed

    return PrunePlan(
        all=all_unused,
        items=_candidates(client, all_unused),
        label=_label,
        remove=remove,
        finish=finish,
        prepare=prepare,
    )


def prune_images(all_unused: bool = False) -> dict:
    """Remove unused images and wait for the result.

    By default only dangling (untagged) images are removed; ``all_unused=True``
    removes every image not referenced by a container (``docker image prune
    -a``). Progress is visible through prune_status() while it runs.
    """
    return _prune.run(lambda: _plan(all_unused))


def start_prune(
    all_unused: bool = False, on_done: Optional[Callable[[dict], None]] = None
) -> dict:
    """Start a prune in the background and return its initial status; the UI
    polls prune_status(). ``on_done`` gets what prune_images() would return."""
    return _prune.start(lambda: _plan(all_unused), on_done)