Docker image version tag (#13097)

In order to keep the BUILD files DRY, support picking up the Docker image version to be built, from other sources rather than require a hard coded value in the BUILD file.

First up is support for looking at the `Dockerfile` FROM commands to glean possible version tags from them.

Example:
```python
docker_image(name="myimage", version="{baseimage.tag}")
```
```Dockerfile
FROM upstream:1.2-feat
...
```

This would then produce `myimage:1.2-feat`.

As more format strings are introduced, so is the potential for users to shot them selves in the foot. So I've added some improved error feedback in case the version or image name template format strings blow up, that should help pinpoint the mistake and present possible solutions.

```python
docker_image(
    name="test-example",
    version="1.2.5-{typo}",
    dependencies=["testprojects/src/python/hello/main"],
)
```
```shell
$ ./pants --no-print-stacktrace package testprojects/src/python/docker:test-example
12:09:17.29 [INFO] Initialization options changed: reinitializing scheduler...
12:09:17.99 [INFO] Scheduler initialized.
12:09:22.38 [ERROR] Exception caught: (pants.engine.internals.scheduler.ExecutionError)
[...]
pants.engine.internals.scheduler.ExecutionError: 1 Exception encountered:

  DockerBuildError: Error in testprojects/src/python/docker:test-example: Invalid format string for the `docker_image(version)` field: '1.2.5-{typo}'.

The key 'typo' is unknown. Try with one of: baseimage, stage0.

```

[ci skip-rust]
[ci skip-build-wheels]

src/python/pants/backend/docker/docker_build.py

@@ -5,9 +5,15 @@
 import logging
 from dataclasses import dataclass
 from os import path
+from typing import cast
 
 from pants.backend.docker.docker_binary import DockerBinary, DockerBinaryRequest
-from pants.backend.docker.docker_build_context import DockerBuildContext, DockerBuildContextRequest
+from pants.backend.docker.docker_build_context import (
+    DockerBuildContext,
+    DockerBuildContextRequest,
+    DockerVersionContextError,
+    DockerVersionContextValue,
+)
 from pants.backend.docker.registries import DockerRegistries
 from pants.backend.docker.subsystem import DockerOptions
 from pants.backend.docker.target_types import (
@@ -23,11 +29,16 @@
 from pants.engine.process import Process, ProcessResult
 from pants.engine.rules import Get, MultiGet, collect_rules, rule
 from pants.engine.unions import UnionRule
+from pants.util.frozendict import FrozenDict
 from pants.util.strutil import bullet_list, pluralize
 
 logger = logging.getLogger(__name__)
 
 
+class DockerNameTemplateError(ValueError):
+    pass
+
+
 @dataclass(frozen=True)
 class BuiltDockerImage(BuiltPackageArtifact):
     tags: tuple[str, ...] = ()
@@ -69,24 +80,59 @@ def dockerfile_path(self) -> str:
         return path.join(self.address.spec_path, self.dockerfile_relpath)
 
     def image_names(
-        self, default_name_template: str, registries: DockerRegistries
+        self,
+        default_name_template: str,
+        registries: DockerRegistries,
+        version_context: FrozenDict[str, DockerVersionContextValue],
     ) -> tuple[str, ...]:
         """This method will always return a non-empty tuple."""
         default_parent = path.basename(path.dirname(self.address.spec_path))
         default_repo = path.basename(self.address.spec_path)
         repo = self.repository.value or default_repo
         name_template = self.name_template.value or default_name_template
-        image_name = name_template.format(
-            name=self.name.value or self.address.target_name,
-            repository=repo,
-            sub_repository="/".join(
-                [default_repo if self.repository.value else default_parent, repo]
-            ),
-        )
-        image_names = tuple(
-            ":".join(s for s in [image_name, tag] if s)
-            for tag in [self.version.value, *(self.tags.value or [])]
-        )
+        try:
+            image_name = name_template.format(
+                name=self.name.value or self.address.target_name,
+                repository=repo,
+                sub_repository="/".join(
+                    [default_repo if self.repository.value else default_parent, repo]
+                ),
+            )
+        except KeyError as e:
+            if self.name_template.value:
+                source = (
+                    "from the `image_name_template` field of the docker_image target "
+                    f"at {self.address}"
+                )
+            else:
+                source = "from the [docker].default_image_name_template configuration option"
+
+            raise DockerNameTemplateError(
+                f"Invalid image name template {source}: {name_template!r}. Unknown key: {e}.\n\n"
+                f"Use any of 'name', 'repository' or 'sub_repository' in the template string."
+            ) from e
+
+        try:
+            image_names = tuple(
+                ":".join(s for s in [image_name, tag] if s)
+                for tag in [
+                    cast(str, self.version.value).format(**version_context),
+                    *(self.tags.value or []),
+                ]
+            )
+        except (KeyError, ValueError) as e:
+            msg = (
+                "Invalid format string for the `version` field of the docker_image target at "
+                f"{self.address}: {self.version.value!r}.\n\n"
+            )
+            if isinstance(e, KeyError):
+                msg += (
+                    f"The key {e} is unknown. Try with one of: "
+                    f'{", ".join(version_context.keys())}.'
+                )
+            else:
+                msg += str(e)
+            raise DockerVersionContextError(msg) from e
 
         registries_options = tuple(registries.get(*(self.registries.value or [])))
         if not registries_options:
@@ -115,7 +161,12 @@ async def build_docker_image(
         ),
     )
 
-    tags = field_set.image_names(options.default_image_name_template, options.registries())
+    tags = field_set.image_names(
+        default_name_template=options.default_image_name_template,
+        registries=options.registries(),
+        version_context=context.version_context,
+    )
+
     result = await Get(
         ProcessResult,
         Process,

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

@@ -5,6 +5,7 @@
 from dataclasses import dataclass
 from itertools import chain
 
+from pants.backend.docker.dockerfile_parser import DockerfileInfo
 from pants.backend.docker.target_types import DockerImageSources
 from pants.core.goals.package import BuiltPackage, PackageFieldSet
 from pants.core.target_types import FilesSources
@@ -22,13 +23,31 @@
     TransitiveTargets,
     TransitiveTargetsRequest,
 )
+from pants.util.frozendict import FrozenDict
 
 logger = logging.getLogger(__name__)
 
 
+class DockerVersionContextError(ValueError):
+    pass
+
+
+class DockerVersionContextValue(FrozenDict[str, str]):
+    """Dict class suitable for use as a format string context object, as it allows to use attribute
+    access rather than item access."""
+
+    def __getattr__(self, attribute: str) -> str:
+        if attribute not in self:
+            raise DockerVersionContextError(
+                f"The key {attribute!r} is unknown. Try with one of: " f'{", ".join(self.keys())}.'
+            )
+        return self[attribute]
+
+
 @dataclass(frozen=True)
 class DockerBuildContext:
     digest: Digest
+    version_context: FrozenDict[str, DockerVersionContextValue]
 
 
 @dataclass(frozen=True)
@@ -71,10 +90,11 @@ async def create_docker_build_context(request: DockerBuildContextRequest) -> Doc
         FieldSetsPerTargetRequest(PackageFieldSet, transitive_targets.dependencies),
     )
 
-    dockerfiles, sources, embedded_pkgs_per_target = await MultiGet(
+    dockerfiles, sources, embedded_pkgs_per_target, dockerfile_info = await MultiGet(
         dockerfiles_request,
         sources_request,
         embedded_pkgs_per_target_request,
+        Get(DockerfileInfo, DockerImageSources, transitive_targets.roots[0][DockerImageSources]),
     )
 
     # Package binary dependencies for build context.
@@ -98,7 +118,18 @@ async def create_docker_build_context(request: DockerBuildContextRequest) -> Doc
         MergeDigests(d for d in all_digests if d),
     )
 
-    return DockerBuildContext(context)
+    version_context: dict[str, DockerVersionContextValue] = {}
+    for stage, tag in [tag.split(maxsplit=1) for tag in dockerfile_info.version_tags]:
+        value = DockerVersionContextValue({"tag": tag})
+        if not version_context:
+            # Refer to the first FROM directive as the "baseimage".
+            version_context["baseimage"] = value
+        version_context[stage] = value
+
+    return DockerBuildContext(
+        digest=context,
+        version_context=FrozenDict(version_context),
+    )
 
 
 def rules():

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

@@ -7,8 +7,13 @@
 
 import pytest
 
-from pants.backend.docker.docker_build_context import DockerBuildContext, DockerBuildContextRequest
+from pants.backend.docker.docker_build_context import (
+    DockerBuildContext,
+    DockerBuildContextRequest,
+    DockerVersionContextValue,
+)
 from pants.backend.docker.docker_build_context import rules as context_rules
+from pants.backend.docker.dockerfile_parser import rules as parser_rules
 from pants.backend.docker.target_types import DockerImage
 from pants.backend.python import target_types_rules
 from pants.backend.python.goals import package_pex_binary
@@ -21,28 +26,33 @@
 from pants.engine.addresses import Address
 from pants.engine.fs import Snapshot
 from pants.testutil.rule_runner import QueryRule, RuleRunner
+from pants.util.frozendict import FrozenDict
 
 
 @pytest.fixture
 def rule_runner() -> RuleRunner:
-    return RuleRunner(
+    rule_runner = RuleRunner(
         rules=[
             *context_rules(),
             *core_target_types_rules(),
             *package_pex_binary.rules(),
+            *parser_rules(),
             *pex_from_targets.rules(),
             *target_types_rules.rules(),
             QueryRule(BuiltPackage, [PexBinaryFieldSet]),
             QueryRule(DockerBuildContext, (DockerBuildContextRequest,)),
         ],
         target_types=[DockerImage, Files, PexBinary],
     )
+    rule_runner.set_options([], env_inherit={"PATH", "PYENV_ROOT", "HOME"})
+    return rule_runner
 
 
 def assert_build_context(
     rule_runner: RuleRunner,
     address: Address,
     expected_files: list[str],
+    expected_version_context: FrozenDict[str, DockerVersionContextValue] | None = None,
 ) -> None:
     context = rule_runner.request(
         DockerBuildContext,
@@ -56,6 +66,8 @@ def assert_build_context(
 
     snapshot = rule_runner.request(Snapshot, [context.digest])
     assert sorted(expected_files) == sorted(snapshot.files)
+    if expected_version_context is not None:
+        assert expected_version_context == context.version_context
 
 
 def test_file_dependencies(rule_runner: RuleRunner) -> None:
@@ -159,7 +171,6 @@ def test_files_out_of_tree(rule_runner: RuleRunner) -> None:
 def test_packaged_pex_path(rule_runner: RuleRunner) -> None:
     # This test is here to ensure that we catch if there is any change in the generated path where
     # built pex binaries go, as we rely on that for dependency inference in the Dockerfile.
-    rule_runner.set_options([], env_inherit={"PATH", "PYENV_ROOT", "HOME"})
     rule_runner.write_files(
         {
             "src/docker/BUILD": """docker_image(dependencies=["src/python/proj/cli:bin"])""",
@@ -174,3 +185,34 @@ def test_packaged_pex_path(rule_runner: RuleRunner) -> None:
         Address("src/docker", target_name="docker"),
         expected_files=["src/docker/Dockerfile", "src.python.proj.cli/bin.pex"],
     )
+
+
+def test_version_context_from_dockerfile(rule_runner: RuleRunner) -> None:
+    rule_runner.write_files(
+        {
+            "src/docker/BUILD": """docker_image()""",
+            "src/docker/Dockerfile": dedent(
+                """\
+                FROM python:3.8
+                FROM alpine as interim
+                FROM interim
+                FROM scratch:1-1 as output
+                """
+            ),
+        }
+    )
+
+    assert_build_context(
+        rule_runner,
+        Address("src/docker"),
+        expected_files=["src/docker/Dockerfile"],
+        expected_version_context=FrozenDict(
+            {
+                "baseimage": DockerVersionContextValue({"tag": "3.8"}),
+                "stage0": DockerVersionContextValue({"tag": "3.8"}),
+                "interim": DockerVersionContextValue({"tag": "latest"}),
+                "stage2": DockerVersionContextValue({"tag": "latest"}),
+                "output": DockerVersionContextValue({"tag": "1-1"}),
+            }
+        ),
+    )