add support for recursively set default value in test mode (#464)

* add support for recursively set default value in test mode * add support for recursively set default value in test mode * add missing change * address comments --------- Co-authored-by: Zihan Li <[email protected]>
linkedin · Nov 1, 2024 · 484211f · 484211f
1 parent e364ee2
commit 484211f
Show file tree

Hide file tree

Showing 6 changed files with 54 additions and 1 deletion.
diff --git a/dao-api/build.gradle b/dao-api/build.gradle
@@ -10,6 +10,7 @@ dependencies {
   compile externalDependency.reflections
   compile externalDependency.commonsLang
   implementation 'com.google.protobuf:protobuf-java:3.21.1'
+  implementation spec.product.pegasus.restliServer
   dataModel project(':core-models')
   dataModel project(':validators')
 

diff --git a/dao-api/src/main/java/com/linkedin/metadata/dao/utils/RecordUtils.java b/dao-api/src/main/java/com/linkedin/metadata/dao/utils/RecordUtils.java
@@ -15,6 +15,8 @@
 import com.linkedin.data.template.UnionTemplate;
 import com.linkedin.metadata.dao.exception.ModelConversionException;
 import com.linkedin.metadata.validator.InvalidSchemaException;
+import com.linkedin.restli.internal.server.response.ResponseUtils;
+import com.linkedin.restli.server.RestLiServiceException;
 import java.io.IOException;
 import java.lang.reflect.Constructor;
 import java.lang.reflect.InvocationTargetException;
@@ -77,6 +79,27 @@ public static String toJsonString(@Nonnull RecordTemplate recordTemplate) {
     }
   }
 
+  /**
+   * Serializes a {@link RecordTemplate} to JSON string.
+   * Also take test mode as input to control the default value fill in strategy
+   * @param recordTemplate the record template to serialize
+   * @return the JSON string serialized using {@link JacksonDataTemplateCodec}.
+   */
+  //Todo: we will remove this method once we verify it works and does not bring too much degrade in test mode.
+  @Nonnull
+  public static String toJsonString(@Nonnull RecordTemplate recordTemplate, boolean isTestMode) {
+    if (isTestMode) {
+      try {
+        DataMap dataWithDefaultValue = (DataMap) ResponseUtils.fillInDataDefault(recordTemplate.schema(), recordTemplate.data());
+        return DATA_TEMPLATE_CODEC.mapToString(dataWithDefaultValue);
+      } catch (RestLiServiceException | IOException e) {
+        throw new ModelConversionException("Failed to serialize RecordTemplate: " + recordTemplate.toString(), e);
+      }
+    } else {
+      return toJsonString(recordTemplate);
+    }
+  }
+
   /**
    * Creates a {@link RecordTemplate} object from a serialized JSON string.
    *

diff --git a/dao-api/src/test/java/com/linkedin/metadata/dao/utils/RecordUtilsTest.java b/dao-api/src/test/java/com/linkedin/metadata/dao/utils/RecordUtilsTest.java
@@ -13,6 +13,7 @@
 import com.linkedin.testing.AspectFoo;
 import com.linkedin.testing.AspectBarArray;
 import com.linkedin.testing.AspectFooArray;
+import com.linkedin.testing.AspectWithDefaultValue;
 import com.linkedin.testing.EntityAspectUnion;
 import com.linkedin.testing.EntityAspectUnionAlias;
 import com.linkedin.testing.EntityAspectUnionComplex;
@@ -22,6 +23,7 @@
 import com.linkedin.testing.PizzaInfo;
 import com.linkedin.testing.StringUnion;
 import com.linkedin.testing.StringUnionArray;
+import com.linkedin.testing.MapValueRecord;
 import com.linkedin.testing.singleaspectentity.EntityValue;
 import com.linkedin.testing.urn.FooUrn;
 import java.io.IOException;
@@ -53,6 +55,17 @@ public void testToJsonString() throws IOException {
     assertEquals(actual, expected);
   }
 
+  @Test
+  public void testToJsonStringWithDefault() throws IOException {
+    AspectWithDefaultValue defaultValueAspect = new AspectWithDefaultValue().setNestedValueWithDefault(new MapValueRecord());
+    String expected =
+        loadJsonFromResource("defaultValueAspect.json").replaceAll("\\s+", "").replaceAll("\\n", "").replaceAll("\\r", "");
+
+    String actual = RecordUtils.toJsonString(defaultValueAspect, true);
+
+    assertEquals(actual, expected);
+  }
+
   @Test
   public void testToRecordTemplate() throws IOException {
     AspectFoo expected = new AspectFoo().setValue("foo");

diff --git a/dao-api/src/test/resources/defaultValueAspect.json b/dao-api/src/test/resources/defaultValueAspect.json
@@ -0,0 +1,6 @@
+{
+  "nestedValueWithDefault":{
+    "mapValueWithDefaultmap":{}
+  },
+  "valueWithDefault": ""
+}
diff --git a/dao-impl/ebean-dao/src/main/java/com/linkedin/metadata/dao/EbeanLocalAccess.java b/dao-impl/ebean-dao/src/main/java/com/linkedin/metadata/dao/EbeanLocalAccess.java
@@ -138,7 +138,7 @@ public <ASPECT extends RecordTemplate> int addWithOptimisticLocking(
     }
 
     AuditedAspect auditedAspect = new AuditedAspect()
-        .setAspect(RecordUtils.toJsonString(newValue))
+        .setAspect(RecordUtils.toJsonString(newValue, isTestMode))
         .setCanonicalName(aspectClass.getCanonicalName())
         .setLastmodifiedby(actor)
         .setLastmodifiedon(new Timestamp(timestamp).toString())

diff --git a/testing/test-models/src/main/pegasus/com/linkedin/testing/AspectWithDefaultValue.pdl b/testing/test-models/src/main/pegasus/com/linkedin/testing/AspectWithDefaultValue.pdl
@@ -0,0 +1,10 @@
+namespace com.linkedin.testing
+
+record AspectWithDefaultValue {
+
+  /**
+   * For unit tests
+   */
+  valueWithDefault: string = ""
+  nestedValueWithDefault: record MapValueRecord {mapValueWithDefaultmap: map[string, string] = { }}
+}