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

Example 1 with ResultSetMetadata

use of org.hypertrace.entity.query.service.v1.ResultSetMetadata in project beam by apache.

the class SpannerChangeStreamErrorTest method mockTableExists.

private void mockTableExists() {
    Statement tableExistsStatement = Statement.of("SELECT t.table_name FROM information_schema.tables AS t WHERE t.table_catalog = '' AND t.table_schema = '' AND t.table_name = 'my-metadata-table'");
    ResultSetMetadata tableExistsResultSetMetadata = ResultSetMetadata.newBuilder().setRowType(StructType.newBuilder().addFields(Field.newBuilder().setName("table_name").setType(Type.newBuilder().setCode(TypeCode.STRING).build()).build()).build()).build();
    ResultSet tableExistsResultSet = ResultSet.newBuilder().addRows(ListValue.newBuilder().addValues(Value.newBuilder().setStringValue(TEST_TABLE).build()).build()).setMetadata(tableExistsResultSetMetadata).build();
    mockSpannerService.putStatementResult(StatementResult.query(tableExistsStatement, tableExistsResultSet));
}

Example 2 with ResultSetMetadata

use of org.hypertrace.entity.query.service.v1.ResultSetMetadata in project java-spanner by googleapis.

the class MockSpannerServiceImpl method returnResultSet.

private void returnResultSet(ResultSet resultSet, ByteString transactionId, TransactionSelector transactionSelector, StreamObserver<ResultSet> responseObserver) {
    ResultSetMetadata metadata = resultSet.getMetadata();
    if (transactionId != null) {
        metadata = metadata.toBuilder().setTransaction(ignoreNextInlineBeginRequest.getAndSet(false) ? Transaction.getDefaultInstance() : Transaction.newBuilder().setId(transactionId).build()).build();
    } else if (transactionSelector.hasBegin() || transactionSelector.hasSingleUse()) {
        Transaction transaction = getTemporaryTransactionOrNull(transactionSelector);
        metadata = metadata.toBuilder().setTransaction(transaction).build();
    }
    resultSet = resultSet.toBuilder().setMetadata(metadata).build();
    responseObserver.onNext(resultSet);
}

Example 3 with ResultSetMetadata

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

Example 4 with ResultSetMetadata

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

Example 5 with ResultSetMetadata

use of org.hypertrace.entity.query.service.v1.ResultSetMetadata in project beam by apache.

the class SpannerChangeStreamErrorTest method mockInvalidChangeStreamRecordReceived.

private void mockInvalidChangeStreamRecordReceived(Timestamp now, Timestamp after3Seconds) {
    Statement changeStreamQueryStatement = Statement.newBuilder("SELECT * FROM READ_my-change-stream(   start_timestamp => @startTimestamp,   end_timestamp => @endTimestamp,   partition_token => @partitionToken,   read_options => null,   heartbeat_milliseconds => @heartbeatMillis)").bind("startTimestamp").to(now).bind("endTimestamp").to(after3Seconds).bind("partitionToken").to((String) null).bind("heartbeatMillis").to(500).build();
    ResultSetMetadata readChangeStreamResultSetMetadata = ResultSetMetadata.newBuilder().setRowType(StructType.newBuilder().addFields(Field.newBuilder().setName("COL1").setType(Type.newBuilder().setCode(TypeCode.ARRAY).setArrayElementType(Type.newBuilder().setCode(TypeCode.STRUCT).setStructType(StructType.newBuilder().addFields(Field.newBuilder().setName("field_name").setType(Type.newBuilder().setCode(TypeCode.STRUCT).setStructType(StructType.newBuilder().addFields(Field.newBuilder().setType(Type.newBuilder().setCode(TypeCode.STRING))))))))))).build();
    ResultSet readChangeStreamResultSet = ResultSet.newBuilder().addRows(ListValue.newBuilder().addValues(Value.newBuilder().setListValue(ListValue.newBuilder().addValues(Value.newBuilder().setListValue(ListValue.newBuilder().addValues(Value.newBuilder().setListValue(ListValue.newBuilder().addValues(Value.newBuilder().setStringValue("bad_value"))))))))).setMetadata(readChangeStreamResultSetMetadata).build();
    mockSpannerService.putStatementResult(StatementResult.query(changeStreamQueryStatement, readChangeStreamResultSet));
}