Updates publish to USES_ENVIRONMENTS (#17313)

This migrates `publish` to create the `BuiltPackage` asset in a correct environment, and run the `publish` executable in the local environment.

Also adds a helper rule, `EnvironmentAwarePackageRequest`, which consolidates the `EnvironmentName` request and `BuiltPackage` request, since `BuiltPackage` is requested all over the place.

See #17129.

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

@@ -91,6 +91,26 @@ def value_or_default(self, *, file_ending: str | None) -> str:
         return os.path.join(self.address.spec_path.replace(os.sep, "."), file_name)
 
 
+@dataclass(frozen=True)
+class EnvironmentAwarePackageRequest:
+    """Request class to request a `BuiltPackage` in an environment-aware fashion."""
+
+    field_set: PackageFieldSet
+
+
+@rule
+async def environment_aware_package(request: EnvironmentAwarePackageRequest) -> BuiltPackage:
+    environment_name = await Get(
+        EnvironmentName,
+        EnvironmentNameRequest,
+        EnvironmentNameRequest.from_field_set(request.field_set),
+    )
+    package = await Get(
+        BuiltPackage, {request.field_set: PackageFieldSet, environment_name: EnvironmentName}
+    )
+    return package
+
+
 class PackageSubsystem(GoalSubsystem):
     name = "package"
     help = "Create a distributable package."
@@ -134,19 +154,9 @@ async def package_asset(workspace: Workspace, dist_dir: DistDir) -> Package:
     if not target_roots_to_field_sets.field_sets:
         return Package(exit_code=0)
 
-    environment_names_per_field_set = await MultiGet(
-        Get(
-            EnvironmentName,
-            EnvironmentNameRequest,
-            EnvironmentNameRequest.from_field_set(field_set),
-        )
-        for field_set in target_roots_to_field_sets.field_sets
-    )
     packages = await MultiGet(
-        Get(BuiltPackage, {field_set: PackageFieldSet, environment_name: EnvironmentName})
-        for field_set, environment_name in zip(
-            target_roots_to_field_sets.field_sets, environment_names_per_field_set
-        )
+        Get(BuiltPackage, EnvironmentAwarePackageRequest(field_set))
+        for field_set in target_roots_to_field_sets.field_sets
     )
 
     merged_digest = await Get(Digest, MergeDigests(pkg.digest for pkg in packages))

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

@@ -29,11 +29,11 @@ async def publish_example(request: PublishToMyRepoRequest, ...) -> PublishProces
 
 from typing_extensions import final
 
-from pants.core.goals.package import BuiltPackage, PackageFieldSet
+from pants.core.goals.package import BuiltPackage, EnvironmentAwarePackageRequest, PackageFieldSet
 from pants.engine.addresses import Address
 from pants.engine.collection import Collection
 from pants.engine.console import Console
-from pants.engine.environment import EnvironmentName
+from pants.engine.environment import ChosenLocalEnvironmentName, EnvironmentName
 from pants.engine.goal import Goal, GoalSubsystem
 from pants.engine.process import InteractiveProcess, InteractiveProcessResult
 from pants.engine.rules import Effect, Get, MultiGet, collect_rules, goal_rule, rule
@@ -180,11 +180,13 @@ def activated(cls, union_membership: UnionMembership) -> bool:
 
 class Publish(Goal):
     subsystem_cls = PublishSubsystem
-    environment_behavior = Goal.EnvironmentBehavior.LOCAL_ONLY  # TODO(#17129) — Migrate this.
+    environment_behavior = Goal.EnvironmentBehavior.USES_ENVIRONMENTS
 
 
 @goal_rule
-async def run_publish(console: Console, publish: PublishSubsystem) -> Publish:
+async def run_publish(
+    console: Console, publish: PublishSubsystem, local_environment: ChosenLocalEnvironmentName
+) -> Publish:
     target_roots_to_package_field_sets, target_roots_to_publish_field_sets = await MultiGet(
         Get(
             TargetRootsToFieldSets,
@@ -242,7 +244,10 @@ async def run_publish(console: Console, publish: PublishSubsystem) -> Publish:
             continue
 
         logger.debug(f"Execute {pub.process}")
-        res = await Effect(InteractiveProcessResult, InteractiveProcess, pub.process)
+        res = await Effect(
+            InteractiveProcessResult,
+            {pub.process: InteractiveProcess, local_environment.val: EnvironmentName},
+        )
         if res.exit_code == 0:
             sigil = console.sigil_succeeded()
             status = "published"
@@ -305,9 +310,12 @@ def default(self, o):
 
 
 @rule
-async def package_for_publish(request: PublishProcessesRequest) -> PublishProcesses:
+async def package_for_publish(
+    request: PublishProcessesRequest, local_environment: ChosenLocalEnvironmentName
+) -> PublishProcesses:
     packages = await MultiGet(
-        Get(BuiltPackage, PackageFieldSet, field_set) for field_set in request.package_field_sets
+        Get(BuiltPackage, EnvironmentAwarePackageRequest(field_set))
+        for field_set in request.package_field_sets
     )
 
     for pkg in packages:
@@ -320,8 +328,10 @@ async def package_for_publish(request: PublishProcessesRequest) -> PublishProces
     publish = await MultiGet(
         Get(
             PublishProcesses,
-            PublishRequest,
-            field_set._request(packages),
+            {
+                field_set._request(packages): PublishRequest,
+                local_environment.val: EnvironmentName,
+            },
         )
         for field_set in request.publish_field_sets
     )