use of com.google.api.ads.admanager.axis.v202205.Row 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 com.google.api.ads.admanager.axis.v202205.Row in project entity-service by hypertrace.
the class EntityQueryServiceTest method testCreateAndGetEntities.
@Test
public void testCreateAndGetEntities() {
// creating an api entity with attributes
Entity.Builder apiEntityBuilder1 = createApiEntity(SERVICE_ID, "api1", API_TYPE);
apiEntityBuilder1.putAttributes(apiAttributesMap.get(API_DISCOVERY_STATE_ATTR), createAttribute("DISCOVERED")).putAttributes(apiAttributesMap.get(API_HTTP_METHOD_ATTR), createAttribute("GET"));
entityDataServiceClient.upsert(apiEntityBuilder1.build());
Entity.Builder apiEntityBuilder2 = createApiEntity(SERVICE_ID, "api2", API_TYPE);
apiEntityBuilder2.putAttributes(apiAttributesMap.get(API_DISCOVERY_STATE_ATTR), createAttribute("UNDER_DISCOVERY")).putAttributes(apiAttributesMap.get(API_HTTP_METHOD_ATTR), createAttribute("GET"));
entityDataServiceClient.upsert(apiEntityBuilder2.build());
// querying the api entity
EntityQueryRequest entityQueryRequest = EntityQueryRequest.newBuilder().setEntityType(EntityType.API.name()).addSelection(createExpression(API_DISCOVERY_STATE_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(2, values.size());
assertEquals("DISCOVERED", values.get(0));
assertEquals("UNDER_DISCOVERY", values.get(1));
}
use of com.google.api.ads.admanager.axis.v202205.Row in project entity-service by hypertrace.
the class EntityQueryServiceImpl method execute.
@Override
public void execute(EntityQueryRequest request, StreamObserver<ResultSetChunk> responseObserver) {
RequestContext requestContext = RequestContext.CURRENT.get();
Optional<String> tenantId = requestContext.getTenantId();
if (tenantId.isEmpty()) {
responseObserver.onError(new ServiceException("Tenant id is missing in the request."));
return;
}
final Iterator<Document> documentIterator;
if (queryAggregationEnabled) {
final Converter<EntityQueryRequest, org.hypertrace.core.documentstore.query.Query> queryConverter = getQueryConverter();
final org.hypertrace.core.documentstore.query.Query query;
try {
query = queryConverter.convert(request, requestContext);
documentIterator = entitiesCollection.aggregate(query);
} catch (final Exception e) {
responseObserver.onError(new ServiceException(e));
return;
}
} else {
// TODO: Optimize this later. For now converting to EDS Query and then again to DocStore
// Query.
Query query = entityQueryConverter.convertToEDSQuery(requestContext, request);
/**
* {@link EntityQueryRequest} selections need to treated differently, since they don't
* transform one to one to {@link org.hypertrace.entity.data.service.v1.EntityDataRequest}
* selections
*/
List<String> docStoreSelections = entityQueryConverter.convertSelectionsToDocStoreSelections(requestContext, request.getSelectionList());
documentIterator = entitiesCollection.search(DocStoreConverter.transform(tenantId.get(), query, docStoreSelections));
}
final DocumentConverter rowConverter = injector.getInstance(DocumentConverter.class);
ResultSetMetadata resultSetMetadata;
try {
resultSetMetadata = this.buildMetadataForSelections(request.getSelectionList());
} catch (final ConversionException e) {
responseObserver.onError(new ServiceException(e));
return;
}
if (!documentIterator.hasNext()) {
ResultSetChunk.Builder resultBuilder = ResultSetChunk.newBuilder();
resultBuilder.setResultSetMetadata(resultSetMetadata);
resultBuilder.setIsLastChunk(true);
resultBuilder.setChunkId(0);
responseObserver.onNext(resultBuilder.build());
responseObserver.onCompleted();
return;
}
boolean isNewChunk = true;
int chunkId = 0, rowCount = 0;
ResultSetChunk.Builder resultBuilder = ResultSetChunk.newBuilder();
while (documentIterator.hasNext()) {
// Set metadata for new chunk
if (isNewChunk) {
resultBuilder.setResultSetMetadata(resultSetMetadata);
isNewChunk = false;
}
try {
final Row row;
if (queryAggregationEnabled) {
row = rowConverter.convertToRow(documentIterator.next(), resultSetMetadata);
resultBuilder.addRow(row);
rowCount++;
} else {
Optional<Entity> entity = DOCUMENT_PARSER.parseOrLog(documentIterator.next(), Entity.newBuilder());
if (entity.isPresent()) {
row = convertToEntityQueryResult(requestContext, entity.get(), request.getSelectionList());
resultBuilder.addRow(row);
rowCount++;
}
}
} catch (final Exception e) {
responseObserver.onError(new ServiceException(e));
return;
}
// current chunk is complete
if (rowCount >= CHUNK_SIZE || !documentIterator.hasNext()) {
resultBuilder.setChunkId(chunkId++);
resultBuilder.setIsLastChunk(!documentIterator.hasNext());
responseObserver.onNext(resultBuilder.build());
resultBuilder = ResultSetChunk.newBuilder();
isNewChunk = true;
rowCount = 0;
}
}
responseObserver.onCompleted();
}
use of com.google.api.ads.admanager.axis.v202205.Row in project entity-service by hypertrace.
the class DocumentConverterImpl method buildRow.
private Row buildRow(final ResultSetMetadata resultSetMetadata, final Map<String, Value> valueMap) {
final Row.Builder builder = Row.newBuilder();
final Value defaultValue = getNullPlaceholderValue();
for (final ColumnMetadata columnMetadata : resultSetMetadata.getColumnMetadataList()) {
final Value value = valueMap.getOrDefault(columnMetadata.getColumnName(), defaultValue);
builder.addColumn(value);
}
return builder.build();
}
use of com.google.api.ads.admanager.axis.v202205.Row in project gateway-service by hypertrace.
the class EntityDataServiceEntityFetcher method getEntities.
@Override
public EntityFetcherResponse getEntities(EntitiesRequestContext requestContext, EntitiesRequest entitiesRequest) {
List<String> entityIdAttributeIds = AttributeMetadataUtil.getIdAttributeIds(attributeMetadataProvider, entityIdColumnsConfigs, requestContext, entitiesRequest.getEntityType());
Map<String, List<String>> requestedAliasesByEntityIdAttributeIds = getExpectedResultNamesForEachAttributeId(entitiesRequest.getSelectionList(), entityIdAttributeIds);
EntityQueryRequest.Builder builder = EntityQueryRequest.newBuilder().setEntityType(entitiesRequest.getEntityType()).setFilter(EntityServiceAndGatewayServiceConverter.convertToEntityServiceFilter(entitiesRequest.getFilter())).addAllSelection(entityIdAttributeIds.stream().map(entityIdAttr -> EntityServiceAndGatewayServiceConverter.createColumnExpression(entityIdAttr).build()).collect(Collectors.toList()));
// add time filter for supported scope
if (!entitiesRequest.getIncludeNonLiveEntities()) {
EntityServiceAndGatewayServiceConverter.addBetweenTimeFilter(entitiesRequest.getStartTimeMillis(), entitiesRequest.getEndTimeMillis(), attributeMetadataProvider, entitiesRequest, builder, requestContext);
}
// Add all expressions in the select that are already not part of the EntityID attributes
entitiesRequest.getSelectionList().stream().filter(ExpressionReader::isAttributeSelection).filter(expression -> ExpressionReader.getAttributeIdFromAttributeSelection(expression).map(attributeId -> !entityIdAttributeIds.contains(attributeId)).orElse(true)).forEach(expression -> builder.addSelection(EntityServiceAndGatewayServiceConverter.convertToEntityServiceExpression(expression)));
int limit = entitiesRequest.getLimit();
if (limit > 0) {
builder.setLimit(limit);
}
int offset = entitiesRequest.getOffset();
if (offset > 0) {
builder.setOffset(offset);
}
if (!entitiesRequest.getOrderByList().isEmpty()) {
builder.addAllOrderBy(EntityServiceAndGatewayServiceConverter.convertToOrderByExpressions(entitiesRequest.getOrderByList()));
}
EntityQueryRequest entityQueryRequest = builder.build();
LOG.debug("Sending Query to EDS ======== \n {}", entityQueryRequest);
Iterator<ResultSetChunk> resultSetChunkIterator = entityQueryServiceClient.execute(builder.build(), requestContext.getHeaders());
Map<String, AttributeMetadata> resultMetadataMap = this.getAttributeMetadataByAlias(requestContext, entitiesRequest);
// We want to retain the order as returned from the respective source. Hence using a
// LinkedHashMap
Map<EntityKey, Builder> entityBuilders = new LinkedHashMap<>();
while (resultSetChunkIterator.hasNext()) {
ResultSetChunk chunk = resultSetChunkIterator.next();
LOG.debug("Received chunk: {}", chunk);
if (chunk.getRowCount() < 1) {
break;
}
for (Row row : chunk.getRowList()) {
// Construct the entity id from the entityIdAttributes columns
EntityKey entityKey = EntityKey.of(IntStream.range(0, entityIdAttributeIds.size()).mapToObj(value -> row.getColumn(value).getString()).toArray(String[]::new));
Builder entityBuilder = entityBuilders.computeIfAbsent(entityKey, k -> Entity.newBuilder());
entityBuilder.setEntityType(entitiesRequest.getEntityType());
entityBuilder.setId(entityKey.toString());
// as post processing.
for (int i = 0; i < entityIdAttributeIds.size(); i++) {
entityBuilder.putAttribute(entityIdAttributeIds.get(i), Value.newBuilder().setString(entityKey.getAttributes().get(i)).setValueType(ValueType.STRING).build());
}
requestedAliasesByEntityIdAttributeIds.forEach((attributeId, requestedAliasList) -> requestedAliasList.forEach(requestedAlias -> entityBuilder.putAttribute(requestedAlias, entityBuilder.getAttributeOrThrow(attributeId))));
for (int i = entityIdAttributeIds.size(); i < chunk.getResultSetMetadata().getColumnMetadataCount(); i++) {
String resultName = chunk.getResultSetMetadata().getColumnMetadata(i).getColumnName();
AttributeMetadata attributeMetadata = resultMetadataMap.get(resultName);
entityBuilder.putAttribute(resultName, EntityServiceAndGatewayServiceConverter.convertQueryValueToGatewayValue(row.getColumn(i), attributeMetadata));
}
}
}
return new EntityFetcherResponse(entityBuilders);
}
Aggregations