Merge pull request #1298 from lisongmin/feat-pull-before-teardown

Feat pull images before teardown containers on up command

Author: p12tic

newsfragments/dont-pull-localhost.bugfix

@@ -0,0 +1,2 @@
+Images that have names starting with `localhost/` will no longer be
+pulled when both image and build sections exist.

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

@@ -30,6 +30,7 @@
 from enum import Enum
 from typing import Any
 from typing import Callable
+from typing import ClassVar
 from typing import Iterable
 from typing import Sequence
 from typing import overload
@@ -2758,11 +2759,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")
@@ -3206,9 +3204,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 PullImageSettings:
+    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
+
+
+def settings_to_pull_args(settings: PullImageSettings) -> list[str]:
+    args = ["--policy", settings.policy]
+    if settings.quiet:
+        args.append("--quiet")
+
+    args.append(settings.image)
+    return args
+
+
+async def pull_image(podman: Podman, settings: PullImageSettings) -> int | None:
+    if settings.policy in ("never", "build"):
+        log.debug("Skipping pull of image %s due to policy %s", settings.image, settings.policy)
+        return 0
+
+    ret = await podman.run([], "pull", settings_to_pull_args(settings))
+    return ret if not settings.ignore_pull_error else 0
+
+
+async def pull_images(
+    podman: Podman,
+    args: argparse.Namespace,
+    services: list[dict[str, Any]],
+) -> int | None:
+    pull_tasks = []
+    settingettings: dict[str, PullImageSettings] = {}
+    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 settingettings:
+                settingettings[image].update_policy(policy)
+            else:
+                settingettings[image] = PullImageSettings(
+                    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.
+                settingettings[image].ignore_pull_error = True
+
+    for s in settingettings.values():
+        pull_tasks.append(pull_image(podman, s))
+
+    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 pull_images(compose.podman, args, pull_services)
+    if err:
+        log.error("Pull image failed")
+        return err
 
     log.info("building images: ...")
 
@@ -3218,8 +3308,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,161 @@
+from argparse import Namespace
+from unittest import IsolatedAsyncioTestCase
+from unittest import mock
+
+from parameterized import parameterized
+
+from podman_compose import PullImageSettings
+from podman_compose import pull_image
+from podman_compose import pull_images
+from podman_compose import settings_to_pull_args
+
+
+class TestPullImageSettings(IsolatedAsyncioTestCase):
+    def test_unsupported_policy_fallback_to_missing(self) -> None:
+        settings = PullImageSettings("localhost/test:1", policy="unsupported")
+        assert settings.policy == "missing"
+
+    def test_update_policy(self) -> None:
+        settings = PullImageSettings("localhost/test:1", policy="never")
+        assert settings.policy == "never"
+
+        # not supported policy
+        settings.update_policy("unsupported")
+        assert settings.policy == "never"
+
+        settings.update_policy("missing")
+        assert settings.policy == "missing"
+
+        settings.update_policy("newer")
+        assert settings.policy == "newer"
+
+        settings.update_policy("always")
+        assert settings.policy == "always"
+
+        # Ensure policy is not downgraded
+        settings.update_policy("build")
+        assert settings.policy == "always"
+
+    def test_pull_args(self) -> None:
+        settings = PullImageSettings("localhost/test:1", policy="always", quiet=True)
+        assert settings_to_pull_args(settings) == [
+            "--policy",
+            "always",
+            "--quiet",
+            "localhost/test:1",
+        ]
+
+        settings.quiet = False
+        assert settings_to_pull_args(settings) == ["--policy", "always", "localhost/test:1"]
+
+    @mock.patch("podman_compose.Podman")
+    async def test_pull_success(self, podman_mock: mock.Mock) -> None:
+        settings = PullImageSettings("localhost/test:1", policy="always", quiet=True)
+
+        run_mock = mock.AsyncMock(return_value=0)
+        podman_mock.run = run_mock
+
+        result = await pull_image(podman_mock, settings)
+        assert result == 0
+        run_mock.assert_called_once_with(
+            [], "pull", ["--policy", "always", "--quiet", "localhost/test:1"]
+        )
+
+    @mock.patch("podman_compose.Podman")
+    async def test_pull_failed(self, podman_mock: mock.Mock) -> None:
+        settings = PullImageSettings(
+            "localhost/test:1",
+            policy="always",
+            quiet=True,
+            ignore_pull_error=True,
+        )
+
+        podman_mock.run = mock.AsyncMock(return_value=1)
+
+        # with ignore_pull_error=True, should return 0 even if pull fails
+        result = await pull_image(podman_mock, settings)
+        assert result == 0
+
+        # with ignore_pull_error=False, should return the actual error code
+        settings.ignore_pull_error = False
+        result = await pull_image(podman_mock, settings)
+        assert result == 1
+
+    @mock.patch("podman_compose.Podman")
+    async def test_pull_with_never_policy(self, podman_mock: mock.Mock) -> None:
+        settings = PullImageSettings(
+            "localhost/test:1",
+            policy="never",
+            quiet=True,
+            ignore_pull_error=True,
+        )
+
+        run_mock = mock.AsyncMock(return_value=1)
+        podman_mock.run = run_mock
+
+        result = await pull_image(podman_mock, settings)
+        assert result == 0
+        assert run_mock.call_count == 0
+
+    @parameterized.expand([
+        (
+            "Local image should not pull",
+            [{"image": "localhost/a:latest"}],
+            [],
+        ),
+        (
+            "Remote image should pull",
+            [{"image": "ghcr.io/a:latest"}],
+            [
+                mock.call([], "pull", ["--policy", "missing", "ghcr.io/a:latest"]),
+            ],
+        ),
+        (
+            "The same image in service should call once",
+            [
+                {"image": "ghcr.io/a:latest"},
+                {"image": "ghcr.io/a:latest"},
+                {"image": "ghcr.io/b:latest"},
+            ],
+            [
+                mock.call([], "pull", ["--policy", "missing", "ghcr.io/a:latest"]),
+                mock.call([], "pull", ["--policy", "missing", "ghcr.io/b:latest"]),
+            ],
+        ),
+    ])
+    @mock.patch("podman_compose.Podman")
+    async def test_pull_image(
+        self,
+        desc: str,
+        services: list[dict],
+        calls: list,
+        podman_mock: mock.Mock,
+    ) -> None:
+        run_mock = mock.AsyncMock(return_value=1)
+        podman_mock.run = run_mock
+
+        assert await pull_images(podman_mock, Namespace(), services) == 0
+        assert run_mock.call_count == len(calls)
+        if calls:
+            run_mock.assert_has_calls(calls, any_order=True)
+
+    @mock.patch("podman_compose.Podman")
+    async def test_pull_image_with_build_section(
+        self,
+        podman_mock: mock.Mock,
+    ) -> None:
+        run_mock = mock.AsyncMock(return_value=1)
+        podman_mock.run = run_mock
+
+        assert (
+            await pull_images(
+                podman_mock,
+                Namespace(),
+                [
+                    {"image": "ghcr.io/a:latest", "build": {"context": "."}},
+                ],
+            )
+            == 0
+        )
+        assert run_mock.call_count == 1
+        run_mock.assert_called_with([], "pull", ["--policy", "missing", "ghcr.io/a:latest"])