java/serving/src/test/java/feast/serving/util/DataGenerator.java
/*
 * SPDX-License-Identifier: Apache-2.0
 * Copyright 2018-2020 The Feast Authors
 *
 * 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
 *
 *     https://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 feast.serving.util;

import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.protobuf.Duration;
import com.google.protobuf.Timestamp;
import feast.proto.core.DataFormatProto.FileFormat;
import feast.proto.core.DataFormatProto.FileFormat.ParquetFormat;
import feast.proto.core.DataFormatProto.StreamFormat;
import feast.proto.core.DataFormatProto.StreamFormat.AvroFormat;
import feast.proto.core.DataFormatProto.StreamFormat.ProtoFormat;
import feast.proto.core.DataSourceProto.DataSource;
import feast.proto.core.DataSourceProto.DataSource.FileOptions;
import feast.proto.core.DataSourceProto.DataSource.KafkaOptions;
import feast.proto.core.DataSourceProto.DataSource.KinesisOptions;
import feast.proto.core.EntityProto;
import feast.proto.core.FeatureProto;
import feast.proto.core.FeatureProto.FeatureSpecV2;
import feast.proto.core.FeatureTableProto.FeatureTableSpec;
import feast.proto.core.StoreProto;
import feast.proto.serving.ServingAPIProto;
import feast.proto.types.ValueProto;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;
import org.apache.commons.lang3.tuple.Triple;

public class DataGenerator {
  // projectName, featureName, exclude
  static Triple<String, String, Boolean> defaultSubscription = Triple.of("*", "*", false);

  static StoreProto.Store defaultStore =
      createStore(
          "test-store", StoreProto.Store.StoreType.REDIS, ImmutableList.of(defaultSubscription));

  public static Triple<String, String, Boolean> getDefaultSubscription() {
    return defaultSubscription;
  }

  public static String valueToString(ValueProto.Value v) {
    String stringRepr;
    switch (v.getValCase()) {
      case STRING_VAL:
        stringRepr = v.getStringVal();
        break;
      case INT64_VAL:
        stringRepr = String.valueOf(v.getInt64Val());
        break;
      case INT32_VAL:
        stringRepr = String.valueOf(v.getInt32Val());
        break;
      case BYTES_VAL:
        stringRepr = v.getBytesVal().toString();
        break;
      default:
        throw new RuntimeException("Type is not supported to be entity");
    }

    return stringRepr;
  }

  public static StoreProto.Store getDefaultStore() {
    return defaultStore;
  }

  public static StoreProto.Store createStore(
      String name,
      StoreProto.Store.StoreType type,
      List<Triple<String, String, Boolean>> subscriptions) {
    StoreProto.Store.Builder builder =
        StoreProto.Store.newBuilder()
            .addAllSubscriptions(
                subscriptions.stream()
                    .map(
                        s ->
                            StoreProto.Store.Subscription.newBuilder()
                                .setProject(s.getLeft())
                                .setName(s.getMiddle())
                                .setExclude(s.getRight())
                                .build())
                    .collect(Collectors.toList()))
            .setName(name)
            .setType(type);

    switch (type) {
      case REDIS:
        StoreProto.Store.RedisConfig redisConfig =
            StoreProto.Store.RedisConfig.newBuilder().build();
        return builder.setRedisConfig(redisConfig).build();
      case REDIS_CLUSTER:
        StoreProto.Store.RedisClusterConfig redisClusterConfig =
            StoreProto.Store.RedisClusterConfig.newBuilder().build();
        return builder.setRedisClusterConfig(redisClusterConfig).build();
      default:
        throw new RuntimeException("Unrecognized Store type");
    }
  }

  public static EntityProto.EntitySpecV2 createEntitySpecV2(
      String name,
      String description,
      ValueProto.ValueType.Enum valueType,
      Map<String, String> tags) {
    return EntityProto.EntitySpecV2.newBuilder()
        .setName(name)
        .setDescription(description)
        .setValueType(valueType)
        .putAllTags(tags)
        .build();
  }

  public static FeatureProto.FeatureSpecV2 createFeatureSpecV2(
      String name, ValueProto.ValueType.Enum valueType, Map<String, String> tags) {
    return FeatureProto.FeatureSpecV2.newBuilder()
        .setName(name)
        .setValueType(valueType)
        .putAllTags(tags)
        .build();
  }

  // Create a Feature Table spec without DataSources configured.
  public static FeatureTableSpec createFeatureTableSpec(
      String name,
      List<String> entities,
      Map<String, ValueProto.ValueType.Enum> features,
      int maxAgeSecs,
      Map<String, String> tags) {

    return FeatureTableSpec.newBuilder()
        .setName(name)
        .addAllEntities(entities)
        .addAllFeatures(
            features.entrySet().stream()
                .map(
                    entry ->
                        FeatureSpecV2.newBuilder()
                            .setName(entry.getKey())
                            .setValueType(entry.getValue())
                            .putAllTags(tags)
                            .build())
                .collect(Collectors.toList()))
        .setMaxAge(Duration.newBuilder().setSeconds(3600).build())
        .setBatchSource(
            DataSource.newBuilder()
                .setTimestampField("ts")
                .setType(DataSource.SourceType.BATCH_FILE)
                .setFileOptions(
                    FileOptions.newBuilder()
                        .setFileFormat(
                            FileFormat.newBuilder()
                                .setParquetFormat(ParquetFormat.newBuilder().build())
                                .build())
                        .setUri("/dev/null")
                        .build())
                .build())
        .putAllLabels(tags)
        .build();
  }

  public static FeatureTableSpec createFeatureTableSpec(
      String name,
      List<String> entities,
      ImmutableMap<String, ValueProto.ValueType.Enum> features,
      int maxAgeSecs,
      Map<String, String> tags) {

    return FeatureTableSpec.newBuilder()
        .setName(name)
        .addAllEntities(entities)
        .addAllFeatures(
            features.entrySet().stream()
                .map(
                    entry ->
                        FeatureSpecV2.newBuilder()
                            .setName(entry.getKey())
                            .setValueType(entry.getValue())
                            .putAllTags(tags)
                            .build())
                .collect(Collectors.toList()))
        .setMaxAge(Duration.newBuilder().setSeconds(maxAgeSecs).build())
        .putAllLabels(tags)
        .build();
  }

  public static DataSource createFileDataSourceSpec(
      String fileURL, String timestampColumn, String datePartitionColumn) {
    return DataSource.newBuilder()
        .setType(DataSource.SourceType.BATCH_FILE)
        .setFileOptions(
            FileOptions.newBuilder().setFileFormat(createParquetFormat()).setUri(fileURL).build())
        .setTimestampField(timestampColumn)
        .setDatePartitionColumn(datePartitionColumn)
        .build();
  }

  public static DataSource createBigQueryDataSourceSpec(
      String bigQueryTable, String timestampColumn, String datePartitionColumn) {
    return DataSource.newBuilder()
        .setType(DataSource.SourceType.BATCH_BIGQUERY)
        .setBigqueryOptions(DataSource.BigQueryOptions.newBuilder().setTable(bigQueryTable).build())
        .setTimestampField(timestampColumn)
        .setDatePartitionColumn(datePartitionColumn)
        .build();
  }

  public static DataSource createKafkaDataSourceSpec(
      String servers, String topic, String classPath, String timestampColumn) {
    return DataSource.newBuilder()
        .setType(DataSource.SourceType.STREAM_KAFKA)
        .setKafkaOptions(
            KafkaOptions.newBuilder()
                .setTopic(topic)
                .setKafkaBootstrapServers(servers)
                .setMessageFormat(createProtoFormat("class.path"))
                .build())
        .setTimestampField(timestampColumn)
        .build();
  }

  public static ValueProto.Value createEmptyValue() {
    return ValueProto.Value.newBuilder().build();
  }

  public static ValueProto.Value createStrValue(String val) {
    return ValueProto.Value.newBuilder().setStringVal(val).build();
  }

  public static ValueProto.Value createDoubleValue(double value) {
    return ValueProto.Value.newBuilder().setDoubleVal(value).build();
  }

  public static ValueProto.Value createInt32Value(int value) {
    return ValueProto.Value.newBuilder().setInt32Val(value).build();
  }

  public static ValueProto.Value createInt64Value(long value) {
    return ValueProto.Value.newBuilder().setInt64Val(value).build();
  }

  public static ServingAPIProto.FeatureReferenceV2 createFeatureReference(
      String featureTableName, String featureName) {
    return ServingAPIProto.FeatureReferenceV2.newBuilder()
        .setFeatureViewName(featureTableName)
        .setFeatureName(featureName)
        .build();
  }

  public static ServingAPIProto.GetOnlineFeaturesRequestV2.EntityRow createEntityRow(
      String entityName, ValueProto.Value entityValue, long seconds) {
    return ServingAPIProto.GetOnlineFeaturesRequestV2.EntityRow.newBuilder()
        .setTimestamp(Timestamp.newBuilder().setSeconds(seconds))
        .putFields(entityName, entityValue)
        .build();
  }

  public static ServingAPIProto.GetOnlineFeaturesRequestV2.EntityRow createCompoundEntityRow(
      ImmutableMap<String, ValueProto.Value> entityNameValues, long seconds) {
    ServingAPIProto.GetOnlineFeaturesRequestV2.EntityRow.Builder entityRow =
        ServingAPIProto.GetOnlineFeaturesRequestV2.EntityRow.newBuilder()
            .setTimestamp(Timestamp.newBuilder().setSeconds(seconds));

    entityNameValues.entrySet().stream()
        .forEach(entry -> entityRow.putFields(entry.getKey(), entry.getValue()));

    return entityRow.build();
  }

  public static DataSource createKinesisDataSourceSpec(
      String region, String streamName, String classPath, String timestampColumn) {
    return DataSource.newBuilder()
        .setType(DataSource.SourceType.STREAM_KINESIS)
        .setKinesisOptions(
            KinesisOptions.newBuilder()
                .setRegion("ap-nowhere1")
                .setStreamName("stream")
                .setRecordFormat(createProtoFormat(classPath))
                .build())
        .setTimestampField(timestampColumn)
        .build();
  }

  public static FileFormat createParquetFormat() {
    return FileFormat.newBuilder().setParquetFormat(ParquetFormat.getDefaultInstance()).build();
  }

  public static StreamFormat createAvroFormat(String schemaJSON) {
    return StreamFormat.newBuilder()
        .setAvroFormat(AvroFormat.newBuilder().setSchemaJson(schemaJSON).build())
        .build();
  }

  public static StreamFormat createProtoFormat(String classPath) {
    return StreamFormat.newBuilder()
        .setProtoFormat(ProtoFormat.newBuilder().setClassPath(classPath).build())
        .build();
  }
}
