use of com.google.datastore.v1.CommitResponse in project google-cloud-java by GoogleCloudPlatform.
the class SessionImplTest method writeAtLeastOnce.
@Test
public void writeAtLeastOnce() throws ParseException {
String timestampString = "2015-10-01T10:54:20.021Z";
ArgumentCaptor<CommitRequest> commit = ArgumentCaptor.forClass(CommitRequest.class);
CommitResponse response = CommitResponse.newBuilder().setCommitTimestamp(Timestamps.parse(timestampString)).build();
Mockito.when(rpc.commit(commit.capture(), Mockito.eq(options))).thenReturn(response);
Timestamp timestamp = session.writeAtLeastOnce(Arrays.asList(Mutation.newInsertBuilder("T").set("C").to("x").build()));
assertThat(timestamp.getSeconds()).isEqualTo(utcTimeSeconds(2015, Calendar.OCTOBER, 1, 10, 54, 20));
assertThat(timestamp.getNanos()).isEqualTo(TimeUnit.MILLISECONDS.toNanos(21));
CommitRequest request = commit.getValue();
assertThat(request.getSingleUseTransaction()).isNotNull();
assertThat(request.getSingleUseTransaction().getReadWrite()).isNotNull();
com.google.spanner.v1.Mutation mutation = com.google.spanner.v1.Mutation.newBuilder().setInsert(Write.newBuilder().setTable("T").addColumns("C").addValues(ListValue.newBuilder().addValues(com.google.protobuf.Value.newBuilder().setStringValue("x")))).build();
assertThat(request.getMutationsList()).containsExactly(mutation);
}
use of com.google.datastore.v1.CommitResponse in project beam by apache.
the class DatastoreV1Test method testDatatoreWriterFnRetriesErrors.
/**
* Tests {@link DatastoreWriterFn} with a failed request which is retried.
*/
@Test
public void testDatatoreWriterFnRetriesErrors() throws Exception {
List<Mutation> mutations = new ArrayList<>();
int numRpcs = 2;
for (int i = 0; i < DatastoreV1.DATASTORE_BATCH_UPDATE_ENTITIES_START * numRpcs; ++i) {
mutations.add(makeUpsert(Entity.newBuilder().setKey(makeKey("key" + i, i + 1)).build()).build());
}
CommitResponse successfulCommit = CommitResponse.getDefaultInstance();
when(mockDatastore.commit(any(CommitRequest.class))).thenReturn(successfulCommit).thenThrow(new DatastoreException("commit", Code.DEADLINE_EXCEEDED, "", null)).thenReturn(successfulCommit);
DatastoreWriterFn datastoreWriter = new DatastoreWriterFn(StaticValueProvider.of(PROJECT_ID), null, mockDatastoreFactory, new FakeWriteBatcher());
DoFnTester<Mutation, Void> doFnTester = DoFnTester.of(datastoreWriter);
doFnTester.setCloningBehavior(CloningBehavior.DO_NOT_CLONE);
doFnTester.processBundle(mutations);
verifyMetricWasSet("BatchDatastoreWrite", "ok", "", 2);
verifyMetricWasSet("BatchDatastoreWrite", "unknown", "", 1);
}
use of com.google.datastore.v1.CommitResponse in project beam by apache.
the class SpannerAccessor method createAndConnect.
private static SpannerAccessor createAndConnect(SpannerConfig spannerConfig) {
SpannerOptions.Builder builder = SpannerOptions.newBuilder();
// Set retryable codes for all API methods
if (spannerConfig.getRetryableCodes() != null) {
builder.getSpannerStubSettingsBuilder().applyToAllUnaryMethods(input -> {
input.setRetryableCodes(spannerConfig.getRetryableCodes());
return null;
});
builder.getSpannerStubSettingsBuilder().executeStreamingSqlSettings().setRetryableCodes(spannerConfig.getRetryableCodes());
}
// Set commit retry settings
UnaryCallSettings.Builder<CommitRequest, CommitResponse> commitSettings = builder.getSpannerStubSettingsBuilder().commitSettings();
ValueProvider<Duration> commitDeadline = spannerConfig.getCommitDeadline();
if (spannerConfig.getCommitRetrySettings() != null) {
commitSettings.setRetrySettings(spannerConfig.getCommitRetrySettings());
} else if (commitDeadline != null && commitDeadline.get().getMillis() > 0) {
// Set the GRPC deadline on the Commit API call.
RetrySettings.Builder commitRetrySettingsBuilder = commitSettings.getRetrySettings().toBuilder();
commitSettings.setRetrySettings(commitRetrySettingsBuilder.setTotalTimeout(org.threeten.bp.Duration.ofMillis(commitDeadline.get().getMillis())).setMaxRpcTimeout(org.threeten.bp.Duration.ofMillis(commitDeadline.get().getMillis())).setInitialRpcTimeout(org.threeten.bp.Duration.ofMillis(commitDeadline.get().getMillis())).build());
}
// Set execute streaming sql retry settings
ServerStreamingCallSettings.Builder<ExecuteSqlRequest, PartialResultSet> executeStreamingSqlSettings = builder.getSpannerStubSettingsBuilder().executeStreamingSqlSettings();
if (spannerConfig.getExecuteStreamingSqlRetrySettings() != null) {
executeStreamingSqlSettings.setRetrySettings(spannerConfig.getExecuteStreamingSqlRetrySettings());
} else {
// Setting the timeout for streaming read to 2 hours. This is 1 hour by default
// after BEAM 2.20.
RetrySettings.Builder executeSqlStreamingRetrySettings = executeStreamingSqlSettings.getRetrySettings().toBuilder();
executeStreamingSqlSettings.setRetrySettings(executeSqlStreamingRetrySettings.setInitialRpcTimeout(org.threeten.bp.Duration.ofMinutes(120)).setMaxRpcTimeout(org.threeten.bp.Duration.ofMinutes(120)).setTotalTimeout(org.threeten.bp.Duration.ofMinutes(120)).build());
}
ValueProvider<String> projectId = spannerConfig.getProjectId();
if (projectId != null) {
builder.setProjectId(projectId.get());
}
ServiceFactory<Spanner, SpannerOptions> serviceFactory = spannerConfig.getServiceFactory();
if (serviceFactory != null) {
builder.setServiceFactory(serviceFactory);
}
ValueProvider<String> host = spannerConfig.getHost();
if (host != null) {
builder.setHost(host.get());
}
ValueProvider<String> emulatorHost = spannerConfig.getEmulatorHost();
if (emulatorHost != null) {
builder.setEmulatorHost(emulatorHost.get());
if (spannerConfig.getIsLocalChannelProvider() != null && spannerConfig.getIsLocalChannelProvider().get()) {
builder.setChannelProvider(LocalChannelProvider.create(emulatorHost.get()));
}
builder.setCredentials(NoCredentials.getInstance());
}
String userAgentString = USER_AGENT_PREFIX + "/" + ReleaseInfo.getReleaseInfo().getVersion();
builder.setHeaderProvider(FixedHeaderProvider.create("user-agent", userAgentString));
SpannerOptions options = builder.build();
Spanner spanner = options.getService();
String instanceId = spannerConfig.getInstanceId().get();
String databaseId = spannerConfig.getDatabaseId().get();
DatabaseClient databaseClient = spanner.getDatabaseClient(DatabaseId.of(options.getProjectId(), instanceId, databaseId));
BatchClient batchClient = spanner.getBatchClient(DatabaseId.of(options.getProjectId(), instanceId, databaseId));
DatabaseAdminClient databaseAdminClient = spanner.getDatabaseAdminClient();
return new SpannerAccessor(spanner, databaseClient, databaseAdminClient, batchClient, spannerConfig);
}
Aggregations