Examples with DatastoreWriterFn org.apache.beam.sdk.io.gcp.datastore.DatastoreV1.DatastoreWriterFn used on opensource projects

Example 1 with DatastoreWriterFn

use of org.apache.beam.sdk.io.gcp.datastore.DatastoreV1.DatastoreWriterFn 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);
}

Example 2 with DatastoreWriterFn

use of org.apache.beam.sdk.io.gcp.datastore.DatastoreV1.DatastoreWriterFn 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 DatastoreWriterFn

use of org.apache.beam.sdk.io.gcp.datastore.DatastoreV1.DatastoreWriterFn in project beam by apache.

the class DatastoreV1Test method testDatastoreWriteFnDisplayData.

@Test
public void testDatastoreWriteFnDisplayData() {
    DatastoreWriterFn datastoreWriter = new DatastoreWriterFn(PROJECT_ID, null);
    DisplayData displayData = DisplayData.from(datastoreWriter);
    assertThat(displayData, hasDisplayItem("projectId", PROJECT_ID));
}

Example 4 with DatastoreWriterFn

use of org.apache.beam.sdk.io.gcp.datastore.DatastoreV1.DatastoreWriterFn 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;
    }
}