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

Example 1 with UpdateOperation

use of org.hypertrace.entity.query.service.v1.UpdateOperation 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));
}

Example 2 with UpdateOperation

use of org.hypertrace.entity.query.service.v1.UpdateOperation 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)));
}

Example 3 with UpdateOperation

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

the class EntityQueryServiceImpl method transformUpdateOperations.

private Map<String, Document> transformUpdateOperations(List<UpdateOperation> updateOperationList, RequestContext requestContext) throws Exception {
    Map<String, Document> documentMap = new HashMap<>();
    for (UpdateOperation updateOperation : updateOperationList) {
        if (!updateOperation.hasSetAttribute()) {
            continue;
        }
        SetAttribute setAttribute = updateOperation.getSetAttribute();
        String attributeId = setAttribute.getAttribute().getColumnName();
        String subDocPath = entityAttributeMapping.getDocStorePathByAttributeId(requestContext, attributeId).orElseThrow(() -> new IllegalArgumentException("Unknown attribute FQN " + attributeId));
        try {
            documentMap.put(subDocPath, convertToJsonDocument(setAttribute.getValue()));
        } catch (Exception e) {
            LOG.error("Failed to put update corresponding to {} in the documentMap", subDocPath, e);
            throw e;
        }
    }
    return Collections.unmodifiableMap(documentMap);
}