sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/utils/QuickCheckGenerators.java

/*
 * 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.extensions.sql.utils;

import static java.util.Arrays.asList;
import static java.util.stream.Collectors.toList;
import static org.apache.beam.sdk.schemas.Schema.toSchema;

import com.pholser.junit.quickcheck.generator.GenerationStatus;
import com.pholser.junit.quickcheck.generator.Generator;
import com.pholser.junit.quickcheck.random.SourceOfRandomness;
import com.pholser.junit.quickcheck.runner.JUnitQuickcheck;
import java.util.List;
import java.util.stream.IntStream;
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;

/** Field type generators invoked by {@link JUnitQuickcheck}. */
public class QuickCheckGenerators {

  private static final FieldTypeGenerator PRIMITIVE_TYPES = new PrimitiveTypes();
  private static final FieldTypeGenerator ARRAYS_OF_ANY = new Arrays();
  private static final FieldTypeGenerator MAPS_OF_ANY = new Maps();
  private static final FieldTypeGenerator ROWS = new Rows();
  private static final FieldTypeGenerator ANY_TYPE = new AnyFieldType();

  /** Generates a primitive {@link TypeName SQL type} randomly. */
  public static class PrimitiveTypes extends FieldTypeGenerator {
    private static final List<FieldType> PRIMITIVE_TYPES =
        java.util.Arrays.stream(TypeName.values())
            .filter(TypeName::isPrimitiveType)
            .map(FieldType::of)
            .collect(toList());

    @Override
    public FieldType generateFieldType(SourceOfRandomness random, GenerationStatus status) {
      return random.choose(PRIMITIVE_TYPES);
    }
  }

  /** Generates {@link TypeName#ARRAY SQL arrays} with random element type. */
  public static class Arrays extends FieldTypeGenerator {
    @Override
    public FieldType generateFieldType(SourceOfRandomness random, GenerationStatus status) {
      return FieldType.array(ANY_TYPE.generate(random, status));
    }
  }

  /**
   * Generates {@link TypeName#MAP SQL maps} with random primitive key type and random value type.
   */
  public static class Maps extends FieldTypeGenerator {
    @Override
    public FieldType generateFieldType(SourceOfRandomness random, GenerationStatus status) {
      return FieldType.map(
          PRIMITIVE_TYPES.generate(random, status), ANY_TYPE.generate(random, status));
    }
  }

  /** Generates {@link TypeName#ROW SQL rows} with random field types. */
  public static class Rows extends FieldTypeGenerator {
    @Override
    public FieldType generateFieldType(SourceOfRandomness random, GenerationStatus status) {
      // stop at 10 levels of nesting to avoid stack overflows
      FieldTypeGenerator rowFieldTypesGenerator =
          (nestingLevel(status) >= 10) ? PRIMITIVE_TYPES : ANY_TYPE;

      return FieldType.row(generateSchema(rowFieldTypesGenerator, random, status));
    }

    private Schema generateSchema(
        FieldTypeGenerator fieldTypeGenerator, SourceOfRandomness random, GenerationStatus status) {

      return IntStream.range(0, status.size() + 1)
          .mapToObj(
              i ->
                  Field.of("field_" + i, fieldTypeGenerator.generate(random, status))
                      .withNullable(true))
          .collect(toSchema());
    }
  }

  /** Generates a {@link FieldType}, randomly delegating to specific type generators. */
  public static class AnyFieldType extends FieldTypeGenerator {
    @Override
    public FieldType generateFieldType(SourceOfRandomness random, GenerationStatus status) {
      return random
          .choose(asList(PRIMITIVE_TYPES, ARRAYS_OF_ANY, MAPS_OF_ANY, ROWS))
          .generate(random, status);
    }
  }

  abstract static class FieldTypeGenerator extends Generator<FieldType> {
    static final GenerationStatus.Key<Integer> NESTING_KEY =
        new GenerationStatus.Key<>("nesting", Integer.class);

    FieldTypeGenerator() {
      super(FieldType.class);
    }

    @Override
    public FieldType generate(SourceOfRandomness random, GenerationStatus status) {
      incrementNesting(status);
      return generateFieldType(random, status);
    }

    protected abstract FieldType generateFieldType(
        SourceOfRandomness random, GenerationStatus status);

    int nestingLevel(GenerationStatus status) {
      return status.valueOf(NESTING_KEY).orElse(-1);
    }

    void incrementNesting(GenerationStatus status) {
      status.setValue(NESTING_KEY, nestingLevel(status) + 1);
    }
  }
}