Search in sources :

Example 11 with PartitionQueryRequest

use of com.google.spanner.v1.PartitionQueryRequest in project java-spanner by googleapis.

the class SpannerClientTest method partitionQueryTest.

@Test
public void partitionQueryTest() throws Exception {
    PartitionResponse expectedResponse = PartitionResponse.newBuilder().addAllPartitions(new ArrayList<Partition>()).setTransaction(Transaction.newBuilder().build()).build();
    mockSpanner.addResponse(expectedResponse);
    PartitionQueryRequest request = PartitionQueryRequest.newBuilder().setSession(SessionName.of("[PROJECT]", "[INSTANCE]", "[DATABASE]", "[SESSION]").toString()).setTransaction(TransactionSelector.newBuilder().build()).setSql("sql114126").setParams(Struct.newBuilder().build()).putAllParamTypes(new HashMap<String, Type>()).setPartitionOptions(PartitionOptions.newBuilder().build()).build();
    PartitionResponse actualResponse = client.partitionQuery(request);
    Assert.assertEquals(expectedResponse, actualResponse);
    List<AbstractMessage> actualRequests = mockSpanner.getRequests();
    Assert.assertEquals(1, actualRequests.size());
    PartitionQueryRequest actualRequest = ((PartitionQueryRequest) actualRequests.get(0));
    Assert.assertEquals(request.getSession(), actualRequest.getSession());
    Assert.assertEquals(request.getTransaction(), actualRequest.getTransaction());
    Assert.assertEquals(request.getSql(), actualRequest.getSql());
    Assert.assertEquals(request.getParams(), actualRequest.getParams());
    Assert.assertEquals(request.getParamTypesMap(), actualRequest.getParamTypesMap());
    Assert.assertEquals(request.getPartitionOptions(), actualRequest.getPartitionOptions());
    Assert.assertTrue(channelProvider.isHeaderSent(ApiClientHeaderProvider.getDefaultApiClientHeaderKey(), GaxGrpcProperties.getDefaultApiClientHeaderPattern()));
}
Also used : PartitionResponse(com.google.spanner.v1.PartitionResponse) Partition(com.google.spanner.v1.Partition) Type(com.google.spanner.v1.Type) AbstractMessage(com.google.protobuf.AbstractMessage) PartitionQueryRequest(com.google.spanner.v1.PartitionQueryRequest) ByteString(com.google.protobuf.ByteString) Test(org.junit.Test)

Example 12 with PartitionQueryRequest

use of com.google.spanner.v1.PartitionQueryRequest in project java-firestore by googleapis.

the class CollectionGroup method getPartitions.

public ApiFuture<List<QueryPartition>> getPartitions(long desiredPartitionCount) {
    if (desiredPartitionCount == 1) {
        // Short circuit if the user only requested a single partition.
        return ApiFutures.immediateFuture(Collections.singletonList(new QueryPartition(partitionQuery, null, null)));
    } else {
        PartitionQueryRequest request = buildRequest(desiredPartitionCount);
        final TraceUtil traceUtil = TraceUtil.getInstance();
        Span span = traceUtil.startSpan(TraceUtil.SPAN_NAME_PARTITIONQUERY);
        try (Scope scope = traceUtil.getTracer().withSpan(span)) {
            return ApiFutures.transform(rpcContext.sendRequest(request, rpcContext.getClient().partitionQueryPagedCallable()), response -> {
                final ImmutableList.Builder<QueryPartition> partitions = ImmutableList.builder();
                consumePartitions(response, queryPartition -> {
                    partitions.add(queryPartition);
                    return null;
                });
                return partitions.build();
            }, MoreExecutors.directExecutor());
        } catch (ApiException exception) {
            span.setStatus(Status.UNKNOWN.withDescription(exception.getMessage()));
            throw FirestoreException.forApiException(exception);
        } finally {
            span.end(TraceUtil.END_SPAN_OPTIONS);
        }
    }
}
Also used : Scope(io.opencensus.common.Scope) PartitionQueryRequest(com.google.firestore.v1.PartitionQueryRequest) ImmutableList(com.google.common.collect.ImmutableList) Span(io.opencensus.trace.Span) ApiException(com.google.api.gax.rpc.ApiException)

Example 13 with PartitionQueryRequest

use of com.google.spanner.v1.PartitionQueryRequest in project java-firestore by googleapis.

the class FirestoreClientTest method partitionQueryExceptionTest.

@Test
public void partitionQueryExceptionTest() throws Exception {
    StatusRuntimeException exception = new StatusRuntimeException(io.grpc.Status.INVALID_ARGUMENT);
    mockFirestore.addException(exception);
    try {
        PartitionQueryRequest request = PartitionQueryRequest.newBuilder().setParent("parent-995424086").setPartitionCount(-1738969222).setPageToken("pageToken873572522").setPageSize(883849137).build();
        client.partitionQuery(request);
        Assert.fail("No exception raised");
    } catch (InvalidArgumentException e) {
    // Expected exception.
    }
}
Also used : InvalidArgumentException(com.google.api.gax.rpc.InvalidArgumentException) PartitionQueryRequest(com.google.firestore.v1.PartitionQueryRequest) StatusRuntimeException(io.grpc.StatusRuntimeException) Test(org.junit.Test)

Example 14 with PartitionQueryRequest

use of com.google.spanner.v1.PartitionQueryRequest in project java-firestore by googleapis.

the class PartitionQuery method requestsOneLessThanDesired.

@Test
public void requestsOneLessThanDesired() throws Exception {
    int desiredPartitionsCount = 2;
    PartitionQueryRequest expectedRequest = PartitionQueryRequest.newBuilder().setParent(DATABASE_NAME).setStructuredQuery(StructuredQuery.newBuilder().addFrom(StructuredQuery.CollectionSelector.newBuilder().setAllDescendants(true).setCollectionId("collectionId")).addOrderBy(StructuredQuery.Order.newBuilder().setField(StructuredQuery.FieldReference.newBuilder().setFieldPath("__name__")).setDirection(StructuredQuery.Direction.ASCENDING))).setPartitionCount(desiredPartitionsCount - 1).build();
    PartitionQueryResponse response = PartitionQueryResponse.newBuilder().addPartitions(CURSOR1).build();
    when(pagedResponse.iterateAll()).thenReturn(ImmutableList.of(CURSOR1));
    when(queryPage.getResponse()).thenReturn(response);
    doReturn(ApiFutures.immediateFuture(pagedResponse)).when(firestoreMock).sendRequest(requestCaptor.capture(), Matchers.<UnaryCallable<PartitionQueryRequest, PartitionQueryPagedResponse>>any());
    firestoreMock.collectionGroup("collectionId").getPartitions(desiredPartitionsCount).get();
    PartitionQueryRequest actualRequest = requestCaptor.getValue();
    assertEquals(actualRequest, expectedRequest);
}
Also used : PartitionQueryRequest(com.google.firestore.v1.PartitionQueryRequest) PartitionQueryPagedResponse(com.google.cloud.firestore.v1.FirestoreClient.PartitionQueryPagedResponse) PartitionQueryResponse(com.google.firestore.v1.PartitionQueryResponse) Test(org.junit.Test)

Example 15 with PartitionQueryRequest

use of com.google.spanner.v1.PartitionQueryRequest in project grpc-gcp-java by GoogleCloudPlatform.

the class SpannerTestCases method testPartitionQuery.

void testPartitionQuery() throws InterruptedException {
    System.out.println("\nTestPartitionQuery");
    ManagedChannel channel = getChannel();
    SpannerBlockingStub stub = getBlockingStub(channel);
    Session session = stub.createSession(CreateSessionRequest.newBuilder().setDatabase(database).build());
    TransactionOptions options = TransactionOptions.newBuilder().setReadOnly(TransactionOptions.ReadOnly.getDefaultInstance()).build();
    TransactionSelector selector = TransactionSelector.newBuilder().setBegin(options).build();
    PartitionQueryRequest request = PartitionQueryRequest.newBuilder().setSession(session.getName()).setSql("select * FROM " + LARGE_TABLE).setTransaction(selector).build();
    BlockingCall<PartitionQueryRequest, PartitionResponse> blockingCall = (PartitionQueryRequest req) -> stub.partitionQuery(req);
    doTestBlocking(channel, request, blockingCall);
    stub.deleteSession(DeleteSessionRequest.newBuilder().setName(session.getName()).build());
    channel.shutdown().awaitTermination(5, TimeUnit.SECONDS);
}
Also used : PartitionResponse(com.google.spanner.v1.PartitionResponse) SpannerBlockingStub(com.google.spanner.v1.SpannerGrpc.SpannerBlockingStub) TransactionOptions(com.google.spanner.v1.TransactionOptions) TransactionSelector(com.google.spanner.v1.TransactionSelector) PartitionQueryRequest(com.google.spanner.v1.PartitionQueryRequest) ManagedChannel(io.grpc.ManagedChannel) GcpManagedChannel(com.google.grpc.gcp.GcpManagedChannel) Session(com.google.spanner.v1.Session)

Aggregations

PartitionQueryRequest (com.google.firestore.v1.PartitionQueryRequest)10 Test (org.junit.Test)9 PartitionQueryResponse (com.google.firestore.v1.PartitionQueryResponse)7 PartitionQueryPagedResponse (com.google.cloud.firestore.v1.FirestoreClient.PartitionQueryPagedResponse)5 PartitionQueryRequest (com.google.spanner.v1.PartitionQueryRequest)5 PartitionResponse (com.google.spanner.v1.PartitionResponse)4 Session (com.google.spanner.v1.Session)3 TransactionOptions (com.google.spanner.v1.TransactionOptions)3 TransactionSelector (com.google.spanner.v1.TransactionSelector)3 PartitionQueryFn (org.apache.beam.sdk.io.gcp.firestore.FirestoreV1ReadFn.PartitionQueryFn)3 PartitionQueryPair (org.apache.beam.sdk.io.gcp.firestore.FirestoreV1ReadFn.PartitionQueryPair)3 ApiException (com.google.api.gax.rpc.ApiException)2 InvalidArgumentException (com.google.api.gax.rpc.InvalidArgumentException)2 GcpManagedChannel (com.google.grpc.gcp.GcpManagedChannel)2 AbstractMessage (com.google.protobuf.AbstractMessage)2 ByteString (com.google.protobuf.ByteString)2 Type (com.google.spanner.v1.Type)2 ManagedChannel (io.grpc.ManagedChannel)2 StatusRuntimeException (io.grpc.StatusRuntimeException)2 Scope (io.opencensus.common.Scope)2