feat: Pull images before teardown containers on up command

Signed-off-by: Songmin Li <lisongmin@protonmail.com>

Author: lisongmin

newsfragments/pull-image-before-teardown.feature

@@ -0,0 +1 @@
+Pull images before teardown containers on compose up, which reduce the downtime of services.

podman_compose.py

@@ -27,9 +27,11 @@
 import tempfile
 import urllib.parse
 from asyncio import Task
+from dataclasses import dataclass
 from enum import Enum
 from typing import Any
 from typing import Callable
+from typing import ClassVar
 from typing import Iterable
 from typing import Union
 from typing import overload
@@ -3131,10 +3133,103 @@ def deps_from_container(args: argparse.Namespace, cnt: dict) -> set:
     return cnt['_deps']
 
 
+@dataclass
+class PullImage:
+    POLICY_PRIORITY: ClassVar[dict[str, int]] = {
+        "always": 3,
+        "newer": 2,
+        "missing": 1,
+        "never": 0,
+        "build": 0,
+    }
+
+    image: str
+    policy: str = "missing"
+    quiet: bool = False
+
+    ignore_pull_error: bool = False
+
+    def __post_init__(self) -> None:
+        if self.policy not in self.POLICY_PRIORITY:
+            log.debug(f"Pull policy {self.policy} is not valid, using 'missing' instead")
+            self.policy = "missing"
+
+    def update_policy(self, new_policy: str) -> None:
+        if new_policy not in self.POLICY_PRIORITY:
+            log.debug(f"Pull policy {new_policy} is not valid, ignoring it")
+            return
+
+        if self.POLICY_PRIORITY[new_policy] > self.POLICY_PRIORITY[self.policy]:
+            self.policy = new_policy
+
+    @property
+    def pull_args(self) -> list[str]:
+        args = ["--policy", self.policy]
+        if self.quiet:
+            args.append("--quiet")
+
+        args.append(self.image)
+        return args
+
+    async def pull(self, podman: Podman) -> int | None:
+        if self.policy in ("never", "build"):
+            log.debug("Skipping pull of image %s due to policy %s", self.image, self.policy)
+            return 0
+
+        ret = await podman.run([], "pull", self.pull_args)
+        return ret if not self.ignore_pull_error else 0
+
+    @classmethod
+    async def pull_images(
+        cls,
+        podman: Podman,
+        args: argparse.Namespace,
+        services: list[dict[str, Any]],
+    ) -> int | None:
+        pull_tasks = []
+        pull_images: dict[str, PullImage] = {}
+        for pull_service in services:
+            if not is_local(pull_service):
+                image = str(pull_service.get("image", ""))
+                policy = getattr(args, "pull", None) or pull_service.get("pull_policy", "missing")
+
+                if image in pull_images:
+                    pull_images[image].update_policy(policy)
+                else:
+                    pull_images[image] = PullImage(
+                        image, policy, getattr(args, "quiet_pull", False)
+                    )
+
+                if "build" in pull_service:
+                    # From https://github.com/compose-spec/compose-spec/blob/main/build.md#using-build-and-image
+                    # When both image and build are specified,
+                    # we should try to pull the image first,
+                    # and then build it if it does not exist.
+                    # we should not stop here if pull fails.
+                    pull_images[image].ignore_pull_error = True
+
+        for pull_image in pull_images.values():
+            pull_tasks.append(pull_image.pull(podman))
+
+        if pull_tasks:
+            ret = await asyncio.gather(*pull_tasks)
+            return next((r for r in ret if not r), 0)
+
+        return 0
+
+
 @cmd_run(podman_compose, "up", "Create and start the entire stack or some of its services")
 async def compose_up(compose: PodmanCompose, args: argparse.Namespace) -> int | None:
     excluded = get_excluded(compose, args)
 
+    log.info("pulling images: ...")
+    pull_services = [v for k, v in compose.services.items() if k not in excluded]
+    err = await PullImage.pull_images(compose.podman, args, pull_services)
+    if err:
+        log.error("Pull image failed")
+        if not args.dry_run:
+            return err
+
     log.info("building images: ...")
 
     if not args.no_build:

tests/unit/test_pull_image.py

@@ -0,0 +1,165 @@
+from argparse import Namespace
+from unittest import IsolatedAsyncioTestCase
+from unittest.mock import AsyncMock
+from unittest.mock import Mock
+from unittest.mock import call
+from unittest.mock import patch
+
+from parameterized import parameterized
+
+from podman_compose import PullImage
+
+
+class TestPullImage(IsolatedAsyncioTestCase):
+    def test_unsupported_policy_fallback_to_missing(self) -> None:
+        pull_image = PullImage("localhost/test:1", policy="unsupported")
+        assert pull_image.policy == "missing"
+
+    def test_update_policy(self) -> None:
+        pull_image = PullImage("localhost/test:1", policy="never")
+        assert pull_image.policy == "never"
+
+        # not supported policy
+        pull_image.update_policy("unsupported")
+        assert pull_image.policy == "never"
+
+        pull_image.update_policy("missing")
+        assert pull_image.policy == "missing"
+
+        pull_image.update_policy("newer")
+        assert pull_image.policy == "newer"
+
+        pull_image.update_policy("always")
+        assert pull_image.policy == "always"
+
+        # Ensure policy is not downgraded
+        pull_image.update_policy("build")
+        assert pull_image.policy == "always"
+
+    def test_pull_args(self) -> None:
+        pull_image = PullImage("localhost/test:1", policy="always", quiet=True)
+        assert pull_image.pull_args == ["--policy", "always", "--quiet", "localhost/test:1"]
+
+        pull_image.quiet = False
+        assert pull_image.pull_args == ["--policy", "always", "localhost/test:1"]
+
+    @patch("podman_compose.Podman")
+    async def test_pull_success(self, podman_mock: Mock) -> None:
+        pull_image = PullImage("localhost/test:1", policy="always", quiet=True)
+
+        run_mock = AsyncMock()
+        run_mock.return_value = 0
+        podman_mock.run = run_mock
+
+        result = await pull_image.pull(podman_mock)
+        assert result == 0
+        run_mock.assert_called_once_with(
+            [], "pull", ["--policy", "always", "--quiet", "localhost/test:1"]
+        )
+
+    @patch("podman_compose.Podman")
+    async def test_pull_failed(self, podman_mock: Mock) -> None:
+        pull_image = PullImage(
+            "localhost/test:1",
+            policy="always",
+            quiet=True,
+            ignore_pull_error=True,
+        )
+
+        run_mock = AsyncMock()
+        run_mock.return_value = 1
+        podman_mock.run = run_mock
+
+        # with ignore_pull_error=True, should return 0 even if pull fails
+        result = await pull_image.pull(podman_mock)
+        assert result == 0
+
+        # with ignore_pull_error=False, should return the actual error code
+        pull_image.ignore_pull_error = False
+        result = await pull_image.pull(podman_mock)
+        assert result == 1
+
+    @patch("podman_compose.Podman")
+    async def test_pull_with_never_policy(self, podman_mock: Mock) -> None:
+        pull_image = PullImage(
+            "localhost/test:1",
+            policy="never",
+            quiet=True,
+            ignore_pull_error=True,
+        )
+
+        run_mock = AsyncMock()
+        run_mock.return_value = 1
+        podman_mock.run = run_mock
+
+        result = await pull_image.pull(podman_mock)
+        assert result == 0
+        assert run_mock.call_count == 0
+
+    @parameterized.expand([
+        (
+            "Local image should not pull",
+            Namespace(),
+            [{"image": "localhost/a:latest"}],
+            0,
+            [],
+        ),
+        (
+            "Remote image should pull",
+            Namespace(),
+            [{"image": "ghcr.io/a:latest"}],
+            1,
+            [
+                call([], "pull", ["--policy", "missing", "ghcr.io/a:latest"]),
+            ],
+        ),
+        (
+            "The same image in service should call once",
+            Namespace(),
+            [
+                {"image": "ghcr.io/a:latest"},
+                {"image": "ghcr.io/a:latest"},
+                {"image": "ghcr.io/b:latest"},
+            ],
+            2,
+            [
+                call([], "pull", ["--policy", "missing", "ghcr.io/a:latest"]),
+                call([], "pull", ["--policy", "missing", "ghcr.io/b:latest"]),
+            ],
+        ),
+    ])
+    @patch("podman_compose.Podman")
+    async def test_pull_images(
+        self,
+        desc: str,
+        args: Namespace,
+        services: list[dict],
+        call_count: int,
+        calls: list,
+        podman_mock: Mock,
+    ) -> None:
+        run_mock = AsyncMock()
+        run_mock.return_value = 0
+        podman_mock.run = run_mock
+
+        assert await PullImage.pull_images(podman_mock, args, services) == 0
+        assert run_mock.call_count == call_count
+        if calls:
+            run_mock.assert_has_calls(calls, any_order=True)
+
+    @patch("podman_compose.Podman")
+    async def test_pull_images_with_build_section(
+        self,
+        podman_mock: Mock,
+    ) -> None:
+        run_mock = AsyncMock()
+        run_mock.return_value = 1
+        podman_mock.run = run_mock
+
+        args: Namespace = Namespace()
+        services: list[dict] = [
+            {"image": "ghcr.io/a:latest", "build": {"context": "."}},
+        ]
+        assert await PullImage.pull_images(podman_mock, args, services) == 0
+        assert run_mock.call_count == 1
+        run_mock.assert_called_with([], "pull", ["--policy", "missing", "ghcr.io/a:latest"])