Update projection mask builder APIs to support updating masks. (#570)

* Update projection mask builder APIs to support updating masks.
* Check if a nested type has ProjectionMask class before generating the fully typesafe builder APIs.
* Added caching in ProjectionMaskApiChecker to avoid loading classes multiple times.

CHANGELOG.md

@@ -15,6 +15,8 @@ and what APIs have changed, if applicable.
 ## [Unreleased]
 - Add fluent client api for subresources
 - Update fluent client APIs to include projection mask as input parameter.
+- Update projection mask builder APIs to support updating the mask objects.
+- Added support for checking if a nested type supports new ProjectionMask API before generating new typesafe APIs for them.
 
 ## [29.17.0] - 2021-03-23
 - Implement D2 cluster subsetting.

generator-test/src/test/java/com/linkedin/pegasus/generator/test/TestMaskMap.java

@@ -4,6 +4,8 @@
 import com.linkedin.data.schema.PathSpecSet;
 import com.linkedin.data.transform.filter.request.MaskCreator;
 import com.linkedin.data.transform.filter.request.MaskTree;
+import com.linkedin.pegasus.generator.test.idl.records.WithCustomRecord;
+import com.linkedin.pegasus.generator.test.pdl.fixtures.CustomRecord;
 import org.testng.Assert;
 import org.testng.annotations.Test;
 
@@ -191,4 +193,145 @@ public void testNonAliasedUnionFieldNestedProjection()
     ));
     Assert.assertEquals(mask.getDataMap(), tree.getDataMap());
   }
+
+  // Tests the case where a partial mask is  created and updated later.
+  @Test
+  public void testReuseMaskSimpleFields()
+  {
+    RecordTest.ProjectionMask mask = RecordTest.createMask()
+        .withBooleanField();
+
+    MaskTree tree = MaskCreator.createPositiveMask(PathSpecSet.of(
+        RecordTest.fields().booleanField()));
+    Assert.assertEquals(mask.getDataMap(), tree.getDataMap());
+
+    // Now update the mask to add new fields.
+    mask.withDoubleField();
+
+    MaskTree tree2 = MaskCreator.createPositiveMask(PathSpecSet.of(
+        RecordTest.fields().booleanField(),
+        RecordTest.fields().doubleField()));
+    Assert.assertEquals(mask.getDataMap(), tree2.getDataMap());
+  }
+
+  @Test
+  public void testReuseMaskNestedRecord()
+  {
+    RecordTest.ProjectionMask mask = RecordTest.createMask()
+        .withBooleanField()
+        .withRecordField(nestedMask -> nestedMask.withLocation());
+
+    MaskTree tree = MaskCreator.createPositiveMask(PathSpecSet.of(
+        RecordTest.fields().booleanField(),
+        RecordTest.fields().recordField().location()));
+    Assert.assertEquals(mask.getDataMap(), tree.getDataMap());
+
+    // Now update the mask to add new fields.
+    mask.withDoubleField()
+        .withRecordField(nestedMask -> nestedMask.withOptionalLocation());
+
+    MaskTree tree2 = MaskCreator.createPositiveMask(PathSpecSet.of(
+        RecordTest.fields().booleanField(),
+        RecordTest.fields().doubleField(),
+        RecordTest.fields().recordField().location(),
+        RecordTest.fields().recordField().optionalLocation()));
+    Assert.assertEquals(mask.getDataMap(), tree2.getDataMap());
+  }
+
+  @Test
+  public void testReuseMaskNestedRecordClearing()
+  {
+    RecordTest.ProjectionMask mask = RecordTest.createMask()
+        .withRecordField(nestedMask -> nestedMask.withLocation());
+
+    MaskTree tree = MaskCreator.createPositiveMask(PathSpecSet.of(
+        RecordTest.fields().recordField().location()));
+    Assert.assertEquals(mask.getDataMap(), tree.getDataMap());
+
+    // Clear the nested mask by projecting the entire field.
+    mask.withRecordField();
+    tree = MaskCreator.createPositiveMask(PathSpecSet.of(
+        RecordTest.fields().recordField()));
+    Assert.assertEquals(mask.getDataMap(), tree.getDataMap());
+
+    // Now update the mask to add new fields.
+    mask.withRecordField(nestedMask -> nestedMask.withOptionalLocation());
+
+    tree = MaskCreator.createPositiveMask(PathSpecSet.of(
+        RecordTest.fields().recordField().optionalLocation()));
+    Assert.assertEquals(mask.getDataMap(), tree.getDataMap());
+  }
+
+  @Test
+  public void testReuseMaskNestedArray()
+  {
+    ArrayTest.ProjectionMask mask = ArrayTest.createMask()
+        .withRecordArray(arrayMask -> arrayMask.withItems(RecordBar.ProjectionMask::withLocation));
+
+    // Now update the mask to add new fields.
+    mask.withRecordArray(arrayMask -> arrayMask.withItems(nestedMask -> nestedMask.withOptionalLocation()));
+
+    MaskTree tree = MaskCreator.createPositiveMask(PathSpecSet.of(
+        ArrayTest.fields().recordArray().items().location(),
+        ArrayTest.fields().recordArray().items().optionalLocation()
+    ));
+    Assert.assertEquals(mask.getDataMap(), tree.getDataMap());
+
+    // Reset the nested mask by projecting all
+    mask.withRecordArray(0, 10);
+    tree = MaskCreator.createPositiveMask(PathSpecSet.of(
+        ArrayTest.fields().recordArray(0, 10)
+    ));
+    Assert.assertEquals(mask.getDataMap(), tree.getDataMap());
+
+    // Apply the nested mask again
+    mask.withRecordArray(arrayMask -> arrayMask.withItems(nestedMask -> nestedMask.withOptionalLocation()));
+    tree = MaskCreator.createPositiveMask(PathSpecSet.of(
+        ArrayTest.fields().recordArray().items().optionalLocation()
+    ));
+    Assert.assertEquals(mask.getDataMap(), tree.getDataMap());
+  }
+
+  @Test
+  public void testReuseUnionFieldNestedProjection()
+  {
+    UnionTest.ProjectionMask mask = UnionTest.createMask()
+        .withUnionWithAliases(unionMask -> unionMask.withMemAnotherInt()
+            .withMemAnotherMap()
+            .withMemMap());
+
+    MaskTree tree = MaskCreator.createPositiveMask(PathSpecSet.of(
+        UnionTest.fields().unionWithAliases().MemAnotherInt(),
+        UnionTest.fields().unionWithAliases().MemAnotherMap(),
+        UnionTest.fields().unionWithAliases().MemMap()
+    ));
+    Assert.assertEquals(mask.getDataMap(), tree.getDataMap());
+
+    mask.withUnionWithAliases(unionMask -> unionMask.withMemInt()
+            .withMemBoolean());
+
+    tree = MaskCreator.createPositiveMask(PathSpecSet.of(
+        UnionTest.fields().unionWithAliases().MemAnotherInt(),
+        UnionTest.fields().unionWithAliases().MemInt(),
+        UnionTest.fields().unionWithAliases().MemBoolean(),
+        UnionTest.fields().unionWithAliases().MemAnotherMap(),
+        UnionTest.fields().unionWithAliases().MemMap()
+    ));
+    Assert.assertEquals(mask.getDataMap(), tree.getDataMap());
+  }
+
+  @Test
+  public void testNestedTypeWithoutProjectionMask()
+  {
+    WithCustomRecord.ProjectionMask mask = WithCustomRecord.createMask()
+        .withCustom(MaskCreator.createPositiveMask(CustomRecord.fields().title()))
+        .withCustomArray(itemsMask -> itemsMask.withItems(
+            MaskCreator.createPositiveMask(CustomRecord.fields().body())));
+
+    MaskTree tree = MaskCreator.createPositiveMask(PathSpecSet.of(
+        WithCustomRecord.fields().custom().title(),
+        WithCustomRecord.fields().customArray().items().body()
+    ));
+    Assert.assertEquals(mask.getDataMap(), tree.getDataMap());
+  }
 }

generator-test/src/test/java/com/linkedin/pegasus/generator/test/pdl/fixtures/CustomRecord.java

@@ -2,9 +2,9 @@
 
 import com.linkedin.data.DataMap;
 import com.linkedin.data.schema.DataSchema;
-import com.linkedin.data.schema.MaskMap;
 import com.linkedin.data.schema.PathSpec;
 import com.linkedin.data.template.DataTemplate;
+import java.util.Collections;
 import java.util.List;
 
 
@@ -60,6 +60,11 @@ public int hashCode()
     return result;
   }
 
+  public static CustomRecord.Fields fields()
+  {
+    return new Fields();
+  }
+
   public static class Fields extends PathSpec
   {
     public Fields(List<String> path, String name)
@@ -71,14 +76,16 @@ public Fields()
     {
       super();
     }
-  }
 
-  public static class ProjectionMask extends MaskMap
-  {
-  }
+    public PathSpec title()
+    {
+      return new PathSpec(getPathComponents(), "title");
+    }
 
-  public static ProjectionMask createMask() {
-    return new ProjectionMask();
+    public PathSpec body()
+    {
+      return new PathSpec(getPathComponents(), "body");
+    }
   }
 
   @Override