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

Example 1 with SetAttribute

use of org.hypertrace.entity.query.service.v1.SetAttribute 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 SetAttribute

use of org.hypertrace.entity.query.service.v1.SetAttribute 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 SetAttribute

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

the class EntityQueryServiceImplTest method testUpdate_success.

@Test
public void testUpdate_success() throws Exception {
    Collection mockEntitiesCollection = mockEntitiesCollection();
    Builder newStatus = LiteralConstant.newBuilder().setValue(Value.newBuilder().setValueType(ValueType.STRING).setString("NEW_STATUS"));
    EntityUpdateRequest updateRequest = EntityUpdateRequest.newBuilder().setEntityType(TEST_ENTITY_TYPE).addEntityIds("entity-id-1").setOperation(UpdateOperation.newBuilder().setSetAttribute(SetAttribute.newBuilder().setAttribute(ColumnIdentifier.newBuilder().setColumnName(ATTRIBUTE_ID2)).setValue(newStatus))).addSelection(Expression.newBuilder().setColumnIdentifier(ColumnIdentifier.newBuilder().setColumnName(ATTRIBUTE_ID1))).addSelection(Expression.newBuilder().setColumnIdentifier(ColumnIdentifier.newBuilder().setColumnName(ATTRIBUTE_ID2))).build();
    StreamObserver<ResultSetChunk> mockResponseObserver = mock(StreamObserver.class);
    Context.current().withValue(RequestContext.CURRENT, mockRequestContextWithTenantId()).call(() -> {
        EntityQueryServiceImpl eqs = new EntityQueryServiceImpl(mockEntitiesCollection, mockMappingForAttributes1And2(), 1, false);
        eqs.update(updateRequest, mockResponseObserver);
        return null;
    });
    verify(mockEntitiesCollection, times(1)).bulkUpdateSubDocs(eq(Map.of(new SingleValueKey("tenant1", "entity-id-1"), Map.of("attributes.status", new JSONDocument(DocStoreJsonFormat.printer().print(newStatus))))));
}

Example 4 with SetAttribute

use of org.hypertrace.entity.query.service.v1.SetAttribute 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 5 with SetAttribute

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

the class EntityQueryServiceImpl method doUpdate.

private void doUpdate(RequestContext requestContext, EntityUpdateRequest request) throws Exception {
    if (request.getOperation().hasSetAttribute()) {
        SetAttribute setAttribute = request.getOperation().getSetAttribute();
        String attributeId = setAttribute.getAttribute().getColumnName();
        String subDocPath = entityAttributeMapping.getDocStorePathByAttributeId(requestContext, attributeId).orElseThrow(() -> new IllegalArgumentException("Unknown attribute FQN " + attributeId));
        JSONDocument jsonDocument = convertToJsonDocument(setAttribute.getValue());
        Map<Key, Map<String, Document>> entitiesUpdateMap = new HashMap<>();
        for (String entityId : request.getEntityIdsList()) {
            SingleValueKey key = new SingleValueKey(requestContext.getTenantId().orElseThrow(), entityId);
            if (entitiesUpdateMap.containsKey(key)) {
                entitiesUpdateMap.get(key).put(subDocPath, jsonDocument);
            } else {
                Map<String, Document> subDocument = new HashMap<>();
                subDocument.put(subDocPath, jsonDocument);
                entitiesUpdateMap.put(key, subDocument);
            }
        }
        try {
            entitiesCollection.bulkUpdateSubDocs(entitiesUpdateMap);
        } catch (Exception e) {
            LOG.error("Failed to update entities {}, subDocPath {}, with new doc {}.", entitiesUpdateMap, subDocPath, jsonDocument, e);
            throw e;
        }
    }
}