Add fetch --repo option

Allow fetching only one repo using the canonical or apparent repo names. Only works with bzlmod.

PiperOrigin-RevId: 570675543
Change-Id: Idf9c412161a2403e8b8b9a4f0ca9e76f63e94de0

Author: SalmaSamy

src/main/java/com/google/devtools/build/lib/bazel/commands/FetchCommand.java

@@ -14,13 +14,15 @@
 package com.google.devtools.build.lib.bazel.commands;
 
 import static com.google.common.primitives.Booleans.countTrue;
+import static java.util.stream.Collectors.joining;
 
 import com.google.common.base.Joiner;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
 import com.google.devtools.build.lib.analysis.NoBuildEvent;
 import com.google.devtools.build.lib.analysis.NoBuildRequestFinishedEvent;
 import com.google.devtools.build.lib.bazel.bzlmod.BazelFetchAllValue;
+import com.google.devtools.build.lib.cmdline.LabelSyntaxException;
 import com.google.devtools.build.lib.cmdline.RepositoryMapping;
 import com.google.devtools.build.lib.cmdline.RepositoryName;
 import com.google.devtools.build.lib.cmdline.TargetPattern;
@@ -38,6 +40,7 @@
 import com.google.devtools.build.lib.query2.engine.QueryExpression;
 import com.google.devtools.build.lib.query2.engine.QuerySyntaxException;
 import com.google.devtools.build.lib.query2.engine.ThreadSafeOutputFormatterCallback;
+import com.google.devtools.build.lib.rules.repository.RepositoryDirectoryValue;
 import com.google.devtools.build.lib.runtime.BlazeCommand;
 import com.google.devtools.build.lib.runtime.BlazeCommandResult;
 import com.google.devtools.build.lib.runtime.Command;
@@ -56,10 +59,13 @@
 import com.google.devtools.build.lib.util.InterruptedFailureDetails;
 import com.google.devtools.build.skyframe.EvaluationContext;
 import com.google.devtools.build.skyframe.EvaluationResult;
+import com.google.devtools.build.skyframe.SkyKey;
 import com.google.devtools.build.skyframe.SkyValue;
 import com.google.devtools.common.options.OptionsParsingResult;
 import java.io.IOException;
 import java.util.EnumSet;
+import java.util.List;
+import net.starlark.java.eval.EvalException;
 
 /** Fetches external repositories. Which is so fetch. */
 @Command(
@@ -90,10 +96,13 @@ public BlazeCommandResult exec(CommandEnvironment env, OptionsParsingResult opti
       return createFailedBlazeCommandResult(Code.OPTIONS_INVALID, errorMessage);
     }
     FetchOptions fetchOptions = options.getOptions(FetchOptions.class);
-    // Validate only one option is provided for fetch
-    boolean moreThanOneOption =
-        countTrue(fetchOptions.all, fetchOptions.configure, !options.getResidue().isEmpty()) > 1;
-    if (moreThanOneOption) {
+    int optionsCount =
+        countTrue(
+            fetchOptions.all,
+            fetchOptions.configure,
+            !fetchOptions.repos.isEmpty(),
+            !options.getResidue().isEmpty());
+    if (optionsCount > 1) {
       String errorMessage = "Only one fetch option should be provided for fetch command.";
       env.getReporter().handle(Event.error(null, errorMessage));
       return createFailedBlazeCommandResult(Code.OPTIONS_INVALID, errorMessage);
@@ -109,9 +118,10 @@ public BlazeCommandResult exec(CommandEnvironment env, OptionsParsingResult opti
                 /* showProgress= */ true,
                 /* id= */ null));
     BlazeCommandResult result;
-
     if (fetchOptions.all || fetchOptions.configure) {
-      return fetchAll(env, options, fetchOptions.configure, threadsOption);
+      result = fetchAll(env, options, threadsOption, fetchOptions.configure);
+    } else if (!fetchOptions.repos.isEmpty()) {
+      result = fetchRepo(env, options, threadsOption, fetchOptions.repos);
     } else {
       result = fetchTarget(env, options, threadsOption);
     }
@@ -125,8 +135,8 @@ public BlazeCommandResult exec(CommandEnvironment env, OptionsParsingResult opti
   private BlazeCommandResult fetchAll(
       CommandEnvironment env,
       OptionsParsingResult options,
-      boolean configureEnabled,
-      LoadingPhaseThreadsOption threadsOption) {
+      LoadingPhaseThreadsOption threadsOption,
+      boolean configureEnabled) {
     if (!options.getOptions(BuildLanguageOptions.class).enableBzlmod) {
       String errorMessage =
           "Bzlmod has to be enabled for fetch --all to work, run with --enable_bzlmod";
@@ -168,6 +178,68 @@ private BlazeCommandResult fetchAll(
     }
   }
 
+  private BlazeCommandResult fetchRepo(
+      CommandEnvironment env,
+      OptionsParsingResult options,
+      LoadingPhaseThreadsOption threadsOption,
+      List<String> repos) {
+    SkyframeExecutor skyframeExecutor = env.getSkyframeExecutor();
+    EvaluationContext evaluationContext =
+        EvaluationContext.newBuilder()
+            .setParallelism(threadsOption.threads)
+            .setEventHandler(env.getReporter())
+            .build();
+    try {
+      env.syncPackageLoading(options);
+      ImmutableSet.Builder<SkyKey> repoDelegatorKeys = ImmutableSet.builder();
+      for (String repo : repos) {
+        RepositoryName repoName = getRepositoryName(env, threadsOption, repo);
+        repoDelegatorKeys.add(RepositoryDirectoryValue.key(repoName));
+      }
+      EvaluationResult<SkyValue> evaluationResult =
+          skyframeExecutor.prepareAndGet(repoDelegatorKeys.build(), evaluationContext);
+      if (evaluationResult.hasError()) {
+        Exception e = evaluationResult.getError().getException();
+        String errorMessage =
+            e != null ? e.getMessage() : "Unexpected error during repository fetching.";
+        env.getReporter().handle(Event.error(errorMessage));
+        return BlazeCommandResult.detailedExitCode(
+            InterruptedFailureDetails.detailedExitCode(errorMessage));
+      }
+      String notFoundRepos =
+          repoDelegatorKeys.build().stream()
+              .filter(
+                  key -> !((RepositoryDirectoryValue) evaluationResult.get(key)).repositoryExists())
+              .map(key -> ((RepositoryDirectoryValue) evaluationResult.get(key)).getErrorMsg())
+              .collect(joining("; "));
+      if (!notFoundRepos.isEmpty()) {
+        String errorMessage = "Fetching repos failed with errors: " + notFoundRepos;
+        env.getReporter().handle(Event.error(errorMessage));
+        return BlazeCommandResult.detailedExitCode(
+            InterruptedFailureDetails.detailedExitCode(errorMessage));
+      }
+
+      // Everything has been fetched successfully!
+      return BlazeCommandResult.success();
+    } catch (AbruptExitException e) {
+      env.getReporter().handle(Event.error(null, "Unknown error: " + e.getMessage()));
+      return BlazeCommandResult.detailedExitCode(e.getDetailedExitCode());
+    } catch (InterruptedException e) {
+      String errorMessage = "Fetch interrupted: " + e.getMessage();
+      env.getReporter().handle(Event.error(errorMessage));
+      return BlazeCommandResult.detailedExitCode(
+          InterruptedFailureDetails.detailedExitCode(errorMessage));
+    } catch (LabelSyntaxException | EvalException | IllegalArgumentException e) {
+      String errorMessage = "Invalid repo name: " + e.getMessage();
+      env.getReporter().handle(Event.error(null, errorMessage));
+      return BlazeCommandResult.detailedExitCode(
+          InterruptedFailureDetails.detailedExitCode(errorMessage));
+    } catch (RepositoryMappingResolutionException e) {
+      env.getReporter().handle(Event.error(e.getMessage()));
+      return BlazeCommandResult.detailedExitCode(e.getDetailedExitCode());
+    }
+  }
+
   private BlazeCommandResult fetchTarget(
       CommandEnvironment env,
       OptionsParsingResult options,
@@ -188,6 +260,7 @@ private BlazeCommandResult fetchTarget(
       RepositoryMapping repoMapping =
           env.getSkyframeExecutor()
               .getMainRepoMapping(keepGoing, threadsOption.threads, env.getReporter());
+
       mainRepoTargetParser =
           new Parser(env.getRelativeWorkingDirectory(), RepositoryName.MAIN, repoMapping);
     } catch (RepositoryMappingResolutionException e) {
@@ -294,6 +367,29 @@ public void processOutput(Iterable<Target> partialResult) {
                 expr));
   }
 
+  private RepositoryName getRepositoryName(
+      CommandEnvironment env, LoadingPhaseThreadsOption threadsOption, String repoName)
+      throws EvalException,
+          LabelSyntaxException,
+          RepositoryMappingResolutionException,
+          InterruptedException {
+    if (repoName.startsWith("@@")) { // canonical RepoName
+      return RepositoryName.create(repoName.substring(2));
+    } else if (repoName.startsWith("@")) { // apparent RepoName
+      RepositoryName.validateUserProvidedRepoName(repoName.substring(1));
+      RepositoryMapping repoMapping =
+          env.getSkyframeExecutor()
+              .getMainRepoMapping(
+                  env.getOptions().getOptions(KeepGoingOption.class).keepGoing,
+                  threadsOption.threads,
+                  env.getReporter());
+      return repoMapping.get(repoName.substring(1));
+    } else {
+      throw new IllegalArgumentException(
+          "The repo value has to be either apparent '@repo' or canonical '@@repo' repo name");
+    }
+  }
+
   private static BlazeCommandResult createFailedBlazeCommandResult(
       Code fetchCommandCode, String message) {
     return BlazeCommandResult.detailedExitCode(

src/main/java/com/google/devtools/build/lib/bazel/commands/FetchOptions.java

@@ -17,6 +17,7 @@
 import com.google.devtools.common.options.OptionDocumentationCategory;
 import com.google.devtools.common.options.OptionEffectTag;
 import com.google.devtools.common.options.OptionsBase;
+import java.util.List;
 
 /** Defines the options specific to Bazel's sync command */
 public class FetchOptions extends OptionsBase {
@@ -36,12 +37,22 @@ public class FetchOptions extends OptionsBase {
       documentationCategory = OptionDocumentationCategory.EXECUTION_STRATEGY,
       effectTags = {OptionEffectTag.CHANGES_INPUTS},
       help =
-          "Only fetch repositories marked as 'configure' for system-configuration purpose. Only"
+          "Only fetches repositories marked as 'configure' for system-configuration purpose. Only"
               + " works when --enable_bzlmod is on.")
   public boolean configure;
 
+  @Option(
+      name = "repo",
+      defaultValue = "null",
+      allowMultiple = true,
+      documentationCategory = OptionDocumentationCategory.EXECUTION_STRATEGY,
+      effectTags = {OptionEffectTag.CHANGES_INPUTS},
+      help =
+          "Only fetches the specified repository, which can be either {@apparent_repo_name} or"
+              + " {@@canonical_repo_name}. Only works when --enable_bzlmod is on.")
+  public List<String> repos;
+
   /*TODO(salmasamy) add more options:
-   * repo: to fetch a specific repo
    * force: to force fetch even if a repo exists
    */
 }

src/main/java/com/google/devtools/build/lib/bazel/commands/fetch.txt

@@ -1,8 +1,10 @@
 
 Usage: %{product} %{command} [<option> ...]
 
-  - --all: Fetches all external dependencies, only works with bzlmod
-  - --configure: Only fetches repositories marked as 'configure', only works with bzlmod
-  - <targets>: Fetches dependencies only for given targets
+  - --all: Fetches all external dependencies, only works with bzlmod.
+  - --configure: Only fetches repositories marked as 'configure', only works with bzlmod.
+  - --repo={repo}: Only fetches the specified repository, which can be either {@apparent_repo_name} or {@@canonical_repo_name}
+, only works with bzlmod.
+  - <targets>: Fetches dependencies only for given targets.
 
 %{options}

src/test/py/bazel/bzlmod/bazel_fetch_test.py

@@ -157,6 +157,67 @@ def testFetchFailsWithMultipleOptions(self):
         'ERROR: Only one fetch option should be provided for fetch command.',
         stderr,
     )
+    exit_code, _, stderr = self.RunBazel(
+        ['fetch', '--all', '--repo=@hello'], allow_failure=True
+    )
+    self.AssertExitCode(exit_code, 2, stderr)
+    self.assertIn(
+        'ERROR: Only one fetch option should be provided for fetch command.',
+        stderr,
+    )
+
+  def testFetchRepo(self):
+    self.main_registry.createCcModule('aaa', '1.0').createCcModule(
+        'bbb', '1.0', {'aaa': '1.0'}
+    ).createCcModule('ccc', '1.0')
+    self.ScratchFile(
+        'MODULE.bazel',
+        [
+            'bazel_dep(name = "bbb", version = "1.0")',
+            'bazel_dep(name = "ccc", version = "1.0", repo_name = "my_repo")',
+            'local_path_override(module_name="bazel_tools", path="tools_mock")',
+            'local_path_override(module_name="local_config_platform", ',
+            'path="platforms_mock")',
+        ],
+    )
+    self.ScratchFile('BUILD')
+    # Test canonical/apparent repo names & multiple repos
+    self.RunBazel(['fetch', '--repo=@@bbb~1.0', '--repo=@my_repo'])
+    _, stdout, _ = self.RunBazel(['info', 'output_base'])
+    repos_fetched = os.listdir(stdout[0] + '/external')
+    self.assertIn('bbb~1.0', repos_fetched)
+    self.assertIn('ccc~1.0', repos_fetched)
+    self.assertNotIn('aaa~1.0', repos_fetched)
+
+  def testFetchInvalidRepo(self):
+    # Invalid repo name (not canonical or apparent)
+    exit_code, _, stderr = self.RunBazel(
+        ['fetch', '--repo=hello'], allow_failure=True
+    )
+    self.AssertExitCode(exit_code, 8, stderr)
+    self.assertIn(
+        'ERROR: Invalid repo name: The repo value has to be either apparent'
+        " '@repo' or canonical '@@repo' repo name",
+        stderr,
+    )
+    # Repo does not exist
+    self.ScratchFile(
+        'MODULE.bazel',
+        [
+            'local_path_override(module_name="bazel_tools", path="tools_mock")',
+            'local_path_override(module_name="local_config_platform", ',
+            'path="platforms_mock")',
+        ],
+    )
+    exit_code, _, stderr = self.RunBazel(
+        ['fetch', '--repo=@@nono', '--repo=@nana'], allow_failure=True
+    )
+    self.AssertExitCode(exit_code, 8, stderr)
+    self.assertIn(
+        "ERROR: Fetching repos failed with errors: Repository '@nono' is not "
+        "defined; No repository visible as '@nana' from main repository",
+        stderr,
+    )
 
 
 if __name__ == '__main__':