diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableWriteSchemaTransformProvider.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableWriteSchemaTransformProvider.java index cc480be6aa7e..480d4199c653 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableWriteSchemaTransformProvider.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableWriteSchemaTransformProvider.java @@ -19,6 +19,7 @@ import static java.util.Optional.ofNullable; import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument; +import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkState; import com.google.auto.service.AutoService; import com.google.auto.value.AutoValue; @@ -34,16 +35,21 @@ import java.util.Map; import org.apache.beam.sdk.io.gcp.bigtable.BigtableWriteSchemaTransformProvider.BigtableWriteSchemaTransformConfiguration; import org.apache.beam.sdk.schemas.AutoValueSchema; +import org.apache.beam.sdk.schemas.Schema; import org.apache.beam.sdk.schemas.annotations.DefaultSchema; import org.apache.beam.sdk.schemas.transforms.SchemaTransform; import org.apache.beam.sdk.schemas.transforms.SchemaTransformProvider; import org.apache.beam.sdk.schemas.transforms.TypedSchemaTransformProvider; +import org.apache.beam.sdk.transforms.GroupByKey; import org.apache.beam.sdk.transforms.MapElements; import org.apache.beam.sdk.transforms.SimpleFunction; +import org.apache.beam.sdk.util.Preconditions; import org.apache.beam.sdk.values.KV; import org.apache.beam.sdk.values.PCollection; import org.apache.beam.sdk.values.PCollectionRowTuple; import org.apache.beam.sdk.values.Row; +import org.apache.beam.sdk.values.TypeDescriptor; +import org.apache.beam.sdk.values.TypeDescriptors; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.primitives.Longs; /** @@ -60,6 +66,13 @@ public class BigtableWriteSchemaTransformProvider private static final String INPUT_TAG = "input"; + private static final Schema BATCHED_MUTATIONS_SCHEMA = + Schema.builder() + .addByteArrayField("key") + .addArrayField( + "mutations", Schema.FieldType.map(Schema.FieldType.STRING, Schema.FieldType.BYTES)) + .build(); + @Override protected SchemaTransform from(BigtableWriteSchemaTransformConfiguration configuration) { return new BigtableWriteSchemaTransform(configuration); @@ -135,18 +148,206 @@ public PCollectionRowTuple expand(PCollectionRowTuple input) { String.format( "Could not find expected input [%s] to %s.", INPUT_TAG, getClass().getSimpleName())); - PCollection beamRowMutations = input.get(INPUT_TAG); - PCollection>> bigtableMutations = - beamRowMutations.apply(MapElements.via(new GetMutationsFromBeamRow())); + Schema inputSchema = input.getSinglePCollection().getSchema(); - bigtableMutations.apply( - BigtableIO.write() - .withTableId(configuration.getTableId()) - .withInstanceId(configuration.getInstanceId()) - .withProjectId(configuration.getProjectId())); + PCollection>> bigtableMutations = null; + if (inputSchema.equals(BATCHED_MUTATIONS_SCHEMA)) { + PCollection beamRowMutations = input.get(INPUT_TAG); + bigtableMutations = + beamRowMutations.apply( + // Original schema inputs gets sent out to the original transform provider mutations + // function + MapElements.via(new GetMutationsFromBeamRow())); + } else if (inputSchema.hasField("type")) { + validateField(inputSchema, "key", Schema.TypeName.BYTES); + validateField(inputSchema, "type", Schema.TypeName.STRING); + if (inputSchema.hasField("value")) { + validateField(inputSchema, "value", Schema.TypeName.BYTES); + } + if (inputSchema.hasField("column_qualifier")) { + validateField(inputSchema, "column_qualifier", Schema.TypeName.BYTES); + } + if (inputSchema.hasField("family_name")) { + validateField(inputSchema, "family_name", Schema.TypeName.BYTES); + } + if (inputSchema.hasField("timestamp_micros")) { + validateField(inputSchema, "timestamp_micros", Schema.TypeName.INT64); + } + if (inputSchema.hasField("start_timestamp_micros")) { + validateField(inputSchema, "start_timestamp_micros", Schema.TypeName.INT64); + } + if (inputSchema.hasField("end_timestamp_micros")) { + validateField(inputSchema, "end_timestamp_micros", Schema.TypeName.INT64); + } + bigtableMutations = changeMutationInput(input); + } else { + throw new RuntimeException( + "Input Schema is invalid: " + + inputSchema + + "\n\nSchema should be formatted in one of two ways:\n " + + "key\": ByteString\n" + + "\"type\": String\n" + + "\"value\": ByteString\n" + + "\"column_qualifier\": ByteString\n" + + "\"family_name\": ByteString\n" + + "\"timestamp_micros\": Long\n" + + "\"start_timestamp_micros\": Long\n" + + "\"end_timestamp_micros\": Long\n" + + "\nOR\n" + + "\n" + + "\"key\": ByteString\n" + + "(\"mutations\", contains map(String, ByteString) of mutations in the mutation schema format"); + } + if (bigtableMutations != null) { + bigtableMutations.apply( + BigtableIO.write() + .withTableId(configuration.getTableId()) + .withInstanceId(configuration.getInstanceId()) + .withProjectId(configuration.getProjectId())); + } else { + throw new RuntimeException( + "Inputted Schema caused mutation error, check error logs and input schema format"); + } return PCollectionRowTuple.empty(input.getPipeline()); } + + private void validateField(Schema inputSchema, String field, Schema.TypeName expectedType) { + Schema.TypeName actualType = inputSchema.getField(field).getType().getTypeName(); + checkState( + actualType.equals(expectedType), + "Schema field '%s' should be of type %s, but was %s.", + field, + expectedType, + actualType); + } + + public PCollection>> changeMutationInput( + PCollectionRowTuple inputR) { + PCollection beamRowMutationsList = inputR.getSinglePCollection(); + // convert all row inputs into KV + PCollection> changedBeamRowMutationsList = + beamRowMutationsList.apply( + MapElements.into( + TypeDescriptors.kvs( + TypeDescriptor.of(ByteString.class), TypeDescriptor.of(Mutation.class))) + .via( + (Row input) -> { + ByteString key = + ByteString.copyFrom( + Preconditions.checkStateNotNull( + input.getBytes("key"), + "Encountered row with null 'key' property.")); + + Mutation bigtableMutation; + String mutationType = + input.getString("type"); // Direct call, can return null + if (mutationType == null) { + throw new IllegalArgumentException("Mutation type cannot be null."); + } + switch (mutationType) { + case "SetCell": + Mutation.SetCell.Builder setMutation = + Mutation.SetCell.newBuilder() + .setValue( + ByteString.copyFrom( + Preconditions.checkStateNotNull( + input.getBytes("value"), + "Encountered SetCell mutation with null 'value' property."))) + .setColumnQualifier( + ByteString.copyFrom( + Preconditions.checkStateNotNull( + input.getBytes("column_qualifier"), + "Encountered SetCell mutation with null 'column_qualifier' property. "))) + .setFamilyNameBytes( + ByteString.copyFrom( + Preconditions.checkStateNotNull( + input.getBytes("family_name"), + "Encountered SetCell mutation with null 'family_name' property."))); + // Use timestamp if provided, else default to -1 (current + // Bigtable + // server time) + // Timestamp (optional, assuming Long type in Row schema) + Long timestampMicros = input.getInt64("timestamp_micros"); + setMutation.setTimestampMicros( + timestampMicros != null ? timestampMicros : -1); + + bigtableMutation = + Mutation.newBuilder().setSetCell(setMutation.build()).build(); + break; + case "DeleteFromColumn": + // set timestamp range if applicable + Mutation.DeleteFromColumn.Builder deleteMutation = + Mutation.DeleteFromColumn.newBuilder() + .setColumnQualifier( + ByteString.copyFrom( + Preconditions.checkStateNotNull( + input.getBytes("column_qualifier"), + "Encountered DeleteFromColumn mutation with null 'column_qualifier' property."))) + .setFamilyNameBytes( + ByteString.copyFrom( + Preconditions.checkStateNotNull( + input.getBytes("family_name"), + "Encountered DeleteFromColumn mutation with null 'family_name' property."))); + + // if start or end timestamp provided + // Timestamp Range (optional, assuming Long type in Row schema) + Long startTimestampMicros = null; + Long endTimestampMicros = null; + + if (input.getSchema().hasField("start_timestamp_micros")) { + startTimestampMicros = input.getInt64("start_timestamp_micros"); + } + if (input.getSchema().hasField("end_timestamp_micros")) { + endTimestampMicros = input.getInt64("end_timestamp_micros"); + } + + if (startTimestampMicros != null || endTimestampMicros != null) { + TimestampRange.Builder timeRange = TimestampRange.newBuilder(); + if (startTimestampMicros != null) { + timeRange.setStartTimestampMicros(startTimestampMicros); + } + if (endTimestampMicros != null) { + timeRange.setEndTimestampMicros(endTimestampMicros); + } + deleteMutation.setTimeRange(timeRange.build()); + } + bigtableMutation = + Mutation.newBuilder() + .setDeleteFromColumn(deleteMutation.build()) + .build(); + break; + case "DeleteFromFamily": + bigtableMutation = + Mutation.newBuilder() + .setDeleteFromFamily( + Mutation.DeleteFromFamily.newBuilder() + .setFamilyNameBytes( + ByteString.copyFrom( + Preconditions.checkStateNotNull( + input.getBytes("family_name"), + "Encountered DeleteFromFamily mutation with null 'family_name' property."))) + .build()) + .build(); + break; + case "DeleteFromRow": + bigtableMutation = + Mutation.newBuilder() + .setDeleteFromRow(Mutation.DeleteFromRow.newBuilder().build()) + .build(); + break; + default: + throw new RuntimeException( + String.format( + "Unexpected mutation type [%s]: Key value is %s", + ((input.getString("type"))), + Arrays.toString(input.getBytes("key")))); + } + return KV.of(key, bigtableMutation); + })); + // now we need to make the KV into a PCollection of KV> + return changedBeamRowMutationsList.apply(GroupByKey.create()); + } } public static class GetMutationsFromBeamRow diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableSimpleWriteSchemaTransformProviderIT.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableSimpleWriteSchemaTransformProviderIT.java new file mode 100644 index 000000000000..7a5dcdc3e999 --- /dev/null +++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableSimpleWriteSchemaTransformProviderIT.java @@ -0,0 +1,655 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.beam.sdk.io.gcp.bigtable; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertThrows; +import static org.junit.Assert.assertTrue; + +import com.google.api.gax.rpc.NotFoundException; +import com.google.cloud.bigtable.admin.v2.BigtableTableAdminClient; +import com.google.cloud.bigtable.admin.v2.BigtableTableAdminSettings; +import com.google.cloud.bigtable.admin.v2.models.CreateTableRequest; +import com.google.cloud.bigtable.data.v2.BigtableDataClient; +import com.google.cloud.bigtable.data.v2.BigtableDataSettings; +import com.google.cloud.bigtable.data.v2.models.Query; +import com.google.cloud.bigtable.data.v2.models.RowCell; +import com.google.cloud.bigtable.data.v2.models.RowMutation; +import java.nio.charset.StandardCharsets; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Date; +import java.util.List; +import java.util.stream.Collectors; +import org.apache.beam.sdk.extensions.gcp.options.GcpOptions; +import org.apache.beam.sdk.io.gcp.bigtable.BigtableWriteSchemaTransformProvider.BigtableWriteSchemaTransformConfiguration; +import org.apache.beam.sdk.schemas.Schema; +import org.apache.beam.sdk.schemas.Schema.FieldType; // Import FieldType +import org.apache.beam.sdk.testing.TestPipeline; +import org.apache.beam.sdk.transforms.Create; +import org.apache.beam.sdk.transforms.PTransform; +import org.apache.beam.sdk.values.PCollection; +import org.apache.beam.sdk.values.PCollectionRowTuple; +import org.apache.beam.sdk.values.Row; +import org.junit.After; +import org.junit.Before; +import org.junit.Rule; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +@RunWith(JUnit4.class) +public class BigtableSimpleWriteSchemaTransformProviderIT { + @Rule public final transient TestPipeline p = TestPipeline.create(); + + private static final String COLUMN_FAMILY_NAME_1 = "test_cf_1"; + private static final String COLUMN_FAMILY_NAME_2 = "test_cf_2"; + private BigtableTableAdminClient tableAdminClient; + private BigtableDataClient dataClient; + private String tableId = String.format("BigtableWriteIT-%tF-% writeTransform; + + @Test + public void testInvalidConfigs() { + // Properties cannot be empty (project, instance, and table) + List invalidConfigs = + Arrays.asList( + BigtableWriteSchemaTransformConfiguration.builder() + .setProjectId("project") + .setInstanceId("instance") + .setTableId(""), + BigtableWriteSchemaTransformConfiguration.builder() + .setProjectId("") + .setInstanceId("instance") + .setTableId("table"), + BigtableWriteSchemaTransformConfiguration.builder() + .setProjectId("project") + .setInstanceId("") + .setTableId("table")); + + for (BigtableWriteSchemaTransformConfiguration.Builder config : invalidConfigs) { + assertThrows( + IllegalArgumentException.class, + () -> { + config.build().validate(); + }); + } + } + + @Before + public void setup() throws Exception { + BigtableTestOptions options = + TestPipeline.testingPipelineOptions().as(BigtableTestOptions.class); + projectId = options.as(GcpOptions.class).getProject(); + instanceId = options.getInstanceId(); + + BigtableDataSettings settings = + BigtableDataSettings.newBuilder().setProjectId(projectId).setInstanceId(instanceId).build(); + // Creates a bigtable data client. + dataClient = BigtableDataClient.create(settings); + + BigtableTableAdminSettings adminSettings = + BigtableTableAdminSettings.newBuilder() + .setProjectId(projectId) + .setInstanceId(instanceId) + .build(); + tableAdminClient = BigtableTableAdminClient.create(adminSettings); + + // set up the table with some pre-written rows to test our mutations on. + // each test is independent of the others + if (!tableAdminClient.exists(tableId)) { + CreateTableRequest createTableRequest = + CreateTableRequest.of(tableId) + .addFamily(COLUMN_FAMILY_NAME_1) + .addFamily(COLUMN_FAMILY_NAME_2); + tableAdminClient.createTable(createTableRequest); + } + + BigtableWriteSchemaTransformConfiguration config = + BigtableWriteSchemaTransformConfiguration.builder() + .setProjectId(projectId) + .setInstanceId(instanceId) + .setTableId(tableId) + .build(); + writeTransform = new BigtableWriteSchemaTransformProvider().from(config); + } + + @After + public void tearDown() { + try { + tableAdminClient.deleteTable(tableId); + System.out.printf("Table %s deleted successfully%n", tableId); + } catch (NotFoundException e) { + System.err.println("Failed to delete a non-existent table: " + e.getMessage()); + } + dataClient.close(); + tableAdminClient.close(); + } + + @Test + public void testSetMutationsExistingColumn() { + RowMutation rowMutation = + RowMutation.create(tableId, "key-1") + .setCell(COLUMN_FAMILY_NAME_1, "col_a", 1000, "val-1-a") + .setCell(COLUMN_FAMILY_NAME_2, "col_c", 1000, "val-1-c"); + dataClient.mutateRow(rowMutation); + Schema testSchema = + Schema.builder() + .addByteArrayField("key") + .addStringField("type") + .addByteArrayField("value") + .addByteArrayField("column_qualifier") + .addByteArrayField("family_name") + .addField("timestamp_micros", FieldType.INT64) // Changed to INT64 + .build(); + + Row mutationRow1 = + Row.withSchema(testSchema) + .withFieldValue("key", "key-1".getBytes(StandardCharsets.UTF_8)) + .withFieldValue("type", "SetCell") + .withFieldValue("value", "new-val-1-a".getBytes(StandardCharsets.UTF_8)) + .withFieldValue("column_qualifier", "col_a".getBytes(StandardCharsets.UTF_8)) + .withFieldValue("family_name", COLUMN_FAMILY_NAME_1.getBytes(StandardCharsets.UTF_8)) + .withFieldValue("timestamp_micros", 2000L) + .build(); + Row mutationRow2 = + Row.withSchema(testSchema) + .withFieldValue("key", "key-1".getBytes(StandardCharsets.UTF_8)) + .withFieldValue("type", "SetCell") + .withFieldValue("value", "new-val-1-c".getBytes(StandardCharsets.UTF_8)) + .withFieldValue("column_qualifier", "col_c".getBytes(StandardCharsets.UTF_8)) + .withFieldValue("family_name", COLUMN_FAMILY_NAME_2.getBytes(StandardCharsets.UTF_8)) + .withFieldValue("timestamp_micros", 2000L) + .build(); + + PCollection inputPCollection = + p.apply(Create.of(Arrays.asList(mutationRow1, mutationRow2))); + inputPCollection.setRowSchema(testSchema); + + PCollectionRowTuple.of("input", inputPCollection) // Use the schema-set PCollection + .apply(writeTransform); + p.run().waitUntilFinish(); + + // get rows from table + List rows = + dataClient.readRows(Query.create(tableId)).stream().collect(Collectors.toList()); + // we should still have only one row with the same key + assertEquals(1, rows.size()); + assertEquals("key-1", rows.get(0).getKey().toStringUtf8()); + + // check that we now have two cells in each column we added to and that + // the last cell in each column has the updated value + com.google.cloud.bigtable.data.v2.models.Row row = rows.get(0); + List cellsColA = + row.getCells(COLUMN_FAMILY_NAME_1, "col_a").stream() + .sorted(RowCell.compareByNative()) + .collect(Collectors.toList()); + List cellsColC = + row.getCells(COLUMN_FAMILY_NAME_2, "col_c").stream() + .sorted(RowCell.compareByNative()) + .collect(Collectors.toList()); + assertEquals(2, cellsColA.size()); + assertEquals(2, cellsColC.size()); + // Bigtable keeps cell history ordered by descending timestamp + assertEquals("new-val-1-a", cellsColA.get(0).getValue().toStringUtf8()); + assertEquals("new-val-1-c", cellsColC.get(0).getValue().toStringUtf8()); + assertEquals("val-1-a", cellsColA.get(1).getValue().toStringUtf8()); + assertEquals("val-1-c", cellsColC.get(1).getValue().toStringUtf8()); + } + + @Test + public void testSetMutationNewColumn() { + RowMutation rowMutation = + RowMutation.create(tableId, "key-1").setCell(COLUMN_FAMILY_NAME_1, "col_a", "val-1-a"); + dataClient.mutateRow(rowMutation); + Schema testSchema = + Schema.builder() + .addByteArrayField("key") + .addStringField("type") + .addByteArrayField("value") + .addByteArrayField("column_qualifier") + .addByteArrayField("family_name") + .addField("timestamp_micros", FieldType.INT64) + .build(); + Row mutationRow = + Row.withSchema(testSchema) + .withFieldValue("key", "key-1".getBytes(StandardCharsets.UTF_8)) + .withFieldValue("type", "SetCell") + .withFieldValue("value", "new-val-1".getBytes(StandardCharsets.UTF_8)) + .withFieldValue("column_qualifier", "new_col".getBytes(StandardCharsets.UTF_8)) + .withFieldValue("family_name", COLUMN_FAMILY_NAME_1.getBytes(StandardCharsets.UTF_8)) + .withFieldValue("timestamp_micros", 999_000L) + .build(); + + PCollection inputPCollection = p.apply(Create.of(Arrays.asList(mutationRow))); + inputPCollection.setRowSchema(testSchema); + + PCollectionRowTuple.of("input", inputPCollection) // Use the schema-set PCollection + .apply(writeTransform); + p.run().waitUntilFinish(); + + // get rows from table + List rows = + dataClient.readRows(Query.create(tableId)).stream().collect(Collectors.toList()); + + // we should still have only one row with the same key + assertEquals(1, rows.size()); + assertEquals("key-1", rows.get(0).getKey().toStringUtf8()); + // check the new column exists with only one cell. + // also check cell value is correct + com.google.cloud.bigtable.data.v2.models.Row row = rows.get(0); + List cellsNewCol = row.getCells(COLUMN_FAMILY_NAME_1, "new_col"); + assertEquals(1, cellsNewCol.size()); + assertEquals("new-val-1", cellsNewCol.get(0).getValue().toStringUtf8()); + } + + @Test + public void testDeleteCellsFromColumn() { + RowMutation rowMutation = + RowMutation.create(tableId, "key-1") + .setCell(COLUMN_FAMILY_NAME_1, "col_a", "val-1-a") + .setCell(COLUMN_FAMILY_NAME_1, "col_b", "val-1-b"); + dataClient.mutateRow(rowMutation); + // write two cells in col_a. both should get deleted + rowMutation = + RowMutation.create(tableId, "key-1").setCell(COLUMN_FAMILY_NAME_1, "col_a", "new-val-1-a"); + dataClient.mutateRow(rowMutation); + Schema testSchema = + Schema.builder() + .addByteArrayField("key") + .addStringField("type") + .addByteArrayField("column_qualifier") + .addByteArrayField("family_name") + .build(); + Row mutationRow = + Row.withSchema(testSchema) + .withFieldValue("key", "key-1".getBytes(StandardCharsets.UTF_8)) + .withFieldValue("type", "DeleteFromColumn") + .withFieldValue("column_qualifier", "col_a".getBytes(StandardCharsets.UTF_8)) + .withFieldValue("family_name", COLUMN_FAMILY_NAME_1.getBytes(StandardCharsets.UTF_8)) + .build(); + + PCollection inputPCollection = p.apply(Create.of(Arrays.asList(mutationRow))); + inputPCollection.setRowSchema(testSchema); + + PCollectionRowTuple.of("input", inputPCollection) // Use the schema-set PCollection + .apply(writeTransform); + p.run().waitUntilFinish(); + + // get rows from table + List rows = + dataClient.readRows(Query.create(tableId)).stream().collect(Collectors.toList()); + + // we should still have one row with the same key + assertEquals(1, rows.size()); + assertEquals("key-1", rows.get(0).getKey().toStringUtf8()); + // get cells from this column family. we started with three cells and deleted two from one + // column. + // we should end up with one cell in the column we didn't touch. + com.google.cloud.bigtable.data.v2.models.Row row = rows.get(0); + List cells = row.getCells(COLUMN_FAMILY_NAME_1); + assertEquals(1, cells.size()); + assertEquals("col_b", cells.get(0).getQualifier().toStringUtf8()); + } + + @Test + public void testDeleteCellsFromColumnWithTimestampRange() { + // write two cells in one column with different timestamps. + RowMutation rowMutation = + RowMutation.create(tableId, "key-1") + .setCell(COLUMN_FAMILY_NAME_1, "col", 100_000_000, "val"); + dataClient.mutateRow(rowMutation); + rowMutation = + RowMutation.create(tableId, "key-1") + .setCell(COLUMN_FAMILY_NAME_1, "col", 200_000_000, "new-val"); + dataClient.mutateRow(rowMutation); + Schema testSchema = + Schema.builder() + .addByteArrayField("key") + .addStringField("type") + .addByteArrayField("column_qualifier") + .addByteArrayField("family_name") + .addField("start_timestamp_micros", FieldType.INT64) + .addField("end_timestamp_micros", FieldType.INT64) + .build(); + Row mutationRow = + Row.withSchema(testSchema) + .withFieldValue("key", "key-1".getBytes(StandardCharsets.UTF_8)) + .withFieldValue("type", "DeleteFromColumn") + .withFieldValue("column_qualifier", "col".getBytes(StandardCharsets.UTF_8)) + .withFieldValue("family_name", COLUMN_FAMILY_NAME_1.getBytes(StandardCharsets.UTF_8)) + .withFieldValue("start_timestamp_micros", 99_990_000L) + .withFieldValue("end_timestamp_micros", 100_000_000L) + .build(); + + PCollection inputPCollection = p.apply(Create.of(Arrays.asList(mutationRow))); + inputPCollection.setRowSchema(testSchema); + + PCollectionRowTuple.of("input", inputPCollection) // Use the schema-set PCollection + .apply(writeTransform); + p.run().waitUntilFinish(); + + // get rows from table + List rows = + dataClient.readRows(Query.create(tableId)).stream().collect(Collectors.toList()); + + // we should still have one row with the same key + assertEquals(1, rows.size()); + assertEquals("key-1", rows.get(0).getKey().toStringUtf8()); + // we had two cells in col_a and deleted the older one. we should be left with the newer cell. + // check cell has correct value and timestamp + com.google.cloud.bigtable.data.v2.models.Row row = rows.get(0); + List cells = row.getCells(COLUMN_FAMILY_NAME_1, "col"); + assertEquals(2, cells.size()); + assertEquals("new-val", cells.get(0).getValue().toStringUtf8()); + assertEquals(200_000_000, cells.get(0).getTimestamp()); + } + + @Test + public void testDeleteColumnFamily() { + RowMutation rowMutation = + RowMutation.create(tableId, "key-1") + .setCell(COLUMN_FAMILY_NAME_1, "col_a", "val") + .setCell(COLUMN_FAMILY_NAME_2, "col_b", "val"); + dataClient.mutateRow(rowMutation); + Schema testSchema = + Schema.builder() + .addByteArrayField("key") + .addStringField("type") + .addByteArrayField("family_name") + .build(); + Row mutationRow = + Row.withSchema(testSchema) + .withFieldValue("key", "key-1".getBytes(StandardCharsets.UTF_8)) + .withFieldValue("type", "DeleteFromFamily") + .withFieldValue("family_name", COLUMN_FAMILY_NAME_1.getBytes(StandardCharsets.UTF_8)) + .build(); + + PCollection inputPCollection = p.apply(Create.of(Arrays.asList(mutationRow))); + inputPCollection.setRowSchema(testSchema); + + PCollectionRowTuple.of("input", inputPCollection) // Use the schema-set PCollection + .apply(writeTransform); + p.run().waitUntilFinish(); + + // get rows from table + List rows = + dataClient.readRows(Query.create(tableId)).stream().collect(Collectors.toList()); + + // we should still have one row with the same key + assertEquals(1, rows.size()); + assertEquals("key-1", rows.get(0).getKey().toStringUtf8()); + // we had one cell in each of two column families. we deleted a column family, so should end up + // with + // one cell in the column family we didn't touch. + com.google.cloud.bigtable.data.v2.models.Row row = rows.get(0); + List cells = row.getCells(); + assertEquals(1, cells.size()); + assertEquals(COLUMN_FAMILY_NAME_2, cells.get(0).getFamily()); + } + + @Test + public void testDeleteRow() { + RowMutation rowMutation = + RowMutation.create(tableId, "key-1").setCell(COLUMN_FAMILY_NAME_1, "col", "val-1"); + dataClient.mutateRow(rowMutation); + rowMutation = + RowMutation.create(tableId, "key-2").setCell(COLUMN_FAMILY_NAME_1, "col", "val-2"); + dataClient.mutateRow(rowMutation); + Schema testSchema = Schema.builder().addByteArrayField("key").addStringField("type").build(); + Row mutationRow = + Row.withSchema(testSchema) + .withFieldValue("key", "key-1".getBytes(StandardCharsets.UTF_8)) + .withFieldValue("type", "DeleteFromRow") + .build(); + + PCollection inputPCollection = p.apply(Create.of(Arrays.asList(mutationRow))); + inputPCollection.setRowSchema(testSchema); + + PCollectionRowTuple.of("input", inputPCollection) // Use the schema-set PCollection + .apply(writeTransform); + p.run().waitUntilFinish(); + + // get rows from table + List rows = + dataClient.readRows(Query.create(tableId)).stream().collect(Collectors.toList()); + + // we created two rows then deleted one, so should end up with the row we didn't touch + assertEquals(1, rows.size()); + assertEquals("key-2", rows.get(0).getKey().toStringUtf8()); + } + + @Test + public void testAllMutations() { + + // --- Initial Setup: Populate the table with diverse data --- + dataClient.mutateRow( + RowMutation.create(tableId, "row-setcell") + .setCell(COLUMN_FAMILY_NAME_1, "col_initial_1", "initial_val_1") + .setCell(COLUMN_FAMILY_NAME_2, "col_initial_2", "initial_val_2")); + + dataClient.mutateRow( + RowMutation.create(tableId, "row-delete-col") + .setCell(COLUMN_FAMILY_NAME_1, "col_to_delete_A", 1000, "val_to_delete_A_old") + .setCell(COLUMN_FAMILY_NAME_1, "col_to_delete_A", 2000, "val_to_delete_A_new") + .setCell(COLUMN_FAMILY_NAME_1, "col_to_keep_B", "val_to_keep_B")); + + dataClient.mutateRow( + RowMutation.create(tableId, "row-delete-col-ts") + .setCell(COLUMN_FAMILY_NAME_1, "ts_col", 1000, "ts_val_old") + .setCell(COLUMN_FAMILY_NAME_1, "ts_col", 2000, "ts_val_new") + .setCell(COLUMN_FAMILY_NAME_2, "ts_col_other_cf", "ts_val_other_cf")); + + dataClient.mutateRow( + RowMutation.create(tableId, "row-delete-family") + .setCell(COLUMN_FAMILY_NAME_1, "col_to_delete_family", "val_delete_family") + .setCell(COLUMN_FAMILY_NAME_2, "col_to_keep_family", "val_keep_family")); + + dataClient.mutateRow( + RowMutation.create(tableId, "row-delete-row") + .setCell(COLUMN_FAMILY_NAME_1, "col", "val_delete_row")); + + dataClient.mutateRow( + RowMutation.create(tableId, "row-final-check") + .setCell(COLUMN_FAMILY_NAME_1, "col_final_1", "val_final_1")); + + // --- Define Schema for various mutation types --- + + Schema uberSchema = + Schema.builder() + .addByteArrayField("key") // Key is always present and non-null + .addStringField( + "type") // Type is always present and non-null (e.g., "SetCell", "DeleteFromRow") + // All other fields are conditional based on the 'type' of mutation, so they must be + // nullable. + .addNullableField("value", FieldType.BYTES) // Used by SetCell + .addNullableField( + "column_qualifier", FieldType.BYTES) // Used by SetCell, DeleteFromColumn + .addNullableField( + "family_name", + FieldType.BYTES) // Used by SetCell, DeleteFromColumn, DeleteFromFamily + .addNullableField("timestamp_micros", FieldType.INT64) // Optional for SetCell + .addNullableField( + "start_timestamp_micros", FieldType.INT64) // Used by DeleteFromColumn with range + .addNullableField( + "end_timestamp_micros", FieldType.INT64) // Used by DeleteFromColumn with range + .build(); + + // --- Create a list of mutation Rows --- + List mutations = new ArrayList<>(); + + // 1. SetCell (Update an existing cell, add a new cell) + // Update "row-setcell", col_initial_1 + mutations.add( + Row.withSchema(uberSchema) + .withFieldValue("key", "row-setcell".getBytes(StandardCharsets.UTF_8)) + .withFieldValue("type", "SetCell") + .withFieldValue("value", "updated_val_1".getBytes(StandardCharsets.UTF_8)) + .withFieldValue("column_qualifier", "col_initial_1".getBytes(StandardCharsets.UTF_8)) + .withFieldValue("family_name", COLUMN_FAMILY_NAME_1.getBytes(StandardCharsets.UTF_8)) + .withFieldValue("timestamp_micros", 3000L) + .build()); + // Add new cell to "row-setcell" + mutations.add( + Row.withSchema(uberSchema) + .withFieldValue("key", "row-setcell".getBytes(StandardCharsets.UTF_8)) + .withFieldValue("type", "SetCell") + .withFieldValue("value", "new_col_val".getBytes(StandardCharsets.UTF_8)) + .withFieldValue("column_qualifier", "new_col_A".getBytes(StandardCharsets.UTF_8)) + .withFieldValue("family_name", COLUMN_FAMILY_NAME_1.getBytes(StandardCharsets.UTF_8)) + .withFieldValue("timestamp_micros", 4000L) + .build()); + + // 2. DeleteFromColumn + // Delete "col_to_delete_A" from "row-delete-col" + mutations.add( + Row.withSchema(uberSchema) + .withFieldValue("key", "row-delete-col".getBytes(StandardCharsets.UTF_8)) + .withFieldValue("type", "DeleteFromColumn") + .withFieldValue("column_qualifier", "col_to_delete_A".getBytes(StandardCharsets.UTF_8)) + .withFieldValue("family_name", COLUMN_FAMILY_NAME_1.getBytes(StandardCharsets.UTF_8)) + .build()); + + // 3. DeleteFromColumn with Timestamp Range + // Delete "ts_col" with timestamp 1000 from "row-delete-col-ts" + mutations.add( + Row.withSchema(uberSchema) + .withFieldValue("key", "row-delete-col-ts".getBytes(StandardCharsets.UTF_8)) + .withFieldValue("type", "DeleteFromColumn") + .withFieldValue("column_qualifier", "ts_col".getBytes(StandardCharsets.UTF_8)) + .withFieldValue("family_name", COLUMN_FAMILY_NAME_1.getBytes(StandardCharsets.UTF_8)) + .withFieldValue("start_timestamp_micros", 999L) // Inclusive + .withFieldValue("end_timestamp_micros", 1001L) // Exclusive + .build()); + + // 4. DeleteFromFamily + // Delete COLUMN_FAMILY_NAME_1 from "row-delete-family" + mutations.add( + Row.withSchema(uberSchema) + .withFieldValue("key", "row-delete-family".getBytes(StandardCharsets.UTF_8)) + .withFieldValue("type", "DeleteFromFamily") + .withFieldValue("family_name", COLUMN_FAMILY_NAME_1.getBytes(StandardCharsets.UTF_8)) + .build()); + + // 5. DeleteFromRow + // Delete "row-delete-row" + mutations.add( + Row.withSchema(uberSchema) + .withFieldValue("key", "row-delete-row".getBytes(StandardCharsets.UTF_8)) + .withFieldValue("type", "DeleteFromRow") + .build()); + + // --- Apply the mutations -- + PCollection inputPCollection = p.apply(Create.of(mutations)); + inputPCollection.setRowSchema(uberSchema); // Set the comprehensive schema for the PCollection + + PCollectionRowTuple.of("input", inputPCollection) // Use the schema-set PCollection + .apply(writeTransform); + p.run().waitUntilFinish(); + + // --- Assertions: Verify the final state of the table --- + + List rows = + dataClient.readRows(Query.create(tableId)).stream().collect(Collectors.toList()); + + assertEquals(5, rows.size()); // Expecting 'row-setcell', 'row-delete-col', 'row-delete-col-ts', + // 'row-final-check' + + // Verify "row-setcell" + com.google.cloud.bigtable.data.v2.models.Row rowSetCell = + rows.stream() + .filter(r -> r.getKey().toStringUtf8().equals("row-setcell")) + .findFirst() + .orElse(null); + assertEquals("row-setcell", rowSetCell.getKey().toStringUtf8()); + List cellsSetCellCol1 = + rowSetCell.getCells(COLUMN_FAMILY_NAME_1, "col_initial_1").stream() + .sorted(RowCell.compareByNative()) + .collect(Collectors.toList()); + assertEquals(2, cellsSetCellCol1.size()); // Original + updated + assertEquals( + "initial_val_1", cellsSetCellCol1.get(0).getValue().toStringUtf8()); // Newest value + assertEquals( + "updated_val_1", cellsSetCellCol1.get(1).getValue().toStringUtf8()); // Oldest value + List cellsSetCellNewCol = rowSetCell.getCells(COLUMN_FAMILY_NAME_1, "new_col_A"); + assertEquals(1, cellsSetCellNewCol.size()); + assertEquals("new_col_val", cellsSetCellNewCol.get(0).getValue().toStringUtf8()); + List cellsSetCellCol2 = rowSetCell.getCells(COLUMN_FAMILY_NAME_2, "col_initial_2"); + assertEquals(1, cellsSetCellCol2.size()); + assertEquals("initial_val_2", cellsSetCellCol2.get(0).getValue().toStringUtf8()); + + // Verify "row-delete-col" + com.google.cloud.bigtable.data.v2.models.Row rowDeleteCol = + rows.stream() + .filter(r -> r.getKey().toStringUtf8().equals("row-delete-col")) + .findFirst() + .orElse(null); + assertEquals("row-delete-col", rowDeleteCol.getKey().toStringUtf8()); + List cellsColToDeleteA = + rowDeleteCol.getCells(COLUMN_FAMILY_NAME_1, "col_to_delete_A"); + assertTrue(cellsColToDeleteA.isEmpty()); // Should be deleted + List cellsColToKeepB = rowDeleteCol.getCells(COLUMN_FAMILY_NAME_1, "col_to_keep_B"); + assertEquals(1, cellsColToKeepB.size()); + assertEquals("val_to_keep_B", cellsColToKeepB.get(0).getValue().toStringUtf8()); + + // Verify "row-delete-col-ts" + com.google.cloud.bigtable.data.v2.models.Row rowDeleteColTs = + rows.stream() + .filter(r -> r.getKey().toStringUtf8().equals("row-delete-col-ts")) + .findFirst() + .orElse(null); + assertEquals("row-delete-col-ts", rowDeleteColTs.getKey().toStringUtf8()); + List cellsTsCol = rowDeleteColTs.getCells(COLUMN_FAMILY_NAME_1, "ts_col"); + assertEquals(1, cellsTsCol.size()); // Only the 2000 timestamp cell should remain + assertEquals("ts_val_new", cellsTsCol.get(0).getValue().toStringUtf8()); + assertEquals(2000, cellsTsCol.get(0).getTimestamp()); + List cellsTsColOtherCf = + rowDeleteColTs.getCells(COLUMN_FAMILY_NAME_2, "ts_col_other_cf"); + assertEquals(1, cellsTsColOtherCf.size()); + assertEquals("ts_val_other_cf", cellsTsColOtherCf.get(0).getValue().toStringUtf8()); + + // Verify "row-delete-family" + com.google.cloud.bigtable.data.v2.models.Row rowDeleteFamily = + rows.stream() + .filter(r -> r.getKey().toStringUtf8().equals("row-delete-family")) + .findFirst() + .orElse(null); + assertEquals("row-delete-family", rowDeleteFamily.getKey().toStringUtf8()); + List cellsCf1 = rowDeleteFamily.getCells(COLUMN_FAMILY_NAME_1); + assertTrue(cellsCf1.isEmpty()); // COLUMN_FAMILY_NAME_1 should be empty + List cellsCf2 = rowDeleteFamily.getCells(COLUMN_FAMILY_NAME_2); + assertEquals(1, cellsCf2.size()); + assertEquals("val_keep_family", cellsCf2.get(0).getValue().toStringUtf8()); + + // Verify "row-delete-row" is gone + assertTrue(rows.stream().noneMatch(r -> r.getKey().toStringUtf8().equals("row-delete-row"))); + + // Verify "row-final-check" still exists + com.google.cloud.bigtable.data.v2.models.Row rowFinalCheck = + rows.stream() + .filter(r -> r.getKey().toStringUtf8().equals("row-final-check")) + .findFirst() + .orElse(null); + assertEquals("row-final-check", rowFinalCheck.getKey().toStringUtf8()); + List cellsFinalCheck = rowFinalCheck.getCells(COLUMN_FAMILY_NAME_1, "col_final_1"); + assertEquals(1, cellsFinalCheck.size()); + assertEquals("val_final_1", cellsFinalCheck.get(0).getValue().toStringUtf8()); + } +} diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableWriteSchemaTransformProviderIT.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableWriteSchemaTransformProviderIT.java index 1a60fe661b52..22159d5fb724 100644 --- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableWriteSchemaTransformProviderIT.java +++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableWriteSchemaTransformProviderIT.java @@ -42,6 +42,7 @@ import org.apache.beam.sdk.testing.TestPipeline; import org.apache.beam.sdk.transforms.Create; import org.apache.beam.sdk.transforms.PTransform; +import org.apache.beam.sdk.values.PCollection; import org.apache.beam.sdk.values.PCollectionRowTuple; import org.apache.beam.sdk.values.Row; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap; @@ -180,7 +181,10 @@ public void testSetMutationsExistingColumn() { .withFieldValue("mutations", mutations) .build(); - PCollectionRowTuple.of("input", p.apply(Create.of(Arrays.asList(mutationRow)))) + PCollection inputPCollection = p.apply(Create.of(Arrays.asList(mutationRow))); + inputPCollection.setRowSchema(SCHEMA); + + PCollectionRowTuple.of("input", inputPCollection) // Use the schema-set PCollection .apply(writeTransform); p.run().waitUntilFinish(); @@ -231,7 +235,10 @@ public void testSetMutationNewColumn() { .withFieldValue("mutations", mutations) .build(); - PCollectionRowTuple.of("input", p.apply(Create.of(Arrays.asList(mutationRow)))) + PCollection inputPCollection = p.apply(Create.of(Arrays.asList(mutationRow))); + inputPCollection.setRowSchema(SCHEMA); + + PCollectionRowTuple.of("input", inputPCollection) // Use the schema-set PCollection .apply(writeTransform); p.run().waitUntilFinish(); @@ -275,7 +282,10 @@ public void testDeleteCellsFromColumn() { .withFieldValue("mutations", mutations) .build(); - PCollectionRowTuple.of("input", p.apply(Create.of(Arrays.asList(mutationRow)))) + PCollection inputPCollection = p.apply(Create.of(Arrays.asList(mutationRow))); + inputPCollection.setRowSchema(SCHEMA); + + PCollectionRowTuple.of("input", inputPCollection) // Use the schema-set PCollection .apply(writeTransform); p.run().waitUntilFinish(); @@ -323,7 +333,10 @@ public void testDeleteCellsFromColumnWithTimestampRange() { .withFieldValue("mutations", mutations) .build(); - PCollectionRowTuple.of("input", p.apply(Create.of(Arrays.asList(mutationRow)))) + PCollection inputPCollection = p.apply(Create.of(Arrays.asList(mutationRow))); + inputPCollection.setRowSchema(SCHEMA); + + PCollectionRowTuple.of("input", inputPCollection) // Use the schema-set PCollection .apply(writeTransform); p.run().waitUntilFinish(); @@ -363,7 +376,10 @@ public void testDeleteColumnFamily() { .withFieldValue("mutations", mutations) .build(); - PCollectionRowTuple.of("input", p.apply(Create.of(Arrays.asList(mutationRow)))) + PCollection inputPCollection = p.apply(Create.of(Arrays.asList(mutationRow))); + inputPCollection.setRowSchema(SCHEMA); + + PCollectionRowTuple.of("input", inputPCollection) // Use the schema-set PCollection .apply(writeTransform); p.run().waitUntilFinish(); @@ -401,7 +417,10 @@ public void testDeleteRow() { .withFieldValue("mutations", mutations) .build(); - PCollectionRowTuple.of("input", p.apply(Create.of(Arrays.asList(mutationRow)))) + PCollection inputPCollection = p.apply(Create.of(Arrays.asList(mutationRow))); + inputPCollection.setRowSchema(SCHEMA); + + PCollectionRowTuple.of("input", inputPCollection) // Use the schema-set PCollection .apply(writeTransform); p.run().waitUntilFinish(); diff --git a/sdks/python/apache_beam/yaml/integration_tests.py b/sdks/python/apache_beam/yaml/integration_tests.py index 4a9ca70cac54..38fa2689268e 100644 --- a/sdks/python/apache_beam/yaml/integration_tests.py +++ b/sdks/python/apache_beam/yaml/integration_tests.py @@ -24,10 +24,13 @@ import logging import os import random +import secrets import sqlite3 import string import unittest import uuid +from datetime import datetime +from datetime import timezone import mock import mysql.connector @@ -36,6 +39,8 @@ import sqlalchemy import yaml from google.cloud import pubsub_v1 +from google.cloud.bigtable import client +from google.cloud.bigtable_admin_v2.types import instance from testcontainers.core.container import DockerContainer from testcontainers.core.waiting_utils import wait_for_logs from testcontainers.google import PubSubContainer @@ -54,6 +59,9 @@ from apache_beam.yaml import yaml_provider from apache_beam.yaml import yaml_transform from apache_beam.yaml.conftest import yaml_test_files_dir +from apitools.base.py.exceptions import HttpError + +_LOGGER = logging.getLogger(__name__) @contextlib.contextmanager @@ -144,6 +152,55 @@ def temp_bigquery_table(project, prefix='yaml_bq_it_'): bigquery_client.client.datasets.Delete(request) +def instance_prefix(instance): + datestr = "".join(filter(str.isdigit, str(datetime.now(timezone.utc).date()))) + instance_id = '%s-%s-%s' % (instance, datestr, secrets.token_hex(4)) + assert len(instance_id) < 34, "instance id length needs to be within [6, 33]" + return instance_id + + +@contextlib.contextmanager +def temp_bigtable_table(project, prefix='yaml_bt_it_'): + INSTANCE = "bt-write-tests" + TABLE_ID = "test-table" + + instance_id = instance_prefix(INSTANCE) + + clientT = client.Client(admin=True, project=project) + # create cluster and instance + instanceT = clientT.instance( + instance_id, + display_name=INSTANCE, + instance_type=instance.Instance.Type.DEVELOPMENT) + cluster = instanceT.cluster("test-cluster", "us-central1-a") + operation = instanceT.create(clusters=[cluster]) + operation.result(timeout=500) + _LOGGER.info("Created instance [%s] in project [%s]", instance_id, project) + + # create table inside instance + table = instanceT.table(TABLE_ID) + table.create() + _LOGGER.info("Created table [%s]", table.table_id) + # in the table that is created, make a new family called cf1 + col_fam = table.column_family('cf1') + col_fam.create() + + # another family called cf2 + col_fam = table.column_family('cf2') + col_fam.create() + + #yielding the tmp table for all the bigTable tests + yield instance_id + + #try catch for deleting table and instance after all tests are ran + try: + _LOGGER.info("Deleting table [%s]", table.table_id) + table.delete() + instanceT.delete() + except HttpError: + _LOGGER.warning("Failed to clean up instance") + + @contextlib.contextmanager def temp_sqlite_database(prefix='yaml_jdbc_it_'): """Context manager to provide a temporary SQLite database via JDBC for @@ -706,7 +763,7 @@ def parse_test_files(filepattern): globals()[suite_name] = type(suite_name, (unittest.TestCase, ), methods) -# Logging setup +# Logging setups logging.getLogger().setLevel(logging.INFO) # Dynamically create test methods from the tests directory. diff --git a/sdks/python/apache_beam/yaml/standard_io.yaml b/sdks/python/apache_beam/yaml/standard_io.yaml index d3d6e8f8d029..80b2e76ace27 100644 --- a/sdks/python/apache_beam/yaml/standard_io.yaml +++ b/sdks/python/apache_beam/yaml/standard_io.yaml @@ -371,3 +371,27 @@ 'WriteToTFRecord': 'beam:schematransform:org.apache.beam:tfrecord_write:v1' config: gradle_target: 'sdks:java:io:expansion-service:shadowJar' + +#BigTable +- type: renaming + transforms: + #'ReadFromBigTable': 'ReadFromBigTable' + 'WriteToBigTable': 'WriteToBigTable' + config: + mappings: + #Temp removing read from bigTable IO +# 'ReadFromBigTable': +# project: 'project_id' +# instance: 'instance_id' +# table: 'table_id' + 'WriteToBigTable': + project: 'project_id' + instance: 'instance_id' + table: 'table_id' + underlying_provider: + type: beamJar + transforms: +# 'ReadFromBigTable': 'beam:schematransform:org.apache.beam:bigtable_read:v1' + 'WriteToBigTable': 'beam:schematransform:org.apache.beam:bigtable_write:v1' + config: + gradle_target: 'sdks:java:io:google-cloud-platform:expansion-service:shadowJar' diff --git a/sdks/python/apache_beam/yaml/tests/bigtable.yaml b/sdks/python/apache_beam/yaml/tests/bigtable.yaml new file mode 100644 index 000000000000..eae5f1bcbb74 --- /dev/null +++ b/sdks/python/apache_beam/yaml/tests/bigtable.yaml @@ -0,0 +1,87 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + + +fixtures: + - name: BT_INSTANCE + type: "apache_beam.yaml.integration_tests.temp_bigtable_table" + config: + project: "apache-beam-testing" + - name: TEMP_DIR + # Need distributed filesystem to be able to read and write from a container. + type: "apache_beam.yaml.integration_tests.gcs_temp_dir" + config: + bucket: "gs://temp-storage-for-end-to-end-tests/temp-it" + + # Tests for BigTable YAML IO + +pipelines: + - pipeline: + type: chain + transforms: + - type: Create + config: + elements: + - {key: 'row1', + type: 'SetCell', + family_name: "cf1", + column_qualifier: "cq1", + value: "value1", + timestamp_micros: -1} + - {key: 'row1', + type: 'SetCell', + family_name: "cf2", + column_qualifier: "cq1", + value: "value2", + timestamp_micros: 1000} + + - type: LogForTesting + - type: MapToFields + name: ConvertStringsToBytes + config: + language: python + fields: + # For 'SetCell' and 'DeleteFromColumn' + key: + callable: | + def convert_to_bytes(row): + return bytes(row.key, 'utf-8') if "key" in row._fields else None + type: + type + family_name: + callable: | + def convert_to_bytes(row): + return bytes(row.family_name, 'utf-8') if 'family_name' in row._fields else None + column_qualifier: + callable: | + def convert_to_bytes(row): + return bytes(row.column_qualifier, 'utf-8') if 'column_qualifier' in row._fields else None + value: + callable: | + def convert_to_bytes(row): + return bytes(row.value, 'utf-8') if 'value' in row._fields else None + timestamp_micros: + timestamp_micros + # The 'type', 'timestamp_micros', 'start_timestamp_micros', 'end_timestamp_micros' + # fields are already of the correct type (String, Long) or are optional. + # We only need to convert fields that are Strings in YAML but need to be Bytes in Java. + + - type: WriteToBigTable + config: + project: 'apache-beam-testing' + instance: "{BT_INSTANCE}" + table: 'test-table'