Merge pull request #982 from github/robertbrignull/read-problems-result-set

Read problems result set if '#select' result set doesn't exist

Author: robertbrignull

dist/query.js

@@ -16,6 +16,7 @@ import {
   Sarif,
   getSarifOutputType,
   QueryMetadata,
+  getBqrsResultCount,
 } from "./codeql";
 
 const test = anyTest as TestFn<{ db: string; tmpDir: string }>;
@@ -260,3 +261,49 @@ test("getting the SARIF output type when the `@kind` metadata is compatible with
 
   t.is(getSarifOutputType(queryMetadata, compatibleQueryKinds), "problem");
 });
+
+test("uses result count from #select result set if it exists", (t) => {
+  const bqrsInfo: BQRSInfo = {
+    resultSets: [{ name: "#select", rows: 3 }],
+    compatibleQueryKinds: [],
+  };
+
+  t.is(getBqrsResultCount(bqrsInfo), 3);
+});
+
+test("uses result count from problems result set if it exists", (t) => {
+  const bqrsInfo: BQRSInfo = {
+    resultSets: [{ name: "problems", rows: 4 }],
+    compatibleQueryKinds: [],
+  };
+
+  t.is(getBqrsResultCount(bqrsInfo), 4);
+});
+
+test("uses result count from #select result set if both #select and problems result sets exist", (t) => {
+  const bqrsInfo: BQRSInfo = {
+    resultSets: [
+      { name: "#select", rows: 3 },
+      { name: "problems", rows: 4 },
+    ],
+    compatibleQueryKinds: [],
+  };
+
+  t.is(getBqrsResultCount(bqrsInfo), 3);
+});
+
+test("throws error if neither #select or problems result sets exist", (t) => {
+  const bqrsInfo: BQRSInfo = {
+    resultSets: [
+      { name: "something", rows: 13 },
+      { name: "unknown", rows: 42 },
+    ],
+    compatibleQueryKinds: [],
+  };
+
+  const error = t.throws(() => getBqrsResultCount(bqrsInfo));
+  t.deepEqual(
+    error?.message,
+    "BQRS does not contain any result sets matching known names. Expected one of #select or problems but found something, unknown",
+  );
+});

src/codeql.test.ts

@@ -449,17 +449,29 @@ export function getSarifResultCount(sarif: Sarif): number {
   return count;
 }
 
+/**
+ * Names of result sets that can be considered the "default" result set
+ * and should be used when calculating number of results and when showing
+ * results to users.
+ * Will check result sets in this order and use the first one that exists.
+ */
+const KNOWN_RESULT_SET_NAMES: string[] = ["#select", "problems"];
+
 /**
  * Gets the number of results in the given BQRS data.
  */
-function getBqrsResultCount(bqrsInfo: BQRSInfo): number {
-  const selectResultSet = bqrsInfo.resultSets.find(
-    (resultSet) => resultSet.name === "#select",
-  );
-  if (!selectResultSet) {
-    throw new Error("No result set named #select");
+export function getBqrsResultCount(bqrsInfo: BQRSInfo): number {
+  for (const name of KNOWN_RESULT_SET_NAMES) {
+    const resultSet = bqrsInfo.resultSets.find((r) => r.name === name);
+    if (resultSet !== undefined) {
+      return resultSet.rows;
+    }
   }
-  return selectResultSet.rows;
+
+  const resultSetNames = bqrsInfo.resultSets.map((r) => r.name);
+  throw new Error(
+    `BQRS does not contain any result sets matching known names. Expected one of ${KNOWN_RESULT_SET_NAMES.join(" or ")} but found ${resultSetNames.join(", ")}`,
+  );
 }
 
 interface DatabaseMetadata {