use of org.hypertrace.entity.query.service.v1.BulkEntityUpdateRequest.EntityUpdateInfo in project entity-service by hypertrace.
the class EntityQueryServiceTest method testBulkUpdate.
@Test
public void testBulkUpdate() throws InterruptedException {
Entity.Builder apiEntityBuilder1 = Entity.newBuilder().setTenantId(TENANT_ID).setEntityType(EntityType.API.name()).setEntityName("api1").putIdentifyingAttributes(EntityConstants.getValue(ServiceAttribute.SERVICE_ATTRIBUTE_ID), createAttribute(SERVICE_ID)).putIdentifyingAttributes(EntityConstants.getValue(ApiAttribute.API_ATTRIBUTE_NAME), createAttribute("api1")).putIdentifyingAttributes(EntityConstants.getValue(ApiAttribute.API_ATTRIBUTE_API_TYPE), createAttribute(API_TYPE));
apiEntityBuilder1.putAttributes(apiAttributesMap.get(API_DISCOVERY_STATE_ATTR), createAttribute("DISCOVERED")).putAttributes(apiAttributesMap.get(API_HTTP_METHOD_ATTR), createAttribute("GET"));
Entity entity1 = entityDataServiceClient.upsert(apiEntityBuilder1.build());
Entity.Builder apiEntityBuilder2 = Entity.newBuilder().setTenantId(TENANT_ID).setEntityType(EntityType.API.name()).setEntityName("api2").putIdentifyingAttributes(EntityConstants.getValue(ServiceAttribute.SERVICE_ATTRIBUTE_ID), createAttribute(SERVICE_ID)).putIdentifyingAttributes(EntityConstants.getValue(ApiAttribute.API_ATTRIBUTE_NAME), createAttribute("api2")).putIdentifyingAttributes(EntityConstants.getValue(ApiAttribute.API_ATTRIBUTE_API_TYPE), createAttribute(API_TYPE));
apiEntityBuilder2.putAttributes(apiAttributesMap.get(API_DISCOVERY_STATE_ATTR), createAttribute("UNDER_DISCOVERY")).putAttributes(apiAttributesMap.get(API_HTTP_METHOD_ATTR), createAttribute("GET"));
Entity entity2 = entityDataServiceClient.upsert(apiEntityBuilder2.build());
// create BulkUpdate request
UpdateOperation update1 = UpdateOperation.newBuilder().setSetAttribute(SetAttribute.newBuilder().setAttribute(ColumnIdentifier.newBuilder().setColumnName(API_DISCOVERY_STATE_ATTR)).setValue(LiteralConstant.newBuilder().setValue(org.hypertrace.entity.query.service.v1.Value.newBuilder().setString("DISCOVERED")))).build();
UpdateOperation update2 = UpdateOperation.newBuilder().setSetAttribute(SetAttribute.newBuilder().setAttribute(ColumnIdentifier.newBuilder().setColumnName(API_HTTP_METHOD_ATTR)).setValue(LiteralConstant.newBuilder().setValue(org.hypertrace.entity.query.service.v1.Value.newBuilder().setString("POST")))).build();
EntityUpdateInfo updateInfo1 = EntityUpdateInfo.newBuilder().addUpdateOperation(update2).build();
EntityUpdateInfo updateInfo2 = EntityUpdateInfo.newBuilder().addUpdateOperation(update1).addUpdateOperation(update2).build();
BulkEntityUpdateRequest bulkUpdateRequest = BulkEntityUpdateRequest.newBuilder().setEntityType(EntityType.API.name()).putEntities(entity1.getEntityId(), updateInfo1).putEntities(entity2.getEntityId(), updateInfo2).build();
Iterator<ResultSetChunk> iterator = GrpcClientRequestContextUtil.executeWithHeadersContext(HEADERS, () -> entityQueryServiceClient.bulkUpdate(bulkUpdateRequest));
while (iterator.hasNext()) {
continue;
}
// Add a small delay for the update to reflect
Thread.sleep(500);
EntityQueryRequest entityQueryRequest = EntityQueryRequest.newBuilder().setEntityType(EntityType.API.name()).addSelection(createExpression(API_DISCOVERY_STATE_ATTR)).addSelection(createExpression(API_HTTP_METHOD_ATTR)).build();
Iterator<ResultSetChunk> resultSetChunkIterator = GrpcClientRequestContextUtil.executeWithHeadersContext(HEADERS, () -> entityQueryServiceClient.execute(entityQueryRequest));
List<String> values = new ArrayList<>();
while (resultSetChunkIterator.hasNext()) {
ResultSetChunk chunk = resultSetChunkIterator.next();
for (Row row : chunk.getRowList()) {
for (int i = 0; i < row.getColumnCount(); i++) {
String value = row.getColumnList().get(i).getString();
values.add(value);
}
}
}
assertEquals(4, values.size());
assertEquals("DISCOVERED", values.get(0));
assertEquals("POST", values.get(1));
assertEquals("DISCOVERED", values.get(2));
assertEquals("POST", values.get(3));
}
use of org.hypertrace.entity.query.service.v1.BulkEntityUpdateRequest.EntityUpdateInfo in project entity-service by hypertrace.
the class EntityQueryServiceTest method testBulkUpdateWithLabels.
@Test
public void testBulkUpdateWithLabels() {
Entity.Builder apiEntityBuilder2 = Entity.newBuilder().setTenantId(TENANT_ID).setEntityType(EntityType.API.name()).setEntityName("api2").putIdentifyingAttributes(EntityConstants.getValue(ServiceAttribute.SERVICE_ATTRIBUTE_ID), createAttribute(SERVICE_ID)).putIdentifyingAttributes(EntityConstants.getValue(ApiAttribute.API_ATTRIBUTE_NAME), createAttribute("api2")).putIdentifyingAttributes(EntityConstants.getValue(ApiAttribute.API_ATTRIBUTE_API_TYPE), createAttribute(API_TYPE));
apiEntityBuilder2.putAttributes(apiAttributesMap.get(API_LABELS_ATTR), createStringArrayAttribute(List.of("Label1")));
Entity entity2 = entityDataServiceClient.upsert(apiEntityBuilder2.build());
UpdateOperation update = UpdateOperation.newBuilder().setSetAttribute(SetAttribute.newBuilder().setAttribute(ColumnIdentifier.newBuilder().setColumnName(API_LABELS_ATTR)).setValue(LiteralConstant.newBuilder().setValue(org.hypertrace.entity.query.service.v1.Value.newBuilder().addAllStringArray(Collections.emptyList()).setValueType(STRING_ARRAY)))).build();
EntityUpdateInfo updateInfo = EntityUpdateInfo.newBuilder().addUpdateOperation(update).build();
BulkEntityUpdateRequest bulkUpdateRequest = BulkEntityUpdateRequest.newBuilder().setEntityType(EntityType.API.name()).putEntities(entity2.getEntityId(), updateInfo).build();
assertDoesNotThrow(() -> GrpcClientRequestContextUtil.executeWithHeadersContext(HEADERS, () -> entityQueryServiceClient.bulkUpdate(bulkUpdateRequest)));
}
use of org.hypertrace.entity.query.service.v1.BulkEntityUpdateRequest.EntityUpdateInfo in project entity-service by hypertrace.
the class EntityQueryServiceImpl method bulkUpdate.
@Override
public void bulkUpdate(BulkEntityUpdateRequest request, StreamObserver<ResultSetChunk> responseObserver) {
// Validations
RequestContext requestContext = RequestContext.CURRENT.get();
Optional<String> maybeTenantId = requestContext.getTenantId();
if (maybeTenantId.isEmpty()) {
responseObserver.onError(new ServiceException("Tenant id is missing in the request."));
return;
}
if (StringUtils.isEmpty(request.getEntityType())) {
responseObserver.onError(new ServiceException("Entity type is missing in the request."));
return;
}
if (request.getEntitiesCount() == 0) {
responseObserver.onError(new ServiceException("Entities are missing in the request."));
}
Map<String, EntityUpdateInfo> entitiesMap = request.getEntitiesMap();
try {
doBulkUpdate(requestContext, entitiesMap);
responseObserver.onCompleted();
} catch (Exception e) {
responseObserver.onError(new ServiceException("Error occurred while executing " + request, e));
}
}
Aggregations