From 53ed1f5a32427f597fefb7c000678503701ab06d Mon Sep 17 00:00:00 2001
From: Yang Yang <yanyang@linkedin.com>
Date: Sat, 28 Sep 2024 20:25:33 -0700
Subject: [PATCH] add getAspectTypesFromAssetType method (#438)

Co-authored-by: Yang Yang <yanyang@yanyang-mn3105.linkedin.biz>
---
 .../metadata/dao/utils/ModelUtils.java        | 33 ++++++++++++++++++-
 .../metadata/dao/utils/ModelUtilsTest.java    | 21 ++++++++++++
 2 files changed, 53 insertions(+), 1 deletion(-)

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
index 57677ce1a..dd4a967c9 100644
--- 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
@@ -333,7 +333,38 @@ public static <SNAPSHOT extends RecordTemplate> List<RecordTemplate> getAspectsF
   }
 
   /**
-   * Extracts the list of aspects in an asset.
+   * Get aspect types from asset type.
+   * @param assetClass asset class
+   * @return A list of aspect classes used in the asset
+   */
+  public static <ASSET extends RecordTemplate> List<Class<? extends RecordTemplate>> getAspectTypesFromAssetType(
+      @Nonnull Class<ASSET> assetClass) {
+    try {
+      final List<Class<? extends RecordTemplate>> aspectTypes = new ArrayList<>();
+      final Field[] assetFields = assetClass.getDeclaredFields();
+      for (final Field assetField : assetFields) {
+        if (assetField.getName().startsWith(FIELD_FIELD_PREFIX)) {
+          final String assetFieldName = assetField.getName().substring(FIELD_FIELD_PREFIX.length());
+          if (assetFieldName.equalsIgnoreCase(URN_FIELD)) {
+            continue;
+          }
+          String methodName = "get" + assetFieldName;
+          Class<?> returnType = assetClass.getMethod(methodName).getReturnType();
+          if (RecordTemplate.class.isAssignableFrom(returnType)) {
+            aspectTypes.add(returnType.asSubclass(RecordTemplate.class));
+          }
+        }
+      }
+      return aspectTypes;
+    } catch (NoSuchMethodException e) {
+      throw new RuntimeException(e);
+    }
+  }
+
+
+
+  /**
+   * Extracts the list of (non-null) aspects in an asset.
    *
    * @param asset the asset to extract aspects from
    * @param <ASSET> must be a valid asset model defined in com.linkedin.metadata.asset
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
index 058eb975d..5019475a1 100644
--- 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
@@ -7,6 +7,7 @@
 import com.linkedin.data.template.UnionTemplate;
 import com.linkedin.metadata.validator.NullFieldException;
 import com.linkedin.testing.AspectAttributes;
+import com.linkedin.testing.AspectFooEvolved;
 import com.linkedin.testing.AspectUnionWithSoftDeletedAspect;
 import com.linkedin.testing.BarAsset;
 import com.linkedin.testing.BarUrnArray;
@@ -328,6 +329,26 @@ public void testGetAspectsFromAsset() {
     assertEquals(aspects.size(), 2);
     assertEquals(aspects.get(0), expectedAspectFoo);
     assertEquals(aspects.get(1), expectedAspectBar);
+
+
+    BarAsset barAsset = new BarAsset();
+    aspects = ModelUtils.getAspectsFromAsset(barAsset);
+    assertTrue(aspects.isEmpty());
+  }
+
+  @Test
+  public void testGetAspectTypesFromAssetType() {
+    List<Class<? extends RecordTemplate>> assetTypes = ModelUtils.getAspectTypesFromAssetType(BarAsset.class);
+    assertEquals(assetTypes.size(), 1);
+    assertEquals(assetTypes.get(0), AspectBar.class);
+
+    assetTypes = ModelUtils.getAspectTypesFromAssetType(EntityAsset.class);
+    assertEquals(assetTypes.size(), 5);
+    assertEquals(assetTypes.get(0), AspectFoo.class);
+    assertEquals(assetTypes.get(1), AspectBar.class);
+    assertEquals(assetTypes.get(2), AspectFooEvolved.class);
+    assertEquals(assetTypes.get(3), AspectFooBar.class);
+    assertEquals(assetTypes.get(4), AspectAttributes.class);
   }
 
   @Test