Search in sources :

Example 6 with RunQueryResponse

use of com.google.firestore.v1.RunQueryResponse in project beam by apache.

the class FirestoreTestingHelper method listDocumentsViaQuery.

Stream<String> listDocumentsViaQuery(String collectionPath) {
    int index = collectionPath.lastIndexOf('/');
    String parent = collectionPath.substring(0, index);
    String collectionId = collectionPath.substring(index + 1);
    FieldReference nameField = FieldReference.newBuilder().setFieldPath("__name__").build();
    RunQueryRequest rqr = RunQueryRequest.newBuilder().setParent(parent).setStructuredQuery(StructuredQuery.newBuilder().addFrom(CollectionSelector.newBuilder().setCollectionId(collectionId)).addOrderBy(Order.newBuilder().setField(nameField).setDirection(Direction.ASCENDING).build()).setSelect(Projection.newBuilder().addFields(nameField).build())).build();
    return StreamSupport.stream(rpc.runQueryCallable().call(rqr).spliterator(), false).filter(RunQueryResponse::hasDocument).map(RunQueryResponse::getDocument).map(Document::getName);
}
Also used : FieldReference(com.google.firestore.v1.StructuredQuery.FieldReference) RunQueryResponse(com.google.firestore.v1.RunQueryResponse) RunQueryRequest(com.google.firestore.v1.RunQueryRequest) Document(com.google.firestore.v1.Document)

Example 7 with RunQueryResponse

use of com.google.firestore.v1.RunQueryResponse in project beam by apache.

the class DatastoreV1Test method testSplitQueryFnWithoutNumSplits.

/**
 * Tests {@link SplitQueryFn} when no query splits is specified.
 */
@Test
public void testSplitQueryFnWithoutNumSplits() throws Exception {
    // Force SplitQueryFn to compute the number of query splits
    int numSplits = 0;
    int expectedNumSplits = 20;
    long entityBytes = expectedNumSplits * DEFAULT_BUNDLE_SIZE_BYTES;
    // In seconds
    long timestamp = 1234L;
    RunQueryRequest latestTimestampRequest = makeRequest(makeLatestTimestampQuery(NAMESPACE), NAMESPACE);
    RunQueryResponse latestTimestampResponse = makeLatestTimestampResponse(timestamp);
    // Per Kind statistics request and response
    RunQueryRequest statRequest = makeRequest(makeStatKindQuery(NAMESPACE, timestamp), NAMESPACE);
    RunQueryResponse statResponse = makeStatKindResponse(entityBytes);
    when(mockDatastore.runQuery(latestTimestampRequest)).thenReturn(latestTimestampResponse);
    when(mockDatastore.runQuery(statRequest)).thenReturn(statResponse);
    when(mockQuerySplitter.getSplits(eq(QUERY), any(PartitionId.class), eq(expectedNumSplits), any(Datastore.class))).thenReturn(splitQuery(QUERY, expectedNumSplits));
    SplitQueryFn splitQueryFn = new SplitQueryFn(V_1_OPTIONS, numSplits, mockDatastoreFactory);
    DoFnTester<Query, Query> doFnTester = DoFnTester.of(splitQueryFn);
    doFnTester.setCloningBehavior(CloningBehavior.DO_NOT_CLONE);
    List<Query> queries = doFnTester.processBundle(QUERY);
    assertEquals(expectedNumSplits, queries.size());
    verify(mockQuerySplitter, times(1)).getSplits(eq(QUERY), any(PartitionId.class), eq(expectedNumSplits), any(Datastore.class));
    verify(mockDatastore, times(1)).runQuery(latestTimestampRequest);
    verify(mockDatastore, times(1)).runQuery(statRequest);
}
Also used : Datastore(com.google.datastore.v1.client.Datastore) GqlQuery(com.google.datastore.v1.GqlQuery) Query(com.google.datastore.v1.Query) RunQueryResponse(com.google.datastore.v1.RunQueryResponse) RunQueryRequest(com.google.datastore.v1.RunQueryRequest) SplitQueryFn(org.apache.beam.sdk.io.gcp.datastore.DatastoreV1.Read.SplitQueryFn) PartitionId(com.google.datastore.v1.PartitionId) Test(org.junit.Test)

Example 8 with RunQueryResponse

use of com.google.firestore.v1.RunQueryResponse in project google-cloud-java by GoogleCloudPlatform.

the class DatastoreTest method buildResponsesForQueryPagination.

private List<RunQueryResponse> buildResponsesForQueryPagination() {
    Entity entity4 = Entity.newBuilder(KEY4).set("value", StringValue.of("value")).build();
    Entity entity5 = Entity.newBuilder(KEY5).set("value", "value").build();
    datastore.add(ENTITY3, entity4, entity5);
    List<RunQueryResponse> responses = new ArrayList<>();
    Query<Key> query = Query.newKeyQueryBuilder().build();
    RunQueryRequest.Builder requestPb = RunQueryRequest.newBuilder();
    query.populatePb(requestPb);
    QueryResultBatch queryResultBatchPb = RunQueryResponse.newBuilder().mergeFrom(((DatastoreImpl) datastore).runQuery(requestPb.build())).getBatch();
    QueryResultBatch queryResultBatchPb1 = QueryResultBatch.newBuilder().mergeFrom(queryResultBatchPb).setMoreResults(QueryResultBatch.MoreResultsType.NOT_FINISHED).clearEntityResults().addAllEntityResults(queryResultBatchPb.getEntityResultsList().subList(0, 1)).setEndCursor(queryResultBatchPb.getEntityResultsList().get(0).getCursor()).build();
    responses.add(RunQueryResponse.newBuilder().setBatch(queryResultBatchPb1).build());
    QueryResultBatch queryResultBatchPb2 = QueryResultBatch.newBuilder().mergeFrom(queryResultBatchPb).setMoreResults(QueryResultBatch.MoreResultsType.NOT_FINISHED).clearEntityResults().addAllEntityResults(queryResultBatchPb.getEntityResultsList().subList(1, 3)).setEndCursor(queryResultBatchPb.getEntityResultsList().get(2).getCursor()).build();
    responses.add(RunQueryResponse.newBuilder().setBatch(queryResultBatchPb2).build());
    QueryResultBatch queryResultBatchPb3 = QueryResultBatch.newBuilder().mergeFrom(queryResultBatchPb).setMoreResults(QueryResultBatch.MoreResultsType.NO_MORE_RESULTS).clearEntityResults().addAllEntityResults(queryResultBatchPb.getEntityResultsList().subList(3, 5)).setEndCursor(queryResultBatchPb.getEntityResultsList().get(4).getCursor()).build();
    responses.add(RunQueryResponse.newBuilder().setBatch(queryResultBatchPb3).build());
    return responses;
}
Also used : QueryResultBatch(com.google.datastore.v1.QueryResultBatch) RunQueryResponse(com.google.datastore.v1.RunQueryResponse) RunQueryRequest(com.google.datastore.v1.RunQueryRequest) ArrayList(java.util.ArrayList)

Example 9 with RunQueryResponse

use of com.google.firestore.v1.RunQueryResponse in project beam by apache.

the class DatastoreV1Test method testEstimatedSizeBytes.

/**
 * Tests {@link DatastoreV1.Read#getEstimatedSizeBytes} to fetch and return estimated size for a
 * query.
 */
@Test
public void testEstimatedSizeBytes() throws Exception {
    long entityBytes = 100L;
    // In seconds
    long timestamp = 1234L;
    RunQueryRequest latestTimestampRequest = makeRequest(makeLatestTimestampQuery(NAMESPACE), NAMESPACE);
    RunQueryResponse latestTimestampResponse = makeLatestTimestampResponse(timestamp);
    // Per Kind statistics request and response
    RunQueryRequest statRequest = makeRequest(makeStatKindQuery(NAMESPACE, timestamp), NAMESPACE);
    RunQueryResponse statResponse = makeStatKindResponse(entityBytes);
    when(mockDatastore.runQuery(latestTimestampRequest)).thenReturn(latestTimestampResponse);
    when(mockDatastore.runQuery(statRequest)).thenReturn(statResponse);
    assertEquals(entityBytes, getEstimatedSizeBytes(mockDatastore, QUERY, NAMESPACE));
    verify(mockDatastore, times(1)).runQuery(latestTimestampRequest);
    verify(mockDatastore, times(1)).runQuery(statRequest);
}
Also used : RunQueryResponse(com.google.datastore.v1.RunQueryResponse) RunQueryRequest(com.google.datastore.v1.RunQueryRequest) Test(org.junit.Test)

Example 10 with RunQueryResponse

use of com.google.firestore.v1.RunQueryResponse in project beam by apache.

the class DatastoreV1Test method makeLatestTimestampResponse.

/**
 * Builds a response of the given timestamp.
 */
private static RunQueryResponse makeLatestTimestampResponse(long timestamp) {
    RunQueryResponse.Builder timestampResponse = RunQueryResponse.newBuilder();
    Entity.Builder entity = Entity.newBuilder();
    entity.setKey(makeKey("dummyKind", "dummyId"));
    entity.putProperties("timestamp", makeValue(new Date(timestamp * 1000)).build());
    EntityResult.Builder entityResult = EntityResult.newBuilder();
    entityResult.setEntity(entity);
    QueryResultBatch.Builder batch = QueryResultBatch.newBuilder();
    batch.addEntityResults(entityResult);
    timestampResponse.setBatch(batch);
    return timestampResponse.build();
}
Also used : Entity(com.google.datastore.v1.Entity) DeleteEntity(org.apache.beam.sdk.io.gcp.datastore.DatastoreV1.DeleteEntity) QueryResultBatch(com.google.datastore.v1.QueryResultBatch) RunQueryResponse(com.google.datastore.v1.RunQueryResponse) EntityResult(com.google.datastore.v1.EntityResult) Date(java.util.Date)

Aggregations

RunQueryResponse (com.google.datastore.v1.RunQueryResponse)6 Test (org.junit.Test)6 RunQueryRequest (com.google.datastore.v1.RunQueryRequest)5 RunQueryResponse (com.google.firestore.v1.RunQueryResponse)4 QueryResultBatch (com.google.datastore.v1.QueryResultBatch)3 RunQueryRequest (com.google.firestore.v1.RunQueryRequest)3 RunQueryFn (org.apache.beam.sdk.io.gcp.firestore.FirestoreV1ReadFn.RunQueryFn)3 Document (com.google.firestore.v1.Document)2 FieldReference (com.google.firestore.v1.StructuredQuery.FieldReference)2 ArrayList (java.util.ArrayList)2 ServerStream (com.google.api.gax.rpc.ServerStream)1 ServerStreamingCallable (com.google.api.gax.rpc.ServerStreamingCallable)1 DatastoreRpc (com.google.cloud.datastore.spi.v1.DatastoreRpc)1 FirestoreStub (com.google.cloud.firestore.v1.stub.FirestoreStub)1 Entity (com.google.datastore.v1.Entity)1 EntityResult (com.google.datastore.v1.EntityResult)1 GqlQuery (com.google.datastore.v1.GqlQuery)1 PartitionId (com.google.datastore.v1.PartitionId)1 Query (com.google.datastore.v1.Query)1 Datastore (com.google.datastore.v1.client.Datastore)1