feat(openapi-v3) enhance scroll API (#14677)

Co-authored-by: jjia <jjia@netflix.com>
Co-authored-by: david-leifker <114954101+david-leifker@users.noreply.github.com>

metadata-service/openapi-servlet/models/src/main/java/io/datahubproject/openapi/v3/models/Criterion.java

@@ -31,7 +31,8 @@ public class Criterion {
     STARTS_WITH,
     ENDS_WITH,
     EXISTS,
-    IN
+    IN,
+    CONTAIN
   }
 
   /** Convert this criterion to its counterpart in RecordTemplate. */

metadata-service/openapi-servlet/models/src/main/java/io/datahubproject/openapi/v3/models/GenericEntityScrollResultV3.java

@@ -11,4 +11,5 @@ public class GenericEntityScrollResultV3
     implements GenericEntityScrollResult<GenericAspectV3, GenericEntityV3> {
   private String scrollId;
   private List<GenericEntityV3> entities;
+  private int totalCount;
 }

metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/controller/GenericEntitiesController.java

@@ -120,7 +120,8 @@ public abstract class GenericEntitiesController<
       Set<String> aspectNames,
       boolean withSystemMetadata,
       @Nullable String scrollId,
-      boolean expandEmpty)
+      boolean expandEmpty,
+      int totalCount)
       throws URISyntaxException;
 
   protected List<E> buildEntityList(
@@ -294,7 +295,8 @@ public abstract class GenericEntitiesController<
             mergedAspects,
             withSystemMetadata,
             result.getScrollId(),
-            true));
+            true,
+            result.getNumEntities()));
   }
 
   @Tag(name = "Generic Entities")

metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/v2/controller/EntityController.java

@@ -129,7 +129,8 @@ public class EntityController
       Set<String> aspectNames,
       boolean withSystemMetadata,
       @Nullable String scrollId,
-      boolean expandEmpty)
+      boolean expandEmpty,
+      int totalCount)
       throws URISyntaxException {
     return GenericEntityScrollResultV2.builder()
         .results(

metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/v3/OpenAPIV3Generator.java

@@ -1403,7 +1403,9 @@ public class OpenAPIV3Generator {
                 "condition",
                 newSchema()
                     .type(TYPE_STRING)
-                    ._enum(Arrays.asList("EQUAL", "STARTS_WITH", "ENDS_WITH", "EXISTS", "IN"))
+                    ._enum(
+                        Arrays.asList(
+                            "EQUAL", "STARTS_WITH", "ENDS_WITH", "EXISTS", "IN", "CONTAIN"))
                     ._default("EQUAL")
                     .description("The condition for the criterion."))
             .addProperties(
@@ -1559,7 +1561,12 @@ public class OpenAPIV3Generator {
                     newSchema()
                         .$ref(
                             String.format(
-                                "#/components/schemas/%s%s", ENTITIES, ENTITY_RESPONSE_SUFFIX))));
+                                "#/components/schemas/%s%s", ENTITIES, ENTITY_RESPONSE_SUFFIX))))
+        .addProperty(
+            "totalCount",
+            newSchema()
+                .type(TYPE_INTEGER)
+                .description("Total number of entities satisfy the criteria."));
   }
 
   private static Schema buildEntityScrollSchema(final EntitySpec entity) {

metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/v3/controller/EntityController.java

@@ -266,7 +266,8 @@ public class EntityController
             entityAspectsBody.getAspects(),
             withSystemMetadata,
             result.getScrollId(),
-            entityAspectsBody.getAspects() != null));
+            entityAspectsBody.getAspects() != null,
+            result.getNumEntities()));
   }
 
   @Tag(name = "EntityVersioning")
@@ -596,13 +597,15 @@ public class EntityController
       Set<String> aspectNames,
       boolean withSystemMetadata,
       @Nullable String scrollId,
-      boolean expandEmpty)
+      boolean expandEmpty,
+      int totalCount)
       throws URISyntaxException {
     return GenericEntityScrollResultV3.builder()
         .entities(
             toRecordTemplates(
                 opContext, searchEntities, aspectNames, withSystemMetadata, expandEmpty))
         .scrollId(scrollId)
+        .totalCount(totalCount)
         .build();
   }
 

metadata-service/openapi-servlet/src/test/java/io/datahubproject/openapi/v3/controller/EntityControllerTest.java

@@ -154,6 +154,7 @@ public class EntityControllerTest extends AbstractTestNGSpringContextTests {
     // Mock scroll ascending/descending results
     ScrollResult expectedResultAscending =
         new ScrollResult()
+            .setNumEntities(3)
             .setEntities(
                 new SearchEntityArray(
                     List.of(
@@ -172,6 +173,7 @@ public class EntityControllerTest extends AbstractTestNGSpringContextTests {
         .thenReturn(expectedResultAscending);
     ScrollResult expectedResultDescending =
         new ScrollResult()
+            .setNumEntities(3)
             .setEntities(
                 new SearchEntityArray(
                     List.of(
@@ -603,6 +605,7 @@ public class EntityControllerTest extends AbstractTestNGSpringContextTests {
 
     ScrollResult expectedResult =
         new ScrollResult()
+            .setNumEntities(2)
             .setEntities(
                 new SearchEntityArray(
                     List.of(
@@ -649,6 +652,7 @@ public class EntityControllerTest extends AbstractTestNGSpringContextTests {
 
     ScrollResult expectedResult =
         new ScrollResult()
+            .setNumEntities(1)
             .setEntities(
                 new SearchEntityArray(List.of(new SearchEntity().setEntity(TEST_URNS.get(0)))))
             .setScrollId("test-scroll-id");
@@ -694,6 +698,7 @@ public class EntityControllerTest extends AbstractTestNGSpringContextTests {
 
     ScrollResult expectedResult =
         new ScrollResult()
+            .setNumEntities(2)
             .setEntities(
                 new SearchEntityArray(
                     List.of(
@@ -768,6 +773,7 @@ public class EntityControllerTest extends AbstractTestNGSpringContextTests {
 
     ScrollResult expectedResult =
         new ScrollResult()
+            .setNumEntities(1)
             .setEntities(
                 new SearchEntityArray(List.of(new SearchEntity().setEntity(TEST_URNS.get(0)))));
 

smoke-test/tests/openapi/v3/entities.json

@@ -113,7 +113,8 @@
     },
     "response": {
       "exclude_regex_paths": [
-        "root\\['scrollId'\\]"
+        "root\\['scrollId'\\]",
+        "root\\['totalCount'\\]"
       ],
       "json": {
         "entities": [
@@ -188,8 +189,64 @@
               }
             }
           }
+        ],
+        "totalCount": 1
+      }
+    }
+  },
+  {
+    "request": {
+      "url": "/openapi/v3/entity/scroll",
+      "params": {
+        "count": "1"
+      },
+      "description": "Generic scroll with filters",
+      "json": {
+        "aspects": [],
+        "entities": [
+          "dataPlatform"
+        ],
+        "filter": {
+          "and": [
+            {
+              "criteria": [
+                {
+                  "field": "name",
+                  "values": ["vert"],
+                  "condition": "CONTAIN"
+                }
               ]
             }
+          ]
+        }
+      }
+    },
+    "response": {
+      "exclude_regex_paths": [
+        "root\\['scrollId'\\]"
+      ],
+      "json": {
+        "entities": [
+          {
+            "urn": "urn:li:dataPlatform:vertexai",
+            "dataPlatformKey": {
+              "value": {
+                "platformName": "vertexai"
+              }
+            },
+            "dataPlatformInfo": {
+              "value": {
+                "name": "vertexai",
+                "datasetNameDelimiter": ".",
+                "type": "OTHERS",
+                "displayName": "vertexai",
+                "logoUrl": "/assets/platforms/vertexai.png"
+              }
+            }
+          }
+        ],
+        "totalCount": 2
+      }
     }
   },
   {
@@ -212,6 +269,7 @@
     "response": {
       "exclude_regex_paths": [
         "root\\['scrollId'\\]",
+        "root\\['totalCount'\\]",
         "root\\['entities'\\]\\[0\\]\\['datasetProperties'\\]\\['systemMetadata'\\]\\['properties'\\]\\['sourceIP'\\]",
         "root\\['entities'\\]\\[0\\]\\['datasetProperties'\\]\\['systemMetadata'\\]\\['properties'\\]\\['telemetryTraceId'\\]",
         "root\\['entities'\\]\\[0\\]\\['datasetProperties'\\]\\['auditStamp'\\]",
@@ -263,7 +321,8 @@
     },
     "response": {
       "exclude_regex_paths": [
-        "root\\['scrollId'\\]"
+        "root\\['scrollId'\\]",
+        "root\\['totalCount'\\]"
       ],
       "json": {
         "entities": [