From 783c04a3e0cf9f3ade12664c169a107154e606ce Mon Sep 17 00:00:00 2001
From: Benjy Weinberger <benjyw@gmail.com>
Date: Mon, 17 Oct 2022 14:03:03 -0700
Subject: [PATCH] Set PEX_ROOT for export post-processing command. (cherry-pick
 of #17177) (#17200)

This is both a leak that foils our recommendations to users for maintenance of
Pants caches and results in unneeded extra work when an export finds an
empty ~/.pex but would have had hits in ~/.cache/pants/named_caches/pex_root
if it had been pointed there.

A reapplication of part of #17177, but not really a straight cherry-pick, as it would not
apply even remotely cleanly.
---
 .../pants/backend/python/goals/export.py      | 20 ++++++++++++++-----
 .../pants/backend/python/goals/export_test.py |  3 ++-
 2 files changed, 17 insertions(+), 6 deletions(-)

diff --git a/src/python/pants/backend/python/goals/export.py b/src/python/pants/backend/python/goals/export.py
index ef488028c30..d1fc3d8fc4b 100644
--- a/src/python/pants/backend/python/goals/export.py
+++ b/src/python/pants/backend/python/goals/export.py
@@ -14,7 +14,7 @@
 from pants.backend.python.util_rules.interpreter_constraints import InterpreterConstraints
 from pants.backend.python.util_rules.pex import Pex, PexRequest
 from pants.backend.python.util_rules.pex_cli import PexPEX
-from pants.backend.python.util_rules.pex_environment import PythonExecutable
+from pants.backend.python.util_rules.pex_environment import PexEnvironment, PythonExecutable
 from pants.backend.python.util_rules.pex_from_targets import RequirementsPexRequest
 from pants.core.goals.export import (
     ExportError,
@@ -79,7 +79,7 @@ def debug_hint(self) -> str | None:
 
 @rule
 async def export_virtualenv(
-    request: _ExportVenvRequest, python_setup: PythonSetup, pex_pex: PexPEX
+    request: _ExportVenvRequest, python_setup: PythonSetup, pex_pex: PexPEX, pex_env: PexEnvironment
 ) -> ExportResult:
     if request.resolve:
         interpreter_constraints = InterpreterConstraints(
@@ -127,6 +127,7 @@ async def export_virtualenv(
     merged_digest = await Get(Digest, MergeDigests([pex_pex.digest, requirements_pex.digest]))
     pex_pex_path = os.path.join("{digest_root}", pex_pex.exe)
     maybe_resolve_str = f"for the resolve '{request.resolve}' " if request.resolve else ""
+    complete_pex_env = pex_env.in_workspace()
     return ExportResult(
         f"virtualenv {maybe_resolve_str}(using Python {py_version})",
         dest,
@@ -143,7 +144,10 @@ async def export_virtualenv(
                     "--remove=all",
                     f"{{digest_root}}/{py_version}",
                 ],
-                {"PEX_MODULE": "pex.tools"},
+                {
+                    **complete_pex_env.environment_dict(python_configured=True),
+                    "PEX_MODULE": "pex.tools",
+                },
             ),
             PostProcessingCommand(["rm", "-f", pex_pex_path]),
         ],
@@ -151,7 +155,9 @@ async def export_virtualenv(
 
 
 @rule
-async def export_tool(request: ExportPythonTool, pex_pex: PexPEX) -> ExportResult:
+async def export_tool(
+    request: ExportPythonTool, pex_pex: PexPEX, pex_env: PexEnvironment
+) -> ExportResult:
     assert request.pex_request is not None
 
     # TODO: Unify export_virtualenv() and export_tool(), since their implementations mostly overlap.
@@ -173,6 +179,7 @@ async def export_tool(request: ExportPythonTool, pex_pex: PexPEX) -> ExportResul
     pex_pex_digest = await Get(Digest, AddPrefix(pex_pex.digest, pex_pex_dir))
 
     merged_digest = await Get(Digest, MergeDigests([pex_pex_digest, pex.digest]))
+    complete_pex_env = pex_env.in_workspace()
     return ExportResult(
         f"virtualenv for the tool '{request.resolve_name}'",
         dest,
@@ -188,7 +195,10 @@ async def export_tool(request: ExportPythonTool, pex_pex: PexPEX) -> ExportResul
                     "--remove=all",
                     f"{{digest_root}}/{request.resolve_name}",
                 ],
-                {"PEX_MODULE": "pex.tools"},
+                {
+                    **complete_pex_env.environment_dict(python_configured=True),
+                    "PEX_MODULE": "pex.tools",
+                },
             ),
             PostProcessingCommand(["rm", "-rf", pex_pex_dest]),
         ],
diff --git a/src/python/pants/backend/python/goals/export_test.py b/src/python/pants/backend/python/goals/export_test.py
index 09ab7cb3cc0..ef32db11a25 100644
--- a/src/python/pants/backend/python/goals/export_test.py
+++ b/src/python/pants/backend/python/goals/export_test.py
@@ -87,7 +87,8 @@ def run(enable_resolves: bool) -> ExportResults:
                 "--remove=all",
                 f"{{digest_root}}/{current_interpreter}",
             )
-            assert ppc0.extra_env == FrozenDict({"PEX_MODULE": "pex.tools"})
+            assert ppc0.extra_env["PEX_MODULE"] == "pex.tools"
+            assert "PEX_ROOT" in ppc0.extra_env
 
             ppc1 = result.post_processing_cmds[1]
             assert ppc1.argv == (