Skip to content

Commit

Permalink
Let repo rules remove env vars for subprocesses
Browse files Browse the repository at this point in the history
RELNOTES: `repository_ctx.execute` can now remove an environment variable when executing a process by associating it with the value `None` in the `environment` argument.
  • Loading branch information
fmeum committed Nov 6, 2024
1 parent 2fd6115 commit c4e87af
Show file tree
Hide file tree
Showing 4 changed files with 75 additions and 5 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -84,9 +84,12 @@
import java.util.Base64;
import java.util.HashMap;
import java.util.HashSet;
import java.util.LinkedHashMap;
import java.util.LinkedHashSet;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.Set;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
Expand Down Expand Up @@ -1835,7 +1838,11 @@ private static String profileArgsDesc(String method, List<String> args) {
name = "environment",
defaultValue = "{}",
named = true,
doc = "Force some environment variables to be set to be passed to the process."),
doc =
"""
Force some environment variables to be set to be passed to the process. The value \
can be <code>None</code> to remove the environment variable.
"""),
@Param(
name = "quiet",
defaultValue = "True",
Expand All @@ -1855,18 +1862,33 @@ private static String profileArgsDesc(String method, List<String> args) {
public StarlarkExecutionResult execute(
Sequence<?> arguments, // <String> or <StarlarkPath> or <Label> expected
StarlarkInt timeoutI,
Dict<?, ?> uncheckedEnvironment, // <String, String> expected
Dict<?, ?> uncheckedEnvironment, // <String, Object> expected
boolean quiet,
String overrideWorkingDirectory,
StarlarkThread thread)
throws EvalException, RepositoryFunctionException, InterruptedException {
validateExecuteArguments(arguments);
int timeout = Starlark.toInt(timeoutI, "timeout");

Map<String, String> forceEnvVariables =
Dict.cast(uncheckedEnvironment, String.class, String.class, "environment");
Map<String, Object> forceEnvVariablesRaw =
Dict.cast(uncheckedEnvironment, String.class, Object.class, "environment");
Map<String, String> forceEnvVariables = new LinkedHashMap<>();
Set<String> removeEnvVariables = new LinkedHashSet<>();
for (Map.Entry<String, Object> entry : forceEnvVariablesRaw.entrySet()) {
String key = entry.getKey();
Object value = entry.getValue();
if (value == Starlark.NONE) {
removeEnvVariables.add(key);
} else if (value instanceof String s) {
forceEnvVariables.put(key, s);
} else {
throw Starlark.errorf("environment values must be strings or None, got %s", value);
}
}

if (canExecuteRemote()) {
// Remote execution only sees the explicitly set environment variables, so removing env vars
// isn't necessary.
return executeRemote(arguments, timeout, forceEnvVariables, quiet, overrideWorkingDirectory);
}

Expand All @@ -1886,7 +1908,7 @@ public StarlarkExecutionResult execute(
WorkspaceRuleEvent.newExecuteEvent(
args,
timeout,
envVariables,
Maps.filterKeys(envVariables, k -> !removeEnvVariables.contains(k)),
forceEnvVariables,
workingDirectory.getPathString(),
quiet,
Expand Down Expand Up @@ -1920,6 +1942,7 @@ public StarlarkExecutionResult execute(
.addArguments(args)
.setDirectory(workingDirectoryPath.getPathFile())
.addEnvironmentVariables(forceEnvVariables)
.removeEnvironmentVariables(removeEnvVariables)
.setTimeout(timeoutMillis)
.setQuiet(quiet)
.execute();
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,7 @@
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.Set;
import net.starlark.java.annot.StarlarkBuiltin;
import net.starlark.java.annot.StarlarkMethod;
import net.starlark.java.eval.EvalException;
Expand Down Expand Up @@ -158,6 +159,13 @@ Builder addEnvironmentVariables(Map<String, String> variables) {
return this;
}

/** Ensure that an environment variable is not passed to the process. */
@CanIgnoreReturnValue
Builder removeEnvironmentVariables(Set<String> removeEnvVariables) {
removeEnvVariables.forEach(envBuilder::remove);
return this;
}

/** Sets the timeout, in milliseconds, after which the executed command will be terminated. */
@CanIgnoreReturnValue
Builder setTimeout(long timeout) {
Expand Down
1 change: 1 addition & 0 deletions src/test/shell/bazel/BUILD
Original file line number Diff line number Diff line change
Expand Up @@ -874,6 +874,7 @@ sh_test(
data = [
":test-deps",
"@bazel_tools//tools/bash/runfiles",
"@local_jdk//:jdk",
],
shard_count = 10,
tags = [
Expand Down
38 changes: 38 additions & 0 deletions src/test/shell/bazel/starlark_repository_test.sh
Original file line number Diff line number Diff line change
Expand Up @@ -3449,4 +3449,42 @@ EOF
expect_log "my_second_repo: @my_first_repo//:foo @@+_repo_rules+my_first_repo//:foo"
}

function test_execute_environment_remove_vars() {
cat >> $(setup_module_dot_bazel) <<'EOF'
my_repo = use_repo_rule("//:repo.bzl", "my_repo")
my_repo(name="repo")
EOF
touch BUILD
cat > repo.bzl <<'EOF'
def _impl(ctx):
st = ctx.execute(
["env"],
environment = {
"CLIENT_ENV_REMOVED": None,
"REPO_ENV_REMOVED": None,
},
)
if st.return_code:
fail("Command did not succeed")
vars = {line.partition("=")[0]: line.partition("=")[-1] for line in st.stdout.split("\n")}
if "CLIENT_ENV_REMOVED" in vars:
fail("CLIENT_ENV_REMOVED should not be in the environment")
if "REPO_ENV_REMOVED" in vars:
fail("REPO_ENV_REMOVED should not be in the environment")
if vars.get("CLIENT_ENV_PRESENT") != "value1":
fail("CLIENT_ENV_PRESENT has wrong value: " + vars.get("CLIENT_ENV_PRESENT"))
if vars.get("REPO_ENV_PRESENT") != "value3":
fail("REPO_ENV_PRESENT has wrong value: " + vars.get("REPO_ENV_PRESENT"))
ctx.file("BUILD", "exports_files(['data.txt'])")
my_repo = repository_rule(_impl)
EOF

CLIENT_ENV_PRESENT=value1 CLIENT_ENV_REMOVED=value2 \
bazel build \
--repo_env=REPO_ENV_PRESENT=value3 --repo_env=REPO_ENV_REMOVED=value4 \
@repo//... &> $TEST_log || fail "expected Bazel to succeed"
}

run_suite "local repository tests"

0 comments on commit c4e87af

Please sign in to comment.