Load environment variables from Docker image

src/python/pants/engine/env_vars.py

@@ -7,8 +7,6 @@
 from dataclasses import dataclass
 from typing import Dict, Optional, Sequence
 
-from pants.engine.internals.session import SessionValues
-from pants.engine.rules import collect_rules, rule
 from pants.util.frozendict import FrozenDict
 from pants.util.meta import frozen_after_init
 from pants.util.ordered_set import FrozenOrderedSet
@@ -93,26 +91,3 @@ class EnvironmentVars(FrozenDict[str, str]):
     `CompleteEnvironmentVars`, as it represents a filtered/relevant subset of the environment, rather
     than the entire unfiltered environment.
     """
-
-
-@rule
-def complete_environment_vars(session_values: SessionValues) -> CompleteEnvironmentVars:
-    return session_values[CompleteEnvironmentVars]
-
-
-@rule
-def environment_vars_subset(
-    session_values: SessionValues, request: EnvironmentVarsRequest
-) -> EnvironmentVars:
-    return EnvironmentVars(
-        session_values[CompleteEnvironmentVars]
-        .get_subset(
-            requested=tuple(request.requested),
-            allowed=(None if request.allowed is None else tuple(request.allowed)),
-        )
-        .items()
-    )
-
-
-def rules():
-    return collect_rules()

src/python/pants/engine/internals/platform_rules.py

@@ -3,9 +3,18 @@
 
 from __future__ import annotations
 
-from pants.core.util_rules.environments import DockerPlatformField, EnvironmentTarget
+from pants.core.util_rules.environments import (
+    DockerImageField,
+    DockerPlatformField,
+    EnvironmentTarget,
+)
+from pants.engine.env_vars import CompleteEnvironmentVars, EnvironmentVars, EnvironmentVarsRequest
+from pants.engine.internals.session import SessionValues
 from pants.engine.platform import Platform
-from pants.engine.rules import collect_rules, rule
+from pants.engine.process import Process, ProcessResult
+from pants.engine.rules import Get, collect_rules, rule
+from pants.util.logging import LogLevel
+from pants.util.strutil import softwrap
 
 
 @rule
@@ -15,5 +24,45 @@ def current_platform(env_tgt: EnvironmentTarget) -> Platform:
     return Platform(env_tgt.val[DockerPlatformField].normalized_value)
 
 
+@rule
+async def complete_environment_vars(
+    session_values: SessionValues, env_tgt: EnvironmentTarget
+) -> CompleteEnvironmentVars:
+    if not env_tgt.val or not env_tgt.val.has_field(DockerImageField):
+        return session_values[CompleteEnvironmentVars]
+    env_process_result = await Get(
+        ProcessResult,
+        Process(
+            ["env", "-0"],
+            description=softwrap(
+                f"""
+                Extract environment variables from the Docker image
+                {env_tgt.val[DockerImageField].value}
+                """
+            ),
+            level=LogLevel.DEBUG,
+        ),
+    )
+    result = {}
+    for line in env_process_result.stdout.decode("utf-8").rstrip().split("\0"):
+        if not line:
+            continue
+        k, v = line.split("=", maxsplit=1)
+        result[k] = v
+    return CompleteEnvironmentVars(result)
+
+
+@rule
+def environment_vars_subset(
+    complete_env_vars: CompleteEnvironmentVars, request: EnvironmentVarsRequest
+) -> EnvironmentVars:
+    return EnvironmentVars(
+        complete_env_vars.get_subset(
+            requested=tuple(request.requested),
+            allowed=(None if request.allowed is None else tuple(request.allowed)),
+        ).items()
+    )
+
+
 def rules():
     return collect_rules()

src/python/pants/engine/internals/platform_rules_test.py

@@ -0,0 +1,42 @@
+# Copyright 2022 Pants project contributors (see CONTRIBUTORS.md).
+# Licensed under the Apache License, Version 2.0 (see LICENSE).
+
+from __future__ import annotations
+
+from textwrap import dedent
+
+from pants.core.util_rules.environments import DockerEnvironmentTarget
+from pants.engine.env_vars import CompleteEnvironmentVars
+from pants.engine.environment import EnvironmentName
+from pants.testutil.rule_runner import QueryRule, RuleRunner
+
+
+def test_docker_complete_env_vars() -> None:
+    rule_runner = RuleRunner(
+        rules=[QueryRule(CompleteEnvironmentVars, [])],
+        target_types=[DockerEnvironmentTarget],
+        singleton_environment=EnvironmentName("docker"),
+    )
+    rule_runner.write_files(
+        {
+            "BUILD": dedent(
+                """\
+                _docker_environment(
+                    name='docker',
+                    image='centos@sha256:a1801b843b1bfaf77c501e7a6d3f709401a1e0c83863037fa3aab063a7fdb9dc',
+                    platform='linux_x86_64',
+                )
+                """
+            )
+        }
+    )
+    rule_runner.set_options(["--environments-preview-names={'docker': '//:docker'}"])
+    result = dict(rule_runner.request(CompleteEnvironmentVars, []))
+
+    # HOSTNAME is not deterministic across machines, so we don't care about the value.
+    assert "HOSTNAME" in result
+    result.pop("HOSTNAME")
+    assert dict(result) == {
+        "PATH": "/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin",
+        "HOME": "/root",
+    }

src/python/pants/init/engine_initializer.py

@@ -16,7 +16,7 @@
 from pants.build_graph.build_configuration import BuildConfiguration
 from pants.core.util_rules import environments, system_binaries
 from pants.core.util_rules.environments import determine_bootstrap_environment
-from pants.engine import desktop, env_vars, fs, process
+from pants.engine import desktop, fs, process
 from pants.engine.console import Console
 from pants.engine.environment import EnvironmentName
 from pants.engine.fs import PathGlobs, Snapshot, Workspace
@@ -262,7 +262,6 @@ def build_root_singleton() -> BuildRoot:
                 *collect_rules(locals()),
                 *build_files.rules(),
                 *fs.rules(),
-                *env_vars.rules(),
                 *desktop.rules(),
                 *git_rules(),
                 *graph.rules(),