From 3a56897087acd1318f3ff2b5a9c448ebc4c30a37 Mon Sep 17 00:00:00 2001
From: Tomasz Pasternak <tomasz.pasternak@jetbrains.com>
Date: Fri, 28 Jun 2024 12:56:20 +0200
Subject: [PATCH] fix: Resolve query failure when there is a large number of
 files in the Working Set

closes #6478
---
 base/src/META-INF/blaze-base.xml              |  4 ++-
 .../BlazeQuerySourceToTargetProvider.java     | 33 +++++++++++++++++--
 2 files changed, 34 insertions(+), 3 deletions(-)

diff --git a/base/src/META-INF/blaze-base.xml b/base/src/META-INF/blaze-base.xml
index f32ea3643c2..d19cd283d42 100644
--- a/base/src/META-INF/blaze-base.xml
+++ b/base/src/META-INF/blaze-base.xml
@@ -409,7 +409,9 @@
     <registryKey defaultValue="true"
                  description="By default Bazel plugin takes control and opens project if .ijwb folder is present. This property can be used to disable this behavior to allow to open .idea or other project models if they are exists."
                  key="bazel.project.auto.open.if.present"/>
-
+    <registryKey defaultValue="false"
+                 description="Bazel Plugin sometimes uses --query_file option to pass the query via a file. By default, the file is deleted right after the query execution is finished. Set this flag to true, to avoid file deletion"
+                 key="bazel.sync.keep.query.files"/>
     <editorNotificationProvider implementation="com.google.idea.blaze.base.wizard2.BazelNotificationProvider"/>
   </extensions>
 
diff --git a/base/src/com/google/idea/blaze/base/dependencies/BlazeQuerySourceToTargetProvider.java b/base/src/com/google/idea/blaze/base/dependencies/BlazeQuerySourceToTargetProvider.java
index c3193007e11..824e173a74c 100644
--- a/base/src/com/google/idea/blaze/base/dependencies/BlazeQuerySourceToTargetProvider.java
+++ b/base/src/com/google/idea/blaze/base/dependencies/BlazeQuerySourceToTargetProvider.java
@@ -41,6 +41,7 @@
 import com.google.idea.blaze.base.sync.workspace.WorkspacePathResolverProvider;
 import com.google.idea.blaze.common.PrintOutput;
 import com.google.idea.blaze.exception.BuildException;
+import com.google.idea.blaze.qsync.project.ProjectProtoTransform;
 import com.google.idea.common.experiments.BoolExperiment;
 import com.intellij.openapi.project.Project;
 import java.io.BufferedReader;
@@ -48,11 +49,18 @@
 import java.io.IOException;
 import java.io.InputStream;
 import java.io.InputStreamReader;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.nio.file.Paths;
+import java.nio.file.StandardOpenOption;
 import java.util.Collection;
 import java.util.List;
 import java.util.Objects;
 import java.util.concurrent.Future;
 import javax.annotation.Nullable;
+
+import com.intellij.openapi.util.registry.Registry;
+import org.jetbrains.annotations.NotNull;
 import org.jetbrains.ide.PooledThreadExecutor;
 
 /**
@@ -185,10 +193,11 @@ private static ImmutableList<TargetInfo> runRecursiveRdepsQuery(
   @Nullable
   private static ImmutableList<TargetInfo> getTargetInfoList(
       Project project, BlazeContext context, ContextType type, String rdepsQuery)
-      throws BlazeQuerySourceToTargetException {
+          throws BlazeQuerySourceToTargetException {
+    Path queryFile = prepareQueryFile(project, rdepsQuery);
     BlazeCommand.Builder command =
         getBlazeCommandBuilder(
-            project, type, rdepsQuery, ImmutableList.of("--output=label_kind"), context);
+            project, type, "--query_file=" + queryFile.toAbsolutePath(), ImmutableList.of("--output=label_kind"), context);
     try (InputStream queryResultStream = runQuery(project, command, context)) {
       BlazeQueryLabelKindParser blazeQueryLabelKindParser =
           new BlazeQueryLabelKindParser(t -> true);
@@ -201,7 +210,27 @@ private static ImmutableList<TargetInfo> getTargetInfoList(
       return blazeQueryLabelKindParser.getTargets();
     } catch (IOException e) {
       throw new BlazeQuerySourceToTargetException("Failed to get target info list", e);
+    } finally {
+      if (!Registry.get("bazel.sync.keep.query.files").asBoolean()) {
+        queryFile.toFile().delete();
+      }
+    }
+  }
+
+  private static @NotNull Path prepareQueryFile(Project project, String rdepsQuery) throws BlazeQuerySourceToTargetException {
+    Path queryFile = null;
+    try {
+      Path queriesDir = Paths.get(project.getBasePath()).resolve("queries");
+      queriesDir.toFile().mkdir();
+      queryFile = Files.createTempFile(queriesDir, "query-", "");
+      Files.writeString(queryFile, rdepsQuery, StandardOpenOption.WRITE);
+    } catch (IOException e) {
+      if (queryFile != null) {
+        queryFile.toFile().delete();
+      }
+      throw new BlazeQuerySourceToTargetException("Couldn't create query file", e);
     }
+    return queryFile;
   }
 
   @Nullable