Resolve Docker addresses relative to the helm_deployment (pantsbuild#…

…19455)

Adds the ability to resolve docker addresses relative to the helm_deployment.

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

@@ -21,10 +21,11 @@
     RenderedHelmFiles,
 )
 from pants.backend.helm.utils.yaml import FrozenYamlIndex, MutableYamlIndex
+from pants.build_graph.address import MaybeAddress
 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.internals.native_engine import AddressInput
+from pants.engine.internals.native_engine import AddressInput, AddressParseException
 from pants.engine.rules import Get, MultiGet, collect_rules, rule
 from pants.engine.target import (
     DependenciesRequest,
@@ -88,12 +89,12 @@ async def analyse_deployment(request: AnalyseHelmDeploymentRequest) -> HelmDeplo
     # Build YAML index of Docker image refs for future processing during depedendecy inference or post-rendering.
     image_refs_index: MutableYamlIndex[str] = MutableYamlIndex()
     for manifest in parsed_manifests:
-        for idx, path, image_ref in manifest.found_image_refs:
+        for entry in manifest.found_image_refs:
             image_refs_index.insert(
                 file_path=PurePath(manifest.filename),
-                document_index=idx,
-                yaml_path=path,
-                item=image_ref,
+                document_index=entry.document_index,
+                yaml_path=entry.path,
+                item=entry.unparsed_image_ref,
             )
 
     return HelmDeploymentReport(
@@ -129,18 +130,47 @@ async def first_party_helm_deployment_mapping(
     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:
-        addr = docker_target_addresses.get(image_ref, None)
-        if addr:
-            return image_ref, addr
-        return None
+    def image_ref_to_address_input(image_ref: str) -> tuple[str, AddressInput] | None:
+        try:
+            return image_ref, AddressInput.parse(
+                image_ref,
+                description_of_origin=f"the helm_deployment at {request.field_set.address}",
+                relative_to=request.field_set.address.spec_path,
+            )
+        except AddressParseException:
+            return None
+
+    indexed_address_inputs = deployment_report.image_refs.transform_values(
+        image_ref_to_address_input
+    )
+    maybe_addresses = await MultiGet(
+        Get(MaybeAddress, AddressInput, ai) for _, ai in indexed_address_inputs.values()
+    )
+
+    docker_target_addresses = {tgt.address for tgt in docker_targets}
+    maybe_addresses_by_ref = {
+        ref: maybe_addr
+        for ((ref, _), maybe_addr) in zip(indexed_address_inputs.values(), maybe_addresses)
+    }
+
+    def image_ref_to_actual_address(
+        image_ref_ai: tuple[str, AddressInput]
+    ) -> tuple[str, Address] | None:
+        image_ref, _ = image_ref_ai
+        maybe_addr = maybe_addresses_by_ref.get(image_ref)
+        if not maybe_addr:
+            return None
+        if not isinstance(maybe_addr.val, Address):
+            return None
+        if maybe_addr.val not in docker_target_addresses:
+            return None
+        return image_ref, maybe_addr.val
 
     return FirstPartyHelmDeploymentMapping(
         address=request.field_set.address,
-        indexed_docker_addresses=deployment_report.image_refs.transform_values(
-            lookup_docker_addreses
+        indexed_docker_addresses=indexed_address_inputs.transform_values(
+            image_ref_to_actual_address
         ),
     )
 

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

@@ -86,7 +86,7 @@ def test_deployment_dependencies_report(rule_runner: RuleRunner) -> None:
                       image: example.com/containers/busybox:1.28
                 """
             ),
-            "src/deployment/BUILD": "helm_deployment(name='foo', dependencies=['//src/mychart'])",
+            "src/deployment/BUILD": "helm_deployment(name='foo', chart='//src/mychart')",
         }
     )
 
@@ -142,6 +142,70 @@ def test_inject_chart_into_deployment_dependencies(rule_runner: RuleRunner) -> N
     assert list(inferred_dependencies.include)[0] == Address("src/mychart")
 
 
+def test_resolve_relative_docker_addresses_to_deployment(rule_runner: RuleRunner) -> None:
+    rule_runner.write_files(
+        {
+            "src/mychart/BUILD": "helm_chart()",
+            "src/mychart/Chart.yaml": HELM_CHART_FILE,
+            "src/mychart/values.yaml": dedent(
+                """\
+                container:
+                  image_ref: docker/image:latest
+                """
+            ),
+            "src/mychart/templates/_helpers.tpl": HELM_TEMPLATE_HELPERS_FILE,
+            "src/mychart/templates/pod.yaml": dedent(
+                """\
+                apiVersion: v1
+                kind: Pod
+                metadata:
+                  name: {{ template "fullname" . }}
+                  labels:
+                    chart: "{{ .Chart.Name }}-{{ .Chart.Version | replace "+" "_" }}"
+                spec:
+                  containers:
+                    - name: myapp-container
+                      image: {{ .Values.container.image_ref }}
+                """
+            ),
+            "src/deployment/BUILD": dedent(
+                """\
+                docker_image(name="myapp")
+
+                helm_deployment(
+                    name="foo",
+                    chart="//src/mychart",
+                    values={
+                        "container.image_ref": ":myapp"
+                    }
+                )
+                """
+            ),
+            "src/deployment/Dockerfile": "FROM busybox:1.28",
+        }
+    )
+
+    source_root_patterns = ("src/*",)
+    rule_runner.set_options(
+        [f"--source-root-patterns={repr(source_root_patterns)}"],
+        env_inherit=PYTHON_BOOTSTRAP_ENV,
+    )
+
+    deployment_addr = Address("src/deployment", target_name="foo")
+    tgt = rule_runner.get_target(deployment_addr)
+    field_set = HelmDeploymentFieldSet.create(tgt)
+
+    expected_image_ref = ":myapp"
+    expected_dependency_addr = Address("src/deployment", target_name="myapp")
+
+    mapping = rule_runner.request(
+        FirstPartyHelmDeploymentMapping, [FirstPartyHelmDeploymentMappingRequest(field_set)]
+    )
+    assert list(mapping.indexed_docker_addresses.values()) == [
+        (expected_image_ref, expected_dependency_addr)
+    ]
+
+
 def test_inject_deployment_dependencies(rule_runner: RuleRunner) -> None:
     rule_runner.write_files(
         {
@@ -163,7 +227,7 @@ def test_inject_deployment_dependencies(rule_runner: RuleRunner) -> None:
                       image: src/image:myapp
                 """
             ),
-            "src/deployment/BUILD": "helm_deployment(name='foo', dependencies=['//src/mychart'])",
+            "src/deployment/BUILD": "helm_deployment(name='foo', chart='//src/mychart')",
             "src/image/BUILD": "docker_image(name='myapp')",
             "src/image/Dockerfile": "FROM busybox:1.28",
         }
@@ -194,8 +258,9 @@ def test_inject_deployment_dependencies(rule_runner: RuleRunner) -> None:
         [InferHelmDeploymentDependenciesRequest(field_set)],
     )
 
-    assert len(inferred_dependencies.include) == 1
-    assert list(inferred_dependencies.include)[0] == expected_dependency_addr
+    # The Helm chart dependency is part of the inferred dependencies
+    assert len(inferred_dependencies.include) == 2
+    assert expected_dependency_addr in inferred_dependencies.include
 
 
 def test_disambiguate_docker_dependency(rule_runner: RuleRunner) -> None:
@@ -223,8 +288,8 @@ def test_disambiguate_docker_dependency(rule_runner: RuleRunner) -> None:
                 """\
                 helm_deployment(
                     name="foo",
+                    chart="//src/mychart",
                     dependencies=[
-                        "//src/mychart",
                         "!//registry/image:latest",
                     ]
                 )
@@ -249,4 +314,6 @@ def test_disambiguate_docker_dependency(rule_runner: RuleRunner) -> None:
         [InferHelmDeploymentDependenciesRequest(HelmDeploymentFieldSet.create(tgt))],
     )
 
-    assert len(inferred_dependencies.include) == 0
+    # Assert only the Helm chart dependency has been inferred
+    assert len(inferred_dependencies.include) == 1
+    assert set(inferred_dependencies.include) == {Address("src/mychart")}

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

@@ -78,10 +78,17 @@ def metadata(self) -> dict[str, Any] | None:
         return {"file": self.file}
 
 
+@dataclass(frozen=True)
+class ParsedImageRefEntry:
+    document_index: int
+    path: YamlPath
+    unparsed_image_ref: str
+
+
 @dataclass(frozen=True)
 class ParsedKubeManifest(EngineAwareReturnType):
     filename: str
-    found_image_refs: tuple[tuple[int, YamlPath, str], ...]
+    found_image_refs: tuple[ParsedImageRefEntry, ...]
 
     def level(self) -> LogLevel | None:
         return LogLevel.DEBUG
@@ -115,7 +122,7 @@ async def parse_kube_manifest(
 
     if result.exit_code == 0:
         output = result.stdout.decode("utf-8").splitlines()
-        image_refs: list[tuple[int, YamlPath, str]] = []
+        image_refs: list[ParsedImageRefEntry] = []
         for line in output:
             parts = line.split(",")
             if len(parts) != 3:
@@ -128,7 +135,13 @@ async def parse_kube_manifest(
                     )
                 )
 
-            image_refs.append((int(parts[0]), YamlPath.parse(parts[1]), parts[2]))
+            image_refs.append(
+                ParsedImageRefEntry(
+                    document_index=int(parts[0]),
+                    path=YamlPath.parse(parts[1]),
+                    unparsed_image_ref=parts[2],
+                )
+            )
 
         return ParsedKubeManifest(filename=request.file.path, found_image_refs=tuple(image_refs))
     else:

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

@@ -9,7 +9,11 @@
 import pytest
 
 from pants.backend.helm.subsystems import k8s_parser
-from pants.backend.helm.subsystems.k8s_parser import ParsedKubeManifest, ParseKubeManifestRequest
+from pants.backend.helm.subsystems.k8s_parser import (
+    ParsedImageRefEntry,
+    ParsedKubeManifest,
+    ParseKubeManifestRequest,
+)
 from pants.backend.helm.testutil import K8S_POD_FILE
 from pants.backend.helm.utils.yaml import YamlPath
 from pants.engine.fs import CreateDigest, Digest, DigestEntries, FileContent, FileEntry
@@ -45,8 +49,8 @@ def test_parser_can_run(rule_runner: RuleRunner) -> None:
     )
 
     expected_image_refs = [
-        (0, YamlPath.parse("/spec/containers/0/image"), "busybox:1.28"),
-        (0, YamlPath.parse("/spec/initContainers/0/image"), "busybox:1.29"),
+        ParsedImageRefEntry(0, YamlPath.parse("/spec/containers/0/image"), "busybox:1.28"),
+        ParsedImageRefEntry(0, YamlPath.parse("/spec/initContainers/0/image"), "busybox:1.29"),
     ]
 
     assert parsed_manifest.found_image_refs == tuple(expected_image_refs)