fix: reflect the internal models to asset models from param and retur…

…n types (#418) * fix: reflect the internal models to asset models from param and return types. * rename variables and add TODO.
linkedin · Sep 7, 2024 · e658496 · e658496
1 parent 1909fd0
commit e658496
Show file tree

Hide file tree

Showing 4 changed files with 47 additions and 36 deletions.
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,31 @@ public static <ASSET extends RecordTemplate, ASPECT_UNION extends UnionTemplate>
       }
       RecordUtils.setRecordTemplatePrimitiveField(asset, URN_FIELD, urn);
 
+      // TODO: cache the asset methods loading
+      final Map<String, String> aspectTypeToAssetSetterMap = new HashMap<>();
+      for (final Method assetMethod : assetClass.getDeclaredMethods()) {
+        if (assetMethod.getName().startsWith("set") && assetMethod.getParameterTypes().length > 0) {
+          aspectTypeToAssetSetterMap.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);
-                }
-              }
+        // TODO: cache the aspect union methods loading
+        final Map<String, String> aspectTypeToAspectUnionGetterMap = new HashMap<>();
+        for (final Method aspectUnionMethod : aspect.getClass().getMethods()) {
+          if (aspectUnionMethod.getName().startsWith("get")) {
+            aspectTypeToAspectUnionGetterMap.put(aspectUnionMethod.getReturnType().getName(),
+                aspectUnionMethod.getName());
+          }
+        }
+        for (final String aspectType : aspectTypeToAssetSetterMap.keySet()) {
+          if (aspectTypeToAspectUnionGetterMap.containsKey(aspectType)) {
+            Object aspectValue =
+                aspect.getClass().getMethod(aspectTypeToAspectUnionGetterMap.get(aspectType)).invoke(aspect);
+            if (aspectValue != null) {
+              asset.getClass()
+                  .getMethod(aspectTypeToAssetSetterMap.get(aspectType), 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/restli-resources/src/test/java/com/linkedin/metadata/restli/BaseEntityResourceTest.java b/restli-resources/src/test/java/com/linkedin/metadata/restli/BaseEntityResourceTest.java
@@ -1458,8 +1458,8 @@ public void testIngestAsset() {
     AspectFoo foo = new AspectFoo().setValue("foo");
     AspectBar bar = new AspectBar().setValue("bar");
     asset.setUrn(urn);
-    asset.setAspectFoo(foo);
-    asset.setAspectBar(bar);
+    asset.setFoo(foo);
+    asset.setBar(bar);
     IngestionTrackingContext trackingContext = new IngestionTrackingContext();
 
     IngestionParams ingestionParams1 = new IngestionParams().setTestMode(true);
@@ -1503,9 +1503,9 @@ public void testGetAsset() {
 
     assertEquals(asset.getUrn(), urn);
 
-    assertEquals(asset.getAspectFoo(), foo);
+    assertEquals(asset.getFoo(), foo);
     assertEquals(asset.getAspectFooEvolved(), fooEvolved);
-    assertEquals(asset.getAspectBar(), bar);
+    assertEquals(asset.getBar(), bar);
     assertEquals(asset.getAspectFooBar(), fooBar);
     assertEquals(asset.getAspectAttributes(), attributes);
   }

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
    */