content: /*
   * Licensed to the Apache Software Foundation (ASF) under one
   * or more contributor license agreements.  See the NOTICE file
   * distributed with this work for additional information
   * regarding copyright ownership.  The ASF licenses this file
   * to you 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 org.apache.beam.sdk.util;
  
  import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull;
  
  import java.io.Serializable;
  import java.util.HashSet;
  import java.util.List;
  import java.util.Set;
  import java.util.regex.Matcher;
  import java.util.regex.Pattern;
  import java.util.stream.Collectors;
  import org.apache.beam.sdk.schemas.Schema;
  import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
  import org.apache.beam.sdk.transforms.windowing.PaneInfo;
  import org.apache.beam.sdk.values.Row;
  import org.apache.beam.sdk.values.ValueInSingleWindow;
  import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.MoreObjects;
  import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions;
  import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Sets;
  import org.checkerframework.checker.nullness.qual.Nullable;
  import org.joda.time.Instant;
  
  /**
   * A utility that interpolates values in a pre-determined {@link String} using an input Beam {@link
   * Row}.
   *
   * <p>The {@link RowStringInterpolator} looks for field names specified inside {curly braces}. For
   * example, if the interpolator is configured with the String {@code "unified {foo} and streaming"},
   * it will look for a field name {@code "foo"} in the input {@link Row} and substitute in that
   * value. If a {@link RowStringInterpolator} is configured with a template String that contains no
   * placeholders (i.e. no curly braces), it will simply return that String, untouched.
   *
   * <p>Nested fields can be specified using dot-notation (e.g. {@code "top.middle.nested"}).
   *
   * <p>Configure a {@link RowStringInterpolator} like so:
   *
   * <pre>{@code
   * String template = "unified {foo} and {bar.baz}!";
   * Row inputRow = {foo: "batch", bar: {baz: "streaming"}, ...};
   *
   * RowStringInterpolator interpolator = new RowStringInterpolator(template, beamSchema);
   * String output = interpolator.interpolate(inputRow, window, paneInfo, timestamp);
   * // output --> "unified batch and streaming!"
   * }</pre>
   *
   * <p>Additionally, {@link #interpolate(Row, BoundedWindow, PaneInfo, Instant)} can be used in
   * streaming scenarios to substitute windowing metadata into the template String. To make use of
   * this, use the relevant placeholder:
   *
   * <ul>
   *   <li>$WINDOW: the window's string representation
   *   <li>$PANE_INDEX: the pane's index
   *   <li>$YYYY: the element timestamp's year
   *   <li>$MM: the element timestamp's month
   *   <li>$DD: the element timestamp's day
   * </ul>
   *
   * <p>For example, your String template can look like:
   *
   * <pre>{@code "unified {foo} and {bar} since {$YYYY}-{$MM}!"}</pre>
   */
  public class RowStringInterpolator implements Serializable {
    private final String template;
    private final Set<String> fieldsToReplace;
    // Represents the string representation of the element's window
    public static final String WINDOW = "$WINDOW";
    public static final String PANE_INDEX = "$PANE_INDEX";
    // Represents the element's pane index
    public static final String YYYY = "$YYYY";
    public static final String MM = "$MM";
    public static final String DD = "$DD";
    private static final Set<String> WINDOWING_METADATA =
        Sets.newHashSet(WINDOW, PANE_INDEX, YYYY, MM, DD);
    private static final Pattern TEMPLATE_PATTERN = Pattern.compile("\\{(.+?)}");
  
    /**
     * @param template a String template, potentially with placeholders in the form of curly braces,
     *     e.g. {@code "my {foo} template"}. During interpolation, these placeholders are replaced
     *     with values in the Beam Row. For more details and examples, refer to the top-level
     *     documentation.
     * @param rowSchema {@link Row}s used for interpolation are expected to be compatible with this
     *     {@link Schema}.
     */
    public RowStringInterpolator(String template, Schema rowSchema) {
      this.template = template;
  
      Matcher m = TEMPLATE_PATTERN.matcher(template);
      fieldsToReplace = new HashSet<>();
      while (m.find()) {
        fieldsToReplace.add(checkStateNotNull(m.group(1)));
      }
  
      List<String> rowFields =
          fieldsToReplace.stream()
              .filter(f -> !WINDOWING_METADATA.contains(f))
              .collect(Collectors.toList());
  
      RowFilter.validateSchemaContainsFields(rowSchema, rowFields, "string interpolation");
    }
  
    /**
     * Performs string interpolation on the template using values from the input {@link Row} and its
     * windowing metadata.
     */
    public String interpolate(ValueInSingleWindow<Row> element) {
      String interpolated = this.template;
      for (String field : fieldsToReplace) {
        Object val;
        Instant timestamp = element.getTimestamp();
        switch (field) {
          case WINDOW:
            val = element.getWindow().toString();
            break;
          case PANE_INDEX:
            val = element.getPaneInfo().getIndex();
            break;
          case YYYY:
            val = timestamp.getChronology().year().get(timestamp.getMillis());
            break;
          case MM:
            val = timestamp.getChronology().monthOfYear().get(timestamp.getMillis());
            break;
          case DD:
            val = timestamp.getChronology().dayOfMonth().get(timestamp.getMillis());
            break;
          default:
            val = MoreObjects.firstNonNull(getValue(element.getValue(), field), "");
            break;
        }
  
        interpolated = interpolated.replace("{" + field + "}", String.valueOf(val));
      }
      return interpolated;
    }
  
    private @Nullable Object getValue(@Nullable Row row, String fieldPath) {
      if (row == null) {
        return null;
      }
      int dotIndex = fieldPath.indexOf('.');
      String field = dotIndex == -1 ? fieldPath : fieldPath.substring(0, dotIndex);
      Preconditions.checkArgument(
          row.getSchema().hasField(field), "Invalid row does not contain field '%s'.", field);
  
      if (dotIndex == -1) {
        return row.getValue(field);
      }
      return getValue(row.getRow(field), fieldPath.substring(dotIndex + 1));
    }
  }
  