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>
This commit is contained in:
Jesse Jia 2025-09-06 07:57:01 -07:00 committed by GitHub
parent 80f206290e
commit 55817d14b8
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
8 changed files with 91 additions and 11 deletions

View File

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

View File

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

View File

@ -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")

View File

@ -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(

View File

@ -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) {

View File

@ -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();
}

View File

@ -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)))));

View File

@ -113,7 +113,8 @@
},
"response": {
"exclude_regex_paths": [
"root\\['scrollId'\\]"
"root\\['scrollId'\\]",
"root\\['totalCount'\\]"
],
"json": {
"entities": [
@ -188,7 +189,63 @@
}
}
}
]
],
"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": [