Final updates for #46 - fixed recommencement bug with RollupRelations…

…hipFieldFinder, and added test proving it works

rollup/main/default/classes/Rollup.cls

@@ -1376,7 +1376,8 @@ global without sharing virtual class Rollup implements Database.Batchable<SObjec
     }
 
     Set<String> objIds = new Set<String>(); // get everything that doesn't have a null Id - a pretty trick
-    Integer amountOfCalcItems = String.isBlank(potentialGrandparentRelationshipFieldPath) ? getCountFromDb(countQuery, objIds, queryWrapper.recordIds) : 0;
+    Set<Id> recordIds = queryWrapper.recordIds; // also used below, bound to the "queryString" variable
+    Integer amountOfCalcItems = getCountFromDb(countQuery, objIds, recordIds);
 
     // emptyRollup used to call "getMaxQueryRows" below, as well as the default rollupControl.BatchChunkSize__c
     Rollup emptyRollup = new Rollup(RollupInvocationPoint.FROM_LWC);
@@ -1395,7 +1396,6 @@ global without sharing virtual class Rollup implements Database.Batchable<SObjec
     queryFields.addAll(whereEval.getRelationshipFieldNames());
     String queryString = getQueryString(calcItemType, new List<String>(queryFields), 'Id', '!=', queryWrapper.getQuery());
     if (amountOfCalcItems < emptyRollup.getMaxQueryRows()) {
-      Set<Id> recordIds = queryWrapper.recordIds;
       List<SObject> calculationItems = Database.query(queryString);
       Rollup thisRollup = getRollup(
         new List<Rollup__mdt>{ rollupInfo },
@@ -1410,7 +1410,7 @@ global without sharing virtual class Rollup implements Database.Batchable<SObjec
     } else {
       // batch to get calc items and then batch to rollup
       return Database.executeBatch(
-        new RollupFullBatchRecalculator(queryString, RollupInvocationPoint.FROM_LWC, rollupInfo, calcItemType),
+        new RollupFullBatchRecalculator(queryString, RollupInvocationPoint.FROM_LWC, rollupInfo, calcItemType, recordIds),
         emptyRollup.rollupControl.BatchChunkSize__c.intValue()
       );
     }
@@ -1901,7 +1901,7 @@ global without sharing virtual class Rollup implements Database.Batchable<SObjec
         BatchChunkSize__c = 2000,
         MaxLookupRowsBeforeBatching__c = Limits.getLimitDmlRows() / 3,
         MaxParentRowsUpdatedAtOnce__c = Limits.getLimitDmlRows() / 2,
-        MaximumRollupRetries__c = 100,
+        MaxRollupRetries__c = 100,
         MaxQueryRows__c = Limits.getLimitQueryRows() / 2,
         ShouldAbortRun__c = false,
         ShouldRunAs__c = 'Queueable'
@@ -1940,9 +1940,11 @@ global without sharing virtual class Rollup implements Database.Batchable<SObjec
         continue;
       }
 
-      Map<String, List<SObject>> calcItemsByLookupField = grandparentRollups.containsKey(rollup.lookupObj)
-        ? grandparentRollups.get(rollup.lookupObj).getParentLookupToRecords()
-        : this.getCalcItemsByLookupField(rollup, this.lookupObjectToUniqueFieldNames.get(rollup.lookupObj));
+      if (grandparentRollups.containsKey(rollup.lookupObj) && rollup.traversal == null) {
+        rollup.traversal = grandparentRollups.get(rollup.lookupObj);
+      }
+
+      Map<String, List<SObject>> calcItemsByLookupField = this.getCalcItemsByLookupField(rollup, this.lookupObjectToUniqueFieldNames.get(rollup.lookupObj));
       // some rollups may not finish retrieving all parent rows the first time around - and that's ok! we can keep
       // trying until all necessary records have been retrieved
       if (rollup.traversal != null && rollup.traversal.getIsFinished() == false) {
@@ -2003,10 +2005,12 @@ global without sharing virtual class Rollup implements Database.Batchable<SObjec
     if (this.deferredRollups.isEmpty() == false && isDeferralAllowed && stackDepth < this.rollupControl?.MaxRollupRetries__c) {
       stackDepth++;
       // tragic, but necessary due to limits on requeueing allowed during testing
-      isDeferralAllowed = Test.isRunningTest() == false && this.rollupControl.MaximumRollupRetries__c > stackDepth;
+      isDeferralAllowed = Test.isRunningTest() == false && this.rollupControl.MaxRollupRetries__c > stackDepth;
+
       this.rollups.clear();
       this.rollups.addAll(this.deferredRollups);
       this.deferredRollups.clear();
+
       if (System.isBatch()) {
         Database.executeBatch(this, this.rollupControl.BatchChunkSize__c.intValue());
       } else {
@@ -2131,16 +2135,21 @@ global without sharing virtual class Rollup implements Database.Batchable<SObjec
   private Integer getLookupRecordsCount(Boolean hasMoreThanOneTarget) {
     // we need to burn a few SOQL calls to consider how many records are going to be queried/updated
     // then, using RollupControl__mdt and/or sensible defaults, we'll decide whether to queue up or batch (or fail - that's always an option)
-    // if there's more than one SObjectType involved (or it's a grandparnet rollup) we bail on retrieving the actual count
-    // because those types of rollups can't be batched (at least not currently)
+    // if there's more than one SObjectType involved we bail on retrieving the actual count
+    // because you can only return one list of SObjects from a batch job's QueryLocator
     SObjectType targetType;
     Map<String, Set<String>> queryCountsToLookupIds = new Map<String, Set<String>>();
     for (Rollup rollup : this.rollups) {
       rollup.isFullRecalc = this.isFullRecalc;
       if (targetType == null) {
         targetType = rollup.lookupObj;
-      } else if (rollup.lookupObj != targetType || String.isNotBlank(rollup.metadata?.GrandparentRelationshipFieldPath__c)) {
+      } else if (rollup.lookupObj != targetType) {
         hasMoreThanOneTarget = true;
+      } else if (String.isNotBlank(rollup.metadata?.GrandparentRelationshipFieldPath__c)) {
+        // getting the count for grandparent (or greater) relationships will be handled further
+        // downstream; for our purposes, it isn't useful to try to get all of the records while
+        // we're still in a sync context
+        continue;
       }
 
       if (hasMoreThanOneTarget) {

rollup/main/default/classes/RollupFullBatchRecalculator.cls

@@ -2,15 +2,19 @@ public class RollupFullBatchRecalculator extends Rollup {
   private final String queryString;
   private final Rollup__mdt rollupInfo;
   private final SObjectType calcItemType;
+  private final Set<Id> recordIds;
 
-  public RollupFullBatchRecalculator(String queryString, RollupInvocationPoint invokePoint, Rollup__mdt rollupInfo, SObjectType calcItemType) {
+  public RollupFullBatchRecalculator(String queryString, RollupInvocationPoint invokePoint, Rollup__mdt rollupInfo, SObjectType calcItemType, Set<Id> recordIds) {
     super(invokePoint);
     this.queryString = queryString;
     this.rollupInfo = rollupInfo;
+    this.recordIds = recordIds;
   }
 
   public override Database.QueryLocator start(Database.BatchableContext bc) {
     Set<Id> objIds = new Set<Id>(); // necessary; there's a bind variable in the query string
+    // note - if the optional where clause was appended to the passed in query string, this.recordIds is also
+    // used as a bind variable
     return Database.getQueryLocator(this.queryString);
   }
 

rollup/main/default/classes/RollupRelationshipFieldFinder.cls

@@ -16,6 +16,7 @@ public without sharing class RollupRelationshipFieldFinder {
   private List<SObject> records;
   private List<String> relationshipParts;
   private Boolean isFirstRun = true;
+  private String currentRelationshipName;
 
   public RollupRelationshipFieldFinder(
     RollupControl__mdt rollupControl,
@@ -32,9 +33,7 @@ public without sharing class RollupRelationshipFieldFinder {
     this.uniqueFinalFieldNames = uniqueFinalFieldNames;
 
     if (this.relationshipParts.size() == 1) {
-      String fieldName = this.relationshipParts[0];
-      this.relationshipParts.set(0, ultimateParent.getDescribe().getName());
-      this.relationshipParts.add(fieldName);
+      this.relationshipParts.add(0, ultimateParent.getDescribe().getName());
     }
     this.originalParts = new List<String>(this.relationshipParts);
   }
@@ -141,6 +140,10 @@ public without sharing class RollupRelationshipFieldFinder {
       Limits.getLimitQueryRows() / 4 < Limits.getQueryRows() ||
       Limits.getLimitHeapSize() / 2 < Limits.getHeapSize()
     ) {
+      // we pop fields off of the list while recursively iterating
+      // which means we need to re-add the last field used if we are stopping
+      // due to limits
+      this.relationshipParts.add(0, this.currentRelationshipName);
       return this.traversal;
     }
 
@@ -198,7 +201,7 @@ public without sharing class RollupRelationshipFieldFinder {
   private Traversal recurseThroughObjectChain(List<SObject> records, SObjectType baseSObjectType) {
     // cache the latest records through in case we need to continue later
     this.recommencementRecords = records;
-    String currentRelationshipName = this.relationshipParts.remove(0);
+    this.currentRelationshipName = this.relationshipParts.remove(0);
     Map<String, SObjectField> fieldMap = baseSObjectType.getDescribe().fields.getMap();
     SObjectField field = this.getField(fieldMap, currentRelationshipName);
 

rollup/main/default/classes/RollupRelationshipFieldFinderTests.cls

@@ -80,7 +80,6 @@ private class RollupRelationshipFieldFinderTests {
     Account intermediateTwo = new Account(Name = 'Intermediate 2');
     insert new List<Account>{ intermediateOne, intermediateTwo };
 
-
     List<Account> updatedAccounts = [SELECT Id, OwnerId, Name FROM Account];
     if (updatedAccounts.size() == 2) {
       // don't run the rest of the test if the org has some kind of ownership assignment going on that would invalidate

rollup/main/default/classes/RollupTests.cls

@@ -2985,6 +2985,45 @@ private class RollupTests {
     System.assertEquals(opps[0].Name + ', ' + opps[1].Name, updatedUser.AboutMe, 'Grandparent rollup should have worked!');
   }
 
+  @isTest
+  static void shouldDeferGrandparentRollupSafelyTillAllParentRecordsAreRetrieved() {
+    Account acc = [SELECT Id, OwnerId FROM Account];
+    List<Opportunity> opps = new List<Opportunity>{
+      new Opportunity(AccountId = acc.Id, StageName = 'Prospecting', Name = 'One', CloseDate = System.today()),
+      new Opportunity(AccountId = acc.Id, StageName = 'Prospecting', Name = 'Two', CloseDate = System.today())
+    };
+    insert opps;
+
+    DMLMock mock = new DMLMock();
+    Rollup.defaultControl = new RollupControl__mdt(MaxQueryRows__c = 2, BatchChunkSize__c = 1, MaxRollupRetries__c = 100);
+    // start as synchronous rollup to allow for one deferral
+    Rollup.specificControl = new RollupControl__mdt(ShouldRunAs__c = 'Synchronous Rollup', MaxQueryRows__c = 2);
+    Rollup.DML = mock;
+    Rollup.shouldRun = true;
+    Rollup.records = opps;
+    Rollup.rollupMetadata = new List<Rollup__mdt>{
+      new Rollup__mdt(
+        CalcItem__c = 'Opportunity',
+        RollupFieldOnCalcItem__c = 'Name',
+        LookupFieldOnCalcItem__c = 'AccountId',
+        LookupObject__c = 'User',
+        LookupFieldOnLookupObject__c = 'Id',
+        RollupFieldOnLookupObject__c = 'AboutMe',
+        RollupOperation__c = 'CONCAT',
+        GrandparentRelationshipFieldPath__c = 'Account.Owner.AboutMe'
+      )
+    };
+    Rollup.apexContext = TriggerOperation.AFTER_INSERT;
+
+    Test.startTest();
+    Rollup.runFromTrigger();
+    Test.stopTest();
+
+    System.assertEquals(1, mock.Records.size(), 'Grandparent record should have been found!');
+    User updatedUser = (User) mock.Records[0];
+    System.assertEquals(opps[0].Name + ', ' + opps[1].Name, updatedUser.AboutMe, 'Grandparent rollup should have worked!');
+  }
+
   //** Helpers */
 
   private static DMLMock loadAccountIdMock(List<SObject> records) {
@@ -3034,4 +3073,4 @@ private class RollupTests {
     String result = String.valueOf(startingNumber++);
     return sObjectType.getDescribe().getKeyPrefix() + '0'.repeat(12 - result.length()) + result;
   }
-}
+}

rollup/main/default/customMetadata/RollupControl.Org_Defaults.md-meta.xml

@@ -22,10 +22,6 @@
         <field>MaxRollupRetries__c</field>
         <value xsi:type="xsd:double">100.0</value>
     </values>
-    <values>
-        <field>MaximumRollupRetries__c</field>
-        <value xsi:nil="true"/>
-    </values>
     <values>
         <field>ShouldAbortRun__c</field>
         <value xsi:type="xsd:boolean">false</value>