feat: Allow environment variables in run configurations (#5885)

* feat: Allow environment variables in run configurations

Allow users to set their own env vars per run-configuration.
It is not respected everywhere (see the list below), but it works on our example projects.

Main features:
- Env vars should work when **run**ing and **debug**ging any binary.
- Env vars should work when **run**ning any test.
- Env vars should work when **debug**ging Go and Java tests. Env vars for C++ test debugging are implemented internally, but they're part of a larger patch and I'll need some more work to untangle that.

base/src/com/google/idea/blaze/base/command/BlazeCommandRunner.java

@@ -25,6 +25,7 @@
 import com.intellij.openapi.project.Project;
 import java.io.InputStream;
 import java.util.List;
+import java.util.Map;
 import java.util.Optional;
 
 /** Runs a blaze command. */
@@ -38,19 +39,21 @@ BlazeBuildOutputs run(
       Project project,
       BlazeCommand.Builder blazeCommandBuilder,
       BuildResultHelper buildResultHelper,
-      BlazeContext context)
+      BlazeContext context,
+      Map<String, String> envVars)
       throws BuildException;
 
   /**
    * Runs a blaze test command, parses the test results into a {@link BlazeTestResults} object using
    * the given {@link BuildResultHelper}.
    */
   BlazeTestResults runTest(
-      Project project,
-      BlazeCommand.Builder blazeCommandBuilder,
-      BuildResultHelper buildResultHelper,
-      BlazeContext context)
-      throws BuildException;
+          Project project,
+          BlazeCommand.Builder blazeCommandBuilder,
+          BuildResultHelper buildResultHelper,
+          BlazeContext context,
+          Map<String, String> envVars)
+          throws BuildException;
 
   /**
    * Runs a blaze query command.

base/src/com/google/idea/blaze/base/command/CommandLineBlazeCommandRunner.java

@@ -52,6 +52,7 @@
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.nio.file.StandardOpenOption;
+import java.util.Map;
 import java.util.Optional;
 import java.util.function.Function;
 
@@ -60,13 +61,14 @@ public class CommandLineBlazeCommandRunner implements BlazeCommandRunner {
 
   @Override
   public BlazeBuildOutputs run(
-      Project project,
-      BlazeCommand.Builder blazeCommandBuilder,
-      BuildResultHelper buildResultHelper,
-      BlazeContext context) {
+          Project project,
+          BlazeCommand.Builder blazeCommandBuilder,
+          BuildResultHelper buildResultHelper,
+          BlazeContext context,
+          Map<String, String> envVars) {
 
     BuildResult buildResult =
-        issueBuild(blazeCommandBuilder, WorkspaceRoot.fromProject(project), context);
+        issueBuild(blazeCommandBuilder, WorkspaceRoot.fromProject(project), envVars, context);
     BuildDepsStatsScope.fromContext(context)
         .ifPresent(stats -> stats.setBazelExitCode(buildResult.exitCode));
     if (buildResult.status == Status.FATAL_ERROR) {
@@ -100,9 +102,14 @@ public BlazeTestResults runTest(
       Project project,
       BlazeCommand.Builder blazeCommandBuilder,
       BuildResultHelper buildResultHelper,
-      BlazeContext context) {
+      BlazeContext context,
+      Map<String, String> envVars) {
+    // For tests, we have to pass the environment variables as `--test_env`, otherwise they don't get forwarded
+    for (Map.Entry<String, String> env: envVars.entrySet()) {
+      blazeCommandBuilder.addBlazeFlags(BlazeFlags.TEST_ENV, String.format("%s=%s", env.getKey(), env.getValue()));
+    }
     BuildResult buildResult =
-        issueBuild(blazeCommandBuilder, WorkspaceRoot.fromProject(project), context);
+        issueBuild(blazeCommandBuilder, WorkspaceRoot.fromProject(project), envVars, context);
     if (buildResult.status == Status.FATAL_ERROR) {
       return BlazeTestResults.NO_RESULTS;
     }
@@ -194,11 +201,12 @@ public InputStream runBlazeInfo(
   }
 
   private BuildResult issueBuild(
-      BlazeCommand.Builder blazeCommandBuilder, WorkspaceRoot workspaceRoot, BlazeContext context) {
+      BlazeCommand.Builder blazeCommandBuilder, WorkspaceRoot workspaceRoot, Map<String, String> envVars, BlazeContext context) {
     blazeCommandBuilder.addBlazeFlags(getExtraBuildFlags(blazeCommandBuilder));
     int retVal =
         ExternalTask.builder(workspaceRoot)
             .addBlazeCommand(blazeCommandBuilder.build())
+            .environmentVars(envVars)
             .context(context)
             .stderr(
                 LineProcessingOutputStream.of(

base/src/com/google/idea/blaze/base/qsync/BazelAppInspectorBuilder.java

@@ -18,6 +18,7 @@
 import static com.google.common.collect.ImmutableList.toImmutableList;
 
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
 import com.google.idea.blaze.base.bazel.BazelExitCodeException;
 import com.google.idea.blaze.base.bazel.BuildSystem;
 import com.google.idea.blaze.base.bazel.BuildSystem.BuildInvoker;
@@ -70,7 +71,7 @@ public AppInspectorInfo buildAppInspector(BlazeContext context, Set<Label> build
               .addBlazeFlags(additionalBlazeFlags);
 
       BlazeBuildOutputs outputs =
-          invoker.getCommandRunner().run(project, builder, buildResultHelper, context);
+          invoker.getCommandRunner().run(project, builder, buildResultHelper, context, ImmutableMap.of());
       BazelExitCodeException.throwIfFailed(builder, outputs.buildResult);
 
       return createAppInspectorInfo(outputs);

base/src/com/google/idea/blaze/base/qsync/BazelDependencyBuilder.java

@@ -23,6 +23,7 @@
 
 import com.google.common.base.Joiner;
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableMultimap;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Lists;
@@ -186,7 +187,7 @@ public OutputInfo build(
       buildDepsStatsBuilder.ifPresent(
           stats -> stats.setBuildFlags(builder.build().toArgumentList()));
       BlazeBuildOutputs outputs =
-          invoker.getCommandRunner().run(project, builder, buildResultHelper, context);
+          invoker.getCommandRunner().run(project, builder, buildResultHelper, context, ImmutableMap.of());
       buildDepsStatsBuilder.ifPresent(
           stats -> {
             stats.setBuildIds(outputs.getBuildIds());

base/src/com/google/idea/blaze/base/qsync/BazelRenderJarBuilder.java

@@ -18,6 +18,7 @@
 import static com.google.common.collect.ImmutableList.toImmutableList;
 
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
 import com.google.idea.blaze.base.bazel.BazelExitCodeException;
 import com.google.idea.blaze.base.bazel.BuildSystem;
 import com.google.idea.blaze.base.bazel.BuildSystem.BuildInvoker;
@@ -76,7 +77,7 @@ public RenderJarInfo buildRenderJar(BlazeContext context, Set<Label> buildTarget
               .addBlazeFlags("--output_groups=intellij-render-resolve-android");
 
       BlazeBuildOutputs outputs =
-          invoker.getCommandRunner().run(project, builder, buildResultHelper, context);
+          invoker.getCommandRunner().run(project, builder, buildResultHelper, context, ImmutableMap.of());
       BazelExitCodeException.throwIfFailed(builder, outputs.buildResult);
 
       return createRenderJarInfo(outputs);

base/src/com/google/idea/blaze/base/run/confighandler/BlazeCommandGenericRunConfigurationRunner.java

@@ -62,6 +62,7 @@
 import com.intellij.execution.ExecutionException;
 import com.intellij.execution.ExecutionResult;
 import com.intellij.execution.Executor;
+import com.intellij.execution.configuration.EnvironmentVariablesData;
 import com.intellij.execution.configurations.CommandLineState;
 import com.intellij.execution.configurations.GeneralCommandLine;
 import com.intellij.execution.configurations.RunProfileState;
@@ -81,6 +82,7 @@
 import java.io.OutputStream;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Map;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
@@ -192,9 +194,16 @@ public OutputStream getProcessInput() {
     private ProcessHandler getScopedProcessHandler(
         Project project, BlazeCommand blazeCommand, WorkspaceRoot workspaceRoot)
         throws ExecutionException {
+      GeneralCommandLine commandLine = new GeneralCommandLine(blazeCommand.toList());
+      EnvironmentVariablesData envVarState = handlerState.getUserEnvVarsState().getData();
+      commandLine.withEnvironment(envVarState.getEnvs());
+      commandLine.withParentEnvironmentType(
+              envVarState.isPassParentEnvs()
+                      ? GeneralCommandLine.ParentEnvironmentType.SYSTEM
+                      : GeneralCommandLine.ParentEnvironmentType.NONE);
       return new ScopedBlazeProcessHandler(
           project,
-          new GeneralCommandLine(blazeCommand.toList()),
+          commandLine,
           workspaceRoot,
           new ScopedBlazeProcessHandler.ScopedProcessHandlerDelegate() {
             @Override
@@ -239,13 +248,14 @@ protected ConsoleView createConsole() {
           });
       addConsoleFilters(consoleFilters.toArray(new Filter[0]));
 
+      @NotNull Map<String, String> envVars = handlerState.getUserEnvVarsState().getData().getEnvs();
       ListenableFuture<BlazeBuildOutputs> blazeBuildOutputsListenableFuture =
           BlazeExecutor.getInstance()
               .submit(
                   () ->
                       invoker
                           .getCommandRunner()
-                          .run(project, blazeCommandBuilder, buildResultHelper, context));
+                          .run(project, blazeCommandBuilder, buildResultHelper, context, envVars));
       Futures.addCallback(
           blazeBuildOutputsListenableFuture,
           new FutureCallback<BlazeBuildOutputs>() {
@@ -311,32 +321,43 @@ protected ConsoleView createConsole() {
         context.addOutputSink(PrintOutput.class, new WritingOutputSink(consoleView));
       }
       addConsoleFilters(consoleFilters.toArray(new Filter[0]));
-      return !invoker.getCommandRunner().canUseCli()
-          ? getCommandRunnerProcessHandler(
+      @NotNull Map<String, String> envVars = handlerState.getUserEnvVarsState().getData().getEnvs();
+
+      if (invoker.getCommandRunner().canUseCli()) {
+        // If we can use the CLI, that means we will run through Bazel (as opposed to a raw process handler)
+        // When running `bazel test`, bazel will not forward the environment to the tests themselves -- we need to use
+        // the --test_env flag for that. Therefore, we convert all the env vars to --test_env flags here.
+        for (Map.Entry<String, String> env: envVars.entrySet()) {
+          blazeCommandBuilder.addBlazeFlags(BlazeFlags.TEST_ENV, String.format("%s=%s", env.getKey(), env.getValue()));
+        }
+
+        return getScopedProcessHandler(project, blazeCommandBuilder.build(), workspaceRoot);
+      }
+      return getCommandRunnerProcessHandlerForTests(
               project,
               invoker,
               buildResultHelper,
               blazeCommandBuilder,
               testResultFinderStrategy,
-              context)
-          : getScopedProcessHandler(project, blazeCommandBuilder.build(), workspaceRoot);
+              context);
     }
 
-    private ProcessHandler getCommandRunnerProcessHandler(
+    private ProcessHandler getCommandRunnerProcessHandlerForTests(
         Project project,
         BuildInvoker invoker,
         BuildResultHelper buildResultHelper,
         BlazeCommand.Builder blazeCommandBuilder,
         BlazeTestResultFinderStrategy testResultFinderStrategy,
         BlazeContext context) {
       ProcessHandler processHandler = getGenericProcessHandler();
+      @NotNull Map<String, String> envVars = handlerState.getUserEnvVarsState().getData().getEnvs();
       ListenableFuture<BlazeTestResults> blazeTestResultsFuture =
           BlazeExecutor.getInstance()
               .submit(
                   () ->
                       invoker
                           .getCommandRunner()
-                          .runTest(project, blazeCommandBuilder, buildResultHelper, context));
+                          .runTest(project, blazeCommandBuilder, buildResultHelper, context, envVars));
       Futures.addCallback(
           blazeTestResultsFuture,
           new FutureCallback<BlazeTestResults>() {

base/src/com/google/idea/blaze/base/run/processhandler/ScopedBlazeProcessHandler.java

@@ -69,7 +69,6 @@ public interface ScopedProcessHandlerDelegate {
    *     context.
    * @throws ExecutionException
    */
-
   public ScopedBlazeProcessHandler(
       Project project,
       GeneralCommandLine command,

base/src/com/google/idea/blaze/base/run/state/BlazeCommandRunConfigurationCommonState.java

@@ -39,18 +39,20 @@ public class BlazeCommandRunConfigurationCommonState extends RunConfigurationCom
   protected final BlazeCommandState command;
   protected final RunConfigurationFlagsState blazeFlags;
   protected final RunConfigurationFlagsState exeFlags;
+  protected final EnvironmentVariablesState envVars;
   protected final BlazeBinaryState blazeBinary;
 
   public BlazeCommandRunConfigurationCommonState(BuildSystemName buildSystemName) {
     command = new BlazeCommandState();
     blazeFlags = new RunConfigurationFlagsState(USER_BLAZE_FLAG_TAG, buildSystemName + " flags:");
     exeFlags = new RunConfigurationFlagsState(USER_EXE_FLAG_TAG, "Executable flags:");
+    envVars = new EnvironmentVariablesState();
     blazeBinary = new BlazeBinaryState();
   }
 
   @Override
   protected ImmutableList<RunConfigurationState> initializeStates() {
-    return ImmutableList.of(command, blazeFlags, exeFlags, blazeBinary);
+    return ImmutableList.of(command, blazeFlags, exeFlags, envVars, blazeBinary);
   }
 
   /** @return The list of blaze flags that the user specified manually. */
@@ -63,6 +65,11 @@ public RunConfigurationFlagsState getExeFlagsState() {
     return exeFlags;
   }
 
+  /** @return The environment variables the user specified manually. */
+  public EnvironmentVariablesState getUserEnvVarsState() {
+    return envVars;
+  }
+
   public BlazeBinaryState getBlazeBinaryState() {
     return blazeBinary;
   }

base/src/com/google/idea/blaze/base/run/state/EnvironmentVariablesState.java

@@ -15,14 +15,19 @@
  */
 package com.google.idea.blaze.base.run.state;
 
+import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
+import com.google.idea.blaze.base.command.BlazeFlags;
 import com.intellij.execution.configuration.EnvironmentVariablesComponent;
 import com.intellij.execution.configuration.EnvironmentVariablesData;
 import com.intellij.openapi.project.Project;
 import com.intellij.openapi.util.InvalidDataException;
 import com.intellij.openapi.util.WriteExternalException;
 import javax.swing.JComponent;
 import org.jdom.Element;
+import org.jetbrains.annotations.VisibleForTesting;
+
+import java.util.Map;
 
 /** State for user-defined environment variables. */
 public class EnvironmentVariablesState implements RunConfigurationState {
@@ -32,10 +37,20 @@ public class EnvironmentVariablesState implements RunConfigurationState {
   private EnvironmentVariablesData data =
       EnvironmentVariablesData.create(ImmutableMap.of(), /* passParentEnvs= */ true);
 
+  public EnvironmentVariablesState() {}
+
   public EnvironmentVariablesData getData() {
     return data;
   }
 
+  public ImmutableList<String> asBlazeTestEnvFlags() {
+    ImmutableList.Builder<String> flags = ImmutableList.builder();
+    for (Map.Entry<String, String> env: data.getEnvs().entrySet()) {
+      flags.add(BlazeFlags.TEST_ENV, String.format("%s=%s", env.getKey(), env.getValue()));
+    }
+    return flags.build();
+  }
+
   @Override
   public void readExternal(Element element) throws InvalidDataException {
     Element child = element.getChild(ELEMENT_TAG);
@@ -57,6 +72,11 @@ public RunConfigurationStateEditor getEditor(Project project) {
     return new Editor();
   }
 
+  @VisibleForTesting
+  public void setEnvVars(Map<String, String> vars) {
+    data = data.with(vars);
+  }
+
   private static class Editor implements RunConfigurationStateEditor {
     private final EnvironmentVariablesComponent component = new EnvironmentVariablesComponent();
 

base/src/com/google/idea/blaze/base/sync/aspects/BlazeIdeInterfaceAspectsImpl.java

@@ -760,7 +760,7 @@ private static BlazeBuildOutputs runBuildForTargets(
       aspectStrategy.addAspectAndOutputGroups(
           builder, outputGroups, activeLanguages, onlyDirectDeps);
 
-      return invoker.getCommandRunner().run(project, builder, buildResultHelper, context);
+      return invoker.getCommandRunner().run(project, builder, buildResultHelper, context, ImmutableMap.of());
     }
   }
 }