public class AvroReaderSplittableDoFn
    extends DoFn<KV<String, ReadableFile>, KV<String, GenericRecord>> {

  public static final Logger LOG = LoggerFactory.getLogger(AvroReaderSplittableDoFn.class);
  private final Counter numberOfAvroRecordsIngested =
      Metrics.counter(AvroReaderSplittableDoFn.class, "numberOfAvroRecordsIngested");
  private final Integer splitSize;
  private final Integer keyRange;

  public AvroReaderSplittableDoFn(Integer keyRange, Integer splitSize) {
    this.keyRange = keyRange;
    this.splitSize = splitSize;
  }

  @ProcessElement
  public void processElement(ProcessContext c, RestrictionTracker<OffsetRange, Long> tracker)
      throws IOException {
    LOG.info(
        "Processing split from {} to {}",
        tracker.currentRestriction().getFrom(),
        tracker.currentRestriction().getTo());
    String fileName = c.element().getKey();

    AvroUtil.AvroSeekableByteChannel channel = AvroUtil.getChannel(c.element().getValue());
    DataFileReader<GenericRecord> fileReader =
        new DataFileReader<>(channel, new GenericDatumReader<>());

    long start = tracker.currentRestriction().getFrom();
    long end = tracker.currentRestriction().getTo();

    // Move the first sync point after the
    fileReader.sync(Math.max(start - SYNC_SIZE, 0));

    // Claim the whole split's range
    if (tracker.tryClaim(end - 1)) {
      // Loop through all records in the split. More precisely from the first
      // sync point after the range's start position until the first sync point
      // after the range's end position.
      while (fileReader.hasNext() && !fileReader.pastSync(end - SYNC_SIZE)) {
        // Read the next Avro record in line
        GenericRecord record = fileReader.next();

        // Output the Avro record
        String outputKey = String.format("%s~%d", fileName, new Random().nextInt(keyRange));
        c.outputWithTimestamp(KV.of(outputKey, record), Instant.now());
        numberOfAvroRecordsIngested.inc();
      }
    }

    fileReader.close();
    channel.close();
  }

  @GetInitialRestriction
  public OffsetRange getInitialRestriction(@Element KV<String, ReadableFile> element)
      throws IOException {
    long totalBytes = element.getValue().getMetadata().sizeBytes();
    LOG.info("Initial Restriction range from {} to {}", 0, totalBytes);
    return new OffsetRange(0, totalBytes);
  }

  @SplitRestriction
  public void splitRestriction(
      @Element KV<String, ReadableFile> file,
      @Restriction OffsetRange range,
      OutputReceiver<OffsetRange> out) {
    List<OffsetRange> splits = range.split(splitSize, splitSize);
    LOG.info("Number of splits: {}", splits.size());
    for (final OffsetRange p : splits) {
      out.output(p);
    }
  }

  @NewTracker
  public OffsetRangeTracker newTracker(@Restriction OffsetRange range) {
    return new OffsetRangeTracker(new OffsetRange(range.getFrom(), range.getTo()));
  }
} public class ConvertAvroRecordToDlpRowDoFn
    extends DoFn<KV<String, GenericRecord>, KV<String, Table.Row>> {

  @ProcessElement
  public void processElement(ProcessContext c) {
    String key = c.element().getKey();
    GenericRecord record = c.element().getValue();
    Table.Row.Builder rowBuilder = Table.Row.newBuilder();
    AvroUtil.getFlattenedValues(
        record,
        (Object value) -> {
          if (value == null) {
            rowBuilder.addValues(Value.newBuilder().setStringValue("").build());
          } else {
            rowBuilder.addValues(Value.newBuilder().setStringValue(value.toString()).build());
          }
        });
    c.output(KV.of(key, rowBuilder.build()));
  }
} public class GenericRecordCoder extends AtomicCoder<GenericRecord> {

  // Cache that holds AvroCoders for the encountered Avro records
  private static final ConcurrentHashMap<String, AvroCoder<GenericRecord>> codersCache =
      new ConcurrentHashMap<>();

  public static GenericRecordCoder of() {
    return new GenericRecordCoder();
  }

  @Override
  public void encode(GenericRecord record, OutputStream outputStream) throws IOException {
    // Add base64-encoded schema to the output stream
    String schema = record.getSchema().toString();
    String schema64 = Base64.getEncoder().encodeToString(schema.getBytes());
    StringUtf8Coder.of().encode(schema64, outputStream);

    // Get AvroCoder associated with the schema
    AvroCoder<GenericRecord> coder =
        codersCache.computeIfAbsent(schema64, key -> AvroCoder.of(record.getSchema()));

    // Add encoded record to the output stream
    coder.encode(record, outputStream);
  }

  @Override
  public GenericRecord decode(InputStream inputStream) throws IOException {
    // Fetch the base64-encoded schema
    String schema64 = StringUtf8Coder.of().decode(inputStream);

    // Get AvroCoder associated with the schema
    AvroCoder<GenericRecord> coder =
        codersCache.computeIfAbsent(
            schema64,
            key ->
                AvroCoder.of(
                    new Schema.Parser()
                        .parse(Arrays.toString(Base64.getDecoder().decode(schema64)))));

    // Decode the Avro record
    return coder.decode(inputStream);
  }
} public class ConvertCSVRecordToDLPRow extends DoFn<KV<String, String>, KV<String, Table.Row>> {

  public static final Logger LOG = LoggerFactory.getLogger(ConvertCSVRecordToDLPRow.class);

  private final Character columnDelimiter;
  private PCollectionView<Map<String, List<String>>> header;

  public ConvertCSVRecordToDLPRow(
      Character columnDelimiter, PCollectionView<Map<String, List<String>>> header) {
    this.columnDelimiter = columnDelimiter;
    this.header = header;
  }

  @ProcessElement
  public void processElement(ProcessContext context) throws IOException {
    Table.Row.Builder rowBuilder = Table.Row.newBuilder();
    String input = Objects.requireNonNull(context.element().getValue());
    String fileName = context.element().getKey();
    Map<String, List<String>> headers = context.sideInput(header);
    List<String> csvHeader = headers.get(fileName);
    if (csvHeader == null) {
      throw new RuntimeException(
          "Unable to find header row for fileName: "
              + fileName
              + ". The side input only contains header for "
              + headers.keySet());
    }

    if (columnDelimiter != null) {

      List<String> values = Util.parseLine(input, columnDelimiter, '"');
      if (values.size() == csvHeader.size()) {
        values.forEach(
            value -> rowBuilder.addValues(Value.newBuilder().setStringValue(value).build()));
        context.output(KV.of(fileName, rowBuilder.build()));

      } else {
        LOG.warn(
            "Rows in {} must have the same number of items {} as there are headers {}",
            fileName,
            values.size(),
            csvHeader.size());
      }
    } else {
      rowBuilder.addValues(Value.newBuilder().setStringValue(input).build());
      context.output(KV.of(fileName, rowBuilder.build()));
    }
  }
} @AutoValue
@SuppressWarnings("serial")
public abstract class BigQueryDynamicWriteTransform
    extends PTransform<PCollection<KV<String, TableRow>>, WriteResult> {

  public static final Logger LOG = LoggerFactory.getLogger(BigQueryDynamicWriteTransform.class);

  public abstract String projectId();

  public abstract String datasetId();

  public static Builder newBuilder() {
    return new AutoValue_BigQueryDynamicWriteTransform.Builder();
  }

  @AutoValue.Builder
  public abstract static class Builder {

    public abstract Builder setDatasetId(String projectId);

    public abstract Builder setProjectId(String datasetId);

    public abstract BigQueryDynamicWriteTransform build();
  }

  @Override
  public WriteResult expand(PCollection<KV<String, TableRow>> input) {

    Write<KV<String, TableRow>> transform = BigQueryIO.<KV<String, TableRow>>write()
        .to(new BQDestination(datasetId(), projectId()))
        .withFormatFunction(KV::getValue)
        .withWriteDisposition(WriteDisposition.WRITE_APPEND)
        .withoutValidation()
        .ignoreInsertIds()
        .withMethod(Write.Method.STREAMING_INSERTS)
        .withCreateDisposition(CreateDisposition.CREATE_IF_NEEDED);

    if (input.getPipeline().getOptions().as(StreamingOptions.class).isStreaming()) {
      transform = transform.withAutoSharding();
    }

    return input.apply("BQ Write", transform);
  }

  public class BQDestination
      extends DynamicDestinations<KV<String, TableRow>, KV<String, TableRow>> {

    private String datasetName;
    private String projectId;

    public BQDestination(String datasetName, String projectId) {
      this.datasetName = datasetName;
      this.projectId = projectId;
    }

    @Override
    public TableDestination getTable(KV<String, TableRow> destination) {
      TableDestination dest =
          new TableDestination(destination.getKey(), "DLP Transformation Storage Table");
      LOG.debug("Table Destination {}", dest.getTableSpec());
      return dest;
    }

    @Override
    public KV<String, TableRow> getDestination(ValueInSingleWindow<KV<String, TableRow>> element) {
      String key = element.getValue().getKey();
      String tableName = String.format("%s:%s.%s", projectId, datasetName, key);
      LOG.debug("Table Name {}", tableName);
      return KV.of(tableName, element.getValue().getValue());
    }

    @Override
    public TableSchema getSchema(KV<String, TableRow> destination) {
      String tableName = destination.getKey().split("\\.")[1];
      LOG.info("Table Name {}", tableName);
      switch (tableName) {
        case "dlp_inspection_result":
          return BigQueryUtils.toTableSchema(Util.dlpInspectionSchema);
        case "error_table":
          return BigQueryUtils.toTableSchema(Util.errorSchema);
        default:
          TableRow bqRow = destination.getValue();
          TableSchema schema = new TableSchema();
          List<TableFieldSchema> fields = new ArrayList<>();
          List<TableCell> cells = bqRow.getF();
          for (int i = 0; i < cells.size(); i++) {
            Map<String, Object> object = cells.get(i);
            String header = object.keySet().iterator().next();
            /* currently all BQ data types are set to String */
            fields.add(
                new TableFieldSchema().setName(Util.checkHeaderName(header)).setType("STRING"));
          }
          schema.setFields(fields);
          return schema;
      }
    }
  }
} @AutoValue
@SuppressWarnings({"deprecation", "serial"})
public abstract class BigQueryReadTransform
    extends PTransform<PBegin, PCollection<KV<String, TableRow>>> {
  public static final Logger LOG = LoggerFactory.getLogger(BigQueryReadTransform.class);

  public abstract String tableRef();

  public abstract Method readMethod();

  public abstract Integer keyRange();

  public abstract String query();

  @AutoValue.Builder
  public abstract static class Builder {

    public abstract Builder setTableRef(String tableRef);

    public abstract Builder setReadMethod(Method readMethod);

    public abstract Builder setKeyRange(Integer keyRange);

    public abstract Builder setQuery(String query);

    public abstract BigQueryReadTransform build();
  }

  public static Builder newBuilder() {
    return new AutoValue_BigQueryReadTransform.Builder();
  }

  @Override
  public PCollection<KV<String, TableRow>> expand(PBegin input) {

    switch (readMethod()) {
      case DEFAULT:
        return input
            .apply(
                "ReadFromBigQuery",
                BigQueryIO.readTableRows()
                    .fromQuery(query())
                    .usingStandardSql()
                    .withMethod(Method.DEFAULT))
            .apply("AddTableNameAsKey", WithKeys.of(tableRef()));
      case EXPORT:
        return input
            .apply(
                "ReadFromBigQuery",
                BigQueryIO.readTableRows()
                    .fromQuery(query())
                    .usingStandardSql()
                    .withMethod(Method.DEFAULT))
            .apply("AddTableNameAsKey", WithKeys.of(tableRef()));

      case DIRECT_READ:
        throw new IllegalArgumentException("Direct read not supported");

      default:
        throw new IllegalArgumentException("Not a valid method");
    }
  }
} @SuppressWarnings("serial")
public class BigQueryTableHeaderDoFn extends DoFn<KV<String, TableRow>, KV<String, List<String>>> {

  public static final Logger LOG = LoggerFactory.getLogger(BigQueryTableHeaderDoFn.class);

  @ProcessElement
  public void processElement(ProcessContext c) {
    List<String> columns = new ArrayList<>();
    c.element().getValue().entrySet().forEach(value -> columns.add(value.getKey()));

    String tableRef = c.element().getKey();
    c.output(KV.of(tableRef, columns));

    LOG.info("Extracted columns for {} :{}", tableRef, columns);
  }
} @SuppressWarnings("serial")
public class CSVFileReaderSplitDoFn extends DoFn<KV<String, ReadableFile>, KV<String, String>> {
  public static final Logger LOG = LoggerFactory.getLogger(CSVFileReaderSplitDoFn.class);
  private final Counter numberOfRowsRead =
      Metrics.counter(CSVFileReaderSplitDoFn.class, "numberOfRowsRead");

  private final String recordDelimiter;
  private final Integer splitSize;

  public CSVFileReaderSplitDoFn(String recordDelimiter, Integer splitSize) {
    this.recordDelimiter = recordDelimiter;
    this.splitSize = splitSize;
  }

  @ProcessElement
  public void processElement(ProcessContext c, RestrictionTracker<OffsetRange, Long> tracker)
      throws IOException {
    String fileName = c.element().getKey();
    try (SeekableByteChannel channel = getReader(c.element().getValue())) {
      FileReader reader =
          new FileReader(
              channel, tracker.currentRestriction().getFrom(), recordDelimiter.getBytes());
      while (tracker.tryClaim(reader.getStartOfNextRecord())) {
        reader.readNextRecord();
        String contents = reader.getCurrent().toStringUtf8();
        String key = fileName;
        numberOfRowsRead.inc();
        c.outputWithTimestamp(KV.of(key, contents), c.timestamp());
      }
    }
  }

  @GetInitialRestriction
  public OffsetRange getInitialRestriction(@Element KV<String, ReadableFile> file)
      throws IOException {
    long totalBytes = file.getValue().getMetadata().sizeBytes();
    LOG.info("Initial Restriction range from 1 to: {}", totalBytes);
    return new OffsetRange(0, totalBytes);
  }

  @SplitRestriction
  public void splitRestriction(
      @Element KV<String, ReadableFile> file,
      @Restriction OffsetRange range,
      OutputReceiver<OffsetRange> out) {
    long totalBytes = file.getValue().getMetadata().sizeBytes();
    List<OffsetRange> splits = range.split(splitSize, splitSize);
    LOG.info("Number of Splits: {} - Total bytes: {}", splits.size(), totalBytes);
    for (final OffsetRange p : splits) {
      out.output(p);
    }
  }

  @NewTracker
  public OffsetRangeTracker newTracker(@Restriction OffsetRange range) {
    return new OffsetRangeTracker(new OffsetRange(range.getFrom(), range.getTo()));
  }

  private static SeekableByteChannel getReader(ReadableFile eventFile) {
    SeekableByteChannel channel = null;
    try {
      channel = eventFile.openSeekable();
    } catch (IOException e) {
      LOG.error("Failed to Open File {}", e.getMessage());
      throw new RuntimeException(e);
    }
    return channel;
  }
} @AutoValue
public abstract class DLPTransform
    extends PTransform<PCollection<KV<String, Table.Row>>, PCollectionTuple> {

  public static final Logger LOG = LoggerFactory.getLogger(DLPTransform.class);

  @Nullable
  public abstract String inspectTemplateName();

  @Nullable
  public abstract String deidTemplateName();

  public abstract Integer batchSize();

  public abstract String projectId();

  public abstract Character columnDelimiter();

  public abstract DLPMethod dlpmethod();

  public abstract String jobName();

  public abstract PCollectionView<Map<String, List<String>>> headers();

  @AutoValue.Builder
  public abstract static class Builder {

    public abstract Builder setInspectTemplateName(String inspectTemplateName);

    public abstract Builder setDeidTemplateName(String inspectTemplateName);

    public abstract Builder setBatchSize(Integer batchSize);

    public abstract Builder setProjectId(String projectId);

    public abstract Builder setHeaders(PCollectionView<Map<String, List<String>>> headers);

    public abstract Builder setColumnDelimiter(Character columnDelimiter);

    public abstract Builder setDlpmethod(DLPMethod method);

    public abstract Builder setJobName(String jobName);

    public abstract DLPTransform build();
  }

  public static Builder newBuilder() {
    return new AutoValue_DLPTransform.Builder();
  }

  @Override
  public PCollectionTuple expand(PCollection<KV<String, Table.Row>> input) {
    switch (dlpmethod()) {
      case INSPECT:
        {
          return input
              .apply(
                  "InspectTransform",
                  DLPInspectText.newBuilder()
                      .setBatchSizeBytes(batchSize())
                      .setColumnDelimiter(columnDelimiter())
                      .setHeaderColumns(headers())
                      .setInspectTemplateName(inspectTemplateName())
                      .setProjectId(projectId())
                      .build())
              .apply(
                  "ConvertInspectResponse",
                  ParDo.of(new ConvertInspectResponse(jobName()))
                      .withOutputTags(
                          Util.inspectOrDeidSuccess, TupleTagList.of(Util.inspectOrDeidFailure)));
        }

      case DEID:
        {
          return input
              .apply(
                  "DeIdTransform",
                  DLPDeidentifyText.newBuilder()
                      .setBatchSizeBytes(batchSize())
                      .setColumnDelimiter(columnDelimiter())
                      .setHeaderColumns(headers())
                      .setInspectTemplateName(inspectTemplateName())
                      .setDeidentifyTemplateName(deidTemplateName())
                      .setProjectId(projectId())
                      .build())
              .apply(
                  "ConvertDeidResponse",
                  ParDo.of(new ConvertDeidResponse())
                      .withOutputTags(
                          Util.inspectOrDeidSuccess, TupleTagList.of(Util.inspectOrDeidFailure)));
        }
      case REID:
        {
          return input
              .apply(
                  "ReIdTransform",
                  DLPReidentifyText.newBuilder()
                      .setBatchSizeBytes(batchSize())
                      .setColumnDelimiter(columnDelimiter())
                      .setHeaderColumns(headers())
                      .setInspectTemplateName(inspectTemplateName())
                      .setReidentifyTemplateName(deidTemplateName())
                      .setProjectId(projectId())
                      .build())
              .apply(
                  "ConvertReidResponse",
                  ParDo.of(new ConvertReidResponse())
                      .withOutputTags(Util.reidSuccess, TupleTagList.of(Util.reidFailure)));
        }
      default:
        {
          throw new IllegalArgumentException("Please validate DLPMethod param!");
        }
    }
  }

  static class ConvertReidResponse
      extends DoFn<KV<String, ReidentifyContentResponse>, KV<String, TableRow>> {

    private final Counter numberOfBytesReidentified =
        Metrics.counter(ConvertDeidResponse.class, "NumberOfBytesReidentified");

    @ProcessElement
    public void processElement(
        @Element KV<String, ReidentifyContentResponse> element, MultiOutputReceiver out) {

      String deidTableName = BigQueryHelpers.parseTableSpec(element.getKey()).getTableId();
      String tableName = String.format("%s_%s", deidTableName, Util.BQ_REID_TABLE_EXT);
      LOG.info("Table Ref {}", tableName);
      Table originalData = element.getValue().getItem().getTable();
      numberOfBytesReidentified.inc(originalData.toByteArray().length);
      List<String> headers =
          originalData.getHeadersList().stream()
              .map(fid -> fid.getName())
              .collect(Collectors.toList());
      List<Table.Row> outputRows = originalData.getRowsList();
      if (outputRows.size() > 0) {
        for (Table.Row outputRow : outputRows) {
          if (outputRow.getValuesCount() != headers.size()) {
            throw new IllegalArgumentException(
                "BigQuery column count must exactly match with data element count");
          }
          out.get(Util.reidSuccess)
              .output(
                  KV.of(
                      tableName,
                      Util.createBqRow(outputRow, headers.toArray(new String[headers.size()]))));
        }
      }
    }
  }

  static class ConvertDeidResponse
      extends DoFn<KV<String, DeidentifyContentResponse>, KV<String, TableRow>> {

    private final Counter numberOfRowDeidentified =
        Metrics.counter(ConvertDeidResponse.class, "numberOfRowDeidentified");

    @ProcessElement
    public void processElement(
        @Element KV<String, DeidentifyContentResponse> element, MultiOutputReceiver out) {

      String fileName = element.getKey();
      Table tokenizedData = element.getValue().getItem().getTable();
      LOG.info("Table Tokenized {}", tokenizedData);
      numberOfRowDeidentified.inc(tokenizedData.getRowsCount());
      List<String> headers =
          tokenizedData.getHeadersList().stream()
              .map(fid -> fid.getName())
              .collect(Collectors.toList());
      List<Table.Row> outputRows = tokenizedData.getRowsList();
      if (outputRows.size() > 0) {
        for (Table.Row outputRow : outputRows) {
          if (outputRow.getValuesCount() != headers.size()) {
            throw new IllegalArgumentException(
                "CSV file's header count must exactly match with data element count");
          }
          out.get(Util.inspectOrDeidSuccess)
              .output(
                  KV.of(
                      fileName,
                      Util.createBqRow(outputRow, headers.toArray(new String[headers.size()]))));
        }
      }
    }
  }

  static class ConvertInspectResponse
      extends DoFn<KV<String, InspectContentResponse>, KV<String, TableRow>> {

    private String jobName;

    public ConvertInspectResponse(String jobName) {
      this.jobName = jobName;
    }

    // Counter to track total number of Inspection Findings fetched from DLP Inspection response
    private final Counter numberOfInspectionFindings =
        Metrics.counter(ConvertInspectResponse.class, "numberOfInspectionFindings");

    // Counter to track total number of times Inspection Findings got truncated in the
    // in the DLP Inspection response
    private final Counter numberOfTimesFindingsTruncated =
        Metrics.counter(ConvertInspectResponse.class, "numberOfTimesFindingsTruncated");

    // Counter to track total number of times Inspection Findings generated in the
    // this should be same number as number of total DLP API calls
    private final Counter numberOfTimesFindingsGenerated =
        Metrics.counter(ConvertInspectResponse.class, "numberOfTimesFindingsGenerated");

    @ProcessElement
    public void processElement(
        @Element KV<String, InspectContentResponse> element, MultiOutputReceiver out) {
      String fileName = element.getKey().split("\\~")[0];
      String timeStamp = Util.getTimeStamp();

      if (element.getValue().getResult().getFindingsTruncated()) {
        numberOfTimesFindingsTruncated.inc();
      }

      numberOfTimesFindingsGenerated.inc();

      element
          .getValue()
          .getResult()
          .getFindingsList()
          .forEach(
              finding -> {
                Row row =
                    Row.withSchema(Util.dlpInspectionSchema)
                        .addValues(
                            jobName,
                            fileName,
                            timeStamp,
                            finding.getQuote(),
                            finding.getInfoType().getName(),
                            finding.getLikelihood().name(),
                            finding.getLocation().getCodepointRange().getStart(),
                            finding.getLocation().getCodepointRange().getEnd(),
                            finding
                                .getLocation()
                                .getContentLocationsList()
                                .get(0)
                                .getRecordLocation()
                                .getFieldId()
                                .getName())
                        .build();
                numberOfInspectionFindings.inc();
                out.get(Util.inspectOrDeidSuccess)
                    .output(KV.of(Util.BQ_DLP_INSPECT_TABLE_NAME, Util.toTableRow(row)));
              });
      element
          .getValue()
          .findInitializationErrors()
          .forEach(
              error -> {
                out.get(Util.inspectOrDeidFailure)
                    .output(
                        KV.of(
                            Util.BQ_ERROR_TABLE_NAME,
                            Util.toTableRow(
                                Row.withSchema(Util.errorSchema)
                                    .addValues(fileName, timeStamp, error.toString(), null)
                                    .build())));
              });
    }
  }
} @SuppressWarnings("serial")
@AutoValue
public abstract class ExtractColumnNamesTransform
    extends PTransform<
        PCollection<KV<String, FileIO.ReadableFile>>, PCollectionView<Map<String, List<String>>>> {

  public abstract FileType fileType();

  @Nullable
  public abstract List<String> headers();

  @AutoValue.Builder
  public abstract static class Builder {

    public abstract ExtractColumnNamesTransform.Builder setFileType(FileType fileType);

    public abstract ExtractColumnNamesTransform.Builder setHeaders(List<String> headers);

    public abstract ExtractColumnNamesTransform build();
  }

  public static ExtractColumnNamesTransform.Builder newBuilder() {
    return new AutoValue_ExtractColumnNamesTransform.Builder();
  }

  @Override
  public PCollectionView<Map<String, List<String>>> expand(
      PCollection<KV<String, FileIO.ReadableFile>> input) {
    PCollection<KV<String, List<String>>> readHeader;
    switch (fileType()) {
      case CSV:
        readHeader = input.apply("ReadHeader", ParDo.of(new CSVColumnNamesDoFn()));
        break;

      case AVRO:
        readHeader = input.apply("ReadHeader", ParDo.of(new AvroColumnNamesDoFn()));
        break;

      case JSON:
        readHeader = input.apply("ReadHeader", ParDo.of(new JsonColumnNameDoFn(headers())));
        break;

      case TXT:
        readHeader = input.apply("ReadHeader", ParDo.of(new TxtColumnNameDoFn(headers())));
        break;

      default:
        throw new IllegalStateException("Unexpected value: " + fileType());
    }
    return readHeader.apply("ViewAsList", View.asMap());
  }
} @AutoValue
public abstract class FilePollingTransform
    extends PTransform<PBegin, PCollection<KV<String, ReadableFile>>> {

  public abstract String filePattern();

  public abstract Duration interval();

  @AutoValue.Builder
  public abstract static class Builder {

    public abstract FilePollingTransform.Builder setFilePattern(String filePattern);

    public abstract FilePollingTransform.Builder setInterval(Duration interval);

    public abstract FilePollingTransform build();
  }

  public static FilePollingTransform.Builder newBuilder() {
    return new AutoValue_FilePollingTransform.Builder();
  }

  @Override
  public PCollection<KV<String, ReadableFile>> expand(PBegin input) {
    return input
        .apply(
            "Poll Input Files",
            FileIO.match()
                .filepattern(filePattern())
                .continuously(interval(), Watch.Growth.never()))
        .apply("Find Pattern Match", FileIO.readMatches().withCompression(Compression.AUTO))
        .apply(ParDo.of(new SanitizeFileNameDoFn()));
  }
} public class MergeBigQueryRowToDlpRow extends DoFn<KV<String, TableRow>, KV<String, Table.Row>> {

  @ProcessElement
  public void processElement(ProcessContext c) {
    TableRow bigqueryRow = c.element().getValue();
    Table.Row.Builder rowBuilder = Table.Row.newBuilder();
    bigqueryRow
        .entrySet()
        .forEach(
            element -> {
              String value = element.getValue().toString();
              rowBuilder.addValues(Value.newBuilder().setStringValue(value).build());
            });
    c.output(KV.of(c.element().getKey(), rowBuilder.build()));
  }
} @SuppressWarnings("serial")
public class PubSubMessageConverts extends DoFn<KV<String, TableRow>, PubsubMessage> {
  private static final Logger LOG = LoggerFactory.getLogger(PubSubMessageConverts.class);

  @ProcessElement
  public void processContext(ProcessContext c) {
    String tableRef = c.element().getKey();
    TableRow row = c.element().getValue();
    String json = Util.gson.toJson(row.getF());
    LOG.info("Json {}", json);
    PubsubMessage message =
        new PubsubMessage(
            json.getBytes(),
            ImmutableMap.<String, String>builder().put("table_name", tableRef).build());
    c.output(message);
  }
} @SuppressWarnings("serial")
public class Util {

  public static final Logger LOG = LoggerFactory.getLogger(Util.class);

  public enum DLPMethod {
    INSPECT,
    DEID,
    REID
  }

  public enum FileType {
    CSV,
    AVRO,
    JSON,
    TXT
  }

  public static final Gson gson = new Gson();
  private static final char DEFAULT_SEPARATOR = ',';
  private static final char DEFAULT_QUOTE = '"';

  private static final DateTimeFormatter BIGQUERY_TIMESTAMP_PRINTER;
  public static final TupleTag<KV<String, String>> agentTranscriptTuple =
      new TupleTag<KV<String, String>>() {};
  public static final TupleTag<KV<String, String>> customerTranscriptTuple =
      new TupleTag<KV<String, String>>() {};
  public static final TupleTag<KV<String, String>> contentTag =
      new TupleTag<KV<String, String>>() {};
  public static final TupleTag<KV<String, ReadableFile>> headerTag =
      new TupleTag<KV<String, ReadableFile>>() {};

  public static final TupleTag<KV<String, TableRow>> inspectOrDeidSuccess =
      new TupleTag<KV<String, TableRow>>() {};
  public static final TupleTag<KV<String, TableRow>> inspectOrDeidFailure =
      new TupleTag<KV<String, TableRow>>() {};

  public static final TupleTag<KV<String, TableRow>> reidSuccess =
      new TupleTag<KV<String, TableRow>>() {};
  public static final TupleTag<KV<String, TableRow>> reidFailure =
      new TupleTag<KV<String, TableRow>>() {};

  public static final TupleTag<KV<String, InspectContentResponse>> inspectApiCallSuccess =
      new TupleTag<KV<String, InspectContentResponse>>() {};
  public static final TupleTag<KV<String, TableRow>> inspectApiCallError =
      new TupleTag<KV<String, TableRow>>() {};

  public static final String BQ_DLP_INSPECT_TABLE_NAME = String.valueOf("dlp_inspection_result");
  public static final String BQ_ERROR_TABLE_NAME = String.valueOf("error_log");
  public static final String BQ_REID_TABLE_EXT = String.valueOf("re_id");

  public static final DateTimeFormatter TIMESTAMP_FORMATTER =
      DateTimeFormat.forPattern("yyyy-MM-dd HH:mm:ss.SSSSSS");

  public static Table.Row convertCsvRowToTableRow(String row) {
    String[] values = row.split(",");
    Table.Row.Builder tableRowBuilder = Table.Row.newBuilder();
    for (String value : values) {
      tableRowBuilder.addValues(Value.newBuilder().setStringValue(value).build());
    }

    return tableRowBuilder.build();
  }

  public static String checkHeaderName(String name) {
    String checkedHeader = name.replaceAll("\\s", "_");
    checkedHeader = checkedHeader.replaceAll("'", "");
    checkedHeader = checkedHeader.replaceAll("/", "");
    checkedHeader = checkedHeader.replaceAll("\\W", "");
    LOG.debug("Name {} checkedHeader {}", name, checkedHeader);
    return checkedHeader;
  }

  public static final Schema dlpInspectionSchema =
      Stream.of(
              Schema.Field.of("job_name", FieldType.STRING).withNullable(true),
              Schema.Field.of("source_file", FieldType.STRING).withNullable(true),
              Schema.Field.of("transaction_time", FieldType.STRING).withNullable(true),
              Schema.Field.of("quote", FieldType.STRING).withNullable(true),
              Schema.Field.of("info_type_name", FieldType.STRING).withNullable(true),
              Schema.Field.of("likelihood", FieldType.STRING).withNullable(true),
              Schema.Field.of("location_start_byte_range", FieldType.INT64).withNullable(true),
              Schema.Field.of("location_end_byte_range", FieldType.INT64).withNullable(true),
              Schema.Field.of("record_loc_field_id", FieldType.STRING).withNullable(true))
          .collect(toSchema());

  public static final Schema errorSchema =
      Stream.of(
              Schema.Field.of("file_name", FieldType.STRING).withNullable(true),
              Schema.Field.of("transaction_timestamp", FieldType.STRING).withNullable(true),
              Schema.Field.of("error_messagee", FieldType.STRING).withNullable(true),
              Schema.Field.of("stack_trace", FieldType.STRING).withNullable(true))
          .collect(toSchema());

  public static String getTimeStamp() {
    return TIMESTAMP_FORMATTER.print(Instant.now().toDateTime(DateTimeZone.UTC));
  }

  public static BufferedReader getReader(ReadableFile csvFile) {
    BufferedReader br = null;
    ReadableByteChannel channel = null;
    /** read the file and create buffered reader */
    try {
      channel = csvFile.openSeekable();

    } catch (IOException e) {
      LOG.error("Failed to Read File {}", e.getMessage());
      throw new RuntimeException(e);
    }

    if (channel != null) {

      br = new BufferedReader(Channels.newReader(channel, Charsets.ISO_8859_1.name()));
    }

    return br;
  }

  public static boolean isDefaultMode(String[] args) {

    for (String arg : args) {
      String[] splitFromEqual = arg.split("=");
      String value = splitFromEqual[1];
      if (value.equals("s3")) {
        return false;
      }
    }
    return true;
  }

  static {
    DateTimeFormatter dateTimePart =
        new DateTimeFormatterBuilder()
            .appendYear(4, 4)
            .appendLiteral('-')
            .appendMonthOfYear(2)
            .appendLiteral('-')
            .appendDayOfMonth(2)
            .appendLiteral(' ')
            .appendHourOfDay(2)
            .appendLiteral(':')
            .appendMinuteOfHour(2)
            .appendLiteral(':')
            .appendSecondOfMinute(2)
            .toFormatter()
            .withZoneUTC();
    BIGQUERY_TIMESTAMP_PRINTER =
        new DateTimeFormatterBuilder()
            .append(dateTimePart)
            .appendLiteral('.')
            .appendFractionOfSecond(3, 3)
            .appendLiteral(" UTC")
            .toFormatter();
  }

  private static Object fromBeamField(FieldType fieldType, Object fieldValue) {
    if (fieldValue == null) {
      if (!fieldType.getNullable()) {
        throw new IllegalArgumentException("Field is not nullable.");
      }
      return null;
    }
    switch (fieldType.getTypeName()) {
      case ARRAY:
      case ITERABLE:
        FieldType elementType = fieldType.getCollectionElementType();
        Iterable<?> items = (Iterable<?>) fieldValue;
        List<Object> convertedItems = Lists.newArrayListWithCapacity(Iterables.size(items));
        for (Object item : items) {
          convertedItems.add(fromBeamField(elementType, item));
        }
        return convertedItems;
      case ROW:
        return toTableRow((Row) fieldValue);
      case DATETIME:
        return ((Instant) fieldValue)
            .toDateTime(DateTimeZone.UTC)
            .toString(BIGQUERY_TIMESTAMP_PRINTER);
      case INT16:
      case INT32:
      case INT64:
      case FLOAT:
      case DOUBLE:
      case STRING:
      case BOOLEAN:
        return fieldValue.toString();
      case DECIMAL:
        return fieldValue.toString();
      case BYTES:
        return BaseEncoding.base64().encode((byte[]) fieldValue);
      default:
        return fieldValue;
    }
  }

  public static TableRow toTableRow(Row row) {
    TableRow output = new TableRow();
    for (int i = 0; i < row.getFieldCount(); i++) {
      Object value = row.getValue(i);
      Field schemaField = row.getSchema().getField(i);
      output = output.set(schemaField.getName(), fromBeamField(schemaField.getType(), value));
    }
    return output;
  }

  public static List<String> getFileHeaders(BufferedReader reader) {
    List<String> headers = new ArrayList<>();
    try {
      CSVRecord csvHeader = CSVFormat.DEFAULT.parse(reader).getRecords().get(0);
      csvHeader.forEach(
          headerValue -> {
            headers.add(headerValue);
          });
    } catch (IOException e) {
      LOG.error("Failed to get csv header values}", e.getMessage());
      throw new RuntimeException(e);
    }
    return headers;
  }

  @SuppressWarnings("null")
  public static List<String> parseLine(String cvsLine, char separators, char customQuote) {

    List<String> result = new ArrayList<>();

    // if empty, return!
    if (cvsLine == null && cvsLine.isEmpty()) {
      return result;
    }

    if (customQuote == ' ') {
      customQuote = DEFAULT_QUOTE;
    }

    if (separators == ' ') {
      separators = DEFAULT_SEPARATOR;
    }

    StringBuffer curVal = new StringBuffer();
    boolean inQuotes = false;
    boolean startCollectChar = false;
    boolean doubleQuotesInColumn = false;

    char[] chars = cvsLine.toCharArray();

    for (char ch : chars) {

      if (inQuotes) {
        startCollectChar = true;
        if (ch == customQuote) {
          inQuotes = false;
          doubleQuotesInColumn = false;
        } else {

          // Fixed : allow "" in custom quote enclosed
          if (ch == '\"') {
            if (!doubleQuotesInColumn) {
              curVal.append(ch);
              doubleQuotesInColumn = true;
            }
          } else {
            curVal.append(ch);
          }
        }
      } else {
        if (ch == customQuote) {

          inQuotes = true;

          // Fixed : allow "" in empty quote enclosed
          if (chars[0] != '"' && customQuote == '\"') {
            curVal.append('"');
          }

          // double quotes in column will hit this!
          if (startCollectChar) {
            curVal.append('"');
          }

        } else if (ch == separators) {

          result.add(curVal.toString());

          curVal = new StringBuffer();
          startCollectChar = false;

        } else if (ch == '\r') {
          // ignore LF characters
          continue;
        } else if (ch == '\n') {
          // the end, break!
          break;
        } else {
          curVal.append(ch);
        }
      }
    }

    result.add(curVal.toString());

    return result;
  }

  public static TableRow createBqRow(Table.Row tokenizedValue, String[] headers) {
    TableRow bqRow = new TableRow();
    AtomicInteger headerIndex = new AtomicInteger(0);
    List<TableCell> cells = new ArrayList<>();
    tokenizedValue
        .getValuesList()
        .forEach(
            value -> {
              String checkedHeaderName =
                  Util.checkHeaderName(headers[headerIndex.getAndIncrement()].toString());
              bqRow.set(checkedHeaderName, value.getStringValue());
              cells.add(new TableCell().set(checkedHeaderName, value.getStringValue()));
            });
    bqRow.setF(cells);
    return bqRow;
  }

  public static String getQueryFromGcs(String gcsPath) {
    GcsPath path = GcsPath.fromUri(URI.create(gcsPath));
    Storage storage = StorageOptions.getDefaultInstance().getService();
    BlobId blobId = BlobId.of(path.getBucket(), path.getObject());
    byte[] content = storage.readAllBytes(blobId);
    String contentString = new String(content, UTF_8);
    LOG.debug("Query: {}", contentString);
    return contentString;
  }
  // index, transcript -> 0 [Customer]: <text> [Agent]:
  public static String parseChatlog(String fileName, String chatLog) {

    String index = chatLog.substring(0, 1);
    String transcript = chatLog.substring(1);
    Random r = new Random();
    // escape header fields
    if (StringUtils.isNumeric(index)) {
      try (Scanner scanner = new Scanner(transcript.trim())) {
        scanner.useDelimiter("\\ ");
        int position = 1;
        StringBuilder agentTranscript = new StringBuilder();
        StringBuilder customerTranscript = new StringBuilder();
        boolean agentScript = true;
        while (scanner.hasNext()) {
          String word = scanner.next().trim();
          if (word.contains("[Customer]:")) {
            customerTranscript.append(StringUtils.remove(word, "[Customer]:"));
            agentScript = false;
            if (agentTranscript.length() > 0) {
              String[] tempValues = {
                String.format("%s%s%s%s%d", fileName, "_", index, "_", r.nextInt(100)),
                "agent",
                agentTranscript.toString(),
                String.valueOf(position),
                "N/A"
              };
              agentTranscript.setLength(0);
              position = position + 1;
              String finalTranscript =
                  Arrays.asList(tempValues).stream().collect(Collectors.joining(","));
              LOG.info(finalTranscript);
              return finalTranscript;
            }

          } else if (word.contains("[Agent]:")) {
            agentTranscript.append(StringUtils.remove(word, "[Agent]:"));
            agentScript = true;
            if (customerTranscript.length() > 0) {
              String[] tempValues = {
                String.format("%s%s%s%s%d", fileName, "_", index, "_", r.nextInt(100)),
                "customer",
                customerTranscript.toString(),
                String.valueOf(position),
                "N/A"
              };
              customerTranscript.setLength(0);
              position = position + 1;
              String finalTranscript =
                  Arrays.asList(tempValues).stream().collect(Collectors.joining(","));
              LOG.info(finalTranscript);
              return finalTranscript;
            }

          } else {
            if (agentScript) {
              agentTranscript.append(StringUtils.SPACE);
              agentTranscript.append(word);
            } else {
              customerTranscript.append(StringUtils.SPACE);
              customerTranscript.append(word);
            }
          }
        }
      }
    }

    return StringUtils.EMPTY;
  }
} @SuppressWarnings("serial")
public class ConvertJsonRecordToDLPRow extends DoFn<KV<String, String>, KV<String, Table.Row>> {
  public static final Logger LOG = LoggerFactory.getLogger(ConvertJsonRecordToDLPRow.class);
  public Gson gson;

  @Setup
  public void setup() {
    gson = new Gson();
  }

  @ProcessElement
  public void processElement(ProcessContext c) {
    String fileName = c.element().getKey();
    JsonObject json = gson.fromJson(c.element().getValue(), JsonObject.class);
    Table.Row.Builder tableRowBuilder = Table.Row.newBuilder();

    json.keySet()
        .forEach(
            key -> {
              String value = json.get(key).getAsString();
              tableRowBuilder.addValues(Value.newBuilder().setStringValue(value));
            });
    Table.Row dlpRow = tableRowBuilder.build();
    LOG.debug("DLP Row {}", dlpRow.toString());
    c.output(KV.of(fileName, dlpRow));
  }
} @SuppressWarnings("serial")
public class JsonReaderSplitDoFn extends DoFn<KV<String, ReadableFile>, KV<String, String>> {
  public static final Logger LOG = LoggerFactory.getLogger(JsonReaderSplitDoFn.class);
  private final Counter numberOfRowsRead =
      Metrics.counter(JsonReaderSplitDoFn.class, "numberOfRowsRead");
  private final Integer keyRange;
  private final String recordDelimiter;
  private final Integer splitSize;

  public JsonReaderSplitDoFn(Integer keyRange, String recordDelimiter, Integer splitSize) {
    this.keyRange = keyRange;
    this.recordDelimiter = recordDelimiter;
    this.splitSize = splitSize;
  }

  @ProcessElement
  public void processElement(ProcessContext c, RestrictionTracker<OffsetRange, Long> tracker)
      throws IOException {
    String fileName = c.element().getKey();
    try (SeekableByteChannel channel = getReader(c.element().getValue())) {
      FileReader reader =
          new FileReader(
              channel, tracker.currentRestriction().getFrom(), recordDelimiter.getBytes());
      while (tracker.tryClaim(reader.getStartOfNextRecord())) {
        reader.readNextRecord();
        String contents = reader.getCurrent().toStringUtf8();
        String key = String.format("%s~%d", fileName, new Random().nextInt(keyRange));
        numberOfRowsRead.inc();
        c.outputWithTimestamp(KV.of(key, contents), Instant.now());
      }
    }
  }

  @GetInitialRestriction
  public OffsetRange getInitialRestriction(@Element KV<String, ReadableFile> file)
      throws IOException {
    long totalBytes = file.getValue().getMetadata().sizeBytes();
    LOG.info("Initial Restriction range from 1 to: {}", totalBytes);
    return new OffsetRange(0, totalBytes);
  }

  @SplitRestriction
  public void splitRestriction(
      @Element KV<String, ReadableFile> file,
      @Restriction OffsetRange range,
      OutputReceiver<OffsetRange> out) {
    long totalBytes = file.getValue().getMetadata().sizeBytes();
    List<OffsetRange> splits = range.split(splitSize, splitSize);
    LOG.info("Number of Splits: {} - Total bytes: {}", splits.size(), totalBytes);
    for (final OffsetRange p : splits) {
      out.output(p);
    }
  }

  @NewTracker
  public OffsetRangeTracker newTracker(@Restriction OffsetRange range) {
    return new OffsetRangeTracker(new OffsetRange(range.getFrom(), range.getTo()));
  }

  private static SeekableByteChannel getReader(ReadableFile eventFile) {
    SeekableByteChannel channel = null;
    try {
      channel = eventFile.openSeekable();
    } catch (IOException e) {
      LOG.error("Failed to Open File {}", e.getMessage());
      throw new RuntimeException(e);
    }
    return channel;
  }
} @SuppressWarnings("serial")
public class ConvertTxtToDLPRow extends DoFn<KV<String, String>, KV<String, Table.Row>> {

  public static final Logger LOG = LoggerFactory.getLogger(ConvertCSVRecordToDLPRow.class);

  private final Character columnDelimiter;
  private PCollectionView<Map<String, List<String>>> header;

  public ConvertTxtToDLPRow(
      Character columnDelimiter, PCollectionView<Map<String, List<String>>> header) {
    this.columnDelimiter = columnDelimiter;
    this.header = header;
  }

  @ProcessElement
  public void processElement(ProcessContext context) throws IOException {
    Table.Row.Builder rowBuilder = Table.Row.newBuilder();
    String input = Objects.requireNonNull(context.element().getValue());
    Map<String, List<String>> headers = context.sideInput(header);
    String fileName = context.element().getKey();
    List<String> csvHeader = headers.get(fileName);
    if (csvHeader == null) {
      throw new RuntimeException(
          "Unable to find header row for fileName: "
              + fileName
              + ". The side input only contains header for "
              + headers.keySet());
    }

    if (columnDelimiter != null) {

      List<String> values = Util.parseLine(input, columnDelimiter, '"');
      if (values.size() == csvHeader.size()) {
        values.forEach(
            value -> rowBuilder.addValues(Value.newBuilder().setStringValue(value).build()));
        context.output(KV.of(fileName, rowBuilder.build()));

      } else {
        LOG.warn(
            "Rows must have the same number of items {} as there are headers {}",
            values.size(),
            csvHeader.size());
      }
    } else {
      rowBuilder.addValues(Value.newBuilder().setStringValue(input).build());
      context.output(KV.of(fileName, rowBuilder.build()));
    }
  }
} @SuppressWarnings("serial")
public class ParseTextLogDoFn extends DoFn<KV<String, String>, KV<String, String>> {

  private static final Logger LOG = LoggerFactory.getLogger(ParseTextLogDoFn.class);

  private Random r;

  @Setup
  public void setup() {
    r = new Random();
  }

  @ProcessElement
  public void processElement(ProcessContext c, MultiOutputReceiver out) {
    String fileName = c.element().getKey();
    String chatLog = c.element().getValue();
    String index = chatLog.substring(0, 1);
    String transcript = chatLog.substring(1);
    // escape header fields
    if (StringUtils.isNumeric(index)) {
      try (Scanner scanner = new Scanner(transcript.trim())) {
        scanner.useDelimiter("\\ ");
        int position = 1;
        StringBuilder agentTranscript = new StringBuilder();
        StringBuilder customerTranscript = new StringBuilder();
        boolean agentScript = true;
        while (scanner.hasNext()) {
          String word = scanner.next().trim();
          if (word.contains("[Customer]:")) {
            customerTranscript.append(StringUtils.remove(word, "[Customer]:"));
            agentScript = false;
            if (agentTranscript.length() > 0) {
              String[] tempValues = {
                String.format("%s%s%s%s%d", fileName, "_", index, "_", r.nextInt(100)),
                "agent",
                agentTranscript.toString(),
                String.valueOf(position),
                "N/A"
              };
              agentTranscript.setLength(0);
              position = position + 1;
              String finalTranscript =
                  Arrays.asList(tempValues).stream().collect(Collectors.joining(","));
              LOG.debug(finalTranscript);
              out.get(Util.agentTranscriptTuple).output(KV.of(fileName, finalTranscript));
            }

          } else if (word.contains("[Agent]:")) {
            agentTranscript.append(StringUtils.remove(word, "[Agent]:"));
            agentScript = true;
            if (customerTranscript.length() > 0) {
              String[] tempValues = {
                String.format("%s%s%s%s%d", fileName, "_", index, "_", r.nextInt(100)),
                "customer",
                customerTranscript.toString(),
                String.valueOf(position),
                "N/A"
              };
              customerTranscript.setLength(0);
              position = position + 1;
              String finalTranscript =
                  Arrays.asList(tempValues).stream().collect(Collectors.joining(","));
              LOG.debug(finalTranscript);
              out.get(Util.customerTranscriptTuple).output(KV.of(fileName, finalTranscript));
            }

          } else {
            if (agentScript) {
              agentTranscript.append(StringUtils.SPACE);
              agentTranscript.append(word);
            } else {
              customerTranscript.append(StringUtils.SPACE);
              customerTranscript.append(word);
            }
          }
        }
      }
    }
  }
} public class TxtReaderSplitDoFn extends DoFn<KV<String, ReadableFile>, KV<String, String>> {
  public static final Logger LOG = LoggerFactory.getLogger(TxtReaderSplitDoFn.class);
  private final Counter numberOfRowsRead =
      Metrics.counter(JsonReaderSplitDoFn.class, "numberOfRowsRead");
  private final Integer keyRange;
  private final String recordDelimiter;
  private final Integer splitSize;

  public TxtReaderSplitDoFn(Integer keyRange, String recordDelimiter, Integer splitSize) {
    this.keyRange = keyRange;
    this.recordDelimiter = recordDelimiter;
    this.splitSize = splitSize;
  }

  @ProcessElement
  public void processElement(ProcessContext c, RestrictionTracker<OffsetRange, Long> tracker)
      throws IOException {
    String fileName = c.element().getKey();
    try (SeekableByteChannel channel = getReader(c.element().getValue())) {
      FileReader reader =
          new FileReader(
              channel, tracker.currentRestriction().getFrom(), recordDelimiter.getBytes());
      while (tracker.tryClaim(reader.getStartOfNextRecord())) {
        reader.readNextRecord();
        String contents = reader.getCurrent().toStringUtf8();
        String key = String.format("%s~%d", fileName, new Random().nextInt(keyRange));
        numberOfRowsRead.inc();
        c.outputWithTimestamp(KV.of(key, contents), Instant.now());
      }
    }
  }

  @GetInitialRestriction
  public OffsetRange getInitialRestriction(@Element KV<String, ReadableFile> file)
      throws IOException {
    long totalBytes = file.getValue().getMetadata().sizeBytes();
    LOG.info("Initial Restriction range from 1 to: {}", totalBytes);
    return new OffsetRange(0, totalBytes);
  }

  @SplitRestriction
  public void splitRestriction(
      @Element KV<String, ReadableFile> file,
      @Restriction OffsetRange range,
      OutputReceiver<OffsetRange> out) {
    long totalBytes = file.getValue().getMetadata().sizeBytes();
    List<OffsetRange> splits = range.split(splitSize, splitSize);
    LOG.info("Number of Splits: {} - Total bytes: {}", splits.size(), totalBytes);
    for (final OffsetRange p : splits) {
      out.output(p);
    }
  }

  @NewTracker
  public OffsetRangeTracker newTracker(@Restriction OffsetRange range) {
    return new OffsetRangeTracker(new OffsetRange(range.getFrom(), range.getTo()));
  }

  private static SeekableByteChannel getReader(ReadableFile eventFile) {
    SeekableByteChannel channel = null;
    try {
      channel = eventFile.openSeekable();
    } catch (IOException e) {
      LOG.error("Failed to Open File {}", e.getMessage());
      throw new RuntimeException(e);
    }
    return channel;
  }
}
import com.google.api.services.bigquery.model.TableRow;
import com.google.privacy.dlp.v2.Table;
import com.google.swarm.tokenization.avro.AvroReaderSplittableDoFn;
import com.google.swarm.tokenization.avro.ConvertAvroRecordToDlpRowDoFn;
import com.google.swarm.tokenization.avro.GenericRecordCoder;
import com.google.swarm.tokenization.beam.ConvertCSVRecordToDLPRow;
import com.google.swarm.tokenization.common.BigQueryDynamicWriteTransform;
import com.google.swarm.tokenization.common.BigQueryReadTransform;
import com.google.swarm.tokenization.common.BigQueryTableHeaderDoFn;
import com.google.swarm.tokenization.common.CSVFileReaderSplitDoFn;
import com.google.swarm.tokenization.common.DLPTransform;
import com.google.swarm.tokenization.common.ExtractColumnNamesTransform;
import com.google.swarm.tokenization.common.FilePollingTransform;
import com.google.swarm.tokenization.common.MergeBigQueryRowToDlpRow;
import com.google.swarm.tokenization.common.PubSubMessageConverts;
import com.google.swarm.tokenization.common.Util;
import com.google.swarm.tokenization.json.ConvertJsonRecordToDLPRow;
import com.google.swarm.tokenization.json.JsonReaderSplitDoFn;
import com.google.swarm.tokenization.txt.ConvertTxtToDLPRow;
import com.google.swarm.tokenization.txt.ParseTextLogDoFn;
import com.google.swarm.tokenization.txt.TxtReaderSplitDoFn;
import java.util.List;
import java.util.Map;
import org.apache.beam.sdk.Pipeline;
import org.apache.beam.sdk.PipelineResult;
import org.apache.beam.sdk.coders.KvCoder;
import org.apache.beam.sdk.coders.StringUtf8Coder;
import org.apache.beam.sdk.io.FileIO.ReadableFile;
import org.apache.beam.sdk.io.gcp.pubsub.PubsubIO;
import org.apache.beam.sdk.options.PipelineOptionsFactory;
import org.apache.beam.sdk.transforms.Flatten;
import org.apache.beam.sdk.transforms.ParDo;
import org.apache.beam.sdk.transforms.Sample;
import org.apache.beam.sdk.transforms.View;
import org.apache.beam.sdk.transforms.windowing.FixedWindows;
import org.apache.beam.sdk.transforms.windowing.Window;
import org.apache.beam.sdk.values.KV;
import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.PCollectionList;
import org.apache.beam.sdk.values.PCollectionTuple;
import org.apache.beam.sdk.values.PCollectionView;
import org.apache.beam.sdk.values.TupleTagList;
import org.joda.time.Duration;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
/*
 * Copyright 2020 Google LLC
 *
 * Licensed 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 com.google.swarm.tokenization;


public class DLPTextToBigQueryStreamingV2 {

  public static final Logger LOG = LoggerFactory.getLogger(DLPTextToBigQueryStreamingV2.class);
  private static final Duration DEFAULT_POLL_INTERVAL = Duration.standardSeconds(3);
  private static final Duration WINDOW_INTERVAL = Duration.standardSeconds(3);
  /** PubSub configuration for default batch size in number of messages */
  public static final Integer PUB_SUB_BATCH_SIZE = 1000;
  /** PubSub configuration for default batch size in bytes */
  public static final Integer PUB_SUB_BATCH_SIZE_BYTES = 10000;

  public static void main(String[] args) {

    DLPTextToBigQueryStreamingV2PipelineOptions options =
        PipelineOptionsFactory.fromArgs(args).as(DLPTextToBigQueryStreamingV2PipelineOptions.class);

    run(options);
  }

  public static PipelineResult run(DLPTextToBigQueryStreamingV2PipelineOptions options) {
    Pipeline p = Pipeline.create(options);

    switch (options.getDLPMethod()) {
      case INSPECT:
      case DEID:
        runInspectAndDeidPipeline(p, options);
        break;

      case REID:
        runReidPipeline(p, options);
        break;

      default:
        throw new IllegalArgumentException("Please validate DLPMethod param!");
    }

    return p.run();
  }

  private static void runInspectAndDeidPipeline(
      Pipeline p, DLPTextToBigQueryStreamingV2PipelineOptions options) {
    PCollection<KV<String, ReadableFile>> inputFiles =
        p.apply(
                FilePollingTransform.newBuilder()
                    .setFilePattern(options.getFilePattern())
                    .setInterval(DEFAULT_POLL_INTERVAL)
                    .build())
            .apply("Fixed Window", Window.into(FixedWindows.of(WINDOW_INTERVAL)));
    final PCollectionView<Map<String, List<String>>> headers =
        inputFiles.apply(
            "Extract Column Names",
            ExtractColumnNamesTransform.newBuilder()
                .setFileType(options.getFileType())
                .setHeaders(options.getHeaders())
                .build());

    PCollection<KV<String, Table.Row>> records;

    switch (options.getFileType()) {
      case AVRO:
        records =
            inputFiles
                .apply(
                    ParDo.of(
                        new AvroReaderSplittableDoFn(
                            options.getKeyRange(), options.getSplitSize())))
                .setCoder(KvCoder.of(StringUtf8Coder.of(), GenericRecordCoder.of()))
                .apply(ParDo.of(new ConvertAvroRecordToDlpRowDoFn()));
        break;
      case CSV:
        records =
            inputFiles
                .apply(
                    "SplitCSVFile",
                    ParDo.of(
                        new CSVFileReaderSplitDoFn(
                            options.getRecordDelimiter(), options.getSplitSize())))
                .apply(
                    "ConvertToDLPRow",
                    ParDo.of(new ConvertCSVRecordToDLPRow(options.getColumnDelimiter(), headers))
                        .withSideInputs(headers));
        break;
      case JSON:
        records =
            inputFiles
                .apply(
                    "SplitJSONFile",
                    ParDo.of(
                        new JsonReaderSplitDoFn(
                            options.getKeyRange(),
                            options.getRecordDelimiter(),
                            options.getSplitSize())))
                .apply("ConvertToDLPRow", ParDo.of(new ConvertJsonRecordToDLPRow()));
        break;
      case TXT:
        PCollectionTuple recordTuple =
            inputFiles
                .apply(
                    "SplitTextFile",
                    ParDo.of(
                        new TxtReaderSplitDoFn(
                            options.getKeyRange(),
                            options.getRecordDelimiter(),
                            options.getSplitSize())))
                .apply(
                    "ParseTextFile",
                    ParDo.of(new ParseTextLogDoFn())
                        .withOutputTags(
                            Util.agentTranscriptTuple,