Add an env attribute to all test and binary rule classes

Fixes bazelbuild#7364.

PiperOrigin-RevId: 345355968

src/main/java/com/google/devtools/build/docgen/PredefinedAttributes.java

@@ -27,16 +27,19 @@
 public class PredefinedAttributes {
 
  /**
- * List of documentation for common attributes of *_test rules, relative to
- * {@link com.google.devtools.build.docgen}.
+ * List of documentation for common attributes of *_test rules, relative to {@link
+ * com.google.devtools.build.docgen}.
  */
- public static final ImmutableList<String> TEST_ATTRIBUTES_DOCFILES = ImmutableList.of(
- "templates/attributes/test/args.html",
- "templates/attributes/test/size.html",
- "templates/attributes/test/timeout.html",
- "templates/attributes/test/flaky.html",
- "templates/attributes/test/shard_count.html",
- "templates/attributes/test/local.html");
+ public static final ImmutableList<String> TEST_ATTRIBUTES_DOCFILES =
+ ImmutableList.of(
+ "templates/attributes/test/args.html",
+ "templates/attributes/test/env.html",
+ "templates/attributes/test/env_inherit.html",
+ "templates/attributes/test/size.html",
+ "templates/attributes/test/timeout.html",
+ "templates/attributes/test/flaky.html",
+ "templates/attributes/test/shard_count.html",
+ "templates/attributes/test/local.html");
 
  /**
  * List of common attributes documentation, relative to {@link com.google.devtools.build.docgen}.
@@ -60,12 +63,13 @@ public class PredefinedAttributes {
  "templates/attributes/common/visibility.html");
 
  /**
- * List of documentation for common attributes of *_binary rules, relative to
- * {@link com.google.devtools.build.docgen}.
+ * List of documentation for common attributes of *_binary rules, relative to {@link
+ * com.google.devtools.build.docgen}.
  */
  public static final ImmutableList<String> BINARY_ATTRIBUTES_DOCFILES =
  ImmutableList.of(
  "templates/attributes/binary/args.html",
+ "templates/attributes/binary/env.html",
  "templates/attributes/binary/output_licenses.html");
 
  private static ImmutableMap<String, RuleDocumentationAttribute> generateAttributeMap(

src/main/java/com/google/devtools/build/docgen/templates/attributes/binary/env.html

@@ -0,0 +1,13 @@
+<p><code>Dictionary of strings; optional; values are subject to
+<a href="${link make-variables#location}">$(location)</a> and
+<a href="${link make-variables}">"Make variable"</a> substitution</code></p>
+
+<p>Specifies additional environment variables to set when the target is
+ executed by <code>bazel run</code>.
+</p>
+
+<p>
+<em class="harmful">NOTE: The environment variables are not set when you run the target
+outside of bazel (for example, by manually executing the binary in
+<code>bazel-bin/</code>).</em>
+</p>

src/main/java/com/google/devtools/build/docgen/templates/attributes/test/env.html

@@ -0,0 +1,8 @@
+<p><code>Dictionary of strings; optional; values are subject to
+<a href="${link make-variables#location}">$(location)</a> and
+<a href="${link make-variables}">"Make variable"</a> substitution</code>
+</p>
+
+<p>Specifies additional environment variables to set when the test is
+ executed by <code>bazel test</code>.
+</p>

src/main/java/com/google/devtools/build/docgen/templates/attributes/test/env_inherit.html

@@ -0,0 +1,5 @@
+<p><code>List of strings; optional</p>
+
+<p>Specifies additional environment variables to inherit from the
+ external environment when the test is executed by <code>bazel test</code>.
+</p>

src/main/java/com/google/devtools/build/lib/analysis/BaseRuleClasses.java

@@ -194,6 +194,8 @@ public Object getDefault(AttributeMap rule) {
  .taggable()
  .nonconfigurable("policy decision: should be consistent across configurations"))
  .add(attr("args", STRING_LIST))
+ .add(attr("env", STRING_DICT))
+ .add(attr("env_inherit", STRING_LIST))
  // Input files for every test action
  .add(
  attr("$test_wrapper", LABEL)
@@ -475,6 +477,7 @@ public static final class BinaryBaseRule implements RuleDefinition {
  public RuleClass build(RuleClass.Builder builder, RuleDefinitionEnvironment env) {
  return builder
  .add(attr("args", STRING_LIST))
+ .add(attr("env", STRING_DICT))
  .add(attr("output_licenses", LICENSE))
  .add(
  attr("$is_executable", BOOLEAN)

src/main/java/com/google/devtools/build/lib/analysis/RunfilesSupport.java

@@ -16,6 +16,8 @@
 
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Preconditions;
+import com.google.common.collect.ImmutableSet;
+import com.google.devtools.build.lib.actions.ActionEnvironment;
 import com.google.devtools.build.lib.actions.Artifact;
 import com.google.devtools.build.lib.actions.CommandLine;
 import com.google.devtools.build.lib.analysis.SourceManifestAction.ManifestType;
@@ -34,9 +36,11 @@
 import com.google.devtools.build.lib.vfs.Path;
 import com.google.devtools.build.lib.vfs.PathFragment;
 import java.util.Collection;
+import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.TreeMap;
 import javax.annotation.Nullable;
 
 /**
@@ -85,6 +89,7 @@ public final class RunfilesSupport {
  private final boolean buildRunfileLinks;
  private final boolean runfilesEnabled;
  private final CommandLine args;
+ private final ActionEnvironment actionEnvironment;
 
  /**
  * Creates the RunfilesSupport helper with the given executable and runfiles.
@@ -94,7 +99,11 @@ public final class RunfilesSupport {
  * @param runfiles the runfiles
  */
  private static RunfilesSupport create(
- RuleContext ruleContext, Artifact executable, Runfiles runfiles, CommandLine args) {
+ RuleContext ruleContext,
+ Artifact executable,
+ Runfiles runfiles,
+ CommandLine args,
+ ActionEnvironment actionEnvironment) {
  Artifact owningExecutable = Preconditions.checkNotNull(executable);
  boolean createManifest = ruleContext.getConfiguration().buildRunfilesManifests();
  boolean buildRunfileLinks = ruleContext.getConfiguration().buildRunfileLinks();
@@ -139,7 +148,8 @@ private static RunfilesSupport create(
  owningExecutable,
  buildRunfileLinks,
  runfilesEnabled,
- args);
+ args,
+ actionEnvironment);
  }
 
  @AutoCodec.Instantiator
@@ -152,7 +162,8 @@ private static RunfilesSupport create(
  Artifact owningExecutable,
  boolean buildRunfileLinks,
  boolean runfilesEnabled,
- CommandLine args) {
+ CommandLine args,
+ ActionEnvironment actionEnvironment) {
  this.runfiles = runfiles;
  this.runfilesInputManifest = runfilesInputManifest;
  this.runfilesManifest = runfilesManifest;
@@ -161,6 +172,7 @@ private static RunfilesSupport create(
  this.buildRunfileLinks = buildRunfileLinks;
  this.runfilesEnabled = runfilesEnabled;
  this.args = args;
+ this.actionEnvironment = actionEnvironment;
  }
 
  /** Returns the executable owning this RunfilesSupport. Only use from Starlark. */
@@ -394,14 +406,23 @@ public CommandLine getArgs() {
  return args;
  }
 
+ /** Returns the immutable environment from the 'env' and 'env_inherit' attribute values. */
+ public ActionEnvironment getActionEnvironment() {
+ return actionEnvironment;
+ }
+
  /**
  * Creates and returns a {@link RunfilesSupport} object for the given rule and executable. Note
  * that this method calls back into the passed in rule to obtain the runfiles.
  */
  public static RunfilesSupport withExecutable(
  RuleContext ruleContext, Runfiles runfiles, Artifact executable) {
  return RunfilesSupport.create(
- ruleContext, executable, runfiles, computeArgs(ruleContext, CommandLine.EMPTY));
+ ruleContext,
+ executable,
+ runfiles,
+ computeArgs(ruleContext, CommandLine.EMPTY),
+ computeActionEnvironment(ruleContext));
  }
 
  /**
@@ -411,7 +432,11 @@ public static RunfilesSupport withExecutable(
  public static RunfilesSupport withExecutable(
  RuleContext ruleContext, Runfiles runfiles, Artifact executable, List<String> appendingArgs) {
  return RunfilesSupport.create(
- ruleContext, executable, runfiles, computeArgs(ruleContext, CommandLine.of(appendingArgs)));
+ ruleContext,
+ executable,
+ runfiles,
+ computeArgs(ruleContext, CommandLine.of(appendingArgs)),
+ computeActionEnvironment(ruleContext));
  }
 
  /**
@@ -421,7 +446,11 @@ public static RunfilesSupport withExecutable(
  public static RunfilesSupport withExecutable(
  RuleContext ruleContext, Runfiles runfiles, Artifact executable, CommandLine appendingArgs) {
  return RunfilesSupport.create(
- ruleContext, executable, runfiles, computeArgs(ruleContext, appendingArgs));
+ ruleContext,
+ executable,
+ runfiles,
+ computeArgs(ruleContext, appendingArgs),
+ computeActionEnvironment(ruleContext));
  }
 
  private static CommandLine computeArgs(RuleContext ruleContext, CommandLine additionalArgs) {
@@ -434,6 +463,30 @@ private static CommandLine computeArgs(RuleContext ruleContext, CommandLine addi
  ruleContext.getExpander().withDataLocations().tokenized("args"), additionalArgs);
  }
 
+ private static ActionEnvironment computeActionEnvironment(RuleContext ruleContext) {
+ if (!ruleContext.getRule().isAttrDefined("env", Type.STRING_DICT)
+ && !ruleContext.getRule().isAttrDefined("env_inherit", Type.STRING_LIST)) {
+ return ActionEnvironment.EMPTY;
+ }
+ TreeMap<String, String> fixedEnv = new TreeMap<>();
+ Set<String> inheritedEnv = new LinkedHashSet<>();
+ if (ruleContext.isAttrDefined("env", Type.STRING_DICT)) {
+ Expander expander = ruleContext.getExpander().withDataLocations();
+ for (Map.Entry<String, String> entry :
+ ruleContext.attributes().get("env", Type.STRING_DICT).entrySet()) {
+ fixedEnv.put(entry.getKey(), expander.expand("env", entry.getValue()));
+ }
+ }
+ if (ruleContext.isAttrDefined("env_inherit", Type.STRING_LIST)) {
+ for (String key : ruleContext.attributes().get("env_inherit", Type.STRING_LIST)) {
+ if (!fixedEnv.containsKey(key)) {
+ inheritedEnv.add(key);
+ }
+ }
+ }
+ return ActionEnvironment.create(fixedEnv, ImmutableSet.copyOf(inheritedEnv));
+ }
+
  /** Returns the path of the input manifest of {@code runfilesDir}. */
  public static Path inputManifestPath(Path runfilesDir) {
  return FileSystemUtils.replaceExtension(runfilesDir, INPUT_MANIFEST_EXT);

src/main/java/com/google/devtools/build/lib/analysis/test/TestActionBuilder.java

@@ -52,7 +52,6 @@
  * Helper class to create test actions.
  */
 public final class TestActionBuilder {
-
  private static final String CC_CODE_COVERAGE_SCRIPT = "CC_CODE_COVERAGE_SCRIPT";
  private static final String LCOV_MERGER = "LCOV_MERGER";
  // The coverage tool Bazel uses to generate a code coverage report for C++.
@@ -396,7 +395,7 @@ private TestParams createTestAction(int shards) throws InterruptedException {
  coverageArtifact,
  coverageDirectory,
  testProperties,
- extraTestEnv,
+ runfilesSupport.getActionEnvironment().addFixedVariables(extraTestEnv),
  executionSettings,
  shard,
  run,

src/main/java/com/google/devtools/build/lib/analysis/test/TestRunnerAction.java

@@ -28,6 +28,7 @@
 import com.google.common.util.concurrent.MoreExecutors;
 import com.google.devtools.build.lib.actions.AbstractAction;
 import com.google.devtools.build.lib.actions.ActionContinuationOrResult;
+import com.google.devtools.build.lib.actions.ActionEnvironment;
 import com.google.devtools.build.lib.actions.ActionExecutionContext;
 import com.google.devtools.build.lib.actions.ActionExecutionException;
 import com.google.devtools.build.lib.actions.ActionInput;
@@ -132,7 +133,7 @@ public class TestRunnerAction extends AbstractAction
  private Boolean unconditionalExecution;
 
  /** Any extra environment variables (and values) added by the rule that created this action. */
- private final ImmutableMap<String, String> extraTestEnv;
+ private final ActionEnvironment extraTestEnv;
 
  /**
  * The set of environment variables that are inherited from the client environment. These are
@@ -176,7 +177,7 @@ private static ImmutableSet<Artifact> nonNullAsSet(Artifact... artifacts) {
  Artifact coverageArtifact,
  @Nullable Artifact coverageDirectory,
  TestTargetProperties testProperties,
- Map<String, String> extraTestEnv,
+ ActionEnvironment extraTestEnv,
  TestTargetExecutionSettings executionSettings,
  int shardNum,
  int runNumber,
@@ -236,12 +237,13 @@ private static ImmutableSet<Artifact> nonNullAsSet(Artifact... artifacts) {
  this.testInfrastructureFailure = baseDir.getChild("test.infrastructure_failure");
  this.workspaceName = workspaceName;
 
- this.extraTestEnv = ImmutableMap.copyOf(extraTestEnv);
+ this.extraTestEnv = extraTestEnv;
  this.requiredClientEnvVariables =
  ImmutableIterable.from(
  Iterables.concat(
  configuration.getActionEnvironment().getInheritedEnv(),
- configuration.getTestActionEnvironment().getInheritedEnv()));
+ configuration.getTestActionEnvironment().getInheritedEnv(),
+ this.extraTestEnv.getInheritedEnv()));
  this.cancelConcurrentTestsOnSuccess = cancelConcurrentTestsOnSuccess;
  this.splitCoveragePostProcessing = splitCoveragePostProcessing;
  this.lcovMergerFilesToRun = lcovMergerFilesToRun;
@@ -378,7 +380,7 @@ protected void computeKey(
  fp.addBoolean(executionSettings.getTestRunnerFailFast());
  RunUnder runUnder = executionSettings.getRunUnder();
  fp.addString(runUnder == null ? "" : runUnder.getValue());
- fp.addStringMap(extraTestEnv);
+ extraTestEnv.addTo(fp);
  // TODO(ulfjack): It might be better for performance to hash the action and test envs in config,
  // and only add a hash here.
  configuration.getActionEnvironment().addTo(fp);
@@ -675,7 +677,7 @@ public Artifact getTestLog() {
  }
 
  /** Returns all environment variables which must be set in order to run this test. */
- public Map<String, String> getExtraTestEnv() {
+ public ActionEnvironment getExtraTestEnv() {
  return extraTestEnv;
  }
 

src/main/java/com/google/devtools/build/lib/exec/TestPolicy.java

@@ -88,7 +88,7 @@ public Map<String, String> computeTestEnvironment(
  }
 
  // Rule-specified test env.
- env.putAll(testAction.getExtraTestEnv());
+ testAction.getExtraTestEnv().resolve(env, clientEnv);
 
  // Overwrite with the environment common to all actions, see --action_env.
  testAction.getConfiguration().getActionEnvironment().resolve(env, clientEnv);

src/main/java/com/google/devtools/build/lib/runtime/commands/RunCommand.java

@@ -471,6 +471,9 @@ public BlazeCommandResult exec(CommandEnvironment env, OptionsParsingResult opti
  }
  } else {
  workingDir = runfilesDir;
+ if (runfilesSupport != null) {
+ runfilesSupport.getActionEnvironment().resolve(runEnvironment, env.getClientEnv());
+ }
  List<String> args = computeArgs(env, targetToRun, commandLineArgs);
  try {
  constructCommandLine(

src/test/shell/integration/run_test.sh

@@ -487,6 +487,37 @@ EOF
  assert_starts_with "${tmpdir}/test_bazel_run_with_custom_tmpdir" "$(cat "${TEST_TMPDIR}/tmpdir_value")"
 }
 
+function test_run_binary_with_env_attribute() {
+ local -r pkg="pkg${LINENO}"
+ mkdir -p ${pkg}
+ cat > $pkg/BUILD <<'EOF'
+sh_binary(
+ name = 't',
+ srcs = [':t.sh'],
+ data = [':t.dat'],
+ env = {
+ "ENV_A": "not_inherited",
+ "ENV_C": "no_surprise",
+ "ENV_DATA": "$(location :t.dat)",
+ },
+)
+EOF
+ cat > $pkg/t.sh <<'EOF'
+#!/bin/sh
+env
+cat $ENV_DATA
+exit 0
+EOF
+ touch $pkg/t.dat
+ chmod +x $pkg/t.sh
+ ENV_B=surprise ENV_C=surprise bazel run //$pkg:t > $TEST_log \
+ || fail "expected test to pass"
+ expect_log "ENV_A=not_inherited"
+ expect_log "ENV_B=surprise"
+ expect_log "ENV_C=no_surprise"
+ expect_log "ENV_DATA=$pkg/t.dat"
+}
+
 # Usage: assert_starts_with PREFIX STRING_TO_CHECK.
 # Asserts that `$1` is a prefix of `$2`.
 function assert_starts_with() {