Add error handling for streaming inserts

VeronicaWasson · VeronicaWasson · commit 293742fea447 · 2023-07-28T21:37:50.000Z
diff --git a/dataflow/snippets/src/main/java/com/example/dataflow/BigQueryStreamExactlyOnce.java b/dataflow/snippets/src/main/java/com/example/dataflow/BigQueryStreamExactlyOnce.java
@@ -30,6 +30,7 @@
 import org.apache.beam.sdk.transforms.MapElements;
 import org.apache.beam.sdk.values.TimestampedValue;
 import org.apache.beam.sdk.values.TypeDescriptor;
+import org.apache.beam.sdk.values.TypeDescriptors;
 import org.joda.time.Duration;
 import org.joda.time.Instant;
 
@@ -45,6 +46,8 @@ private static TestStream<String> createEventSource() {
             TimestampedValue.of("Bob,30",
                 startTime.plus(Duration.standardSeconds(1))),
             TimestampedValue.of("Charles,40",
+                startTime.plus(Duration.standardSeconds(2))),
+            TimestampedValue.of("Dylan,Invalid value",
                 startTime.plus(Duration.standardSeconds(2))))
         .advanceWatermarkToInfinity();
   }
@@ -80,10 +83,19 @@ public static PipelineResult main(String[] args) {
             .withCreateDisposition(CreateDisposition.CREATE_NEVER)
             .withWriteDisposition(WriteDisposition.WRITE_APPEND)
             .withMethod(Write.Method.STORAGE_WRITE_API)
-            // For exactly-once processing, set the number of Write API streams and the triggering
-            // frequency.
-            .withNumStorageWriteApiStreams(1)
-            .withTriggeringFrequency(Duration.standardSeconds(5)));
+            // For exactly-once processing, set the triggering frequency.
+            .withTriggeringFrequency(Duration.standardSeconds(5)))
+        // Get the collection of write errors.
+        .getFailedStorageApiInserts()
+        .apply(MapElements.into(TypeDescriptors.strings())
+            // Process each error. In production systems, it's useful to write the errors to
+            // another destination, such as a dead-letter table or queue.
+            .via(
+                x -> {
+                  System.out.println("Failed insert: " + x.getErrorMessage());
+                  System.out.println("Row: " + x.getRow());
+                  return "";
+                }));
     return pipeline.run();
   }
 }
diff --git a/dataflow/snippets/src/test/java/com/example/dataflow/BigQueryWriteIT.java b/dataflow/snippets/src/test/java/com/example/dataflow/BigQueryWriteIT.java
@@ -17,6 +17,7 @@
 package com.example.dataflow;
 
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
 
 import com.google.cloud.bigquery.BigQuery;
 import com.google.cloud.bigquery.BigQuery.DatasetDeleteOption;
@@ -136,5 +137,8 @@ public void streamExactlyOnce() throws Exception {
         QueryJobConfiguration.newBuilder(query).setDefaultDataset(datasetName).build();
     TableResult result = bigquery.query(queryConfig);
     assertEquals(3, result.getTotalRows());
+    // Verify that the bad data was written to the error collection.
+    String got = bout.toString();
+    assertTrue(got.contains("Failed insert: "));
   }
 }

Original file line number	Diff line number	Diff line change
`@@ -17,6 +17,7 @@`
`17`	`17`	`package com.example.dataflow;`
`18`	`18`
`19`	`19`	`import static org.junit.Assert.assertEquals;`
	`20`	`+import static org.junit.Assert.assertTrue;`
`20`	`21`
`21`	`22`	`import com.google.cloud.bigquery.BigQuery;`
`22`	`23`	`import com.google.cloud.bigquery.BigQuery.DatasetDeleteOption;`
`@@ -136,5 +137,8 @@ public void streamExactlyOnce() throws Exception {`
`136`	`137`	`QueryJobConfiguration.newBuilder(query).setDefaultDataset(datasetName).build();`
`137`	`138`	`TableResult result = bigquery.query(queryConfig);`
`138`	`139`	`assertEquals(3, result.getTotalRows());`
	`140`	`+ // Verify that the bad data was written to the error collection.`
	`141`	`+ String got = bout.toString();`
	`142`	`+ assertTrue(got.contains("Failed insert: "));`
`139`	`143`	`}`
`140`	`144`	`}`