Examples with Mutation com.google.spanner.v1.Mutation used on opensource projects

Example 71 with Mutation

use of com.google.spanner.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 {
    when(mockBigtableSource.getTableId()).thenReturn(StaticValueProvider.of(TABLE_ID));
    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();
}

Example 72 with Mutation

use of com.google.spanner.v1.Mutation in project DataflowTemplates by GoogleCloudPlatform.

the class BeamRowToBigtableFn method processElement.

@ProcessElement
@SuppressWarnings("unused")
public void processElement(@Element Row row, OutputReceiver<KV<ByteString, Iterable<Mutation>>> out) {
    // Generate the Bigtable Rowkey. This key will be used for all Cells for this row.
    ByteString rowkey = generateRowKey(row);
    // Number of mutations accumulated to be outputted.
    int cellsProcessed = 0;
    // DoFn return value.
    ImmutableList.Builder<Mutation> mutations = ImmutableList.builder();
    // Retrieve all fields that are not part of the rowkey. All fields that are part of the
    // rowkey must have a metadata key set to key_order. All other fields should become their own
    // cells in Bigtable.
    List<Field> nonKeyColumns = row.getSchema().getFields().stream().filter(f -> f.getType().getMetadataString(CassandraRowMapperFn.KEY_ORDER_METADATA_KEY).isEmpty()).collect(Collectors.toList());
    // Iterate over all the fields with three cases. Primitives, collections and maps.
    for (Field field : nonKeyColumns) {
        List<Mutation> mutationsToAdd;
        TypeName type = field.getType().getTypeName();
        if (type.isPrimitiveType()) {
            ByteString value = primitiveFieldToBytes(type, row.getValue(field.getName()));
            SetCell cell = createCell(defaultColumnFamily.get(), field.getName(), value);
            mutationsToAdd = Arrays.asList(Mutation.newBuilder().setSetCell(cell).build());
        } else if (type.isCollectionType()) {
            mutationsToAdd = createCollectionMutations(row, field);
        } else if (type.isMapType()) {
            mutationsToAdd = createMapMutations(row, field);
        } else {
            throw new UnsupportedOperationException("Mapper does not support type:" + field.getType().getTypeName().toString());
        }
        for (Mutation mutationToAdd : mutationsToAdd) {
            mutations.add(mutationToAdd);
            cellsProcessed++;
            if (this.splitLargeRows && cellsProcessed % maxMutationsPerRequest == 0) {
                // Send a MutateRow request when we have accumulated max mutations per request.
                out.output(KV.of(rowkey, mutations.build()));
                mutations = ImmutableList.builder();
            }
        }
    }
    // Flush any remaining mutations.
    ImmutableList remainingMutations = mutations.build();
    if (!remainingMutations.isEmpty()) {
        out.output(KV.of(rowkey, remainingMutations));
    }
}

Example 73 with Mutation

use of com.google.spanner.v1.Mutation in project DataflowTemplates by GoogleCloudPlatform.

the class ParquetToBigtableTest method applyParquettoBigtableFnWithSplitLargeRows.

/**
 * Test whether {@link ParquetToBigtable} correctly maps a GenericRecord to a KV with
 * SplitLargeRows enabled..
 */
@Test
public void applyParquettoBigtableFnWithSplitLargeRows() throws Exception {
    byte[] rowKey1 = "row1".getBytes();
    ByteBuffer key1 = ByteBuffer.wrap(rowKey1);
    List<BigtableCell> cells1 = new ArrayList<>();
    addParquetCell(cells1, "family1", "column1", 1, "10");
    addParquetCell(cells1, "family1", "column1", 2, "20");
    addParquetCell(cells1, "family1", "column2", 1, "30");
    addParquetCell(cells1, "family2", "column1", 1, "40");
    GenericRecord parquetRow1 = new GenericRecordBuilder(BigtableRow.getClassSchema()).set("key", key1).set("cells", cells1).build();
    byte[] rowKey2 = "row2".getBytes();
    ByteBuffer key2 = ByteBuffer.wrap(rowKey2);
    List<BigtableCell> cells2 = new ArrayList<>();
    addParquetCell(cells2, "family2", "column2", 2, "40");
    GenericRecord parquetRow2 = new GenericRecordBuilder(BigtableRow.getClassSchema()).set("key", key2).set("cells", cells2).build();
    final List<GenericRecord> parquetRows = ImmutableList.of(parquetRow1, parquetRow2);
    KV<ByteString, Iterable<Mutation>> rowMutations1 = createBigtableRowMutations("row1");
    addBigtableMutation(rowMutations1, "family1", "column1", 1, "10");
    addBigtableMutation(rowMutations1, "family1", "column1", 2, "20");
    KV<ByteString, Iterable<Mutation>> rowMutations2 = createBigtableRowMutations("row1");
    addBigtableMutation(rowMutations2, "family1", "column2", 1, "30");
    addBigtableMutation(rowMutations2, "family2", "column1", 1, "40");
    KV<ByteString, Iterable<Mutation>> rowMutations3 = createBigtableRowMutations("row2");
    addBigtableMutation(rowMutations3, "family2", "column2", 2, "40");
    final List<KV<ByteString, Iterable<Mutation>>> expectedBigtableRows = ImmutableList.of(rowMutations1, rowMutations2, rowMutations3);
    PCollection<KV<ByteString, Iterable<Mutation>>> bigtableRows = pipeline.apply("Create", Create.of(parquetRows).withCoder(AvroCoder.of(GenericRecord.class, BigtableRow.getClassSchema()))).apply("TransformToBigtable", ParDo.of(ParquetToBigtableFn.createWithSplitLargeRows(StaticValueProvider.of(true), 2)));
    PAssert.that(bigtableRows).containsInAnyOrder(expectedBigtableRows);
    pipeline.run().waitUntilFinish();
}

Example 74 with Mutation

use of com.google.spanner.v1.Mutation in project DataflowTemplates by GoogleCloudPlatform.

the class TestUtils method addBigtableMutation.

static void addBigtableMutation(KV<ByteString, Iterable<Mutation>> rowMutations, String family, String qualifier, long timestamp, String value) {
    SetCell setCell = SetCell.newBuilder().setFamilyName(family).setColumnQualifier(toByteString(qualifier)).setTimestampMicros(timestamp).setValue(toByteString(value)).build();
    Mutation mutation = Mutation.newBuilder().setSetCell(setCell).build();
    ((List) rowMutations.getValue()).add(mutation);
}

Example 75 with Mutation

use of com.google.spanner.v1.Mutation in project DataflowTemplates by GoogleCloudPlatform.

the class BeamRowToBigtableFnTest method processElementWithMapColumn.

@Test
public void processElementWithMapColumn() {
    String columnFamily = "default";
    String rowKeyValue = "rowkeyvalue";
    String rowKeyColumnName = "rowkey";
    String listColumnName = "mapColumnName";
    Map<Integer, String> mapValue = new HashMap<>();
    mapValue.put(0, "first");
    mapValue.put(1, "second");
    mapValue.put(2, "third");
    Schema schema = Schema.builder().addField(Schema.Field.of(rowKeyColumnName, FieldType.STRING.withMetadata(CassandraRowMapperFn.KEY_ORDER_METADATA_KEY, "0"))).addField(Schema.Field.of(listColumnName, FieldType.map(FieldType.INT32, FieldType.STRING))).build();
    Row input = Row.withSchema(schema).addValue(rowKeyValue).addValue(mapValue).build();
    final List<Row> rows = Collections.singletonList(input);
    List<Mutation> mutations = new ArrayList<>();
    mutations.add(createMutation(columnFamily, "mapColumnName[0].key", ByteString.copyFrom(Bytes.toBytes(0))));
    mutations.add(createMutation(columnFamily, "mapColumnName[0].value", ByteString.copyFrom(Bytes.toBytes(mapValue.get(0)))));
    mutations.add(createMutation(columnFamily, "mapColumnName[1].key", ByteString.copyFrom(Bytes.toBytes(1))));
    mutations.add(createMutation(columnFamily, "mapColumnName[1].value", ByteString.copyFrom(Bytes.toBytes(mapValue.get(1)))));
    mutations.add(createMutation(columnFamily, "mapColumnName[2].key", ByteString.copyFrom(Bytes.toBytes(2))));
    mutations.add(createMutation(columnFamily, "mapColumnName[2].value", ByteString.copyFrom(Bytes.toBytes(mapValue.get(2)))));
    final List<KV<ByteString, Iterable<Mutation>>> expectedBigtableRows = ImmutableList.of(KV.of(ByteString.copyFrom(Bytes.toBytes("rowkeyvalue")), mutations));
    PCollection<KV<ByteString, Iterable<Mutation>>> bigtableRows = pipeline.apply("Create", Create.of(rows)).apply("Transform to Bigtable", ParDo.of(BeamRowToBigtableFn.create(ValueProvider.StaticValueProvider.of("#"), ValueProvider.StaticValueProvider.of("default"))));
    PAssert.that(bigtableRows).containsInAnyOrder(expectedBigtableRows);
    pipeline.run();
}