[feat] add findEntitiesV2 to support logical expression filter (#568)

Author: ybz1013

dao-impl/ebean-dao/src/main/java/com/linkedin/metadata/dao/EbeanLocalRelationshipQueryDAO.java

@@ -110,10 +110,15 @@ public void setSchemaConfig(EbeanLocalDAO.SchemaConfig schemaConfig) {
   @Nonnull
   public <SNAPSHOT extends RecordTemplate> List<SNAPSHOT> findEntities(@Nonnull Class<SNAPSHOT> snapshotClass,
       @Nonnull LocalRelationshipFilter filter, int offset, int count) throws OperationNotSupportedException {
+    return findEntitiesCore(snapshotClass, filter, offset, count, false);
+  }
+
+  private <SNAPSHOT extends RecordTemplate> List<SNAPSHOT> findEntitiesCore(@Nonnull Class<SNAPSHOT> snapshotClass,
+      @Nonnull LocalRelationshipFilter filter, int offset, int count, boolean logicalExpressionFilterEnabled) throws OperationNotSupportedException {
     if (_schemaConfig == EbeanLocalDAO.SchemaConfig.OLD_SCHEMA_ONLY) {
       throw new OperationNotSupportedException("findEntities is not supported in OLD_SCHEMA_MODE");
     }
-    validateEntityFilter(filter, snapshotClass);
+    validateEntityFilter(filter, snapshotClass, logicalExpressionFilterEnabled);
 
     final String tableName = SQLSchemaUtils.getTableName(ModelUtils.getUrnTypeFromSnapshot(snapshotClass));
     final StringBuilder sqlBuilder = new StringBuilder();
@@ -129,6 +134,25 @@ public <SNAPSHOT extends RecordTemplate> List<SNAPSHOT> findEntities(@Nonnull Cl
         .collect(Collectors.toList());
   }
 
+  /**
+   * Finds a list of entities of a specific type based on the given filter on the entity.
+   * Similar to {@link #findEntities(Class, LocalRelationshipFilter, int, int)},
+   * but this method uses the LogicalExpressionLocalRelationshipCriterion in LocalRelationshipFilter.
+   * The SNAPSHOT class must be defined within com.linkedin.metadata.snapshot package in metadata-models.
+   * This method is not supported in OLD_SCHEMA_ONLY mode.
+   * @param snapshotClass the snapshot class to query.
+   * @param filter the filter to apply when querying. Uses `logicalExpressionCriteria` instead of `criteria`.
+   * @param offset the offset the query should start at. Ignored if set to a negative value.
+   * @param count the maximum number of entities to return. Ignored if set to a non-positive value.
+   * @return A list of entity records of class SNAPSHOT.
+   * @throws OperationNotSupportedException when called in OLD_SCHEMA_ONLY mode. This exception must be explicitly handled by the caller.
+   */
+  @Nonnull
+  public <SNAPSHOT extends RecordTemplate> List<SNAPSHOT> findEntitiesV2(@Nonnull Class<SNAPSHOT> snapshotClass,
+      @Nonnull LocalRelationshipFilter filter, int offset, int count) throws OperationNotSupportedException {
+    return findEntitiesCore(snapshotClass, filter, offset, count, true);
+  }
+
   /**
    * Finds a list of entities of a specific type based on the given source, destination, and relationship filters.
    * Every SNAPSHOT class must be defined within com.linkedin.metadata.snapshot package in metadata-models.
@@ -449,18 +473,15 @@ private String getMgEntityTableName(@Nullable String entityType) {
     return SQLSchemaUtils.getTableName(entityType);
   }
 
-  /**
-   * Validate:
-   * 1. The entity filter only contains supported conditions.
-   * 2. if the entity class is null, then the filter should be empty.
-   * If any of above is violated, throw an IllegalArgumentException.
-   */
   private <ENTITY extends RecordTemplate> void validateEntityFilter(@Nonnull LocalRelationshipFilter filter, @Nullable Class<ENTITY> entityClass) {
-    if (entityClass == null && filter.hasCriteria() && filter.getCriteria().size() > 0) {
-      throw new IllegalArgumentException("Entity class is null but filter is not empty.");
-    }
+    validateEntityFilter(filter, entityClass, false);
+  }
 
-    validateFilterCriteria(filter, false);
+  private <ENTITY extends RecordTemplate> void validateEntityFilter(@Nonnull LocalRelationshipFilter filter, @Nullable Class<ENTITY> entityClass,
+      boolean logicalExpressionFilterEnabled) {
+    validateEntityTypeAndFilter(filter,
+        entityClass != null ? ModelUtils.getUrnTypeFromSnapshot(entityClass) : null,
+        logicalExpressionFilterEnabled);
   }
 
   /**
@@ -498,10 +519,10 @@ private static void validateLogicalExpressionFilter(@Nonnull LocalRelationshipFi
 
     if (logicalExpressionFilterEnabled && filter.hasCriteria()) {
         throw new IllegalArgumentException(
-            "Please do not use the 'criteria' field and use the 'logicalExpressionCriteria' field instead for findRelationshipsV4 API.");
+            "Please do not use the 'criteria' field and use the 'logicalExpressionCriteria' field instead for this API.");
     } else if (!logicalExpressionFilterEnabled && filter.hasLogicalExpressionCriteria()) {
         throw new IllegalArgumentException(
-            "Please do not use the 'logicalExpressionCriteria' field and use the 'criteria' field instead for findRelationshipsV2/V3 API.");
+            "Please do not use the 'logicalExpressionCriteria' field and use the 'criteria' field instead for this API.");
     }
   }
 

dao-impl/ebean-dao/src/test/java/com/linkedin/metadata/dao/localrelationship/EbeanLocalRelationshipQueryDAOTest.java

@@ -175,6 +175,68 @@ public void testFindOneEntityTwoAspects() throws URISyntaxException, OperationNo
     assertEquals(fooSnapshotList.get(0).getAspects(), expected);
   }
 
+  @Test
+  public void testFindEntitiesV2WithV1Filter() throws URISyntaxException, OperationNotSupportedException {
+    // Ingest data
+    _fooUrnEBeanLocalAccess.add(new FooUrn(1), new AspectFoo().setValue("foo"), AspectFoo.class, new AuditStamp(), null, false);
+
+    // Prepare filter
+    LocalRelationshipCriterion filterCriterion = EBeanDAOUtils.buildRelationshipFieldCriterion(LocalRelationshipValue.create("foo"),
+        Condition.EQUAL,
+        new AspectField().setAspect(AspectFoo.class.getCanonicalName()).setPath("/value"));
+
+    LocalRelationshipFilter filter = new LocalRelationshipFilter().setCriteria(new LocalRelationshipCriterionArray(filterCriterion));
+
+    assertThrows(IllegalArgumentException.class, () -> _localRelationshipQueryDAO.findEntitiesV2(FooSnapshot.class, filter, 0, 10));
+  }
+
+  @Test
+  public void testFindOneEntityV2() throws URISyntaxException, OperationNotSupportedException {
+    // Ingest data
+    _fooUrnEBeanLocalAccess.add(new FooUrn(1), new AspectFoo().setValue("foo"), AspectFoo.class, new AuditStamp(), null, false);
+
+    // Prepare filter
+    LocalRelationshipCriterion filterCriterion = EBeanDAOUtils.buildRelationshipFieldCriterion(LocalRelationshipValue.create("foo"),
+        Condition.EQUAL,
+        new AspectField().setAspect(AspectFoo.class.getCanonicalName()).setPath("/value"));
+
+    LocalRelationshipFilter filter = new LocalRelationshipFilter().setLogicalExpressionCriteria(
+        wrapCriterionAsLogicalExpression(filterCriterion));
+    List<FooSnapshot> fooSnapshotList = _localRelationshipQueryDAO.findEntitiesV2(FooSnapshot.class, filter, 0, 10);
+
+    assertEquals(fooSnapshotList.size(), 1);
+    assertEquals(fooSnapshotList.get(0).getAspects().size(), 1);
+    assertEquals(fooSnapshotList.get(0).getAspects().get(0).getAspectFoo(), new AspectFoo().setValue("foo"));
+  }
+
+  @Test
+  public void testFindOneEntityTwoAspectsV2() throws URISyntaxException, OperationNotSupportedException {
+    // Ingest data
+    _fooUrnEBeanLocalAccess.add(new FooUrn(1), new AspectFoo().setValue("foo"), AspectFoo.class, new AuditStamp(), null, false);
+    _fooUrnEBeanLocalAccess.add(new FooUrn(1), new AspectBar().setValue("bar"), AspectBar.class, new AuditStamp(), null, false);
+
+    // Prepare filter
+    LocalRelationshipCriterion filterCriterion = EBeanDAOUtils.buildRelationshipFieldCriterion(LocalRelationshipValue.create("foo"),
+        Condition.EQUAL,
+        new AspectField().setAspect(AspectFoo.class.getCanonicalName()).setPath("/value"));
+
+    LocalRelationshipFilter filter = new LocalRelationshipFilter().setLogicalExpressionCriteria(
+        wrapCriterionAsLogicalExpression(filterCriterion));
+
+    List<FooSnapshot> fooSnapshotList = _localRelationshipQueryDAO.findEntitiesV2(FooSnapshot.class, filter, 0, 10);
+
+    assertEquals(fooSnapshotList.size(), 1);
+    assertEquals(fooSnapshotList.get(0).getAspects().size(), 2);
+    EntityAspectUnion fooAspectUnion = new EntityAspectUnion();
+    fooAspectUnion.setAspectFoo(new AspectFoo().setValue("foo"));
+    EntityAspectUnion barAspectUnion = new EntityAspectUnion();
+    barAspectUnion.setAspectBar(new AspectBar().setValue("bar"));
+
+    EntityAspectUnionArray expected = new EntityAspectUnionArray(fooAspectUnion, barAspectUnion);
+
+    assertEquals(fooSnapshotList.get(0).getAspects(), expected);
+  }
+
   @DataProvider(name = "schemaConfig")
   public static Object[][] schemaConfig() {
     return new Object[][] {