Examples with CommitRequest com.google.spanner.v1.CommitRequest used on opensource projects

Example 1 with CommitRequest

use of com.google.spanner.v1.CommitRequest 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);
}

Example 2 with CommitRequest

use of com.google.spanner.v1.CommitRequest in project beam by apache.

the class DatastoreV1Test method datastoreWriterFnTest.

// A helper method to test DatastoreWriterFn for various batch sizes.
private void datastoreWriterFnTest(int numMutations) throws Exception {
    // Create the requested number of mutations.
    List<Mutation> mutations = new ArrayList<>(numMutations);
    for (int i = 0; i < numMutations; ++i) {
        mutations.add(makeUpsert(Entity.newBuilder().setKey(makeKey("key" + i, i + 1)).build()).build());
    }
    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);
    int start = 0;
    while (start < numMutations) {
        int end = Math.min(numMutations, start + DatastoreV1.DATASTORE_BATCH_UPDATE_ENTITIES_START);
        CommitRequest.Builder commitRequest = CommitRequest.newBuilder();
        commitRequest.setMode(CommitRequest.Mode.NON_TRANSACTIONAL);
        commitRequest.addAllMutations(mutations.subList(start, end));
        // Verify all the batch requests were made with the expected mutations.
        verify(mockDatastore, times(1)).commit(commitRequest.build());
        start = end;
    }
}

Example 3 with CommitRequest

use of com.google.spanner.v1.CommitRequest in project beam by apache.

the class DatastoreV1Test method testDatatoreWriterFnWithLargeEntities.

/**
 * Tests {@link DatastoreWriterFn} with large entities that need to be split into more batches.
 */
@Test
public void testDatatoreWriterFnWithLargeEntities() throws Exception {
    List<Mutation> mutations = new ArrayList<>();
    int entitySize = 0;
    for (int i = 0; i < 12; ++i) {
        Entity entity = Entity.newBuilder().setKey(makeKey("key" + i, i + 1)).putProperties("long", makeValue(new String(new char[900_000])).setExcludeFromIndexes(true).build()).build();
        // Take the size of any one entity.
        entitySize = entity.getSerializedSize();
        mutations.add(makeUpsert(entity).build());
    }
    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);
    // This test is over-specific currently; it requires that we split the 12 entity writes into 3
    // requests, but we only need each CommitRequest to be less than 10MB in size.
    int entitiesPerRpc = DATASTORE_BATCH_UPDATE_BYTES_LIMIT / entitySize;
    int start = 0;
    while (start < mutations.size()) {
        int end = Math.min(mutations.size(), start + entitiesPerRpc);
        CommitRequest.Builder commitRequest = CommitRequest.newBuilder();
        commitRequest.setMode(CommitRequest.Mode.NON_TRANSACTIONAL);
        commitRequest.addAllMutations(mutations.subList(start, end));
        // Verify all the batch requests were made with the expected mutations.
        verify(mockDatastore).commit(commitRequest.build());
        start = end;
    }
}

Example 4 with CommitRequest

use of com.google.spanner.v1.CommitRequest 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);
}