Perform address mapping on a per-deployment basis (pantsbuild#16492)

Improvement of Helm deployment dependency inference

src/python/pants/backend/helm/dependency_inference/deployment.py

@@ -6,16 +6,13 @@
 import logging
 from dataclasses import dataclass
 from pathlib import PurePath
+from typing import Any
 
 from pants.backend.docker.target_types import AllDockerImageTargets
 from pants.backend.docker.target_types import rules as docker_target_types_rules
 from pants.backend.helm.subsystems import k8s_parser
 from pants.backend.helm.subsystems.k8s_parser import ParsedKubeManifest, ParseKubeManifestRequest
-from pants.backend.helm.target_types import (
-    AllHelmDeploymentTargets,
-    HelmDeploymentDependenciesField,
-    HelmDeploymentFieldSet,
-)
+from pants.backend.helm.target_types import HelmDeploymentFieldSet
 from pants.backend.helm.target_types import rules as helm_target_types_rules
 from pants.backend.helm.util_rules import renderer
 from pants.backend.helm.util_rules.renderer import (
@@ -25,17 +22,16 @@
 )
 from pants.backend.helm.utils.yaml import FrozenYamlIndex, MutableYamlIndex
 from pants.engine.addresses import Address
+from pants.engine.engine_aware import EngineAwareParameter, EngineAwareReturnType
 from pants.engine.fs import Digest, DigestEntries, FileEntry
 from pants.engine.rules import Get, MultiGet, collect_rules, rule
 from pants.engine.target import (
     DependenciesRequest,
     ExplicitlyProvidedDependencies,
-    FieldSet,
     InferDependenciesRequest,
     InferredDependencies,
 )
 from pants.engine.unions import UnionRule
-from pants.util.frozendict import FrozenDict
 from pants.util.logging import LogLevel
 from pants.util.ordered_set import FrozenOrderedSet, OrderedSet
 from pants.util.strutil import pluralize, softwrap
@@ -44,23 +40,37 @@
 
 
 @dataclass(frozen=True)
-class HelmDeploymentReport:
+class AnalyseHelmDeploymentRequest(EngineAwareParameter):
+    field_set: HelmDeploymentFieldSet
+
+    def debug_hint(self) -> str | None:
+        return self.field_set.address.spec
+
+
+@dataclass(frozen=True)
+class HelmDeploymentReport(EngineAwareReturnType):
     address: Address
     image_refs: FrozenYamlIndex[str]
 
     @property
     def all_image_refs(self) -> FrozenOrderedSet[str]:
         return FrozenOrderedSet(self.image_refs.values())
 
+    def level(self) -> LogLevel | None:
+        return LogLevel.DEBUG
+
+    def metadata(self) -> dict[str, Any] | None:
+        return {"address": self.address, "image_refs": self.image_refs}
+
 
 @rule(desc="Analyse Helm deployment", level=LogLevel.DEBUG)
-async def analyse_deployment(field_set: HelmDeploymentFieldSet) -> HelmDeploymentReport:
+async def analyse_deployment(request: AnalyseHelmDeploymentRequest) -> HelmDeploymentReport:
     rendered_deployment = await Get(
         RenderedHelmFiles,
         HelmDeploymentRequest(
             cmd=HelmDeploymentCmd.RENDER,
-            field_set=field_set,
-            description=f"Rendering Helm deployment {field_set.address}",
+            field_set=request.field_set,
+            description=f"Rendering Helm deployment {request.field_set.address}",
         ),
     )
 
@@ -85,35 +95,39 @@ async def analyse_deployment(field_set: HelmDeploymentFieldSet) -> HelmDeploymen
                 item=image_ref,
             )
 
-    return HelmDeploymentReport(address=field_set.address, image_refs=image_refs_index.frozen())
+    return HelmDeploymentReport(
+        address=request.field_set.address, image_refs=image_refs_index.frozen()
+    )
+
+
+@dataclass(frozen=True)
+class FirstPartyHelmDeploymentMappingRequest(EngineAwareParameter):
+    field_set: HelmDeploymentFieldSet
+
+    def debug_hint(self) -> str | None:
+        return self.field_set.address.spec
 
 
 @dataclass(frozen=True)
-class FirstPartyHelmDeploymentMappings:
+class FirstPartyHelmDeploymentMapping:
     """A mapping between `helm_deployment` target addresses and tuples made up of a Docker image
     reference and a `docker_image` target address.
 
     The tuples of Docker image references and addresses are stored in a YAML index so we can track
     the locations in which the Docker image refs appear in the deployment files.
     """
 
-    deployment_to_docker_addresses: FrozenDict[Address, FrozenYamlIndex[tuple[str, Address]]]
-
-    def docker_addresses_referenced_by(self, address: Address) -> list[tuple[str, Address]]:
-        if address not in self.deployment_to_docker_addresses:
-            return []
-        return list(self.deployment_to_docker_addresses[address].values())
+    address: Address
+    indexed_docker_addresses: FrozenYamlIndex[tuple[str, Address]]
 
 
 @rule
-async def first_party_helm_deployment_mappings(
-    deployment_targets: AllHelmDeploymentTargets, docker_targets: AllDockerImageTargets
-) -> FirstPartyHelmDeploymentMappings:
-    field_sets = [HelmDeploymentFieldSet.create(tgt) for tgt in deployment_targets]
-    all_deployments_info = await MultiGet(
-        Get(HelmDeploymentReport, HelmDeploymentFieldSet, field_set) for field_set in field_sets
+async def first_party_helm_deployment_mapping(
+    request: FirstPartyHelmDeploymentMappingRequest, docker_targets: AllDockerImageTargets
+) -> FirstPartyHelmDeploymentMapping:
+    deployment_report = await Get(
+        HelmDeploymentReport, AnalyseHelmDeploymentRequest(request.field_set)
     )
-
     docker_target_addresses = {tgt.address.spec: tgt.address for tgt in docker_targets}
 
     def lookup_docker_addreses(image_ref: str) -> tuple[str, Address] | None:
@@ -122,38 +136,32 @@ def lookup_docker_addreses(image_ref: str) -> tuple[str, Address] | None:
             return image_ref, addr
         return None
 
-    # Builds a mapping between `helm_deployment` addresses and a YAML index of `docker_image` addresses.
-    address_mapping = {
-        fs.address: info.image_refs.transform_values(lookup_docker_addreses)
-        for fs, info in zip(field_sets, all_deployments_info)
-    }
-    return FirstPartyHelmDeploymentMappings(
-        deployment_to_docker_addresses=FrozenDict(address_mapping)
+    return FirstPartyHelmDeploymentMapping(
+        address=request.field_set.address,
+        indexed_docker_addresses=deployment_report.image_refs.transform_values(
+            lookup_docker_addreses
+        ),
     )
 
 
-@dataclass(frozen=True)
-class HelmDeploymentDependenciesInferenceFieldSet(FieldSet):
-    required_fields = (HelmDeploymentDependenciesField,)
-
-    dependencies: HelmDeploymentDependenciesField
-
-
 class InferHelmDeploymentDependenciesRequest(InferDependenciesRequest):
-    infer_from = HelmDeploymentDependenciesInferenceFieldSet
+    infer_from = HelmDeploymentFieldSet
 
 
 @rule(desc="Find the dependencies needed by a Helm deployment")
 async def inject_deployment_dependencies(
-    request: InferHelmDeploymentDependenciesRequest, mapping: FirstPartyHelmDeploymentMappings
+    request: InferHelmDeploymentDependenciesRequest,
 ) -> InferredDependencies:
-    explicitly_provided_deps = await Get(
-        ExplicitlyProvidedDependencies, DependenciesRequest(request.field_set.dependencies)
+    explicitly_provided_deps, mapping = await MultiGet(
+        Get(ExplicitlyProvidedDependencies, DependenciesRequest(request.field_set.dependencies)),
+        Get(
+            FirstPartyHelmDeploymentMapping,
+            FirstPartyHelmDeploymentMappingRequest(request.field_set),
+        ),
     )
-    candidate_docker_addresses = mapping.docker_addresses_referenced_by(request.field_set.address)
 
     dependencies: OrderedSet[Address] = OrderedSet()
-    for imager_ref, candidate_address in candidate_docker_addresses:
+    for imager_ref, candidate_address in mapping.indexed_docker_addresses.values():
         matches = frozenset([candidate_address]).difference(explicitly_provided_deps.includes)
         explicitly_provided_deps.maybe_warn_of_ambiguous_dependency_inference(
             matches,

src/python/pants/backend/helm/dependency_inference/deployment_test.py

@@ -10,8 +10,9 @@
 from pants.backend.docker.target_types import DockerImageTarget
 from pants.backend.helm.dependency_inference import deployment
 from pants.backend.helm.dependency_inference.deployment import (
-    FirstPartyHelmDeploymentMappings,
-    HelmDeploymentDependenciesInferenceFieldSet,
+    AnalyseHelmDeploymentRequest,
+    FirstPartyHelmDeploymentMapping,
+    FirstPartyHelmDeploymentMappingRequest,
     HelmDeploymentReport,
     InferHelmDeploymentDependenciesRequest,
 )
@@ -51,8 +52,8 @@ def rule_runner() -> RuleRunner:
             *process.rules(),
             *stripped_source_files.rules(),
             *tool.rules(),
-            QueryRule(FirstPartyHelmDeploymentMappings, ()),
-            QueryRule(HelmDeploymentReport, (HelmDeploymentFieldSet,)),
+            QueryRule(FirstPartyHelmDeploymentMapping, (FirstPartyHelmDeploymentMappingRequest,)),
+            QueryRule(HelmDeploymentReport, (AnalyseHelmDeploymentRequest,)),
             QueryRule(InferredDependencies, (InferHelmDeploymentDependenciesRequest,)),
         ],
     )
@@ -97,7 +98,9 @@ def test_deployment_dependencies_report(rule_runner: RuleRunner) -> None:
     target = rule_runner.get_target(Address("src/deployment", target_name="foo"))
     field_set = HelmDeploymentFieldSet.create(target)
 
-    dependencies_report = rule_runner.request(HelmDeploymentReport, [field_set])
+    dependencies_report = rule_runner.request(
+        HelmDeploymentReport, [AnalyseHelmDeploymentRequest(field_set)]
+    )
 
     expected_container_refs = [
         "busybox:1.28",
@@ -144,22 +147,21 @@ def test_inject_deployment_dependencies(rule_runner: RuleRunner) -> None:
 
     deployment_addr = Address("src/deployment", target_name="foo")
     tgt = rule_runner.get_target(deployment_addr)
+    field_set = HelmDeploymentFieldSet.create(tgt)
 
     expected_image_ref = "src/image:myapp"
     expected_dependency_addr = Address("src/image", target_name="myapp")
 
-    mappings = rule_runner.request(FirstPartyHelmDeploymentMappings, [])
-    assert mappings.docker_addresses_referenced_by(deployment_addr) == [
+    mapping = rule_runner.request(
+        FirstPartyHelmDeploymentMapping, [FirstPartyHelmDeploymentMappingRequest(field_set)]
+    )
+    assert list(mapping.indexed_docker_addresses.values()) == [
         (expected_image_ref, expected_dependency_addr)
     ]
 
     inferred_dependencies = rule_runner.request(
         InferredDependencies,
-        [
-            InferHelmDeploymentDependenciesRequest(
-                HelmDeploymentDependenciesInferenceFieldSet.create(tgt)
-            )
-        ],
+        [InferHelmDeploymentDependenciesRequest(field_set)],
     )
 
     assert len(inferred_dependencies.include) == 1
@@ -214,11 +216,7 @@ def test_disambiguate_docker_dependency(rule_runner: RuleRunner) -> None:
 
     inferred_dependencies = rule_runner.request(
         InferredDependencies,
-        [
-            InferHelmDeploymentDependenciesRequest(
-                HelmDeploymentDependenciesInferenceFieldSet.create(tgt)
-            )
-        ],
+        [InferHelmDeploymentDependenciesRequest(HelmDeploymentFieldSet.create(tgt))],
     )
 
     assert len(inferred_dependencies.include) == 0

src/python/pants/backend/helm/subsystems/k8s_parser.py

@@ -7,6 +7,7 @@
 import pkgutil
 from dataclasses import dataclass
 from pathlib import PurePath
+from typing import Any
 
 from pants.backend.helm.utils.yaml import YamlPath
 from pants.backend.python.goals import lockfile
@@ -17,14 +18,14 @@
 from pants.backend.python.util_rules import pex
 from pants.backend.python.util_rules.pex import PexRequest, VenvPex, VenvPexProcess
 from pants.core.goals.generate_lockfiles import GenerateToolLockfileSentinel
-from pants.engine.engine_aware import EngineAwareParameter
+from pants.engine.engine_aware import EngineAwareParameter, EngineAwareReturnType
 from pants.engine.fs import CreateDigest, Digest, FileContent, FileEntry
 from pants.engine.process import FallibleProcessResult
 from pants.engine.rules import Get, collect_rules, rule
 from pants.engine.unions import UnionRule
 from pants.util.docutil import git_url
 from pants.util.logging import LogLevel
-from pants.util.strutil import softwrap
+from pants.util.strutil import pluralize, softwrap
 
 logger = logging.getLogger(__name__)
 
@@ -99,12 +100,27 @@ class ParseKubeManifestRequest(EngineAwareParameter):
     def debug_hint(self) -> str | None:
         return self.file.path
 
+    def metadata(self) -> dict[str, Any] | None:
+        return {"file": self.file}
+
 
 @dataclass(frozen=True)
-class ParsedKubeManifest:
+class ParsedKubeManifest(EngineAwareReturnType):
     filename: str
     found_image_refs: tuple[tuple[int, YamlPath, str], ...]
 
+    def level(self) -> LogLevel | None:
+        return LogLevel.DEBUG
+
+    def message(self) -> str | None:
+        return f"Found {pluralize(len(self.found_image_refs), 'image reference')} in file {self.filename}"
+
+    def metadata(self) -> dict[str, Any] | None:
+        return {
+            "filename": self.filename,
+            "found_image_refs": self.found_image_refs,
+        }
+
 
 @rule(desc="Parse Kubernetes resource manifest")
 async def parse_kube_manifest(
@@ -150,7 +166,6 @@ async def parse_kube_manifest(
             softwrap(
                 f"""
                 Could not parse Kubernetes manifests in file: {request.file.path}.
-
                 {parser_error}
                 """
             )

src/python/pants/backend/helm/subsystems/k8s_parser_test.py

@@ -40,7 +40,8 @@ def test_parser_can_run(rule_runner: RuleRunner) -> None:
     file_entries = rule_runner.request(DigestEntries, [file_digest])
 
     parsed_manifest = rule_runner.request(
-        ParsedKubeManifest, [ParseKubeManifestRequest(cast(FileEntry, file_entries[0]))]
+        ParsedKubeManifest,
+        [ParseKubeManifestRequest(file=cast(FileEntry, file_entries[0]))],
     )
 
     expected_image_refs = [
@@ -70,7 +71,8 @@ def test_parser_returns_no_image_refs(rule_runner: RuleRunner) -> None:
     file_entries = rule_runner.request(DigestEntries, [file_digest])
 
     parsed_manifest = rule_runner.request(
-        ParsedKubeManifest, [ParseKubeManifestRequest(cast(FileEntry, file_entries[0]))]
+        ParsedKubeManifest,
+        [ParseKubeManifestRequest(file=cast(FileEntry, file_entries[0]))],
     )
 
     assert len(parsed_manifest.found_image_refs) == 0

src/python/pants/backend/helm/subsystems/post_renderer_main.py

@@ -120,6 +120,13 @@ def main(args: list[str]) -> None:
 
             output_manifest_map[source_filename].append(dump_yaml_data(yaml, manifest_yaml))
 
+    # Include in the output the files that didn't require any replacement
+    remaining_files = set(input_manifest_map.keys()) - set(output_manifest_map.keys())
+    for remaining_file in remaining_files:
+        input_file = input_manifest_map.get(remaining_file)
+        if input_file:
+            output_manifest_map[remaining_file] = input_file
+
     print_manifests(output_manifest_map)
 
 

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

@@ -20,7 +20,10 @@
     DockerBuildContext,
     DockerBuildContextRequest,
 )
-from pants.backend.helm.dependency_inference.deployment import FirstPartyHelmDeploymentMappings
+from pants.backend.helm.dependency_inference.deployment import (
+    FirstPartyHelmDeploymentMapping,
+    FirstPartyHelmDeploymentMappingRequest,
+)
 from pants.backend.helm.subsystems import post_renderer
 from pants.backend.helm.subsystems.post_renderer import HelmPostRenderer, SetupHelmPostRenderer
 from pants.backend.helm.target_types import HelmDeploymentFieldSet
@@ -45,11 +48,13 @@ def debug_hint(self) -> str | None:
 @rule(desc="Prepare Helm deployment post-renderer", level=LogLevel.DEBUG)
 async def prepare_post_renderer_for_helm_deployment(
     request: HelmDeploymentPostRendererRequest,
-    mappings: FirstPartyHelmDeploymentMappings,
     docker_options: DockerOptions,
 ) -> HelmPostRenderer:
-    docker_addr_index = mappings.deployment_to_docker_addresses[request.field_set.address]
-    docker_addresses = [addr for _, addr in docker_addr_index.values()]
+    mapping = await Get(
+        FirstPartyHelmDeploymentMapping, FirstPartyHelmDeploymentMappingRequest(request.field_set)
+    )
+
+    docker_addresses = [addr for _, addr in mapping.indexed_docker_addresses.values()]
 
     logger.debug(
         softwrap(
@@ -113,7 +118,7 @@ def find_replacement(value: tuple[str, Address]) -> str | None:
         _, addr = value
         return docker_addr_ref_mapping.get(addr)
 
-    replacements = docker_addr_index.transform_values(find_replacement)
+    replacements = mapping.indexed_docker_addresses.transform_values(find_replacement)
 
     return await Get(
         HelmPostRenderer,