Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Let repo rules remove env vars for subprocesses #24221

Closed
wants to merge 1 commit into from
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
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"
Loading