fix(noCode): Improving efficiency of EntityService "listLatestAspects…

…" API (#2711)
datahub-project · Jun 18, 2021 · 550a9de · 550a9de
1 parent 1efc08d
commit 550a9de
Show file tree

Hide file tree

Showing 5 changed files with 14 additions and 6 deletions.
diff --git a/gms/impl/src/main/java/com/linkedin/metadata/resources/dataplatform/DataPlatforms.java b/gms/impl/src/main/java/com/linkedin/metadata/resources/dataplatform/DataPlatforms.java
@@ -95,6 +95,7 @@ public Task<DataPlatform> get(
   @RestMethod.GetAll
   public Task<List<DataPlatform>> getAllDataPlatforms(@Nonnull @PagingContextParam(defaultCount = 100) PagingContext pagingContext) {
     return Task.value(_entityService.listLatestAspects(
+        "dataPlatform",
         "dataPlatformInfo",
         pagingContext.getStart(),
         pagingContext.getCount())

diff --git a/metadata-io/src/main/java/com/linkedin/metadata/entity/EntityService.java b/metadata-io/src/main/java/com/linkedin/metadata/entity/EntityService.java
@@ -110,17 +110,19 @@ public abstract VersionedAspect getVersionedAspect(
       long version);
 
   /**
-   * Retrieves a list of all persisted aspects with a specific name, sorted by corresponding urn.
+   * Retrieves a list of all aspects belonging to an entity of a particular type, sorted by urn.
    *
    * Note that once we drop support for legacy 'getAllDataPlatforms' endpoint,
    * we can drop support for this unless otherwise required. Only visible for backwards compatibility.
    *
-   * @param aspectName name of the aspect requested
+   * @param entityName name of the entity type the aspect belongs to, e.g. 'dataset'
+   * @param aspectName name of the aspect requested, e.g. 'ownership'
    * @param start the starting index of the returned aspects, used in pagination
    * @param count the count of the aspects to be returned, used in pagination
    * @return a {@link ListResult} of {@link RecordTemplate}s representing the requested aspect.
    */
   public abstract ListResult<RecordTemplate> listLatestAspects(
+      @Nonnull final String entityName,
       @Nonnull final String aspectName,
       final int start,
       int count);

diff --git a/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/EbeanAspectDao.java b/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/EbeanAspectDao.java
@@ -396,23 +396,27 @@ public ListResult<String> listAspectMetadata(
 
   @Nonnull
   public ListResult<String> listLatestAspectMetadata(
+      @Nonnull final String entityName,
       @Nonnull final String aspectName,
       final int start,
       final int pageSize) {
-    return listAspectMetadata(aspectName, LATEST_ASPECT_VERSION, start, pageSize);
+    return listAspectMetadata(entityName, aspectName, LATEST_ASPECT_VERSION, start, pageSize);
   }
 
   @Nonnull
   public ListResult<String> listAspectMetadata(
+      @Nonnull final String entityName,
       @Nonnull final String aspectName,
       final long version,
       final int start,
       final int pageSize) {
     validateConnection();
 
+    final String urnPrefixMatcher = "urn:li:" + entityName + ":%";
     final PagedList<EbeanAspectV2> pagedList = _server.find(EbeanAspectV2.class)
         .select(EbeanAspectV2.ALL_COLUMNS)
         .where()
+        .like(EbeanAspectV2.URN_COLUMN, urnPrefixMatcher)
         .eq(EbeanAspectV2.ASPECT_COLUMN, aspectName)
         .eq(EbeanAspectV2.VERSION_COLUMN, version)
         .setFirstRow(start)

diff --git a/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/EbeanEntityService.java b/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/EbeanEntityService.java
@@ -138,11 +138,12 @@ public VersionedAspect getVersionedAspect(@Nonnull Urn urn, @Nonnull String aspe
   @Override
   @Nonnull
   public ListResult<RecordTemplate> listLatestAspects(
+      @Nonnull final String entityName,
       @Nonnull final String aspectName,
       final int start,
       int count) {
 
-    final ListResult<String> aspectMetadataList = _entityDao.listLatestAspectMetadata(aspectName, start, count);
+    final ListResult<String> aspectMetadataList = _entityDao.listLatestAspectMetadata(entityName, aspectName, start, count);
 
     final List<RecordTemplate> aspects = new ArrayList<>();
     for (int i = 0; i < aspectMetadataList.getValues().size(); i++) {

diff --git a/metadata-io/src/test/java/com/linkedin/metadata/entity/EbeanEntityServiceTest.java b/metadata-io/src/test/java/com/linkedin/metadata/entity/EbeanEntityServiceTest.java
@@ -201,7 +201,7 @@ public void testIngestListLatestAspects() throws Exception {
     _entityService.ingestAspect(entityUrn3, aspectName, writeAspect3, TEST_AUDIT_STAMP);
 
     // List aspects
-    ListResult<RecordTemplate> batch1 = _entityService.listLatestAspects(aspectName, 0, 2);
+    ListResult<RecordTemplate> batch1 = _entityService.listLatestAspects(entityUrn1.getEntityType(), aspectName, 0, 2);
 
     assertEquals(2, batch1.getNextStart());
     assertEquals(2, batch1.getPageSize());
@@ -211,7 +211,7 @@ public void testIngestListLatestAspects() throws Exception {
     assertTrue(DataTemplateUtil.areEqual(writeAspect1,  batch1.getValues().get(0)));
     assertTrue(DataTemplateUtil.areEqual(writeAspect2,  batch1.getValues().get(1)));
 
-    ListResult<RecordTemplate> batch2 = _entityService.listLatestAspects(aspectName, 2, 2);
+    ListResult<RecordTemplate> batch2 = _entityService.listLatestAspects(entityUrn1.getEntityType(), aspectName, 2, 2);
     assertEquals(1, batch2.getValues().size());
     assertTrue(DataTemplateUtil.areEqual(writeAspect3,  batch2.getValues().get(0)));
   }