Migrates generate-lockfiles to USES_ENVIRONMENTS.

src/python/pants/core/goals/generate_lockfiles.py

@@ -11,7 +11,7 @@
 from typing import Callable, ClassVar, Iterable, Sequence
 
 from pants.engine.collection import Collection
-from pants.engine.environment import EnvironmentName
+from pants.engine.environment import ChosenLocalEnvironmentName, EnvironmentName
 from pants.engine.fs import Digest, MergeDigests, Workspace
 from pants.engine.goal import Goal, GoalSubsystem
 from pants.engine.internals.selectors import Get, MultiGet
@@ -51,6 +51,14 @@ class GenerateLockfile:
     lockfile_dest: str
 
 
+@dataclass(frozen=True)
+class GenerateLockfileWithEnvironments(GenerateLockfile):
+    """Allows a `GenerateLockfile` subclass to specify which environments the request is compatible
+    with, if the relevant backend supports environments."""
+
+    environments: tuple[EnvironmentName, ...]
+
+
 @dataclass(frozen=True)
 class WrappedGenerateLockfile:
     request: GenerateLockfile
@@ -359,14 +367,15 @@ def activated(cls, union_membership: UnionMembership) -> bool:
 
 class GenerateLockfilesGoal(Goal):
     subsystem_cls = GenerateLockfilesSubsystem
-    environment_behavior = Goal.EnvironmentBehavior.LOCAL_ONLY  # TODO(#17129) — Migrate this.
+    environment_behavior = Goal.EnvironmentBehavior.USES_ENVIRONMENTS
 
 
 @goal_rule
 async def generate_lockfiles_goal(
     workspace: Workspace,
     union_membership: UnionMembership,
     generate_lockfiles_subsystem: GenerateLockfilesSubsystem,
+    local_environment: ChosenLocalEnvironmentName,
 ) -> GenerateLockfilesGoal:
     known_user_resolve_names = await MultiGet(
         Get(KnownUserResolveNames, KnownUserResolveNamesRequest, request())
@@ -378,27 +387,45 @@ async def generate_lockfiles_goal(
         set(generate_lockfiles_subsystem.resolve),
     )
 
+    # This is the "planning" phase of lockfile generation. Currently this is all done in the local
+    # environment, since there's not currently a clear mechanism to prescribe an environment.
     all_specified_user_requests = await MultiGet(
-        Get(UserGenerateLockfiles, RequestedUserResolveNames, resolve_names)
+        Get(
+            UserGenerateLockfiles,
+            {resolve_names: RequestedUserResolveNames, local_environment.val: EnvironmentName},
+        )
         for resolve_names in requested_user_resolve_names
     )
     specified_tool_requests = await MultiGet(
-        Get(WrappedGenerateLockfile, GenerateToolLockfileSentinel, sentinel())
+        Get(
+            WrappedGenerateLockfile,
+            {sentinel(): GenerateToolLockfileSentinel, local_environment.val: EnvironmentName},
+        )
         for sentinel in requested_tool_sentinels
     )
     applicable_tool_requests = filter_tool_lockfile_requests(
         specified_tool_requests,
         resolve_specified=bool(generate_lockfiles_subsystem.resolve),
     )
 
+    # Execute the actual lockfile generation in each request's environment.
+    # Currently, since resolves specify a single filename for output, we pick a resonable
+    # environment to execute the request in. Currently we warn if multiple environments are
+    # specified.
+    all_requests = itertools.chain(*all_specified_user_requests, applicable_tool_requests)
     results = await MultiGet(
-        Get(GenerateLockfileResult, GenerateLockfile, req)
-        for req in (
-            *(req for reqs in all_specified_user_requests for req in reqs),
-            *applicable_tool_requests,
+        Get(
+            GenerateLockfileResult,
+            {
+                req: GenerateLockfile,
+                _preferred_environment(req, local_environment.val): EnvironmentName,
+            },
         )
+        for req in all_requests
     )
 
+    # Lockfiles are actually written here. This would be an acceptable place to handle conflict
+    # resolution behaviour if we start executing requests in multiple environments.
     merged_digest = await Get(Digest, MergeDigests(res.digest for res in results))
     workspace.write_digest(merged_digest)
     for result in results:
@@ -407,6 +434,25 @@ async def generate_lockfiles_goal(
     return GenerateLockfilesGoal(exit_code=0)
 
 
+def _preferred_environment(request: GenerateLockfile, default: EnvironmentName) -> EnvironmentName:
+
+    if not isinstance(request, GenerateLockfileWithEnvironments):
+        return default  # This request has not been migrated to use environments.
+
+    if len(request.environments) == 1:
+        return request.environments[0]
+
+    ret = default if default in request.environments else request.environments[0]
+
+    logger.warning(
+        f"The `{request.__class__.__name__}` for resolve `{request.resolve_name}` specifies more "
+        "than one environment. Pants will generate the lockfile using only the environment "
+        f"`{ret.val}`, which may have unintended effects when executing in the other environments."
+    )
+
+    return ret
+
+
 # -----------------------------------------------------------------------------------------------
 # Helpers for determining the resolve
 # -----------------------------------------------------------------------------------------------

src/python/pants/core/goals/generate_lockfiles_test.py

@@ -3,6 +3,8 @@
 
 from __future__ import annotations
 
+from typing import Iterable
+
 import pytest
 
 from pants.backend.python.subsystems.setup import PythonSetup
@@ -19,15 +21,18 @@
     NO_TOOL_LOCKFILE,
     AmbiguousResolveNamesError,
     GenerateLockfile,
+    GenerateLockfileWithEnvironments,
     GenerateToolLockfileSentinel,
     KnownUserResolveNames,
     NoCompatibleResolveException,
     RequestedUserResolveNames,
     UnrecognizedResolveNamesError,
     WrappedGenerateLockfile,
+    _preferred_environment,
     determine_resolves_to_generate,
     filter_tool_lockfile_requests,
 )
+from pants.engine.environment import EnvironmentName
 from pants.engine.target import Dependencies, Target
 from pants.testutil.option_util import create_subsystem
 from pants.util.strutil import softwrap
@@ -266,3 +271,41 @@ def maybe_get_resolve(t: Target) -> str | None:
             """
         )
     )
+
+
+_default = "_default"
+
+
+@pytest.mark.parametrize(
+    ("env_names", "expected", "in_output", "not_in_output"),
+    (
+        ((), _default, None, "anything"),
+        (("jeremy", "derek"), "jeremy", "`jeremy`, which may have ", None),
+        (("jeremy", "_default"), "_default", "`_default`, which may have ", None),
+        (("_default", "jeremy"), "_default", "`_default`, which may have ", None),
+        (("jeremy",), "jeremy", None, "anything"),
+    ),
+)
+def test_preferred_environment(
+    env_names: Iterable[str],
+    expected: str,
+    in_output: str | None,
+    not_in_output: str | None,
+    caplog,
+):
+
+    resolve_name = "boop"
+    resolve_dest = "beep"
+    if not env_names:
+        request = GenerateLockfile(resolve_name, resolve_dest)
+    else:
+        envs = tuple(EnvironmentName(name) for name in env_names)
+        request = GenerateLockfileWithEnvironments(resolve_name, resolve_dest, envs)
+
+    default = EnvironmentName(_default)
+
+    preferred_env = _preferred_environment(request, default)
+
+    assert preferred_env.val == expected
+    assert in_output is None or in_output in caplog.text
+    assert not_in_output is None or not_in_output not in caplog.text