Search in sources :

Example 11 with Mutation

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;
    }
}
Also used : CommitRequest(com.google.datastore.v1.CommitRequest) ArrayList(java.util.ArrayList) Mutation(com.google.datastore.v1.Mutation) DatastoreWriterFn(org.apache.beam.sdk.io.gcp.datastore.DatastoreV1.DatastoreWriterFn)

Example 12 with Mutation

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));
}
Also used : DeleteKeyFn(org.apache.beam.sdk.io.gcp.datastore.DatastoreV1.DeleteKeyFn) Mutation(com.google.datastore.v1.Mutation) DeleteKey(org.apache.beam.sdk.io.gcp.datastore.DatastoreV1.DeleteKey) Key(com.google.datastore.v1.Key) DatastoreHelper.makeKey(com.google.datastore.v1.client.DatastoreHelper.makeKey) DatastoreV1.isValidKey(org.apache.beam.sdk.io.gcp.datastore.DatastoreV1.isValidKey) Test(org.junit.Test)

Example 13 with Mutation

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()));
}
Also used : Table(com.google.bigtable.admin.v2.Table) ByteString(com.google.protobuf.ByteString) KV(org.apache.beam.sdk.values.KV) Mutation(com.google.bigtable.v2.Mutation) TestPipeline(org.apache.beam.sdk.testing.TestPipeline) Pipeline(org.apache.beam.sdk.Pipeline) Test(org.junit.Test)

Example 14 with Mutation

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));
}
Also used : Entity(com.google.datastore.v1.Entity) DeleteEntity(org.apache.beam.sdk.io.gcp.datastore.DatastoreV1.DeleteEntity) UpsertFn(org.apache.beam.sdk.io.gcp.datastore.DatastoreV1.UpsertFn) Mutation(com.google.datastore.v1.Mutation) DeleteKey(org.apache.beam.sdk.io.gcp.datastore.DatastoreV1.DeleteKey) Key(com.google.datastore.v1.Key) DatastoreHelper.makeKey(com.google.datastore.v1.client.DatastoreHelper.makeKey) DatastoreV1.isValidKey(org.apache.beam.sdk.io.gcp.datastore.DatastoreV1.isValidKey) Test(org.junit.Test)

Example 15 with Mutation

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());
}
Also used : Mutation(com.google.bigtable.v2.Mutation) Test(org.junit.Test)

Aggregations

Test (org.junit.Test)19 Mutation (com.google.bigtable.v2.Mutation)13 Mutation (com.google.datastore.v1.Mutation)7 Mutation (com.airbnb.jitney.event.spinaltap.v1.Mutation)5 ByteString (com.google.protobuf.ByteString)5 ArrayList (java.util.ArrayList)5 MysqlUpdateMutation (com.airbnb.spinaltap.mysql.mutation.MysqlUpdateMutation)4 Key (com.google.datastore.v1.Key)4 DatastoreHelper.makeKey (com.google.datastore.v1.client.DatastoreHelper.makeKey)4 MysqlDeleteMutation (com.airbnb.spinaltap.mysql.mutation.MysqlDeleteMutation)3 MysqlInsertMutation (com.airbnb.spinaltap.mysql.mutation.MysqlInsertMutation)3 MysqlMutation (com.airbnb.spinaltap.mysql.mutation.MysqlMutation)3 CommitRequest (com.google.datastore.v1.CommitRequest)3 Entity (com.google.datastore.v1.Entity)3 DatastoreException (com.google.datastore.v1.client.DatastoreException)3 DatastoreWriterFn (org.apache.beam.sdk.io.gcp.datastore.DatastoreV1.DatastoreWriterFn)3 DeleteEntity (org.apache.beam.sdk.io.gcp.datastore.DatastoreV1.DeleteEntity)3 DeleteKey (org.apache.beam.sdk.io.gcp.datastore.DatastoreV1.DeleteKey)3 DatastoreV1.isValidKey (org.apache.beam.sdk.io.gcp.datastore.DatastoreV1.isValidKey)3 KV (org.apache.beam.sdk.values.KV)3