Enable env vars in node_build_script (pantsbuild#19100)

## Context
Many javascript and typescript build scripts use environment variables
at build time. Adding this feature to pants would allow users to keep
using this workflow as well as taking advantage of the cache system
whenever specific environment variables change.

## Feature
This PR adds `extra_env_vars` to `node_build_script` just like in other
backends. You can either set a variable to a specific value or the
variable name so pants will use the current environments' exported
variable.
```
package_json(
    scripts=[
        node_build_script(
            entry_point="build", 
            extra_env_vars=["FOO=BAR", "ENVIRONMENT_VAR"]
         )
     ]
)
```

Closes pantsbuild#19068

src/python/pants/backend/javascript/package/rules.py

@@ -16,6 +16,7 @@
     NodeBuildScript,
     NodeBuildScriptEntryPointField,
     NodeBuildScriptExtraCaches,
+    NodeBuildScriptExtraEnvVarsField,
     NodeBuildScriptOutputDirectoriesField,
     NodeBuildScriptOutputFilesField,
     NodeBuildScriptSourcesField,
@@ -32,6 +33,7 @@
     PackageFieldSet,
 )
 from pants.core.target_types import ResourceSourceField
+from pants.engine.env_vars import EnvironmentVars, EnvironmentVarsRequest
 from pants.engine.internals.native_engine import AddPrefix, Digest, Snapshot
 from pants.engine.internals.selectors import Get
 from pants.engine.process import ProcessResult
@@ -66,6 +68,7 @@ class NodeBuildScriptPackageFieldSet(PackageFieldSet):
     output_directories: NodeBuildScriptOutputDirectoriesField
     output_files: NodeBuildScriptOutputFilesField
     extra_caches: NodeBuildScriptExtraCaches
+    extra_env_vars: NodeBuildScriptExtraEnvVarsField
 
 
 @dataclass(frozen=True)
@@ -128,6 +131,7 @@ class NodeBuildScriptRequest:
     output_directories: tuple[str, ...]
     script_name: str
     extra_caches: tuple[str, ...]
+    extra_env_vars: tuple[str, ...]
 
     def __post_init__(self) -> None:
         if not (self.output_directories or self.output_files):
@@ -154,6 +158,7 @@ def from_generate_request(
             or (),
             script_name=req.protocol_target[NodeBuildScriptEntryPointField].value,
             extra_caches=req.protocol_target[NodeBuildScriptExtraCaches].value or (),
+            extra_env_vars=req.protocol_target[NodeBuildScriptExtraEnvVarsField].value or (),
         )
 
     @classmethod
@@ -164,6 +169,7 @@ def from_package_request(cls, req: NodeBuildScriptPackageFieldSet) -> NodeBuildS
             output_directories=req.output_directories.value or (),
             script_name=req.script_name.value,
             extra_caches=req.extra_caches.value or (),
+            extra_env_vars=req.extra_env_vars.value or (),
         )
 
     def get_paths(self) -> Iterable[str]:
@@ -180,12 +186,14 @@ async def run_node_build_script(req: NodeBuildScriptRequest) -> NodeBuildScriptR
     output_dirs = req.output_directories
     script_name = req.script_name
     extra_caches = req.extra_caches
+    extra_env_vars = req.extra_env_vars
 
     def cache_name(cache_path: str) -> str:
         parts = (installation.project_env.package_dir(), script_name, cache_path)
         return "_".join(_NOT_ALPHANUMERIC.sub("_", part) for part in parts if part)
 
     args = ("run", script_name)
+    target_env_vars = await Get(EnvironmentVars, EnvironmentVarsRequest(extra_env_vars))
     result = await Get(
         ProcessResult,
         NodeJsProjectEnvironmentProcess(
@@ -204,6 +212,7 @@ def cache_name(cache_path: str) -> str:
             per_package_caches=FrozenDict(
                 {cache_name(extra_cache): extra_cache for extra_cache in extra_caches or ()}
             ),
+            extra_env=target_env_vars,
         ),
     )
 

src/python/pants/backend/javascript/package/rules_test.py

@@ -234,3 +234,42 @@ def test_packages_files_as_resource_in_workspace(
     rule_runner.write_digest(result.snapshot.digest)
     with open(os.path.join(rule_runner.build_root, "src/js/a/dist/index.cjs")) as f:
         assert f.read() == "blarb\n"
+
+
+def test_extra_envs(rule_runner: RuleRunner) -> None:
+    rule_runner.write_files(
+        {
+            "src/js/BUILD": dedent(
+                """\
+                package_json(
+                    scripts=[
+                        node_build_script(entry_point="build", extra_env_vars=["FOO=BAR"], output_files=["dist/index.cjs"])
+                    ]
+                )
+                """
+            ),
+            "src/js/package.json": json.dumps(
+                {
+                    "name": "ham",
+                    "version": "0.0.1",
+                    "browser": "lib/index.mjs",
+                    "scripts": {"build": "mkdir dist && echo $FOO >> dist/index.cjs"},
+                }
+            ),
+            "src/js/package-lock.json": json.dumps({}),
+            "src/js/lib/BUILD": dedent(
+                """\
+                javascript_sources()
+                """
+            ),
+            "src/js/lib/index.mjs": "",
+        }
+    )
+    tgt = rule_runner.get_target(Address("src/js", generated_name="build"))
+    snapshot = rule_runner.request(Snapshot, (EMPTY_DIGEST,))
+    result = rule_runner.request(
+        GeneratedSources, [GenerateResourcesFromNodeBuildScriptRequest(snapshot, tgt)]
+    )
+    rule_runner.write_digest(result.snapshot.digest)
+    with open(os.path.join(rule_runner.build_root, "src/js/dist/index.cjs")) as f:
+        assert f.read() == "BAR\n"

src/python/pants/backend/javascript/package_json.py

@@ -88,6 +88,7 @@ class NodeBuildScript(NodeScript):
     output_directories: tuple[str, ...] = ()
     output_files: tuple[str, ...] = ()
     extra_caches: tuple[str, ...] = ()
+    extra_env_vars: tuple[str, ...] = ()
 
     alias: ClassVar[str] = "node_build_script"
 
@@ -98,6 +99,7 @@ def create(
         output_directories: Iterable[str] = (),
         output_files: Iterable[str] = (),
         extra_caches: Iterable[str] = (),
+        extra_env_vars: Iterable[str] = (),
     ) -> NodeBuildScript:
         """A build script, mapped from the `scripts` section of a package.json file.
 
@@ -110,6 +112,7 @@ def create(
             output_directories=tuple(output_directories),
             output_files=tuple(output_files),
             extra_caches=tuple(extra_caches),
+            extra_env_vars=tuple(extra_env_vars),
         )
 
 
@@ -410,6 +413,20 @@ class NodeBuildScriptOutputDirectoriesField(StringSequenceField):
     )
 
 
+class NodeBuildScriptExtraEnvVarsField(StringSequenceField):
+    alias = "extra_env_vars"
+    required = False
+    default = ()
+    help = help_text(
+        """
+        Additional environment variables to include in environment when running a build script process.
+
+        Entries are strings in the form `ENV_VAR=value` to use explicitly; or just
+        `ENV_VAR` to copy the value of a variable in Pants's own environment.
+        """
+    )
+
+
 class NodeBuildScriptExtraCaches(StringSequenceField):
     alias = "extra_caches"
     required = False
@@ -452,6 +469,7 @@ class NodeBuildScriptTarget(Target):
         NodeBuildScriptOutputFilesField,
         NodeBuildScriptSourcesField,
         NodeBuildScriptExtraCaches,
+        NodeBuildScriptExtraEnvVarsField,
         NodePackageDependenciesField,
         OutputPathField,
     )
@@ -923,6 +941,7 @@ async def generate_node_package_targets(
                         NodeBuildScriptEntryPointField.alias: build_script.entry_point,
                         NodeBuildScriptOutputDirectoriesField.alias: build_script.output_directories,
                         NodeBuildScriptOutputFilesField.alias: build_script.output_files,
+                        NodeBuildScriptExtraEnvVarsField.alias: build_script.extra_env_vars,
                         NodeBuildScriptExtraCaches.alias: build_script.extra_caches,
                         NodePackageDependenciesField.alias: [
                             file_tgt.address.spec,

src/python/pants/backend/javascript/run/rules.py

@@ -13,10 +13,12 @@
 from pants.backend.javascript.nodejs_project_environment import NodeJsProjectEnvironmentProcess
 from pants.backend.javascript.package_json import (
     NodeBuildScriptEntryPointField,
+    NodeBuildScriptExtraEnvVarsField,
     NodePackageDependenciesField,
 )
 from pants.core.goals.run import RunFieldSet, RunInSandboxBehavior, RunRequest
 from pants.core.util_rules.environments import EnvironmentField
+from pants.engine.env_vars import EnvironmentVars, EnvironmentVarsRequest
 from pants.engine.internals.selectors import Get
 from pants.engine.process import Process
 from pants.engine.rules import Rule, collect_rules, rule
@@ -29,6 +31,7 @@ class RunNodeBuildScriptFieldSet(RunFieldSet):
     run_in_sandbox_behavior = RunInSandboxBehavior.RUN_REQUEST_HERMETIC
 
     entry_point: NodeBuildScriptEntryPointField
+    extra_env_vars: NodeBuildScriptExtraEnvVarsField
     environment: EnvironmentField
 
 
@@ -39,13 +42,17 @@ async def run_node_build_script(
     installation = await Get(
         InstalledNodePackageWithSource, InstalledNodePackageRequest(field_set.address)
     )
+    target_env_vars = await Get(
+        EnvironmentVars, EnvironmentVarsRequest(field_set.extra_env_vars.value or ())
+    )
     process = await Get(
         Process,
         NodeJsProjectEnvironmentProcess(
             installation.project_env,
             args=("--prefix", "{chroot}", "run", str(field_set.entry_point.value)),
             description=f"Running {str(field_set.entry_point.value)}.",
             input_digest=installation.digest,
+            extra_env=target_env_vars,
         ),
     )
 

src/python/pants/backend/javascript/run/rules_test.py

@@ -75,3 +75,38 @@ def test_creates_run_requests_package_json_scripts(rule_runner: RuleRunner) -> N
         result = rule_runner.request(RunRequest, [RunNodeBuildScriptFieldSet.create(tgt)])
 
         assert result.args == ("npm", "--prefix", "{chroot}", "run", script)
+
+
+def test_extra_envs(rule_runner: RuleRunner) -> None:
+    rule_runner.write_files(
+        {
+            "src/js/BUILD": dedent(
+                """\
+                package_json(
+                    scripts=[
+                        node_build_script(entry_point="build", extra_env_vars=["FOO=BAR"], output_files=["dist/index.cjs"])
+                    ]
+                )
+                """
+            ),
+            "src/js/package.json": json.dumps(
+                {
+                    "name": "ham",
+                    "version": "0.0.1",
+                    "browser": "lib/index.mjs",
+                    "scripts": {"build": "mkdir dist && echo $FOO >> dist/index.cjs"},
+                }
+            ),
+            "src/js/package-lock.json": json.dumps({}),
+            "src/js/lib/BUILD": dedent(
+                """\
+                javascript_sources()
+                """
+            ),
+            "src/js/lib/index.mjs": "",
+        }
+    )
+    tgt = rule_runner.get_target(Address("src/js", generated_name="build"))
+
+    result = rule_runner.request(RunRequest, [RunNodeBuildScriptFieldSet.create(tgt)])
+    assert result.extra_env.get("FOO") == "BAR"