fix: reflect the internal models to asset models from param and return types.

RealChrisL · RealChrisL · commit b677196a0099 · 2024-09-06T16:27:48.000-07:00
diff --git a/dao-api/src/main/java/com/linkedin/metadata/dao/utils/ModelUtils.java b/dao-api/src/main/java/com/linkedin/metadata/dao/utils/ModelUtils.java
@@ -498,23 +498,27 @@ public static <ASSET extends RecordTemplate, ASPECT_UNION extends UnionTemplate>
       }
       RecordUtils.setRecordTemplatePrimitiveField(asset, URN_FIELD, urn);
 
+      final Map<String, String> paramToAssetMethod = new HashMap<>();
+      for (final Method assetMethod : assetClass.getDeclaredMethods()) {
+        if (assetMethod.getName().startsWith("set") && assetMethod.getParameterTypes().length > 0) {
+          paramToAssetMethod.put(assetMethod.getParameterTypes()[0].getName(), assetMethod.getName());
+        }
+      }
+
       for (final ASPECT_UNION aspect : aspects) {
-        final Field[] aspectUnionFields = aspect.getClass().getDeclaredFields();
-        final Field[] assetFields = asset.getClass().getDeclaredFields();
-        for (final Field aspectUnionField : aspectUnionFields) {
-          if (aspectUnionField.getName().startsWith(MEMBER_FIELD_PREFIX)) {
-            final String aspectFieldName = aspectUnionField.getName().substring(MEMBER_FIELD_PREFIX.length());
-            for (final Field assetField : assetFields) {
-              if (assetField.getName().startsWith(FIELD_FIELD_PREFIX) && assetField.getName()
-                  .substring(FIELD_FIELD_PREFIX.length())
-                  .equals(aspectFieldName)) {
-                final Object aspectValue = aspect.getClass().getMethod("get" + aspectFieldName).invoke(aspect);
-                if (aspectValue != null) {
-                  asset.getClass()
-                      .getMethod("set" + aspectFieldName, aspectValue.getClass())
-                      .invoke(asset, aspectValue);
-                }
-              }
+        final Map<String, String> returnToUnionMethod = new HashMap<>();
+        for (final Method aspectUnionMethod : aspect.getClass().getMethods()) {
+          if (aspectUnionMethod.getName().startsWith("get")) {
+            returnToUnionMethod.put(aspectUnionMethod.getReturnType().getName(), aspectUnionMethod.getName());
+          }
+        }
+        for (final String param : paramToAssetMethod.keySet()) {
+          if (returnToUnionMethod.containsKey(param)) {
+            Object aspectValue = aspect.getClass().getMethod(returnToUnionMethod.get(param)).invoke(aspect);
+            if (aspectValue != null) {
+              asset.getClass()
+                  .getMethod(paramToAssetMethod.get(param), aspectValue.getClass())
+                  .invoke(asset, aspectValue);
             }
           }
         }
diff --git a/dao-api/src/test/java/com/linkedin/metadata/dao/utils/ModelUtilsTest.java b/dao-api/src/test/java/com/linkedin/metadata/dao/utils/ModelUtilsTest.java
@@ -319,8 +319,8 @@ public void testGetAspectsFromAsset() {
     EntityAsset asset = new EntityAsset();
     AspectFoo expectedAspectFoo = new AspectFoo().setValue("foo");
     AspectBar expectedAspectBar = new AspectBar().setValue("bar");
-    asset.setAspectFoo(expectedAspectFoo);
-    asset.setAspectBar(expectedAspectBar);
+    asset.setFoo(expectedAspectFoo);
+    asset.setBar(expectedAspectBar);
 
     List<? extends RecordTemplate> aspects = ModelUtils.getAspectsFromAsset(asset);
 
@@ -601,10 +601,10 @@ public void testNewAsset() {
         ModelUtils.newAsset(EntityAsset.class, expectedUrn, Lists.newArrayList(aspectUnion1, aspectUnion2));
 
     assertEquals(asset.getUrn(), expectedUrn);
-    assertTrue(asset.hasAspectFoo());
-    assertTrue(asset.hasAspectBar());
-    assertEquals(asset.getAspectFoo(), expectedFoo);
-    assertEquals(asset.getAspectBar(), expectedBar);
+    assertTrue(asset.hasFoo());
+    assertTrue(asset.hasBar());
+    assertEquals(asset.getFoo(), expectedFoo);
+    assertEquals(asset.getBar(), expectedBar);
     assertFalse(asset.hasAspectAttributes());
   }
 
@@ -626,10 +626,10 @@ public void testConvertSnapshotToAsset() {
     EntityAsset entityAsset = ModelUtils.convertSnapshotToAsset(EntityAsset.class, snapshot);
 
     assertEquals(entityAsset.getUrn(), expectedUrn);
-    assertTrue(entityAsset.hasAspectFoo());
-    assertTrue(entityAsset.hasAspectBar());
-    assertEquals(entityAsset.getAspectFoo(), expectedFoo);
-    assertEquals(entityAsset.getAspectBar(), expectedBar);
+    assertTrue(entityAsset.hasFoo());
+    assertTrue(entityAsset.hasBar());
+    assertEquals(entityAsset.getFoo(), expectedFoo);
+    assertEquals(entityAsset.getBar(), expectedBar);
     assertFalse(entityAsset.hasAspectAttributes());
   }
 
@@ -651,9 +651,9 @@ public void testConvertInternalSnapshotToAsset() {
     EntityAsset entityAsset = ModelUtils.convertSnapshotToAsset(EntityAsset.class, internalEntitySnapshot);
 
     assertEquals(entityAsset.getUrn(), expectedUrn);
-    assertEquals(entityAsset.getAspectFoo(), expectedFoo);
+    assertEquals(entityAsset.getFoo(), expectedFoo);
     assertEquals(entityAsset.getAspectFooBar(), expectedFooBar);
-    assertFalse(entityAsset.hasAspectBar());
+    assertFalse(entityAsset.hasBar());
     assertFalse(entityAsset.hasAspectAttributes());
   }
 
@@ -665,8 +665,8 @@ public void testConvertAssetToInternalSnapshot() {
     AspectBar expectedBar = new AspectBar().setValue("bar");
     AspectFooBar expectedFooBar = new AspectFooBar().setBars(new BarUrnArray(new BarUrn(2)));
     asset.setUrn(expectedUrn);
-    asset.setAspectFoo(expectedFoo);
-    asset.setAspectBar(expectedBar);
+    asset.setFoo(expectedFoo);
+    asset.setBar(expectedBar);
     asset.setAspectFooBar(expectedFooBar);
 
     InternalEntitySnapshot internalEntitySnapshot =
diff --git a/testing/test-models/src/main/pegasus/com/linkedin/testing/EntityAsset.pdl b/testing/test-models/src/main/pegasus/com/linkedin/testing/EntityAsset.pdl
@@ -15,20 +15,23 @@ record EntityAsset {
   /**
    * For unit tests
    */
-  aspectFoo: optional AspectFoo,
+  foo: optional AspectFoo,
+
   /**
    * For unit tests
    */
-  aspectBar: optional AspectBar,
+  bar: optional AspectBar,
 
   /**
    * For unit tests
    */
   aspectFooEvolved: optional AspectFooEvolved,
+
   /**
    * For unit tests
    */
   aspectFooBar: optional AspectFooBar,
+
   /**
    * For unit tests
    */
