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()));
}
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);
}
}
}
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.
}
}
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);
}
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);
}
Aggregations