docker: Allow overriding FROM args when determining upstream image de…

…pendencies (pantsbuild#18009)

This solution allows `docker_build_args` to override the 
```dockerfile
ARG VAR=<address>
FROM $VAR
```
inference by assuming a supplied docker build arg option is preferred to the "default value" in the Dockerfile when inferring dependencies and assembling docker build context.


Fixes pantsbuild#18004.

src/python/pants/backend/docker/util_rules/dependencies.py

@@ -8,10 +8,14 @@
 
 from pants.backend.docker.subsystems.dockerfile_parser import DockerfileInfo, DockerfileInfoRequest
 from pants.backend.docker.target_types import DockerImageDependenciesField
+from pants.backend.docker.util_rules.docker_build_args import (
+    DockerBuildArgs,
+    DockerBuildArgsRequest,
+)
 from pants.core.goals.package import AllPackageableTargets, OutputPathField
 from pants.engine.addresses import Addresses, UnparsedAddressInputs
 from pants.engine.rules import Get, collect_rules, rule
-from pants.engine.target import FieldSet, InferDependenciesRequest, InferredDependencies
+from pants.engine.target import FieldSet, InferDependenciesRequest, InferredDependencies, Targets
 from pants.engine.unions import UnionRule
 from pants.util.strutil import softwrap
 
@@ -33,12 +37,19 @@ async def infer_docker_dependencies(
 ) -> InferredDependencies:
     """Inspects the Dockerfile for references to known packagable targets."""
     dockerfile_info = await Get(DockerfileInfo, DockerfileInfoRequest(request.field_set.address))
+    targets = await Get(Targets, Addresses([request.field_set.address]))
+    build_args = await Get(DockerBuildArgs, DockerBuildArgsRequest(targets.expect_single()))
+    merged_from_build_args = {
+        k: build_args.to_dict().get(k, v)
+        for k, v in dockerfile_info.from_image_build_args.to_dict().items()
+    }
+    dockerfile_build_args = {k: v for k, v in merged_from_build_args.items() if v}
 
     putative_image_addresses = set(
         await Get(
             Addresses,
             UnparsedAddressInputs(
-                (v for v in dockerfile_info.from_image_build_args.to_dict().values() if v),
+                dockerfile_build_args.values(),
                 owning_address=dockerfile_info.address,
                 description_of_origin=softwrap(
                     f"""

src/python/pants/backend/docker/util_rules/dependencies_test.py

@@ -8,7 +8,7 @@
 from pants.backend.docker.goals import package_image
 from pants.backend.docker.subsystems import dockerfile_parser
 from pants.backend.docker.target_types import DockerImageDependenciesField, DockerImageTarget
-from pants.backend.docker.util_rules import dockerfile
+from pants.backend.docker.util_rules import docker_build_args, dockerfile
 from pants.backend.docker.util_rules.dependencies import (
     InferDockerDependencies,
     infer_docker_dependencies,
@@ -32,6 +32,7 @@ def rule_runner() -> RuleRunner:
         rules=[
             *dockerfile.rules(),
             *dockerfile_parser.rules(),
+            *docker_build_args.rules(),
             package.find_all_packageable_targets,
             *package_image.rules(),
             *package_pex_binary.rules(),
@@ -135,3 +136,40 @@ def test_infer_docker_dependencies(files, rule_runner: RuleRunner) -> None:
             Address("project/hello/main/go", target_name="go_bin"),
         ]
     )
+
+
+def test_does_not_infer_dependency_when_docker_build_arg_overwrites(
+    rule_runner: RuleRunner,
+) -> None:
+    rule_runner.write_files(
+        {
+            "src/upstream/BUILD": dedent(
+                """\
+                docker_image(
+                  name="image",
+                  repository="upstream/{name}",
+                  image_tags=["1.0"],
+                  instructions=["FROM alpine:3.16.1"],
+                )
+                """
+            ),
+            "src/downstream/BUILD": "docker_image(name='image')",
+            "src/downstream/Dockerfile": dedent(
+                """\
+                ARG BASE_IMAGE=src/upstream:image
+                FROM $BASE_IMAGE
+                """
+            ),
+        }
+    )
+
+    tgt = rule_runner.get_target(Address("src/downstream", target_name="image"))
+    rule_runner.set_options(
+        ["--docker-build-args=BASE_IMAGE=alpine:3.17.0"],
+        env_inherit={"PATH", "PYENV_ROOT", "HOME"},
+    )
+    inferred = rule_runner.request(
+        InferredDependencies,
+        [InferDockerDependencies(tgt[DockerImageDependenciesField])],
+    )
+    assert inferred == InferredDependencies([])

src/python/pants/backend/docker/util_rules/docker_build_context.py

@@ -331,11 +331,12 @@ async def create_docker_build_context(request: DockerBuildContextRequest) -> Doc
     if request.build_upstream_images:
         # Update build arg values for FROM image build args.
 
-        # Get the FROM image build args with defined values in the Dockerfile.
-        dockerfile_build_args = {
-            k: v for k, v in dockerfile_info.from_image_build_args.to_dict().items() if v
+        # Get the FROM image build args with defined values in the Dockerfile & build args.
+        merged_from_build_args = {
+            k: build_args.to_dict().get(k, v)
+            for k, v in dockerfile_info.from_image_build_args.to_dict().items()
         }
-
+        dockerfile_build_args = {k: v for k, v in merged_from_build_args.items() if v}
         # Parse the build args values into Address instances.
         from_image_addresses = await Get(
             Addresses,

src/python/pants/backend/docker/util_rules/docker_build_context_test.py

@@ -256,6 +256,48 @@ def test_from_image_build_arg_dependency(rule_runner: RuleRunner) -> None:
     )
 
 
+def test_from_image_build_arg_dependency_overwritten(rule_runner: RuleRunner) -> None:
+    rule_runner.write_files(
+        {
+            "src/upstream/BUILD": dedent(
+                """\
+                docker_image(
+                  name="image",
+                  repository="upstream/{name}",
+                  image_tags=["1.0"],
+                  instructions=["FROM alpine:3.16.1"],
+                )
+                """
+            ),
+            "src/downstream/BUILD": "docker_image(name='image')",
+            "src/downstream/Dockerfile": dedent(
+                """\
+                ARG BASE_IMAGE=src/upstream:image
+                FROM $BASE_IMAGE
+                """
+            ),
+        }
+    )
+
+    assert_build_context(
+        rule_runner,
+        Address("src/downstream", target_name="image"),
+        expected_files=["src/downstream/Dockerfile"],
+        build_upstream_images=True,
+        expected_interpolation_context={
+            "tags": {
+                "baseimage": "3.10-slim",
+                "stage0": "3.10-slim",
+            },
+            "build_args": {
+                "BASE_IMAGE": "python:3.10-slim",
+            },
+        },
+        expected_num_upstream_images=0,
+        pants_args=["--docker-build-args=BASE_IMAGE=python:3.10-slim"],
+    )
+
+
 def test_from_image_build_arg_not_in_repo_issue_15585(rule_runner: RuleRunner) -> None:
     rule_runner.write_files(
         {