use of com.google.datastore.v1.Mutation 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 com.google.datastore.v1.Mutation in project beam by apache.
the class DatastoreV1Test method testDeleteKeys.
/**
* Test that valid keys are transformed to delete mutations.
*/
@Test
public void testDeleteKeys() {
Key key = makeKey("bird", "finch").build();
DeleteKeyFn deleteKeyFn = new DeleteKeyFn();
Mutation expectedMutation = makeDelete(key).build();
assertEquals(expectedMutation, deleteKeyFn.apply(key));
}
use of com.google.datastore.v1.Mutation in project beam by apache.
the class BigtableWriteIT method testE2EBigtableWrite.
@Test
public void testE2EBigtableWrite() throws Exception {
final String tableName = bigtableOptions.getInstanceName().toTableNameStr(tableId);
final String instanceName = bigtableOptions.getInstanceName().toString();
final int numRows = 1000;
final List<KV<ByteString, ByteString>> testData = generateTableData(numRows);
createEmptyTable(instanceName, tableId);
Pipeline p = Pipeline.create(options);
p.apply(GenerateSequence.from(0).to(numRows)).apply(ParDo.of(new DoFn<Long, KV<ByteString, Iterable<Mutation>>>() {
@ProcessElement
public void processElement(ProcessContext c) {
int index = c.element().intValue();
Iterable<Mutation> mutations = ImmutableList.of(Mutation.newBuilder().setSetCell(Mutation.SetCell.newBuilder().setValue(testData.get(index).getValue()).setFamilyName(COLUMN_FAMILY_NAME)).build());
c.output(KV.of(testData.get(index).getKey(), mutations));
}
})).apply(BigtableIO.write().withBigtableOptions(bigtableOptions).withTableId(tableId));
p.run();
// Test number of column families and column family name equality
Table table = getTable(tableName);
assertThat(table.getColumnFamiliesMap().keySet(), Matchers.hasSize(1));
assertThat(table.getColumnFamiliesMap(), Matchers.hasKey(COLUMN_FAMILY_NAME));
// Test table data equality
List<KV<ByteString, ByteString>> tableData = getTableData(tableName);
assertThat(tableData, Matchers.containsInAnyOrder(testData.toArray()));
}
use of com.google.datastore.v1.Mutation in project beam by apache.
the class DatastoreV1Test method testAddEntities.
@Test
public /**
* Test that entities with valid keys are transformed to upsert mutations.
*/
void testAddEntities() throws Exception {
Key key = makeKey("bird", "finch").build();
Entity entity = Entity.newBuilder().setKey(key).build();
UpsertFn upsertFn = new UpsertFn();
Mutation expectedMutation = makeUpsert(entity).build();
assertEquals(expectedMutation, upsertFn.apply(entity));
}
use of com.google.datastore.v1.Mutation in project simple-bigtable by spotify.
the class BigtableMutationImplTest method testDeleteColumnFamilyAndQualifier.
@Test
public void testDeleteColumnFamilyAndQualifier() throws Exception {
final BigtableMutationImpl bigtableMutationImpl = (BigtableMutationImpl) this.bigtableMutation.deleteColumn("family", "qualifier");
verifyGetMutateRowRequest();
bigtableMutationImpl.execute();
bigtableMutationImpl.executeAsync();
assertEquals(1, bigtableMutationImpl.getMutateRowRequest().getMutationsCount());
final Mutation mutation = bigtableMutationImpl.getMutateRowRequest().getMutations(0);
assertEquals(Mutation.MutationCase.DELETE_FROM_COLUMN, mutation.getMutationCase());
assertEquals("family", mutation.getDeleteFromColumn().getFamilyName());
assertEquals("qualifier", mutation.getDeleteFromColumn().getColumnQualifier().toStringUtf8());
assertEquals(Mutation.DeleteFromRow.getDefaultInstance(), mutation.getDeleteFromRow());
assertEquals(Mutation.DeleteFromFamily.getDefaultInstance(), mutation.getDeleteFromFamily());
assertEquals(Mutation.SetCell.getDefaultInstance(), mutation.getSetCell());
verify(bigtableMock.getMockedDataClient()).mutateRow(bigtableMutation.getMutateRowRequest().build());
verify(bigtableMock.getMockedDataClient()).mutateRowAsync(bigtableMutation.getMutateRowRequest().build());
verifyNoMoreInteractions(bigtableMock.getMockedDataClient());
}
Aggregations