Examples with BulkEntityArrayAttributeUpdateRequest org.hypertrace.entity.query.service.v1.BulkEntityArrayAttributeUpdateRequest used on opensource projects

Example 1 with BulkEntityArrayAttributeUpdateRequest

use of org.hypertrace.entity.query.service.v1.BulkEntityArrayAttributeUpdateRequest in project entity-service by hypertrace.

the class EntityQueryServiceImplTest method testBulkUpdateEntityArrayAttribute.

@Test
void testBulkUpdateEntityArrayAttribute() throws Exception {
    List<String> entityIds = IntStream.rangeClosed(1, 5).mapToObj(i -> UUID.randomUUID().toString()).collect(Collectors.toList());
    BulkEntityArrayAttributeUpdateRequest request = BulkEntityArrayAttributeUpdateRequest.newBuilder().addAllEntityIds(entityIds).setAttribute(ColumnIdentifier.newBuilder().setColumnName(ATTRIBUTE_ID3).build()).setEntityType(TEST_ENTITY_TYPE).setOperation(BulkEntityArrayAttributeUpdateRequest.Operation.OPERATION_ADD).addAllValues(List.of(LiteralConstant.newBuilder().setValue(Value.newBuilder().setString("Label1")).build(), LiteralConstant.newBuilder().setValue(Value.newBuilder().setString("Label2")).build())).build();
    when(mockAttributeMapping.getDocStorePathByAttributeId(requestContext, ATTRIBUTE_ID3)).thenReturn(Optional.of(EDS_COLUMN_NAME3));
    StreamObserver<BulkEntityArrayAttributeUpdateResponse> mockResponseObserver = mock(StreamObserver.class);
    Context.current().withValue(RequestContext.CURRENT, mockRequestContextWithTenantId()).call(() -> {
        EntityQueryServiceImpl eqs = new EntityQueryServiceImpl(entitiesCollection, mockAttributeMapping, 1, false);
        eqs.bulkUpdateEntityArrayAttribute(request, mockResponseObserver);
        return null;
    });
    ArgumentCaptor<BulkArrayValueUpdateRequest> argumentCaptor = ArgumentCaptor.forClass(BulkArrayValueUpdateRequest.class);
    verify(entitiesCollection, times(1)).bulkOperationOnArrayValue(argumentCaptor.capture());
    BulkArrayValueUpdateRequest bulkArrayValueUpdateRequest = argumentCaptor.getValue();
    assertEquals(entityIds.stream().map(entityId -> new SingleValueKey("tenant1", entityId)).collect(Collectors.toCollection(LinkedHashSet::new)), bulkArrayValueUpdateRequest.getKeys());
    assertEquals(BulkArrayValueUpdateRequest.Operation.ADD, bulkArrayValueUpdateRequest.getOperation());
    assertEquals(EDS_COLUMN_NAME3 + ".valueList.values", bulkArrayValueUpdateRequest.getSubDocPath());
    List<Document> subDocuments = bulkArrayValueUpdateRequest.getSubDocuments();
    assertEquals(2, subDocuments.size());
    assertEquals("{\"value\":{\"string\":\"Label1\"}}", subDocuments.get(0).toString());
    assertEquals("{\"value\":{\"string\":\"Label2\"}}", subDocuments.get(1).toString());
}

Example 2 with BulkEntityArrayAttributeUpdateRequest

use of org.hypertrace.entity.query.service.v1.BulkEntityArrayAttributeUpdateRequest in project gateway-service by hypertrace.

the class EdsEntityUpdater method bulkUpdateEntities.

public BulkUpdateEntitiesResponse bulkUpdateEntities(BulkUpdateEntitiesRequest request, UpdateExecutionContext updateExecutionContext) {
    MultiValuedAttributeOperation multiValuedAttributeOperation = request.getOperation().getMultiValuedAttributeOperation();
    org.hypertrace.entity.query.service.v1.BulkEntityArrayAttributeUpdateRequest updateRequest = org.hypertrace.entity.query.service.v1.BulkEntityArrayAttributeUpdateRequest.newBuilder().setEntityType(request.getEntityType()).addAllEntityIds(request.getEntityIdsList()).setOperation(EntityServiceAndGatewayServiceConverter.convertToBulkEntityArrayAttributeUpdateOperation(multiValuedAttributeOperation.getType())).setAttribute(EntityServiceAndGatewayServiceConverter.convertToQueryColumnIdentifier(multiValuedAttributeOperation.getAttribute())).addAllValues(multiValuedAttributeOperation.getValuesList().stream().map(value -> EntityServiceAndGatewayServiceConverter.convertToQueryLiteral(value)).collect(Collectors.toList())).build();
    eqsClient.bulkUpdateEntityArrayAttribute(updateRequest, updateExecutionContext.getRequestHeaders());
    return BulkUpdateEntitiesResponse.newBuilder().build();
}

Example 3 with BulkEntityArrayAttributeUpdateRequest

use of org.hypertrace.entity.query.service.v1.BulkEntityArrayAttributeUpdateRequest in project entity-service by hypertrace.

the class EntityQueryServiceImpl method bulkUpdateEntityArrayAttribute.

@Override
public void bulkUpdateEntityArrayAttribute(BulkEntityArrayAttributeUpdateRequest request, StreamObserver<BulkEntityArrayAttributeUpdateResponse> responseObserver) {
    RequestContext requestContext = RequestContext.CURRENT.get();
    String tenantId = requestContext.getTenantId().orElse(null);
    if (isNull(tenantId)) {
        responseObserver.onError(new ServiceException("Tenant id is missing in the request."));
        return;
    }
    try {
        Set<Key> keys = request.getEntityIdsList().stream().map(entityId -> new SingleValueKey(tenantId, entityId)).collect(Collectors.toCollection(LinkedHashSet::new));
        String attributeId = request.getAttribute().getColumnName();
        String subDocPath = entityAttributeMapping.getDocStorePathByAttributeId(requestContext, attributeId).orElseThrow(() -> new IllegalArgumentException("Unknown attribute " + attributeId));
        List<Document> subDocuments = request.getValuesList().stream().map(this::convertToJsonDocument).collect(toUnmodifiableList());
        BulkArrayValueUpdateRequest bulkArrayValueUpdateRequest = new BulkArrayValueUpdateRequest(keys, subDocPath + ARRAY_VALUE_PATH_SUFFIX, getMatchingOperation(request.getOperation()), subDocuments);
        entitiesCollection.bulkOperationOnArrayValue(bulkArrayValueUpdateRequest);
        responseObserver.onNext(BulkEntityArrayAttributeUpdateResponse.newBuilder().build());
        responseObserver.onCompleted();
    } catch (Exception e) {
        responseObserver.onError(e);
    }
}

Example 4 with BulkEntityArrayAttributeUpdateRequest

use of org.hypertrace.entity.query.service.v1.BulkEntityArrayAttributeUpdateRequest in project entity-service by hypertrace.

the class EntityQueryServiceTest method testBulkArrayValueUpdateWithLabels.

@Test
public void testBulkArrayValueUpdateWithLabels() {
    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_LABELS_ATTR), createStringArrayAttribute(List.of("Label1")));
    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_LABELS_ATTR), createStringArrayAttribute(List.of("Label2")));
    Entity entity2 = entityDataServiceClient.upsert(apiEntityBuilder2.build());
    BulkEntityArrayAttributeUpdateRequest request = BulkEntityArrayAttributeUpdateRequest.newBuilder().build().newBuilder().setEntityType(EntityType.API.name()).addAllEntityIds(Set.of(entity1.getEntityId(), entity2.getEntityId())).setAttribute(ColumnIdentifier.newBuilder().setColumnName(API_LABELS_ATTR)).setOperation(BulkEntityArrayAttributeUpdateRequest.Operation.OPERATION_ADD).addAllValues(List.of(LiteralConstant.newBuilder().setValue(org.hypertrace.entity.query.service.v1.Value.newBuilder().setString("Label3").build()).build(), LiteralConstant.newBuilder().setValue(org.hypertrace.entity.query.service.v1.Value.newBuilder().setString("Label4").build()).build())).build();
    assertDoesNotThrow(() -> GrpcClientRequestContextUtil.executeWithHeadersContext(HEADERS, () -> entityQueryServiceClient.bulkUpdateEntityArrayAttribute(request)));
    EntityQueryRequest entityQueryRequest = EntityQueryRequest.newBuilder().setEntityType(EntityType.API.name()).addSelection(createExpression(API_ID_ATTR)).addSelection(createExpression(API_LABELS_ATTR)).build();
    Iterator<ResultSetChunk> resultSetChunkIterator = GrpcClientRequestContextUtil.executeWithHeadersContext(HEADERS, () -> entityQueryServiceClient.execute(entityQueryRequest));
    Map<String, List<String>> labelsMap = new HashMap<>();
    while (resultSetChunkIterator.hasNext()) {
        ResultSetChunk chunk = resultSetChunkIterator.next();
        for (Row row : chunk.getRowList()) {
            assertEquals(2, row.getColumnCount());
            String apiName = row.getColumn(0).getString();
            ProtocolStringList labelsValue = row.getColumn(1).getStringArrayList();
            labelsMap.put(apiName, labelsValue);
        }
    }
    assertEquals(2, labelsMap.size());
    assertEquals(List.of("Label1", "Label3", "Label4"), labelsMap.get(entity1.getEntityId()));
    assertEquals(List.of("Label2", "Label3", "Label4"), labelsMap.get(entity2.getEntityId()));
}