alluxio-master-2.x/table/client/src/main/java/alluxio/client/table/RetryHandlingTableMasterClient.java
/*
 * The Alluxio Open Foundation licenses this work under the Apache License, version 2.0
 * (the "License"). You may not use this work except in compliance with the License, which is
 * available at www.apache.org/licenses/LICENSE-2.0
 *
 * This software is distributed on an "AS IS" basis, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND,
 * either express or implied, as more fully set forth in the License.
 *
 * See the NOTICE file distributed with this work for information regarding copyright ownership.
 */

package alluxio.client.table;

import alluxio.AbstractMasterClient;
import alluxio.Constants;
import alluxio.exception.status.AlluxioStatusException;
import alluxio.grpc.ServiceType;
import alluxio.grpc.table.AttachDatabasePRequest;
import alluxio.grpc.table.ColumnStatisticsInfo;
import alluxio.grpc.table.Constraint;
import alluxio.grpc.table.Database;
import alluxio.grpc.table.DetachDatabasePRequest;
import alluxio.grpc.table.GetAllDatabasesPRequest;
import alluxio.grpc.table.GetAllTablesPRequest;
import alluxio.grpc.table.GetDatabasePRequest;
import alluxio.grpc.table.GetPartitionColumnStatisticsPRequest;
import alluxio.grpc.table.GetTableColumnStatisticsPRequest;
import alluxio.grpc.table.GetTablePRequest;
import alluxio.grpc.table.GetTransformJobInfoPRequest;
import alluxio.grpc.table.Partition;
import alluxio.grpc.table.ReadTablePRequest;
import alluxio.grpc.table.SyncDatabasePRequest;
import alluxio.grpc.table.SyncStatus;
import alluxio.grpc.table.TableInfo;
import alluxio.grpc.table.TableMasterClientServiceGrpc;
import alluxio.grpc.table.TransformJobInfo;
import alluxio.grpc.table.TransformTablePRequest;
import alluxio.master.MasterClientContext;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;
import javax.annotation.concurrent.ThreadSafe;

/**
 * A wrapper for the gRPC client to interact with the table master.
 */
@ThreadSafe
public final class RetryHandlingTableMasterClient extends AbstractMasterClient
    implements TableMasterClient {
  private static final Logger RPC_LOG = LoggerFactory.getLogger(TableMasterClient.class);
  private TableMasterClientServiceGrpc.TableMasterClientServiceBlockingStub mClient = null;

  /**
   * Creates a new block master client.
   *
   * @param conf master client configuration
   */
  public RetryHandlingTableMasterClient(MasterClientContext conf) {
    super(conf);
  }

  @Override
  protected ServiceType getRemoteServiceType() {
    return ServiceType.TABLE_MASTER_CLIENT_SERVICE;
  }

  @Override
  protected String getServiceName() {
    return Constants.TABLE_MASTER_CLIENT_SERVICE_NAME;
  }

  @Override
  protected long getServiceVersion() {
    return Constants.TABLE_MASTER_CLIENT_SERVICE_VERSION;
  }

  @Override
  protected void afterConnect() {
    mClient = TableMasterClientServiceGrpc.newBlockingStub(mChannel);
  }

  @Override
  public List<String> getAllDatabases() throws AlluxioStatusException {
    return retryRPC(() -> mClient.getAllDatabases(
        GetAllDatabasesPRequest.newBuilder().build()).getDatabaseList(),
        RPC_LOG, "GetAllDatabases", "");
  }

  @Override
  public Database getDatabase(String databaseName) throws AlluxioStatusException {
    return retryRPC(() -> mClient.getDatabase(GetDatabasePRequest.newBuilder()
        .setDbName(databaseName).build()),
        RPC_LOG, "GetDatabase", "databaseName=%s", databaseName).getDb();
  }

  @Override
  public List<String> getAllTables(String databaseName) throws AlluxioStatusException {
    return retryRPC(() -> mClient.getAllTables(
        GetAllTablesPRequest.newBuilder().setDatabase(databaseName).build()).getTableList(),
        RPC_LOG, "GetAllTables", "databaseName=%s", databaseName);
  }

  @Override
  public TableInfo getTable(String databaseName, String tableName) throws AlluxioStatusException {
    return retryRPC(() -> mClient.getTable(
        GetTablePRequest.newBuilder().setDbName(databaseName).setTableName(tableName).build())
        .getTableInfo(), RPC_LOG, "GetTable", "databaseName=%s,tableName=%s",
        databaseName, tableName);
  }

  @Override
  public SyncStatus attachDatabase(String udbType, String udbConnectionUri, String udbDbName,
      String dbName, Map<String, String> configuration, boolean ignoreSyncErrors)
      throws AlluxioStatusException {
    return retryRPC(() -> mClient.attachDatabase(
        AttachDatabasePRequest.newBuilder().setUdbType(udbType)
            .setUdbConnectionUri(udbConnectionUri).setUdbDbName(udbDbName).setDbName(dbName)
            .putAllOptions(configuration).setIgnoreSyncErrors(ignoreSyncErrors).build())
        .getSyncStatus(),
        RPC_LOG, "AttachDatabase", "udbType=%s,udbConnectionUri=%s,udbDbName=%s,dbName=%s,"
            + "configuration=%s,ignoreSyncErrors=%s",
        udbType, udbConnectionUri, udbDbName, dbName, configuration, ignoreSyncErrors);
  }

  @Override
  public boolean detachDatabase(String dbName)
      throws AlluxioStatusException {
    return retryRPC(() -> mClient.detachDatabase(
        DetachDatabasePRequest.newBuilder().setDbName(dbName).build()).getSuccess(),
        RPC_LOG, "DetachDatabase", "dbName=%s", dbName);
  }

  @Override
  public SyncStatus syncDatabase(String dbName) throws AlluxioStatusException {
    return retryRPC(() -> mClient.syncDatabase(
        SyncDatabasePRequest.newBuilder().setDbName(dbName).build()).getStatus(),
        RPC_LOG, "SyncDatabase", "dbName=%s", dbName);
  }

  @Override
  public List<Partition> readTable(String databaseName, String tableName, Constraint constraint)
      throws AlluxioStatusException {
    return retryRPC(() -> mClient.readTable(
        ReadTablePRequest.newBuilder().setDbName(databaseName).setTableName(tableName)
            .setConstraint(constraint).build()).getPartitionsList(),
        RPC_LOG, "ReadTable", "databaseName=%s,tableName=%s,constraint=%s", databaseName, tableName,
        constraint);
  }

  @Override
  public List<ColumnStatisticsInfo> getTableColumnStatistics(
          String databaseName,
          String tableName,
          List<String> columnNames) throws AlluxioStatusException {
    return retryRPC(() -> mClient.getTableColumnStatistics(
        GetTableColumnStatisticsPRequest.newBuilder().setDbName(databaseName)
            .setTableName(tableName).addAllColNames(columnNames).build()).getStatisticsList(),
        RPC_LOG, "GetTableColumnStatistics",
        "databaseName=%s,tableName=%s,columnNames=%s", databaseName, tableName, columnNames);
  }

  @Override
  public List<String> getPartitionNames(
          String databaseName,
          String tableName) throws AlluxioStatusException {
    return null;
  }

  @Override
  public Map<String, List<ColumnStatisticsInfo>> getPartitionColumnStatistics(
          String databaseName,
          String tableName,
          List<String> partitionNames,
          List<String> columnNames) throws AlluxioStatusException {
    return retryRPC(() -> mClient.getPartitionColumnStatistics(
        GetPartitionColumnStatisticsPRequest.newBuilder().setDbName(databaseName)
            .setTableName(tableName).addAllColNames(columnNames)
            .addAllPartNames(partitionNames).build()).getPartitionStatisticsMap(),
        RPC_LOG, "GetPartitionColumnStatistics",
        "databaseName=%s,tableName=%s,partitionNames=%s,columnNames=%s",
        databaseName, tableName, partitionNames, columnNames)
        .entrySet().stream().collect(Collectors.toMap(Map.Entry::getKey,
            e -> e.getValue().getStatisticsList(), (e1, e2) -> e1));
  }

  @Override
  public long transformTable(String dbName, String tableName, String definition)
      throws AlluxioStatusException {
    return retryRPC(() -> mClient.transformTable(
        TransformTablePRequest.newBuilder()
            .setDbName(dbName)
            .setTableName(tableName)
            .setDefinition(definition)
            .build()).getJobId(),
        RPC_LOG, "TransformTable", "dbName=%s,tableName=%s,definition=%s",
        dbName, tableName, definition);
  }

  @Override
  public TransformJobInfo getTransformJobInfo(long jobId) throws AlluxioStatusException {
    return retryRPC(() -> mClient.getTransformJobInfo(
        GetTransformJobInfoPRequest.newBuilder()
            .setJobId(jobId)
            .build()).getInfo(0),
        RPC_LOG, "GetTransformJobInfo", "jobId=%d", jobId);
  }

  @Override
  public List<TransformJobInfo> getAllTransformJobInfo() throws AlluxioStatusException {
    return retryRPC(() -> mClient.getTransformJobInfo(
        GetTransformJobInfoPRequest.newBuilder().build()).getInfoList(),
        RPC_LOG, "GetAllTransformJobInfo", "");
  }
}
