Support Docker custom tags via plugin hook

src/python/pants/backend/helm/util_rules/post_renderer.py

@@ -5,10 +5,12 @@
 
 import logging
 from dataclasses import dataclass
+from itertools import chain
 
 from pants.backend.docker.goals.package_image import DockerFieldSet
 from pants.backend.docker.subsystems import dockerfile_parser
 from pants.backend.docker.subsystems.docker_options import DockerOptions
+from pants.backend.docker.target_types import DockerImageTags, DockerImageTagsRequest
 from pants.backend.docker.util_rules import (
     docker_binary,
     docker_build_args,
@@ -29,8 +31,9 @@
 from pants.backend.helm.target_types import HelmDeploymentFieldSet
 from pants.engine.addresses import Address, Addresses
 from pants.engine.engine_aware import EngineAwareParameter
-from pants.engine.rules import Get, MultiGet, collect_rules, rule
-from pants.engine.target import Targets
+from pants.engine.rules import Get, MultiGet, collect_rules, rule, rule_helper
+from pants.engine.target import Targets, WrappedTarget, WrappedTargetRequest
+from pants.engine.unions import UnionMembership
 from pants.util.logging import LogLevel
 from pants.util.strutil import bullet_list, softwrap
 
@@ -45,10 +48,29 @@ def debug_hint(self) -> str | None:
         return self.field_set.address.spec
 
 
+@rule_helper
+async def _obtain_custom_image_tags(
+    address: Address, union_membership: UnionMembership
+) -> DockerImageTags:
+    wrapped_target = await Get(
+        WrappedTarget, WrappedTargetRequest(address, description_of_origin="<infalible>")
+    )
+
+    image_tags_requests = union_membership.get(DockerImageTagsRequest)
+    found_image_tags = await MultiGet(
+        Get(DockerImageTags, DockerImageTagsRequest, image_tags_request_cls(wrapped_target.target))
+        for image_tags_request_cls in image_tags_requests
+        if image_tags_request_cls.is_applicable(wrapped_target.target)
+    )
+
+    return DockerImageTags(chain.from_iterable(found_image_tags))
+
+
 @rule(desc="Prepare Helm deployment post-renderer", level=LogLevel.DEBUG)
 async def prepare_post_renderer_for_helm_deployment(
     request: HelmDeploymentPostRendererRequest,
     docker_options: DockerOptions,
+    union_membership: UnionMembership,
 ) -> HelmPostRenderer:
     mapping = await Get(
         FirstPartyHelmDeploymentMapping, FirstPartyHelmDeploymentMappingRequest(request.field_set)
@@ -79,16 +101,19 @@ async def prepare_post_renderer_for_helm_deployment(
     docker_targets = await Get(Targets, Addresses(docker_addresses))
     field_sets = [DockerFieldSet.create(tgt) for tgt in docker_targets]
 
-    def resolve_docker_image_ref(address: Address, context: DockerBuildContext) -> str | None:
+    async def resolve_docker_image_ref(address: Address, context: DockerBuildContext) -> str | None:
         docker_field_sets = [fs for fs in field_sets if fs.address == address]
         if not docker_field_sets:
             return None
 
+        additional_image_tags = await _obtain_custom_image_tags(address, union_membership)
+
         docker_field_set = docker_field_sets[0]
         image_refs = docker_field_set.image_refs(
             default_repository=docker_options.default_repository,
             registries=docker_options.registries(),
             interpolation_context=context.interpolation_context,
+            additional_tags=tuple(additional_image_tags),
         )
 
         # Choose first non-latest image reference found, or fallback to 'latest'.
@@ -110,7 +135,7 @@ def resolve_docker_image_ref(address: Address, context: DockerBuildContext) -> s
         return resolved_ref
 
     docker_addr_ref_mapping = {
-        addr: resolve_docker_image_ref(addr, ctx)
+        addr: await resolve_docker_image_ref(addr, ctx)
         for addr, ctx in zip(docker_addresses, docker_contexts)
     }
 

src/python/pants/backend/helm/util_rules/post_renderer_test.py

@@ -7,7 +7,11 @@
 
 import pytest
 
-from pants.backend.docker.target_types import DockerImageTarget
+from pants.backend.docker.target_types import (
+    DockerImageTags,
+    DockerImageTagsRequest,
+    DockerImageTarget,
+)
 from pants.backend.helm.dependency_inference import deployment as infer_deployment
 from pants.backend.helm.subsystems.post_renderer import (
     HELM_POST_RENDERER_CFG_FILENAME,
@@ -31,10 +35,23 @@
 from pants.core.util_rules import source_files
 from pants.engine.addresses import Address
 from pants.engine.process import ProcessResult
-from pants.engine.rules import QueryRule
+from pants.engine.rules import QueryRule, rule
+from pants.engine.target import Target
+from pants.engine.unions import UnionRule
 from pants.testutil.rule_runner import PYTHON_BOOTSTRAP_ENV, RuleRunner
 
 
+class CustomTestImageTagRequest(DockerImageTagsRequest):
+    @classmethod
+    def is_applicable(cls, target: Target) -> bool:
+        return "bar" in target.address.target_name
+
+
+@rule
+async def custom_test_image_tags(_: CustomTestImageTagRequest) -> DockerImageTags:
+    return DockerImageTags(["custom-tag"])
+
+
 @pytest.fixture
 def rule_runner() -> RuleRunner:
     return RuleRunner(
@@ -43,6 +60,8 @@ def rule_runner() -> RuleRunner:
             *infer_deployment.rules(),
             *source_files.rules(),
             *post_renderer.rules(),
+            custom_test_image_tags,
+            UnionRule(DockerImageTagsRequest, CustomTestImageTagRequest),
             QueryRule(HelmPostRenderer, (HelmDeploymentPostRendererRequest,)),
             QueryRule(RenderedHelmFiles, (HelmDeploymentRequest,)),
             QueryRule(ProcessResult, (HelmProcess,)),
@@ -134,8 +153,8 @@ def test_can_prepare_post_renderer(rule_runner: RuleRunner) -> None:
             /spec/containers/1/image: app_foo:latest
             /spec/initContainers/0/image: init_foo:latest
         - paths:
-            /spec/containers/1/image: app_bar:latest
-            /spec/initContainers/0/image: init_bar:latest
+            /spec/containers/1/image: app_bar:custom-tag
+            /spec/initContainers/0/image: init_bar:custom-tag
         """
     )
 
@@ -169,12 +188,12 @@ def test_can_prepare_post_renderer(rule_runner: RuleRunner) -> None:
         spec:
           initContainers:
             - name: myapp-init-container
-              image: init_bar:latest
+              image: init_bar:custom-tag
           containers:
             - name: busy
               image: busybox:1.29
             - name: myapp-container
-              image: app_bar:latest
+              image: app_bar:custom-tag
         """
     )