code: /*
   * SPDX-License-Identifier: Apache-2.0
   * Copyright 2018-2019 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 dev.feast;
  
  import com.google.common.collect.Lists;
  import feast.proto.serving.ServingAPIProto;
  import feast.proto.serving.ServingAPIProto.GetFeastServingInfoRequest;
  import feast.proto.serving.ServingAPIProto.GetFeastServingInfoResponse;
  import feast.proto.serving.ServingAPIProto.GetOnlineFeaturesRequest;
  import feast.proto.serving.ServingAPIProto.GetOnlineFeaturesResponse;
  import feast.proto.serving.ServingServiceGrpc;
  import feast.proto.serving.ServingServiceGrpc.ServingServiceBlockingStub;
  import feast.proto.types.ValueProto;
  import io.grpc.CallCredentials;
  import io.grpc.ManagedChannel;
  import io.grpc.ManagedChannelBuilder;
  import io.grpc.netty.shaded.io.grpc.netty.GrpcSslContexts;
  import io.grpc.netty.shaded.io.grpc.netty.NettyChannelBuilder;
  import io.opentracing.contrib.grpc.TracingClientInterceptor;
  import io.opentracing.util.GlobalTracer;
  import java.io.File;
  import java.time.Instant;
  import java.util.*;
  import java.util.concurrent.TimeUnit;
  import java.util.stream.Collectors;
  import javax.net.ssl.SSLException;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  @SuppressWarnings("WeakerAccess")
  public class FeastClient implements AutoCloseable {
    Logger logger = LoggerFactory.getLogger(FeastClient.class);
  
    private static final int CHANNEL_SHUTDOWN_TIMEOUT_SEC = 5;
  
    private final ManagedChannel channel;
    private final ServingServiceBlockingStub stub;
  
    /**
     * Create a client to access Feast Serving.
     *
     * @param host hostname or ip address of Feast serving GRPC server
     * @param port port number of Feast serving GRPC server
     * @return {@link FeastClient}
     */
    public static FeastClient create(String host, int port) {
      // configure client with no security config.
      return FeastClient.createSecure(host, port, SecurityConfig.newBuilder().build());
    }
  
    /**
     * Create a authenticated client that can access Feast serving with authentication enabled.
     *
     * @param host hostname or ip address of Feast serving GRPC server
     * @param port port number of Feast serving GRPC server
     * @param securityConfig security options to configure the Feast client. See {@link
     *     SecurityConfig} for options.
     * @return {@link FeastClient}
     */
    public static FeastClient createSecure(String host, int port, SecurityConfig securityConfig) {
      // Configure client TLS
      ManagedChannel channel = null;
      if (securityConfig.isTLSEnabled()) {
        if (securityConfig.getCertificatePath().isPresent()) {
          String certificatePath = securityConfig.getCertificatePath().get();
          // Use custom certificate for TLS
          File certificateFile = new File(certificatePath);
          try {
            channel =
                NettyChannelBuilder.forAddress(host, port)
                    .useTransportSecurity()
                    .sslContext(GrpcSslContexts.forClient().trustManager(certificateFile).build())
                    .build();
          } catch (SSLException e) {
            throw new IllegalArgumentException(
                String.format("Invalid Certificate provided at path: %s", certificatePath), e);
          }
        } else {
          // Use system certificates for TLS
          channel = ManagedChannelBuilder.forAddress(host, port).useTransportSecurity().build();
        }
      } else {
        // Disable TLS
        channel = ManagedChannelBuilder.forAddress(host, port).usePlaintext().build();
      }
  
      return new FeastClient(channel, securityConfig.getCredentials());
    }
  
    /**
     * Obtain info about Feast Serving.
     *
     * @return {@link GetFeastServingInfoResponse} containing Feast version, Serving type etc.
     */
    public GetFeastServingInfoResponse getFeastServingInfo() {
      return stub.getFeastServingInfo(GetFeastServingInfoRequest.newBuilder().build());
    }
  
    /**
     * Get online features from Feast, without indicating project, will use `default`.
     *
     * <p>See {@link #getOnlineFeatures(List, List, String)}
     *
     * @param featureRefs list of string feature references to retrieve in the following format
     *     featureTable:feature, where 'featureTable' and 'feature' refer to the FeatureTable and
     *     Feature names respectively. Only the Feature name is required.
     * @param entities list of {@link Row} to select the entities to retrieve the features for.
     * @return list of {@link Row} containing retrieved data fields.
     */
    public List<Row> getOnlineFeatures(List<String> featureRefs, List<Row> entities) {
      GetOnlineFeaturesRequest.Builder requestBuilder = GetOnlineFeaturesRequest.newBuilder();
  
      requestBuilder.setFeatures(
          ServingAPIProto.FeatureList.newBuilder().addAllVal(featureRefs).build());
  
      requestBuilder.putAllEntities(getEntityValuesMap(entities));
  
      GetOnlineFeaturesResponse response = stub.getOnlineFeatures(requestBuilder.build());
  
      List<Row> results = Lists.newArrayList();
      if (response.getResultsCount() == 0) {
        return results;
      }
  
      for (int rowIdx = 0; rowIdx < response.getResults(0).getValuesCount(); rowIdx++) {
        Row row = Row.create();
        for (int featureIdx = 0; featureIdx < response.getResultsCount(); featureIdx++) {
          row.set(
              response.getMetadata().getFeatureNames().getVal(featureIdx),
              response.getResults(featureIdx).getValues(rowIdx),
              response.getResults(featureIdx).getStatuses(rowIdx));
  
          row.setEntityTimestamp(
              Instant.ofEpochSecond(
                  response.getResults(featureIdx).getEventTimestamps(rowIdx).getSeconds()));
        }
        for (Map.Entry<String, ValueProto.Value> entry :
            entities.get(rowIdx).getFields().entrySet()) {
          row.set(entry.getKey(), entry.getValue());
        }
  
        results.add(row);
      }
      return results;
    }
  
    private Map<String, ValueProto.RepeatedValue> getEntityValuesMap(List<Row> entities) {
      Map<String, ValueProto.RepeatedValue.Builder> columnarEntities = new HashMap<>();
      for (Row row : entities) {
        for (Map.Entry<String, ValueProto.Value> field : row.getFields().entrySet()) {
          if (!columnarEntities.containsKey(field.getKey())) {
            columnarEntities.put(field.getKey(), ValueProto.RepeatedValue.newBuilder());
          }
          columnarEntities.get(field.getKey()).addVal(field.getValue());
        }
      }
  
      return columnarEntities.entrySet().stream()
          .map((e) -> Map.entry(e.getKey(), e.getValue().build()))
          .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
    }
  
    /**
     * Get online features from Feast.
     *
     * <p>Example of retrieving online features for the driver FeatureTable, with features driver_id
     * and driver_name
     *
     * <pre>{@code
     * FeastClient client = FeastClient.create("localhost", 6566);
     * List<String> requestedFeatureIds = Arrays.asList("driver:driver_id", "driver:driver_name");
     * List<Row> requestedRows =
     *         Arrays.asList(Row.create().set("driver_id", 123), Row.create().set("driver_id", 456));
     * List<Row> retrievedFeatures = client.getOnlineFeatures(requestedFeatureIds, requestedRows);
     * retrievedFeatures.forEach(System.out::println);
     * }</pre>
     *
     * @param featureRefs list of string feature references to retrieve in the following format
     *     featureTable:feature, where 'featureTable' and 'feature' refer to the FeatureTable and
     *     Feature names respectively. Only the Feature name is required.
     * @param rows list of {@link Row} to select the entities to retrieve the features for
     * @param project {@link String} Specifies the project override. If specified uses the project for
     *     retrieval. Overrides the projects set in Feature References if also specified.
     * @return list of {@link Row} containing retrieved data fields.
     */
    public List<Row> getOnlineFeatures(List<String> featureRefs, List<Row> rows, String project) {
      return getOnlineFeatures(featureRefs, rows);
    }
  
    protected FeastClient(ManagedChannel channel, Optional<CallCredentials> credentials) {
      this.channel = channel;
      TracingClientInterceptor tracingInterceptor =
          TracingClientInterceptor.newBuilder().withTracer(GlobalTracer.get()).build();
  
      ServingServiceBlockingStub servingStub =
          ServingServiceGrpc.newBlockingStub(tracingInterceptor.intercept(channel));
  
      if (credentials.isPresent()) {
        servingStub = servingStub.withCallCredentials(credentials.get());
      }
  
      this.stub = servingStub;
    }
  
    public void close() throws Exception {
      if (channel != null) {
        channel.shutdown().awaitTermination(CHANNEL_SHUTDOWN_TIMEOUT_SEC, TimeUnit.SECONDS);
      }
    }
  }
  