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.schemas.Schema.TypeName.BOOLEAN;
  import static org.apache.beam.sdk.schemas.Schema.TypeName.BYTE;
  import static org.apache.beam.sdk.schemas.Schema.TypeName.DATETIME;
  import static org.apache.beam.sdk.schemas.Schema.TypeName.DECIMAL;
  import static org.apache.beam.sdk.schemas.Schema.TypeName.DOUBLE;
  import static org.apache.beam.sdk.schemas.Schema.TypeName.FLOAT;
  import static org.apache.beam.sdk.schemas.Schema.TypeName.INT16;
  import static org.apache.beam.sdk.schemas.Schema.TypeName.INT32;
  import static org.apache.beam.sdk.schemas.Schema.TypeName.INT64;
  import static org.apache.beam.sdk.schemas.Schema.TypeName.STRING;
  import static org.apache.beam.sdk.util.RowJsonValueExtractors.booleanValueExtractor;
  import static org.apache.beam.sdk.util.RowJsonValueExtractors.byteValueExtractor;
  import static org.apache.beam.sdk.util.RowJsonValueExtractors.dateValueExtractor;
  import static org.apache.beam.sdk.util.RowJsonValueExtractors.datetimeValueExtractor;
  import static org.apache.beam.sdk.util.RowJsonValueExtractors.decimalValueExtractor;
  import static org.apache.beam.sdk.util.RowJsonValueExtractors.doubleValueExtractor;
  import static org.apache.beam.sdk.util.RowJsonValueExtractors.floatValueExtractor;
  import static org.apache.beam.sdk.util.RowJsonValueExtractors.intValueExtractor;
  import static org.apache.beam.sdk.util.RowJsonValueExtractors.localDatetimeValueExtractor;
  import static org.apache.beam.sdk.util.RowJsonValueExtractors.longValueExtractor;
  import static org.apache.beam.sdk.util.RowJsonValueExtractors.shortValueExtractor;
  import static org.apache.beam.sdk.util.RowJsonValueExtractors.stringValueExtractor;
  import static org.apache.beam.sdk.util.RowJsonValueExtractors.timeValueExtractor;
  import static org.apache.beam.sdk.values.Row.toRow;
  import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList.toImmutableList;
  
  import com.fasterxml.jackson.core.JsonGenerator;
  import com.fasterxml.jackson.core.JsonParser;
  import com.fasterxml.jackson.databind.DeserializationContext;
  import com.fasterxml.jackson.databind.JsonNode;
  import com.fasterxml.jackson.databind.SerializerProvider;
  import com.fasterxml.jackson.databind.deser.std.StdDeserializer;
  import com.fasterxml.jackson.databind.node.JsonNodeType;
  import com.fasterxml.jackson.databind.ser.std.StdSerializer;
  import com.google.auto.value.AutoValue;
  import java.io.IOException;
  import java.math.BigDecimal;
  import java.time.LocalDate;
  import java.time.LocalDateTime;
  import java.time.LocalTime;
  import java.util.stream.Stream;
  import java.util.stream.StreamSupport;
  import org.apache.beam.sdk.schemas.Schema;
  import org.apache.beam.sdk.schemas.Schema.Field;
  import org.apache.beam.sdk.schemas.Schema.FieldType;
  import org.apache.beam.sdk.schemas.Schema.TypeName;
  import org.apache.beam.sdk.schemas.logicaltypes.SqlTypes;
  import org.apache.beam.sdk.util.RowJsonValueExtractors.ValueExtractor;
  import org.apache.beam.sdk.values.Row;
  import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList;
  import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap;
  import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableSet;
  import org.checkerframework.checker.nullness.qual.Nullable;
  import org.joda.time.ReadableInstant;
  
  /**
   * Jackson serializer and deserializer for {@link Row Rows}.
   *
   * <p>Supports converting between JSON primitive types and:
   *
   * <ul>
   *   <li>{@link Schema.TypeName#BYTE}
   *   <li>{@link Schema.TypeName#INT16}
   *   <li>{@link Schema.TypeName#INT32}
   *   <li>{@link Schema.TypeName#INT64}
   *   <li>{@link Schema.TypeName#FLOAT}
   *   <li>{@link Schema.TypeName#DOUBLE}
   *   <li>{@link Schema.TypeName#BOOLEAN}
   *   <li>{@link Schema.TypeName#STRING}
   *   <li>{@link Schema.TypeName#DECIMAL}
   *   <li>{@link Schema.TypeName#DATETIME}
   * </ul>
   */
  @SuppressWarnings({
    "nullness" // TODO(https://github.com/apache/beam/issues/20497)
  })
  public class RowJson {
    private static final ImmutableSet<TypeName> SUPPORTED_TYPES =
        ImmutableSet.of(BYTE, INT16, INT32, INT64, FLOAT, DOUBLE, BOOLEAN, STRING, DECIMAL, DATETIME);
    private static final ImmutableSet<String> KNOWN_LOGICAL_TYPE_IDENTIFIERS =
        ImmutableSet.of(
            SqlTypes.DATE.getIdentifier(),
            SqlTypes.TIME.getIdentifier(),
            SqlTypes.DATETIME.getIdentifier());
  
    /**
     * Throws {@link UnsupportedRowJsonException} if {@code schema} contains an unsupported field
     * type.
     */
    public static void verifySchemaSupported(Schema schema) {
      ImmutableList<UnsupportedField> unsupportedFields = findUnsupportedFields(schema);
      if (!unsupportedFields.isEmpty()) {
        throw new UnsupportedRowJsonException(
            String.format(
                "Field type%s %s not supported when converting between JSON and Rows. Supported types are: %s",
                unsupportedFields.size() > 1 ? "s" : "",
                unsupportedFields.toString(),
                SUPPORTED_TYPES.toString()));
      }
    }
  
    private static class UnsupportedField {
      final String descriptor;
      final TypeName typeName;
  
      UnsupportedField(String descriptor, TypeName typeName) {
        this.descriptor = descriptor;
        this.typeName = typeName;
      }
  
      @Override
      public String toString() {
        return this.descriptor + "=" + this.typeName;
      }
    }
  
    private static ImmutableList<UnsupportedField> findUnsupportedFields(Schema schema) {
      return schema.getFields().stream()
          .flatMap((field) -> findUnsupportedFields(field).stream())
          .collect(toImmutableList());
    }
  
    private static ImmutableList<UnsupportedField> findUnsupportedFields(Field field) {
      return findUnsupportedFields(field.getType(), field.getName());
    }
  
    private static ImmutableList<UnsupportedField> findUnsupportedFields(
        FieldType fieldType, String fieldName) {
      TypeName fieldTypeName = fieldType.getTypeName();
  
      if (fieldTypeName.isCompositeType()) {
        return fieldType.getRowSchema().getFields().stream()
            .flatMap(
                (field) ->
                    findUnsupportedFields(field.getType(), fieldName + "." + field.getName())
                        .stream())
            .collect(toImmutableList());
      }
  
      if (fieldTypeName.isCollectionType()) {
        return findUnsupportedFields(fieldType.getCollectionElementType(), fieldName + "[]");
      }
  
      if (fieldTypeName.isLogicalType()) {
        if (KNOWN_LOGICAL_TYPE_IDENTIFIERS.contains(fieldType.getLogicalType().getIdentifier())) {
          return ImmutableList.of();
        } else {
          return findUnsupportedFields(fieldType.getLogicalType().getBaseType(), fieldName);
        }
      }
  
      if (!SUPPORTED_TYPES.contains(fieldTypeName)) {
        return ImmutableList.of(new UnsupportedField(fieldName, fieldTypeName));
      }
  
      return ImmutableList.of();
    }
  
    /** Jackson deserializer for parsing JSON into {@link Row Rows}. */
    public static class RowJsonDeserializer extends StdDeserializer<Row> {
  
      private static final boolean SEQUENTIAL = false;
  
      /**
       * An enumeration type for specifying how {@link RowJsonDeserializer} should expect null values
       * to be represented.
       *
       * <p>For example, when parsing JSON for the Schema {@code (str: REQUIRED STRING, int: NULLABLE
       * INT64)}:
       *
       * <ul>
       *   <li>If configured with {@code REQUIRE_NULL}, {@code {"str": "foo", "int": null}} would be
       *       accepted.
       *   <li>If configured with {@code REQUIRE_MISSING}, {@code {"str": "bar"}} would be accepted,
       *       and would yield a {@link Row} with {@code null} for the {@code int} field.
       *   <li>If configured with {@code ALLOW_MISSING_OR_NULL}, either JSON string would be accepted.
       * </ul>
       */
      public enum NullBehavior {
        /**
         * Specifies that a null value may be represented as either a missing field or a null value in
         * the input JSON.
         */
        ACCEPT_MISSING_OR_NULL,
        /**
         * Specifies that a null value must be represented with a null value in JSON. If the field is
         * missing an {@link UnsupportedRowJsonException} will be thrown.
         */
        REQUIRE_NULL,
        /**
         * Specifies that a null value must be represented with a missing field in JSON. If the field
         * has a null value an {@link UnsupportedRowJsonException} will be thrown.
         */
        REQUIRE_MISSING,
      }
  
      private static final ImmutableMap<TypeName, ValueExtractor<?>> JSON_VALUE_GETTERS =
          ImmutableMap.<TypeName, ValueExtractor<?>>builder()
              .put(BYTE, byteValueExtractor())
              .put(INT16, shortValueExtractor())
              .put(INT32, intValueExtractor())
              .put(INT64, longValueExtractor())
              .put(FLOAT, floatValueExtractor())
              .put(DOUBLE, doubleValueExtractor())
              .put(BOOLEAN, booleanValueExtractor())
              .put(STRING, stringValueExtractor())
              .put(DECIMAL, decimalValueExtractor())
              .put(DATETIME, datetimeValueExtractor())
              .build();
  
      private final Schema schema;
      private NullBehavior nullBehavior = NullBehavior.ACCEPT_MISSING_OR_NULL;
  
      /** Creates a deserializer for a {@link Row} {@link Schema}. */
      public static RowJsonDeserializer forSchema(Schema schema) {
        verifySchemaSupported(schema);
        return new RowJsonDeserializer(schema);
      }
  
      private RowJsonDeserializer(Schema schema) {
        super(Row.class);
        this.schema = schema;
      }
  
      /**
       * Sets the behaviour of the deserializer when retrieving null values in the input JSON. See
       * {@link NullBehavior} for a description of the options. Default value is {@code
       * ACCEPT_MISSING_OR_NULL}.
       */
      public RowJsonDeserializer withNullBehavior(NullBehavior behavior) {
        this.nullBehavior = behavior;
        return this;
      }
  
      @Override
      public Row deserialize(JsonParser jsonParser, DeserializationContext deserializationContext)
          throws IOException {
  
        // Parse and convert the root object to Row as if it's a nested field with name 'root'
        return (Row)
            extractJsonNodeValue(
                FieldValue.of("root", FieldType.row(schema), jsonParser.readValueAsTree()));
      }
  
      private Object extractJsonNodeValue(FieldValue fieldValue) {
        if (fieldValue.type().getNullable()) {
          if (!fieldValue.isJsonValuePresent()) {
            switch (this.nullBehavior) {
              case ACCEPT_MISSING_OR_NULL:
              case REQUIRE_MISSING:
                return null;
              case REQUIRE_NULL:
                throw new UnsupportedRowJsonException(
                    "Field '" + fieldValue.name() + "' is not present in the JSON object.");
            }
          }
  
          if (fieldValue.isJsonNull()) {
            switch (this.nullBehavior) {
              case ACCEPT_MISSING_OR_NULL:
              case REQUIRE_NULL:
                return null;
              case REQUIRE_MISSING:
                throw new UnsupportedRowJsonException(
                    "Field '" + fieldValue.name() + "' has a null value in the JSON object.");
            }
          }
        } else {
          // field is not nullable
          if (!fieldValue.isJsonValuePresent()) {
            throw new UnsupportedRowJsonException(
                "Non-nullable field '" + fieldValue.name() + "' is not present in the JSON object.");
          } else if (fieldValue.isJsonNull()) {
            throw new UnsupportedRowJsonException(
                "Non-nullable field '" + fieldValue.name() + "' has value null in the JSON object.");
          }
        }
  
        if (fieldValue.isRowType()) {
          return jsonObjectToRow(fieldValue);
        }
  
        if (fieldValue.isArrayType()) {
          return jsonArrayToList(fieldValue);
        }
  
        if (fieldValue.typeName().isLogicalType()) {
          String identifier = fieldValue.type().getLogicalType().getIdentifier();
          if (SqlTypes.DATE.getIdentifier().equals(identifier)) {
            return dateValueExtractor().extractValue(fieldValue.jsonValue());
          } else if (SqlTypes.TIME.getIdentifier().equals(identifier)) {
            return timeValueExtractor().extractValue(fieldValue.jsonValue());
          } else if (SqlTypes.DATETIME.getIdentifier().equals(identifier)) {
            return localDatetimeValueExtractor().extractValue(fieldValue.jsonValue());
          } else {
            return extractJsonNodeValue(
                FieldValue.of(
                    fieldValue.name(),
                    fieldValue.type().getLogicalType().getBaseType(),
                    fieldValue.jsonValue()));
          }
        }
  
        return extractJsonPrimitiveValue(fieldValue);
      }
  
      private Row jsonObjectToRow(FieldValue rowFieldValue) {
        if (!rowFieldValue.isJsonObject()) {
          throw new UnsupportedRowJsonException(
              "Expected JSON object for field '"
                  + rowFieldValue.name()
                  + "'. Unable to convert '"
                  + rowFieldValue.jsonValue().asText()
                  + "' to Beam Row, it is not a JSON object. Currently only JSON objects can be parsed to Beam Rows");
        }
  
        return rowFieldValue.rowSchema().getFields().stream()
            .map(
                schemaField ->
                    extractJsonNodeValue(
                        FieldValue.of(
                            schemaField.getName(),
                            schemaField.getType(),
                            rowFieldValue.jsonFieldValue(schemaField.getName()))))
            .collect(toRow(rowFieldValue.rowSchema()));
      }
  
      private Object jsonArrayToList(FieldValue arrayFieldValue) {
        if (!arrayFieldValue.isJsonArray()) {
          throw new UnsupportedRowJsonException(
              "Expected JSON array for field '"
                  + arrayFieldValue.name()
                  + "'. Instead got "
                  + arrayFieldValue.jsonNodeType().name());
        }
  
        return arrayFieldValue
            .jsonArrayElements()
            .map(
                jsonArrayElement ->
                    extractJsonNodeValue(
                        FieldValue.of(
                            arrayFieldValue.name() + "[]",
                            arrayFieldValue.arrayElementType(),
                            jsonArrayElement)))
            .collect(toImmutableList());
      }
  
      private static Object extractJsonPrimitiveValue(FieldValue fieldValue) {
        try {
          return JSON_VALUE_GETTERS.get(fieldValue.typeName()).extractValue(fieldValue.jsonValue());
        } catch (RuntimeException e) {
          throw new UnsupportedRowJsonException(
              "Unable to get value from field '"
                  + fieldValue.name()
                  + "'. Schema type '"
                  + fieldValue.typeName()
                  + "'. JSON node type "
                  + fieldValue.jsonNodeType().name(),
              e);
        }
      }
  
      /**
       * Helper class to keep track of schema field type, name, and actual json value for the field.
       */
      @AutoValue
      abstract static class FieldValue {
        abstract String name();
  
        abstract FieldType type();
  
        abstract @Nullable JsonNode jsonValue();
  
        TypeName typeName() {
          return type().getTypeName();
        }
  
        boolean isJsonValuePresent() {
          return jsonValue() != null;
        }
  
        boolean isJsonNull() {
          return jsonValue().isNull();
        }
  
        JsonNodeType jsonNodeType() {
          return jsonValue().getNodeType();
        }
  
        boolean isJsonArray() {
          return jsonValue().isArray();
        }
  
        Stream<JsonNode> jsonArrayElements() {
          return StreamSupport.stream(jsonValue().spliterator(), SEQUENTIAL);
        }
  
        boolean isArrayType() {
          return TypeName.ARRAY.equals(type().getTypeName())
              || TypeName.ITERABLE.equals(type().getTypeName());
        }
  
        FieldType arrayElementType() {
          return type().getCollectionElementType();
        }
  
        boolean isJsonObject() {
          return jsonValue().isObject();
        }
  
        JsonNode jsonFieldValue(String fieldName) {
          return jsonValue().get(fieldName);
        }
  
        boolean isRowType() {
          return TypeName.ROW.equals(type().getTypeName());
        }
  
        Schema rowSchema() {
          return type().getRowSchema();
        }
  
        static FieldValue of(String name, FieldType type, JsonNode jsonValue) {
          return new AutoValue_RowJson_RowJsonDeserializer_FieldValue(name, type, jsonValue);
        }
      }
    }
  
    /** Jackson serializer for converting {@link Row Rows} to JSON. */
    public static class RowJsonSerializer extends StdSerializer<Row> {
  
      private final Schema schema;
      private Boolean dropNullsOnWrite = false;
  
      /** Creates a serializer for a {@link Row} {@link Schema}. */
      public static RowJsonSerializer forSchema(Schema schema) {
        verifySchemaSupported(schema);
        return new RowJsonSerializer(schema);
      }
  
      private RowJsonSerializer(Schema schema) {
        super(Row.class);
        this.schema = schema;
      }
  
      /** Serializer drops nulls on write if set to true instead of writing fieldName: null. */
      public RowJsonSerializer withDropNullsOnWrite(Boolean dropNullsOnWrite) {
        this.dropNullsOnWrite = dropNullsOnWrite;
        return this;
      }
  
      @Override
      public void serialize(Row value, JsonGenerator gen, SerializerProvider provider)
          throws IOException {
        writeRow(value, this.schema, gen);
      }
  
      // TODO: ByteBuddy generate based on schema?
      private void writeRow(Row row, Schema schema, JsonGenerator gen) throws IOException {
        gen.writeStartObject();
        for (int i = 0; i < schema.getFieldCount(); ++i) {
          Field field = schema.getField(i);
          Object value = row.getValue(i);
          if (dropNullsOnWrite && value == null && field.getType().getNullable()) {
            continue;
          }
          gen.writeFieldName(field.getName());
          if (field.getType().getNullable() && value == null) {
            gen.writeNull();
            continue;
          }
          writeValue(gen, field.getType(), value);
        }
        gen.writeEndObject();
      }
  
      private void writeValue(JsonGenerator gen, FieldType type, Object value) throws IOException {
        switch (type.getTypeName()) {
          case BOOLEAN:
            gen.writeBoolean((boolean) value);
            break;
          case STRING:
            gen.writeString((String) value);
            break;
          case BYTE:
            gen.writeNumber((byte) value);
            break;
          case DOUBLE:
            gen.writeNumber((double) value);
            break;
          case FLOAT:
            gen.writeNumber((float) value);
            break;
          case INT16:
            gen.writeNumber((short) value);
            break;
          case INT32:
            gen.writeNumber((int) value);
            break;
          case INT64:
            gen.writeNumber((long) value);
            break;
          case DECIMAL:
            gen.writeNumber((BigDecimal) value);
            break;
          case DATETIME:
            // ReadableInstant.toString produces ISO8601 format string
            gen.writeString(((ReadableInstant) value).toString());
            break;
          case ARRAY:
          case ITERABLE:
            gen.writeStartArray();
            for (Object element : (Iterable<Object>) value) {
              writeValue(gen, type.getCollectionElementType(), element);
            }
            gen.writeEndArray();
            break;
          case ROW:
            writeRow((Row) value, type.getRowSchema(), gen);
            break;
          case LOGICAL_TYPE:
            String identifier = type.getLogicalType().getIdentifier();
            if (SqlTypes.DATE.getIdentifier().equals(identifier)) {
              gen.writeString(((LocalDate) value).toString()); // ISO 8601 format
            } else if (SqlTypes.TIME.getIdentifier().equals(identifier)) {
              gen.writeString(((LocalTime) value).toString()); // ISO 8601 format
            } else if (SqlTypes.DATETIME.getIdentifier().equals(identifier)) {
              gen.writeString(((LocalDateTime) value).toString()); // ISO 8601 format
            } else {
              writeValue(gen, type.getLogicalType().getBaseType(), value);
            }
            break;
          default:
            throw new IllegalArgumentException("Unsupported field type: " + type);
        }
      }
    }
  
    /** Gets thrown when Row parsing or serialization fails for any reason. */
    public static class UnsupportedRowJsonException extends RuntimeException {
  
      UnsupportedRowJsonException(String message, Throwable reason) {
        super(message, reason);
      }
  
      UnsupportedRowJsonException(String message) {
        super(message);
      }
    }
  }
  