From 6b9d7f4f429a9a92bacd4090f164f3b257a56c7e Mon Sep 17 00:00:00 2001
From: Chirag Ramani <chirag.ramani7@gmail.com>
Date: Sun, 23 Apr 2023 17:50:56 -0700
Subject: [PATCH] Add jsonproto option to query --output flag

---
 .../build/lib/query2/query/output/BUILD       |  1 +
 .../query2/query/output/OutputFormatters.java |  1 +
 .../output/ProtoJSONOutputFormatter.java      | 48 +++++++++++++++++++
 .../shell/integration/bazel_query_test.sh     | 23 +++++++++
 4 files changed, 73 insertions(+)
 create mode 100644 src/main/java/com/google/devtools/build/lib/query2/query/output/ProtoJSONOutputFormatter.java

diff --git a/src/main/java/com/google/devtools/build/lib/query2/query/output/BUILD b/src/main/java/com/google/devtools/build/lib/query2/query/output/BUILD
index 044d91f6360a24..dca8afa3ed44b4 100644
--- a/src/main/java/com/google/devtools/build/lib/query2/query/output/BUILD
+++ b/src/main/java/com/google/devtools/build/lib/query2/query/output/BUILD
@@ -38,5 +38,6 @@ java_library(
         "//third_party:guava",
         "//third_party:jsr305",
         "//third_party/protobuf:protobuf_java",
+        "//third_party/protobuf:protobuf_java_util",
     ],
 )
diff --git a/src/main/java/com/google/devtools/build/lib/query2/query/output/OutputFormatters.java b/src/main/java/com/google/devtools/build/lib/query2/query/output/OutputFormatters.java
index 752cc6520e49aa..d90e1d8d33d46c 100644
--- a/src/main/java/com/google/devtools/build/lib/query2/query/output/OutputFormatters.java
+++ b/src/main/java/com/google/devtools/build/lib/query2/query/output/OutputFormatters.java
@@ -39,6 +39,7 @@ public static ImmutableList<OutputFormatter> getDefaultFormatters() {
         new GraphOutputFormatter(),
         new XmlOutputFormatter(),
         new ProtoOutputFormatter(),
+        new ProtoJSONOutputFormatter(),
         new StreamedProtoOutputFormatter());
   }
 
diff --git a/src/main/java/com/google/devtools/build/lib/query2/query/output/ProtoJSONOutputFormatter.java b/src/main/java/com/google/devtools/build/lib/query2/query/output/ProtoJSONOutputFormatter.java
new file mode 100644
index 00000000000000..cfbaca7d016ec6
--- /dev/null
+++ b/src/main/java/com/google/devtools/build/lib/query2/query/output/ProtoJSONOutputFormatter.java
@@ -0,0 +1,48 @@
+// Copyright 2023 The Bazel Authors. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+package com.google.devtools.build.lib.query2.query.output;
+
+import com.google.devtools.build.lib.cmdline.RepositoryMapping;
+import com.google.devtools.build.lib.packages.Target;
+import com.google.devtools.build.lib.query2.engine.OutputFormatterCallback;
+import java.io.IOException;
+import java.io.OutputStream;
+import java.nio.charset.StandardCharsets;
+import com.google.protobuf.util.JsonFormat;
+
+/**
+ * An output formatter that outputs a protocol buffer json representation of a query result and outputs
+ * the json to the output print stream. 
+ */
+public class ProtoJSONOutputFormatter extends ProtoOutputFormatter {
+    @Override
+    public String getName() {
+      return "jsonproto";
+    }
+    private final JsonFormat.Printer jsonPrinter = JsonFormat.printer();
+  
+    @Override
+    public OutputFormatterCallback<Target> createPostFactoStreamCallback(
+        final OutputStream out, final QueryOptions options, RepositoryMapping mainRepoMapping) {
+      return new OutputFormatterCallback<Target>() {
+        @Override
+        public void processOutput(Iterable<Target> partialResult)
+            throws IOException, InterruptedException {
+          for (Target target : partialResult) {
+            out.write(jsonPrinter.print(toTargetProtoBuffer(target)).getBytes(StandardCharsets.UTF_8));
+          }
+        }
+      };
+    }
+}
diff --git a/src/test/shell/integration/bazel_query_test.sh b/src/test/shell/integration/bazel_query_test.sh
index d97b5ceff4e3fd..93c03b7fbb92d2 100755
--- a/src/test/shell/integration/bazel_query_test.sh
+++ b/src/test/shell/integration/bazel_query_test.sh
@@ -1098,4 +1098,27 @@ EOF
   expect_log "//pkg3:t4"
 }
 
+function test_basic_query_jsonproto() {
+  local pkg="${FUNCNAME[0]}"
+  mkdir -p "$pkg" || fail "mkdir -p $pkg"
+  cat > "$pkg/BUILD" <<'EOF'
+genrule(
+    name = "bar",
+    srcs = ["dummy.txt"],
+    outs = ["bar_out.txt"],
+    cmd = "echo unused > $(OUTS)",
+)
+EOF
+  bazel query --output=jsonproto --noimplicit_deps "//$pkg:bar" > output 2> "$TEST_log" \
+    || fail "Expected success"
+  cat output >> "$TEST_log"
+
+  # Verify that the appropriate action was included.
+  assert_contains "\"ruleClass\": \"genrule\"" output
+  assert_contains "\"name\": \"//$pkg:bar\"" output
+  assert_contains "\"ruleInput\": \[\"//$pkg:dummy.txt\"\]" output
+  assert_contains "\"ruleOutput\": \[\"//$pkg:bar_out.txt\"\]" output
+  assert_contains "echo unused" output
+}
+
 run_suite "${PRODUCT_NAME} query tests"