Adding error handler for SpannerReadSchemaTransformProvider and missi… (#35241)

* Adding error handler for SpannerReadSchemaTransformProvider and missing tests for SpannerSchemaTransformProvider

* Removed not used logging

* Spotless Apply

* Spotless Apply

* Spotless Apply

* Typo correction

Author: TanuSharma2511

sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerReadSchemaTransformProvider.java

@@ -17,28 +17,40 @@
  */
 package org.apache.beam.sdk.io.gcp.spanner;
 
+import static org.apache.beam.sdk.util.Preconditions.checkArgumentNotNull;
 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.checkNotNull;
 
 import com.google.auto.service.AutoService;
 import com.google.auto.value.AutoValue;
 import com.google.cloud.spanner.Struct;
 import java.io.Serializable;
+import java.util.Arrays;
 import java.util.Collections;
 import java.util.List;
 import javax.annotation.Nullable;
+import org.apache.beam.sdk.metrics.Counter;
+import org.apache.beam.sdk.metrics.Metrics;
 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.annotations.SchemaFieldDescription;
 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.MapElements;
+import org.apache.beam.sdk.schemas.transforms.providers.ErrorHandling;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.DoFn.FinishBundle;
+import org.apache.beam.sdk.transforms.DoFn.FinishBundleContext;
+import org.apache.beam.sdk.transforms.DoFn.MultiOutputReceiver;
+import org.apache.beam.sdk.transforms.DoFn.ProcessElement;
+import org.apache.beam.sdk.transforms.ParDo;
 import org.apache.beam.sdk.values.PCollection;
 import org.apache.beam.sdk.values.PCollectionRowTuple;
+import org.apache.beam.sdk.values.PCollectionTuple;
 import org.apache.beam.sdk.values.Row;
-import org.apache.beam.sdk.values.TypeDescriptor;
+import org.apache.beam.sdk.values.TupleTag;
+import org.apache.beam.sdk.values.TupleTagList;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Strings;
 
 /** A provider for reading from Cloud Spanner using a Schema Transform Provider. */
@@ -61,6 +73,11 @@ public class SpannerReadSchemaTransformProvider
     extends TypedSchemaTransformProvider<
         SpannerReadSchemaTransformProvider.SpannerReadSchemaTransformConfiguration> {
 
+  public static final TupleTag<Row> OUTPUT_TAG = new TupleTag<Row>() {};
+  public static final TupleTag<Row> ERROR_TAG = new TupleTag<Row>() {};
+  public static final Schema ERROR_SCHEMA =
+      Schema.builder().addStringField("error").addStringField("row").build();
+
   @Override
   public String identifier() {
     return "beam:schematransform:org.apache.beam:spanner_read:v1";
@@ -133,6 +150,7 @@ static class SpannerSchemaTransformRead extends SchemaTransform implements Seria
     @Override
     public PCollectionRowTuple expand(PCollectionRowTuple input) {
       checkNotNull(input, "Input to SpannerReadSchemaTransform cannot be null.");
+      boolean handleErrors = ErrorHandling.hasOutput(configuration.getErrorHandling());
       SpannerIO.Read read =
           SpannerIO.readWithSchema()
               .withProjectId(configuration.getProjectId())
@@ -152,12 +170,66 @@ public PCollectionRowTuple expand(PCollectionRowTuple input) {
       }
       PCollection<Struct> spannerRows = input.getPipeline().apply(read);
       Schema schema = spannerRows.getSchema();
-      PCollection<Row> rows =
+
+      PCollectionTuple outputTuple =
           spannerRows.apply(
-              MapElements.into(TypeDescriptor.of(Row.class))
-                  .via((Struct struct) -> StructUtils.structToBeamRow(struct, schema)));
+              ParDo.of(
+                      new ErrorFn("spanner-read-error-counter", ERROR_SCHEMA, schema, handleErrors))
+                  .withOutputTags(OUTPUT_TAG, TupleTagList.of(ERROR_TAG)));
+
+      PCollectionRowTuple outputRows =
+          PCollectionRowTuple.of("output", outputTuple.get(OUTPUT_TAG).setRowSchema(schema));
+
+      // Error handling
+      PCollection<Row> errorOutput = outputTuple.get(ERROR_TAG).setRowSchema(ERROR_SCHEMA);
+      if (handleErrors) {
+        outputRows =
+            outputRows.and(
+                checkArgumentNotNull(configuration.getErrorHandling()).getOutput(), errorOutput);
+      }
 
-      return PCollectionRowTuple.of("output", rows.setRowSchema(schema));
+      return outputRows;
+    }
+  }
+
+  public static class ErrorFn extends DoFn<Struct, Row> {
+    private final Counter errorCounter;
+    private Long errorsInBundle = 0L;
+    private final boolean handleErrors;
+    private final Schema errorSchema;
+    private final Schema schema;
+
+    public ErrorFn(String name, Schema errorSchema, Schema schema, boolean handleErrors) {
+      this.errorCounter = Metrics.counter(SpannerReadSchemaTransformProvider.class, name);
+      this.handleErrors = handleErrors;
+      this.errorSchema = errorSchema;
+      this.schema = schema;
+    }
+
+    @ProcessElement
+    public void processElement(@DoFn.Element Struct struct, MultiOutputReceiver receiver) {
+      Row mappedRow = null;
+      try {
+        mappedRow = StructUtils.structToBeamRow(struct, schema);
+      } catch (Exception e) {
+        if (!handleErrors) {
+          throw new RuntimeException(e);
+        }
+        errorsInBundle += 1;
+        receiver
+            .get(ERROR_TAG)
+            .output(
+                Row.withSchema(errorSchema).addValues(e.getMessage(), struct.toString()).build());
+      }
+      if (mappedRow != null) {
+        receiver.get(OUTPUT_TAG).output(mappedRow);
+      }
+    }
+
+    @FinishBundle
+    public void finish(FinishBundleContext c) {
+      errorCounter.inc(errorsInBundle);
+      errorsInBundle = 0L;
     }
   }
 
@@ -168,7 +240,7 @@ public List<String> inputCollectionNames() {
 
   @Override
   public List<String> outputCollectionNames() {
-    return Collections.singletonList("output");
+    return Arrays.asList("output", "errors");
   }
 
   @DefaultSchema(AutoValueSchema.class)
@@ -193,6 +265,8 @@ public abstract static class Builder {
 
       public abstract Builder setBatching(Boolean batching);
 
+      public abstract Builder setErrorHandling(ErrorHandling errorHandling);
+
       public abstract SpannerReadSchemaTransformConfiguration build();
     }
 
@@ -261,6 +335,10 @@ public static Builder builder() {
         "Set to false to disable batching. Useful when using a query that is not compatible with the PartitionQuery API. Defaults to true.")
     @Nullable
     public abstract Boolean getBatching();
+
+    @SchemaFieldDescription("This option specifies whether and where to output unwritable rows.")
+    @Nullable
+    public abstract ErrorHandling getErrorHandling();
   }
 
   @Override