Feat pull images before teardown containers on up command

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
@@ -2688,11 +2690,8 @@ def is_local(container: dict) -> bool:
     * prefixed with localhost/
     * has a build section and is not prefixed
     """
-    return (
-        "/" not in container["image"]
-        if "build" in container
-        else container["image"].startswith("localhost/")
-    )
+    image = container.get("image", "")
+    return image.startswith("localhost/") or ("build" in container and "/" not in image)
 
 
 @cmd_run(podman_compose, "wait", "wait running containers to stop")
@@ -3131,9 +3130,101 @@ def deps_from_container(args: argparse.Namespace, cnt: dict) -> set:
     return cnt['_deps']
 
 
-@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)
+@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("Pull policy %s is not valid, using 'missing' instead", self.policy)
+            self.policy = "missing"
+
+    def update_policy(self, new_policy: str) -> None:
+        if new_policy not in self.POLICY_PRIORITY:
+            log.debug("Pull policy %s is not valid, ignoring it", new_policy)
+            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
+
+
+async def prepare_images(
+    compose: PodmanCompose, args: argparse.Namespace, excluded: set[str]
+) -> int | None:
+    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")
+        return err
 
     log.info("building images: ...")
 
@@ -3143,8 +3234,20 @@ async def compose_up(compose: PodmanCompose, args: argparse.Namespace) -> int |
         build_exit_code = await compose.commands["build"](compose, build_args)
         if build_exit_code != 0:
             log.error("Build command failed")
-            if not args.dry_run:
-                return build_exit_code
+            return build_exit_code
+
+    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)
+
+    exit_code = await prepare_images(compose, args, excluded)
+    if exit_code != 0:
+        log.error("Prepare images failed")
+        if not args.dry_run:
+            return exit_code
 
     # if needed, tear down existing containers
 

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"])