feat(yaml): Add BigTable write connector

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<Row> beamRowMutations = input.get(INPUT_TAG);
-      PCollection<KV<ByteString, Iterable<Mutation>>> 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<KV<ByteString, Iterable<Mutation>>> bigtableMutations = null;
+      if (inputSchema.equals(BATCHED_MUTATIONS_SCHEMA)) {
+        PCollection<Row> 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<KV<ByteString, Iterable<Mutation>>> changeMutationInput(
+        PCollectionRowTuple inputR) {
+      PCollection<Row> beamRowMutationsList = inputR.getSinglePCollection();
+      // convert all row inputs into KV<ByteString, Mutation>
+      PCollection<KV<ByteString, Mutation>> 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<ByteString, Iterable<Mutation>>
+      return changedBeamRowMutationsList.apply(GroupByKey.create());
+    }
   }
 
   public static class GetMutationsFromBeamRow