Search in sources :

Example 21 with Mutation

use of com.airbnb.jitney.event.spinaltap.v1.Mutation in project SpinalTap by airbnb.

the class KafkaDestinationTest method getMutation.

private Mutation getMutation(byte[] payload) throws Exception {
    Mutation mutation = new Mutation();
    deserializer.get().deserialize(mutation, payload);
    return mutation;
}
Also used : MysqlUpdateMutation(com.airbnb.spinaltap.mysql.mutation.MysqlUpdateMutation) Mutation(com.airbnb.jitney.event.spinaltap.v1.Mutation) MysqlMutation(com.airbnb.spinaltap.mysql.mutation.MysqlMutation) MysqlDeleteMutation(com.airbnb.spinaltap.mysql.mutation.MysqlDeleteMutation) MysqlInsertMutation(com.airbnb.spinaltap.mysql.mutation.MysqlInsertMutation)

Example 22 with Mutation

use of com.airbnb.jitney.event.spinaltap.v1.Mutation in project beam by apache.

the class DatastoreV1Test method testDeleteEntities.

/**
 * Test that entities with valid keys are transformed to delete mutations.
 */
@Test
public void testDeleteEntities() throws Exception {
    Key key = makeKey("bird", "finch").build();
    Entity entity = Entity.newBuilder().setKey(key).build();
    DeleteEntityFn deleteEntityFn = new DeleteEntityFn();
    Mutation expectedMutation = makeDelete(entity.getKey()).build();
    assertEquals(expectedMutation, deleteEntityFn.apply(entity));
}
Also used : Entity(com.google.datastore.v1.Entity) DeleteEntity(org.apache.beam.sdk.io.gcp.datastore.DatastoreV1.DeleteEntity) DeleteEntityFn(org.apache.beam.sdk.io.gcp.datastore.DatastoreV1.DeleteEntityFn) 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 23 with Mutation

use of com.airbnb.jitney.event.spinaltap.v1.Mutation 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;
    }
}
Also used : CommitRequest(com.google.datastore.v1.CommitRequest) Entity(com.google.datastore.v1.Entity) DeleteEntity(org.apache.beam.sdk.io.gcp.datastore.DatastoreV1.DeleteEntity) ArrayList(java.util.ArrayList) Mutation(com.google.datastore.v1.Mutation) DatastoreWriterFn(org.apache.beam.sdk.io.gcp.datastore.DatastoreV1.DatastoreWriterFn) Test(org.junit.Test)

Example 24 with Mutation

use of com.airbnb.jitney.event.spinaltap.v1.Mutation in project beam by apache.

the class BigtableIOTest method testWritingFailsTableDoesNotExist.

/**
 * Tests that when writing to a non-existent table, the write fails.
 */
@Test
public void testWritingFailsTableDoesNotExist() throws Exception {
    final String table = "TEST-TABLE";
    PCollection<KV<ByteString, Iterable<Mutation>>> emptyInput = p.apply(Create.empty(KvCoder.of(ByteStringCoder.of(), IterableCoder.of(ProtoCoder.of(Mutation.class)))));
    // Exception will be thrown by write.validate() when writeToDynamic is applied.
    thrown.expect(IllegalArgumentException.class);
    thrown.expectMessage(String.format("Table %s does not exist", table));
    emptyInput.apply("write", defaultWrite.withTableId(table));
    p.run();
}
Also used : ByteString(com.google.protobuf.ByteString) KV(org.apache.beam.sdk.values.KV) Mutation(com.google.bigtable.v2.Mutation) Test(org.junit.Test)

Example 25 with Mutation

use of com.airbnb.jitney.event.spinaltap.v1.Mutation in project beam by apache.

the class BigtableServiceImplTest method testWrite.

/**
 * This test ensures that protobuf creation and interactions with {@link BulkMutation} work as
 * expected.
 *
 * @throws IOException
 * @throws InterruptedException
 */
@Test
public void testWrite() throws IOException, InterruptedException {
    BigtableService.Writer underTest = new BigtableServiceImpl.BigtableWriterImpl(mockSession, TABLE_NAME);
    Mutation mutation = Mutation.newBuilder().setSetCell(SetCell.newBuilder().setFamilyName("Family").build()).build();
    ByteString key = ByteString.copyFromUtf8("key");
    SettableFuture<MutateRowResponse> fakeResponse = SettableFuture.create();
    when(mockBulkMutation.add(any(MutateRowsRequest.Entry.class))).thenReturn(fakeResponse);
    underTest.writeRecord(KV.of(key, ImmutableList.of(mutation)));
    Entry expected = MutateRowsRequest.Entry.newBuilder().setRowKey(key).addMutations(mutation).build();
    verify(mockBulkMutation, times(1)).add(expected);
    underTest.close();
    verify(mockBulkMutation, times(1)).flush();
}
Also used : MutateRowResponse(com.google.bigtable.v2.MutateRowResponse) Entry(com.google.bigtable.v2.MutateRowsRequest.Entry) ByteString(com.google.protobuf.ByteString) Mutation(com.google.bigtable.v2.Mutation) BulkMutation(com.google.cloud.bigtable.grpc.async.BulkMutation) 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 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 MysqlMutationMetadata (com.airbnb.spinaltap.mysql.mutation.MysqlMutationMetadata)2