error handling

Author: santhh

src/main/java/com/google/swarm/tokenization/DLPS3ScannerPipeline.java

@@ -27,7 +27,6 @@
 import org.apache.beam.sdk.io.gcp.bigquery.BigQueryIO;
 import org.apache.beam.sdk.io.gcp.pubsub.PubsubIO;
 import org.apache.beam.sdk.options.PipelineOptionsFactory;
-import org.apache.beam.sdk.values.KV;
 import org.apache.beam.sdk.values.PCollection;
 import org.apache.beam.sdk.values.PCollectionTuple;
 import org.apache.beam.sdk.values.Row;
@@ -46,18 +45,20 @@ public static void main(String[] args) {
   public static PipelineResult run(S3ReaderOptions options) {
     Pipeline p = Pipeline.create(options);
 
-    PCollection<KV<String, String>> nonInspectedContents =
+    PCollectionTuple nonInspectedContents =
         p.apply(
             "File Read Transforrm",
             FileReaderTransform.newBuilder().setSubscriber(options.getSubscriber()).build());
 
     PCollectionTuple inspectedData =
-        nonInspectedContents.apply(
-            "DLPScanner",
-            DLPTransform.newBuilder()
-                .setInspectTemplateName(options.getInspectTemplateName())
-                .setProjectId(options.getProject())
-                .build());
+        nonInspectedContents
+            .get(Util.readRowSuccess)
+            .apply(
+                "DLPScanner",
+                DLPTransform.newBuilder()
+                    .setInspectTemplateName(options.getInspectTemplateName())
+                    .setProjectId(options.getProject())
+                    .build());
 
     PCollection<Row> inspectedContents =
         inspectedData.get(Util.inspectData).setRowSchema(Util.bqDataSchema);

src/main/java/com/google/swarm/tokenization/common/FileReaderSplitDoFn.java

@@ -51,18 +51,18 @@ public void processElement(ProcessContext c, RestrictionTracker<OffsetRange, Lon
         readBuffer.clear();
         LOG.debug(
             "Current Restriction {}, Content Size{}", tracker.currentRestriction(), buffer.size());
-        c.output(KV.of(fileName, buffer.toStringUtf8().trim()));
+        c.output(Util.readRowSuccess, KV.of(fileName, buffer.toStringUtf8().trim()));
       }
     } catch (Exception e) {
-
-      LOG.error(e.getMessage());
+      c.output(Util.readRowFailure, KV.of(fileName, e.getMessage()));
     }
   }
 
   @GetInitialRestriction
   public OffsetRange getInitialRestriction(@Element KV<String, ReadableFile> file)
       throws IOException {
-    long totalBytes = file.getValue().getMetadata().sizeBytes();
+    
+	long totalBytes = file.getValue().getMetadata().sizeBytes();
     long totalSplit = 0;
     if (totalBytes < BATCH_SIZE) {
       totalSplit = 2;
@@ -97,14 +97,10 @@ public OffsetRangeTracker newTracker(@Restriction OffsetRange range) {
     return new OffsetRangeTracker(new OffsetRange(range.getFrom(), range.getTo()));
   }
 
-  private static SeekableByteChannel getReader(ReadableFile eventFile) {
+  private static SeekableByteChannel getReader(ReadableFile eventFile) throws IOException {
     SeekableByteChannel channel = null;
-    try {
-      channel = eventFile.openSeekable();
-    } catch (IOException e) {
-      LOG.error("Failed to Open File {}", e.getMessage());
-      throw new RuntimeException(e);
-    }
+    LOG.info("event File Channel {}",eventFile.getMetadata().resourceId().getFilename());
+    channel = eventFile.openSeekable();
     return channel;
   }
 }

src/main/java/com/google/swarm/tokenization/common/FileReaderTransform.java

@@ -19,20 +19,20 @@
 import com.google.swarm.tokenization.common.CSVFileReaderTransform.Builder;
 import org.apache.beam.sdk.extensions.gcp.util.gcsfs.GcsPath;
 import org.apache.beam.sdk.io.FileIO;
+import org.apache.beam.sdk.io.fs.EmptyMatchTreatment;
 import org.apache.beam.sdk.io.gcp.pubsub.PubsubIO;
 import org.apache.beam.sdk.io.gcp.pubsub.PubsubMessage;
 import org.apache.beam.sdk.transforms.DoFn;
 import org.apache.beam.sdk.transforms.PTransform;
 import org.apache.beam.sdk.transforms.ParDo;
-import org.apache.beam.sdk.values.KV;
 import org.apache.beam.sdk.values.PBegin;
-import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.PCollectionTuple;
+import org.apache.beam.sdk.values.TupleTagList;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 @AutoValue
-public abstract class FileReaderTransform
-    extends PTransform<PBegin, PCollection<KV<String, String>>> {
+public abstract class FileReaderTransform extends PTransform<PBegin, PCollectionTuple> {
 
   public static final Logger LOG = LoggerFactory.getLogger(FileReaderTransform.class);
 
@@ -50,17 +50,20 @@ public static Builder newBuilder() {
   }
 
   @Override
-  public PCollection<KV<String, String>> expand(PBegin input) {
+  public PCollectionTuple expand(PBegin input) {
 
     return input
         .apply(
             "ReadFileMetadata",
             PubsubIO.readMessagesWithAttributes().fromSubscription(subscriber()))
         .apply("ConvertToGCSUri", ParDo.of(new MapPubSubMessage()))
-        .apply("FindFile", FileIO.matchAll())
+        .apply("FindFile", FileIO.matchAll().withEmptyMatchTreatment(EmptyMatchTreatment.ALLOW))
         .apply(FileIO.readMatches())
         .apply("AddFileNameAsKey", ParDo.of(new FileSourceDoFn()))
-        .apply("ReadFile", ParDo.of(new FileReaderSplitDoFn()));
+        .apply(
+            "ReadFile",
+            ParDo.of(new FileReaderSplitDoFn())
+                .withOutputTags(Util.readRowSuccess, TupleTagList.of(Util.readRowFailure)));
   }
 
   public class MapPubSubMessage extends DoFn<PubsubMessage, String> {

src/main/java/com/google/swarm/tokenization/common/Util.java

@@ -39,6 +39,7 @@
 import org.apache.beam.sdk.options.ValueProvider;
 import org.apache.beam.sdk.schemas.Schema;
 import org.apache.beam.sdk.schemas.Schema.FieldType;
+import org.apache.beam.sdk.values.KV;
 import org.apache.beam.sdk.values.Row;
 import org.apache.beam.sdk.values.TupleTag;
 import org.apache.commons.csv.CSVFormat;
@@ -67,6 +68,8 @@ public class Util {
   public static TupleTag<Row> inspectData = new TupleTag<Row>() {};
   public static TupleTag<Row> auditData = new TupleTag<Row>() {};
   public static TupleTag<Row> errorData = new TupleTag<Row>() {};
+  public static TupleTag<KV<String, String>> readRowSuccess = new TupleTag<KV<String, String>>() {};
+  public static TupleTag<KV<String, String>> readRowFailure = new TupleTag<KV<String, String>>() {};
 
   public static String parseBucketName(String value) {
     return value.substring(5, value.length() - 1);