fix: modify large queries use query more to fetch all needed records

src/i18n/i18n.ts

@@ -91,5 +91,6 @@ export const messages = {
   covIdFormatErr: 'Cannot specify code coverage with a TestRunId result',
   startHandshake: 'Attempting StreamingClient handshake',
   finishHandshake: 'Finished StreamingClient handshake',
-  subscribeStarted: 'Subscribing to ApexLog events'
+  subscribeStarted: 'Subscribing to ApexLog events',
+  invalidCountQueryResult: 'Invalid query result for count query ("%s")'
 };

src/tests/asyncTests.ts

@@ -28,11 +28,16 @@ import {
   TestResult,
   TestRunIdResult
 } from './types';
-import { calculatePercentage, isValidTestRunID } from './utils';
+import {
+  calculatePercentage,
+  isValidTestRunID,
+  verifyCountQueries
+} from './utils';
 import * as util from 'util';
 import { QUERY_RECORD_LIMIT } from './constants';
 import { CodeCoverage } from './codeCoverage';
 import { HttpRequest } from 'jsforce';
+import { OrgConfigProperties } from '@salesforce/core/lib/org/orgConfigProperties';
 
 export class AsyncTests {
   public readonly connection: Connection;
@@ -298,11 +303,13 @@ export class AsyncTests {
     apexTestResultQuery +=
       'ApexClass.Id, ApexClass.Name, ApexClass.NamespacePrefix ';
     apexTestResultQuery += 'FROM ApexTestResult WHERE QueueItemId IN (%s)';
-
+    const apexTestResultQueryCount =
+      'Select count(id) from ApexTestResult where QueueItemId IN (%s)';
     const apexResultIds = testQueueResult.records.map(record => record.Id);
 
     // iterate thru ids, create query with id, & compare query length to char limit
     const queries: string[] = [];
+    const countQueries: string[] = [];
     for (let i = 0; i < apexResultIds.length; i += QUERY_RECORD_LIMIT) {
       const recordSet: string[] = apexResultIds
         .slice(i, i + QUERY_RECORD_LIMIT)
@@ -311,12 +318,29 @@ export class AsyncTests {
         apexTestResultQuery,
         recordSet.join(',')
       );
+      const countQuery: string = util.format(
+        apexTestResultQueryCount,
+        recordSet.join(',')
+      );
+      countQueries.push(countQuery);
       queries.push(query);
     }
 
-    const queryPromises = queries.map(query => {
+    const countQueryPromises = countQueries.map(query => {
+      return this.connection.singleRecordQuery<{
+        expr0: number;
+        tooling: true;
+      }>(query);
+    });
+
+    const countQueryResult = await Promise.all(countQueryPromises);
+
+    verifyCountQueries(countQueryResult, countQueries);
+
+    const queryPromises = queries.map((query, index) => {
       return this.connection.tooling.query<ApexTestResultRecord>(query, {
-        autoFetch: true
+        autoFetch: true,
+        maxFetch: countQueryResult[index].expr0
       });
     });
     const apexTestResults = await Promise.all(queryPromises);

src/tests/codeCoverage.ts

@@ -5,7 +5,7 @@
  * For full license text, see LICENSE.txt file in the repo root or https://opensource.org/licenses/BSD-3-Clause
  */
 
-import { Connection } from '@salesforce/core';
+import { ConfigAggregator, Connection } from '@salesforce/core';
 import {
   ApexCodeCoverage,
   ApexCodeCoverageAggregate,
@@ -16,8 +16,10 @@ import {
   PerClassCoverage
 } from './types';
 import * as util from 'util';
-import { calculatePercentage } from './utils';
+import { calculatePercentage, verifyCountQueries } from './utils';
 import { QUERY_RECORD_LIMIT } from './constants';
+import { OrgConfigProperties } from '@salesforce/core/lib/org/orgConfigProperties';
+import { nls } from '../i18n';
 
 export class CodeCoverage {
   public readonly connection: Connection;
@@ -159,31 +161,58 @@ export class CodeCoverage {
   ): Promise<ApexCodeCoverage[]> {
     const perClassCodeCovQuery =
       'SELECT ApexTestClassId, ApexClassOrTrigger.Id, ApexClassOrTrigger.Name, TestMethodName, NumLinesCovered, NumLinesUncovered, Coverage FROM ApexCodeCoverage WHERE ApexTestClassId IN (%s)';
-    return this.fetchResults(apexTestClassSet, perClassCodeCovQuery);
+    const countPerClassCodeCovQuery =
+      'SELECT count(ApexTestClassId) FROM ApexCodeCoverage WHERE ApexTestClassId IN (%s)';
+    return this.fetchResults(
+      apexTestClassSet,
+      perClassCodeCovQuery,
+      countPerClassCodeCovQuery
+    );
   }
 
   private async queryAggregateCodeCov(
     apexClassIdSet: Set<string>
   ): Promise<ApexCodeCoverageAggregate[]> {
     const codeCoverageQuery =
       'SELECT ApexClassOrTrigger.Id, ApexClassOrTrigger.Name, NumLinesCovered, NumLinesUncovered, Coverage FROM ApexCodeCoverageAggregate WHERE ApexClassorTriggerId IN (%s)';
-    return this.fetchResults(apexClassIdSet, codeCoverageQuery);
+    const countCodeCoverageQuery =
+      'SELECT count(ApexClassOrTrigger.Id) FROM ApexCodeCoverageAggregate WHERE ApexClassorTriggerId IN (%s)';
+    return this.fetchResults(
+      apexClassIdSet,
+      codeCoverageQuery,
+      countCodeCoverageQuery
+    );
   }
 
   private async fetchResults<
     T extends ApexCodeCoverage | ApexCodeCoverageAggregate
-  >(idSet: Set<string>, selectQuery: string): Promise<T[]> {
+  >(idSet: Set<string>, selectQuery: string, countQuery: string): Promise<T[]> {
+    const config = await ConfigAggregator.create();
+    const countQueries = this.createQueries(countQuery, idSet);
     const queries = this.createQueries(selectQuery, idSet);
 
-    const queryPromises = queries.map(query => {
+    const countQueryPromises = countQueries.map(query =>
+      this.connection.singleRecordQuery<{ expr0: number }>(query, {
+        tooling: true
+      })
+    );
+
+    const countQueryResult = await Promise.all(countQueryPromises);
+
+    verifyCountQueries(countQueryResult, countQueries);
+
+    const queryPromises = queries.map((query, index) => {
       // The query method returns a type QueryResult from jsforce
       // that has takes a type that extends the jsforce Record.
       // ApexCodeCoverageRecord and ApexCodeCoverageAggregateRecord
       // are the Records compatible types defined in this project.
       return this.connection.tooling.query<
         ApexCodeCoverageRecord | ApexCodeCoverageAggregateRecord
       >(query, {
-        autoFetch: true
+        autoFetch: true,
+        maxFetch:
+          countQueryResult[index]?.expr0 ??
+          config.getPropertyValue(OrgConfigProperties.ORG_MAX_QUERY_LIMIT)
       });
     });
 

src/tests/utils.ts

@@ -8,6 +8,7 @@
 import { Connection } from '@salesforce/core';
 import { CLASS_ID_PREFIX, TEST_RUN_ID_PREFIX } from './constants';
 import { NamespaceInfo } from './types';
+import { nls } from '../i18n';
 
 export function isValidTestRunID(testRunId: string): boolean {
   return (
@@ -54,3 +55,23 @@ export async function queryNamespaces(
 
   return [...orgNamespaces, ...installedNamespaces];
 }
+
+export function verifyCountQueries(
+  countQueryResult: { expr0: number }[],
+  countQueries: string[]
+) {
+  // check count query results to ensure that there are results that can be used to set maxFetch
+  countQueryResult.forEach((result, index) => {
+    if (!result?.expr0 || result?.expr0 <= 0) {
+      throw new Error(
+        nls.localize(
+          'invalidCountQueryResult',
+          countQueries[index].substring(
+            0,
+            countQueries[index].indexOf('IN (') + 30
+          )
+        )
+      );
+    }
+  });
+}