use of org.apache.beam.sdk.io.gcp.datastore.DatastoreV1 in project beam by apache.
the class DataStoreReadWriteIT method testWriteRead_viaCoreBeamIO.
@Test
public void testWriteRead_viaCoreBeamIO() {
String projectId = options.getProject();
Key ancestor = makeKey(KIND, UUID.randomUUID().toString()).build();
Key itemKey = makeKey(ancestor, KIND, UUID.randomUUID().toString()).setPartitionId(PartitionId.newBuilder().setProjectId(projectId).build()).build();
Row testWriteRow = Row.withSchema(SOURCE_SCHEMA).addValues(itemKey.toByteArray(), "4000").build();
writePipeline.apply(Create.of(testWriteRow).withRowSchema(SOURCE_SCHEMA)).apply(RowToEntity.create("__key__", KIND)).apply(DatastoreIO.v1().write().withProjectId(projectId));
writePipeline.run().waitUntilFinish();
Query.Builder query = Query.newBuilder();
query.addKindBuilder().setName(KIND);
query.setFilter(makeFilter("__key__", Operator.EQUAL, makeValue(itemKey)));
DatastoreV1.Read read = DatastoreIO.v1().read().withProjectId(projectId).withQuery(query.build());
PCollection<Row> rowsRead = readPipeline.apply(read).apply(EntityToRow.create(SOURCE_SCHEMA, "__key__"));
PAssert.that(rowsRead).containsInAnyOrder(testWriteRow);
readPipeline.run().waitUntilFinish();
}
use of org.apache.beam.sdk.io.gcp.datastore.DatastoreV1 in project beam by apache.
the class DatastoreV1Test method testSplitQueryFnWithQueryLimit.
/**
* Tests {@link DatastoreV1.Read.SplitQueryFn} when the query has a user specified limit.
*/
@Test
public void testSplitQueryFnWithQueryLimit() throws Exception {
Query queryWithLimit = QUERY.toBuilder().setLimit(Int32Value.newBuilder().setValue(1)).build();
SplitQueryFn splitQueryFn = new SplitQueryFn(V_1_OPTIONS, 10, mockDatastoreFactory);
DoFnTester<Query, Query> doFnTester = DoFnTester.of(splitQueryFn);
doFnTester.setCloningBehavior(CloningBehavior.DO_NOT_CLONE);
List<Query> queries = doFnTester.processBundle(queryWithLimit);
assertEquals(1, queries.size());
verifyNoMoreInteractions(mockDatastore);
verifyNoMoreInteractions(mockQuerySplitter);
}
use of org.apache.beam.sdk.io.gcp.datastore.DatastoreV1 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 org.apache.beam.sdk.io.gcp.datastore.DatastoreV1 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;
}
}
use of org.apache.beam.sdk.io.gcp.datastore.DatastoreV1 in project beam by apache.
the class DatastoreV1Test method testBuildWrite.
/**
* Test building a Write using builder methods.
*/
@Test
public void testBuildWrite() throws Exception {
DatastoreV1.Write write = DatastoreIO.v1().write().withProjectId(PROJECT_ID);
assertEquals(PROJECT_ID, write.getProjectId());
}
Aggregations