alluxio-master-2.x/table/server/underdb/glue/src/main/java/alluxio/table/under/glue/GlueDatabase.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.table.under.glue;

import alluxio.AlluxioURI;
import alluxio.exception.AlluxioException;
import alluxio.exception.status.NotFoundException;
import alluxio.grpc.table.ColumnStatisticsInfo;
import alluxio.grpc.table.Layout;
import alluxio.grpc.table.layout.hive.PartitionInfo;
import alluxio.master.table.DatabaseInfo;
import alluxio.table.common.UdbPartition;
import alluxio.table.common.layout.HiveLayout;
import alluxio.table.common.udb.PathTranslator;
import alluxio.table.common.udb.UdbBypassSpec;
import alluxio.table.common.udb.UdbConfiguration;
import alluxio.table.common.udb.UdbContext;
import alluxio.table.common.udb.UdbTable;
import alluxio.table.common.udb.UdbUtils;
import alluxio.table.common.udb.UnderDatabase;
import alluxio.util.io.PathUtils;

import com.amazonaws.AmazonClientException;
import com.amazonaws.ClientConfiguration;
import com.amazonaws.Protocol;
import com.amazonaws.auth.AWSCredentialsProvider;
import com.amazonaws.auth.AWSStaticCredentialsProvider;
import com.amazonaws.auth.BasicAWSCredentials;
import com.amazonaws.auth.DefaultAWSCredentialsProviderChain;
import com.amazonaws.services.glue.AWSGlueAsync;
import com.amazonaws.services.glue.AWSGlueAsyncClientBuilder;
import com.amazonaws.services.glue.model.AWSGlueException;
import com.amazonaws.services.glue.model.Column;
import com.amazonaws.services.glue.model.Database;
import com.amazonaws.services.glue.model.EntityNotFoundException;
import com.amazonaws.services.glue.model.GetColumnStatisticsForPartitionRequest;
import com.amazonaws.services.glue.model.GetColumnStatisticsForTableRequest;
import com.amazonaws.services.glue.model.GetDatabaseRequest;
import com.amazonaws.services.glue.model.GetDatabaseResult;
import com.amazonaws.services.glue.model.GetPartitionsRequest;
import com.amazonaws.services.glue.model.GetPartitionsResult;
import com.amazonaws.services.glue.model.GetTableRequest;
import com.amazonaws.services.glue.model.GetTablesRequest;
import com.amazonaws.services.glue.model.GetTablesResult;
import com.amazonaws.services.glue.model.GlueEncryptionException;
import com.amazonaws.services.glue.model.Partition;
import com.amazonaws.services.glue.model.Table;
import com.amazonaws.services.glue.model.ValidationException;
import com.google.common.annotations.VisibleForTesting;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

/**
 * Glue database implementation.
 */
public class GlueDatabase implements UnderDatabase {
  private static final Logger LOG = LoggerFactory.getLogger(GlueDatabase.class);

  private final UdbContext mUdbContext;
  private final AWSGlueAsync mGlueClient;
  private final UdbConfiguration mGlueConfiguration;
  /** the name of the glue db. */
  private final String mGlueDbName;

  /** the owner name of glue database, which is an fake place holder. */
  private final String mOwnerName = "PUBLIC_OWNER";
  private final alluxio.grpc.table.PrincipalType mOwnerType = alluxio.grpc.table.PrincipalType.ROLE;

  @VisibleForTesting
  protected GlueDatabase(UdbContext udbContext, UdbConfiguration glueConfig, String glueDbName) {
    mUdbContext = udbContext;
    mGlueConfiguration = glueConfig;
    mGlueClient = createAsyncGlueClient(glueConfig);
    mGlueDbName = glueDbName;
  }

  /**
   * Create an instance of the Glue database UDB.
   *
   * @param udbContext the db context
   * @param configuration the configuration
   * @return the new instance
   */
  public static GlueDatabase create(UdbContext udbContext, UdbConfiguration configuration) {
    String glueDbName = udbContext.getUdbDbName();
    if (glueDbName == null || glueDbName.isEmpty()) {
      throw new IllegalArgumentException(
          "Glue database name cannot be empty: " + glueDbName);
    } else if (configuration.get(Property.GLUE_REGION) == null) {
      throw new IllegalArgumentException("GlueUdb Error: Please setup aws region.");
    }

    return new GlueDatabase(udbContext, configuration, glueDbName);
  }

  @Override
  public UdbContext getUdbContext() {
    return mUdbContext;
  }

  @Override
  public DatabaseInfo getDatabaseInfo() throws IOException {
    try {
      GetDatabaseRequest dbRequest = new GetDatabaseRequest()
          .withCatalogId(mGlueConfiguration.get(Property.CATALOG_ID))
          .withName(mGlueDbName);
      GetDatabaseResult dbResult = mGlueClient.getDatabase(dbRequest);
      Database glueDatabase = dbResult.getDatabase();
      // Glue database location, description and parameters could be null
      String glueDbLocation = glueDatabase.getLocationUri() == null
          ? "" : glueDatabase.getLocationUri();
      String glueDbDescription = glueDatabase.getDescription() == null
          ? "" : glueDatabase.getDescription();
      Map<String, String> glueParameters = new HashMap<>();
      if (glueDatabase.getParameters() != null) {
        glueParameters.putAll(glueDatabase.getParameters());
      }
      return new DatabaseInfo(
          glueDbLocation,
          mOwnerName,
          mOwnerType,
          glueDbDescription,
          glueParameters);
    } catch (EntityNotFoundException e) {
      throw new IOException("Cannot find glue database: " + mGlueDbName
          + "Catalog ID: " + mGlueConfiguration.get(Property.CATALOG_ID)
          + ". " + e.getMessage(), e);
    }
  }

  /**
   * This method allows user to test udb glue client with remote glue server.
   *
   * @param config udbconfiguration
   * @return glue async client
   */
  protected static AWSGlueAsync createAsyncGlueClient(UdbConfiguration config) {
    ClientConfiguration clientConfig = new ClientConfiguration()
        .withMaxConnections(config.getInt(Property.MAX_GLUE_CONNECTION));

    if (!config.get(Property.AWS_PROXY_HOST).isEmpty()) {
      clientConfig.withProxyProtocol(getProtocol(config.get(Property.AWS_PROXY_PROTOCOL)))
          .withProxyHost(config.get(Property.AWS_PROXY_HOST))
          .withProxyPort(config.getInt(Property.AWS_PROXY_PORT))
          .withProxyUsername(config.get(Property.AWS_PROXY_USER_NAME))
          .withProxyPassword(config.get(Property.AWS_PROXY_PASSWORD));
    }

    AWSGlueAsyncClientBuilder asyncClientBuilder = AWSGlueAsyncClientBuilder
        .standard()
        .withClientConfiguration(clientConfig);

    if (!config.get(Property.GLUE_REGION).isEmpty()) {
      LOG.info("Set Glue region: {}.", config.get(Property.GLUE_REGION));
      asyncClientBuilder.setRegion(config.get(Property.GLUE_REGION));
    } else {
      LOG.warn("GlueDatabase: Please setup the AWS region.");
    }

    asyncClientBuilder.setCredentials(getAWSCredentialsProvider(config));

    return asyncClientBuilder.build();
  }

  private static AWSCredentialsProvider getAWSCredentialsProvider(UdbConfiguration config) {
    //TODO(shouwei): add compelete authentication method for glue udb
    if (!config.get(Property.AWS_GLUE_ACCESS_KEY).isEmpty()
        && !config.get(Property.AWS_GLUE_SECRET_KEY).isEmpty()) {
      return new AWSStaticCredentialsProvider(
          new BasicAWSCredentials(
          config.get(Property.AWS_GLUE_ACCESS_KEY),
          config.get(Property.AWS_GLUE_SECRET_KEY)));
    }
    return DefaultAWSCredentialsProviderChain.getInstance();
  }

  private static Protocol getProtocol(String protocol) {
    if (protocol.equals("HTTP")) {
      return Protocol.HTTP;
    } else if (protocol.equals("HTTPS")) {
      return Protocol.HTTPS;
    } else {
      LOG.warn("Invalid protocol type {}."
          + "Avaiable proxy protocol type HTTP and HTTPS.", protocol);
    }
    return null;
  }

  @Override
  public String getType() {
    return GlueDatabaseFactory.TYPE;
  }

  @Override
  public String getName() {
    return mGlueDbName;
  }

  @Override
  public List<String> getTableNames() throws IOException {
    try {
      String nextToken = null;
      List<String> tableNames = new ArrayList<>();
      do {
        GetTablesRequest tablesRequest =
            new GetTablesRequest()
                .withCatalogId(mGlueConfiguration.get(Property.CATALOG_ID))
                .withDatabaseName(mGlueDbName)
                .withNextToken(nextToken);
        GetTablesResult tablesResult = mGlueClient.getTables(tablesRequest);
        tablesResult.getTableList().forEach(table -> tableNames.add(table.getName()));
        nextToken = tablesResult.getNextToken();
      } while (nextToken != null);
      return tableNames;
    } catch (EntityNotFoundException e) {
      throw new IOException("Failed to get glue tables: " + e.getMessage()
          + " in Database: " + mGlueDbName
          + "; with Catalog ID: " + mGlueConfiguration.get(Property.CATALOG_ID) + ".", e);
    }
  }

  @VisibleForTesting
  private PathTranslator mountAlluxioPaths(Table table, List<Partition> partitions,
      UdbBypassSpec bypassSpec)
      throws IOException {
    String tableName = table.getName();
    AlluxioURI ufsUri;
    AlluxioURI alluxioUri = mUdbContext.getTableLocation(tableName);
    String glueUfsUri = table.getStorageDescriptor().getLocation();

    try {
      PathTranslator pathTranslator = new PathTranslator();
      if (bypassSpec.hasFullTable(tableName)) {
        pathTranslator.addMapping(glueUfsUri, glueUfsUri);
        return pathTranslator;
      }
      ufsUri = new AlluxioURI(table.getStorageDescriptor().getLocation());
      pathTranslator.addMapping(
          UdbUtils.mountAlluxioPath(
              tableName,
              ufsUri,
              alluxioUri,
              mUdbContext,
              mGlueConfiguration),
          glueUfsUri);

      for (Partition partition : partitions) {
        AlluxioURI partitionUri;
        String partitionName;
        if (partition.getStorageDescriptor() != null
            && partition.getStorageDescriptor().getLocation() != null
            && ufsUri.isAncestorOf(
            partitionUri = new AlluxioURI(
                partition.getStorageDescriptor().getLocation()))) {
          glueUfsUri = partition.getStorageDescriptor().getLocation();
          partitionName = partition.getValues().toString();
          try {
            partitionName = GlueUtils.makePartitionName(
                table.getPartitionKeys(),
                partition.getValues());
          } catch (IOException e) {
            LOG.warn("Error making partition name for table {},"
                    + " partition {} in database {} with CatalogID {}.",
                tableName,
                partition.getValues().toString(),
                mGlueDbName,
                mGlueConfiguration.get(Property.CATALOG_ID));
          }
          if (bypassSpec.hasPartition(tableName, partitionName)) {
            pathTranslator.addMapping(partitionUri.getPath(), partitionUri.getPath());
            continue;
          }
          alluxioUri = new AlluxioURI(
              PathUtils.concatPath(
                  mUdbContext.getTableLocation(tableName).getPath(),
                  partitionName));
          // mount partition path if it is not already mounted as part of the table path mount
          pathTranslator
              .addMapping(
                  UdbUtils.mountAlluxioPath(
                      tableName,
                      partitionUri,
                      alluxioUri,
                      mUdbContext,
                      mGlueConfiguration),
                  glueUfsUri);
        }
      }
      return pathTranslator;
    } catch (AlluxioException e) {
      throw new IOException(
          "Failed to mount table location. tableName: " + tableName
              + " glueUfsLocation: " + glueUfsUri
              + " AlluxioLocation: " + alluxioUri + " error: " + e.getMessage(), e);
    }
  }

  private List<ColumnStatisticsInfo> getTableColumnStatistics(String dbName, String tableName,
      GetColumnStatisticsForTableRequest getColumnStatisticsForTableRequest) {
    // TODO(shouwei): Add Async support for table column statistics
    try {
      return getClient().getColumnStatisticsForTable(getColumnStatisticsForTableRequest)
          .getColumnStatisticsList().stream().map(GlueUtils::toProto).collect(Collectors.toList());
    } catch (AmazonClientException e) {
      LOG.warn("Cannot get the table column statistics info for table {}.{} with error {}.",
          dbName, tableName, e.toString());
    }
    return Collections.emptyList();
  }

  private List<ColumnStatisticsInfo> getPartitionColumnStatistics(String dbName, String tableName,
      GetColumnStatisticsForPartitionRequest getColumnStatisticsForPartitionRequest) {
    // TODO(shouwei): Add Async support for partition column statistics
    try {
      List<ColumnStatisticsInfo> partColumnStatistic = getClient()
          .getColumnStatisticsForPartition(getColumnStatisticsForPartitionRequest)
          .getColumnStatisticsList().stream().map(GlueUtils::toProto).collect(Collectors.toList());
      return partColumnStatistic;
    } catch (AmazonClientException e) {
      LOG.warn("Cannot get the partition column statistics info for table {}.{} with error {}.",
          dbName, tableName, e.toString());
    }
    return Collections.emptyList();
  }

  @Override
  public UdbTable getTable(String tableName, UdbBypassSpec bypassSpec) throws IOException {
    Table table;
    List<Partition> partitions;
    try {
      GetTableRequest tableRequest = new GetTableRequest()
          .withCatalogId(mGlueConfiguration.get(Property.CATALOG_ID))
          .withDatabaseName(mGlueDbName)
          .withName(tableName);
      table = getClient().getTable(tableRequest).getTable();

      partitions = batchGetPartitions(getClient(), tableName);
      PathTranslator pathTranslator = mountAlluxioPaths(table, partitions, bypassSpec);

      List<Column> partitionColumns;
      if (table.getPartitionKeys() == null) {
        partitionColumns = Collections.emptyList();
      } else {
        partitionColumns = table.getPartitionKeys();
      }

      // Get table parameters
      Map<String, String> tableParameters = table.getParameters() == null
          ? Collections.emptyMap() : table.getParameters();

      // Get column statistics info for table
      List<String> columnNames = table.getStorageDescriptor()
          .getColumns().stream().map(Column::getName).collect(Collectors.toList());
      GetColumnStatisticsForTableRequest getColumnStatisticsForTableRequest =
          new GetColumnStatisticsForTableRequest()
              .withCatalogId(mGlueConfiguration.get(Property.CATALOG_ID))
              .withDatabaseName(mGlueDbName)
              .withTableName(tableName)
              .withColumnNames(columnNames);
      List<ColumnStatisticsInfo> columnStatisticsTableData = new ArrayList<>();
      if (mGlueConfiguration.getBoolean(Property.TABLE_COLUMN_STATISTICS_ENABLE)) {
        columnStatisticsTableData = getTableColumnStatistics(
            mGlueDbName, tableName, getColumnStatisticsForTableRequest);
      }

      // Get column statistics info for partitions
      // potential expensive call
      Map<String, List<ColumnStatisticsInfo>> statsMap = new HashMap<>();
      if (mGlueConfiguration.getBoolean(Property.PARTITION_COLUMN_STATISTICS_ENABLE)) {
        for (Partition partition : partitions) {
          List<String> partitionValue = partition.getValues();
          if (partitionValue != null) {
            GetColumnStatisticsForPartitionRequest getColumnStatisticsForPartitionRequest =
                new GetColumnStatisticsForPartitionRequest()
                    .withCatalogId(mGlueConfiguration.get(Property.CATALOG_ID))
                    .withDatabaseName(mGlueDbName)
                    .withTableName(tableName)
                    .withColumnNames(columnNames)
                    .withPartitionValues(partitionValue);
            String partName = GlueUtils.makePartitionName(partitionColumns, partition.getValues());
            statsMap.put(partName, getPartitionColumnStatistics(
                mGlueDbName, tableName, getColumnStatisticsForPartitionRequest));
          }
        }
      }

      PartitionInfo partitionInfo = PartitionInfo.newBuilder()
          // Database name is not required for glue table, use mGlueDbName
          .setDbName(mGlueDbName)
          .setTableName(tableName)
          .addAllDataCols(GlueUtils.toProto(table.getStorageDescriptor().getColumns()))
          .setStorage(GlueUtils.toProto(table.getStorageDescriptor(), pathTranslator))
          .putAllParameters(tableParameters)
          .build();

      Layout layout = Layout.newBuilder()
          .setLayoutType(HiveLayout.TYPE)
          .setLayoutData(partitionInfo.toByteString())
          .build();

      List<UdbPartition> udbPartitions = new ArrayList<>();
      if (partitionColumns.isEmpty()) {
        PartitionInfo.Builder partitionInfoBuilder = PartitionInfo.newBuilder()
            .setDbName(mGlueDbName)
            .setTableName(tableName)
            .addAllDataCols(GlueUtils.toProto(table.getStorageDescriptor().getColumns()))
            .setStorage(GlueUtils.toProto(table.getStorageDescriptor(), pathTranslator))
            .setPartitionName(tableName)
            .putAllParameters(tableParameters);
        udbPartitions.add(new GluePartition(
            new HiveLayout(partitionInfoBuilder.build(), Collections.emptyList())));
      } else {
        for (Partition partition : partitions) {
          String partName = GlueUtils.makePartitionName(partitionColumns, partition.getValues());
          PartitionInfo.Builder partitionInfoBuilder = PartitionInfo.newBuilder()
              .setDbName(mGlueDbName)
              .setTableName(tableName)
              .addAllDataCols(GlueUtils.toProto(partition.getStorageDescriptor().getColumns()))
              .setStorage(GlueUtils.toProto(partition.getStorageDescriptor(), pathTranslator))
              .setPartitionName(partName)
              .putAllParameters(partition.getParameters() == null
                  ? Collections.emptyMap() : partition.getParameters());
          if (partition.getValues() != null) {
            partitionInfoBuilder.addAllValues(partition.getValues());
          }
          udbPartitions.add(new GluePartition(new HiveLayout(partitionInfoBuilder.build(),
              statsMap.getOrDefault(partName, Collections.emptyList()))));
        }
      }

      return new GlueTable(this,
          pathTranslator,
          tableName,
          GlueUtils.toProtoSchema(table.getStorageDescriptor().getColumns()),
          columnStatisticsTableData,
          // Glue does not provide FieldSchema from API directly
          // Get FieldSchema from partition keys
          GlueUtils.toProto(table.getPartitionKeys()),
          udbPartitions,
          layout,
          table);
    } catch (EntityNotFoundException e) {
      throw new NotFoundException("Table " + tableName
          + " does not exist in Database: " + mGlueDbName
          + "; Catalog ID: " + mGlueConfiguration.get(Property.CATALOG_ID)
          + ".", e);
    } catch (ValidationException e) {
      throw new IOException("Failed to get table: " + tableName
          + " in Database: " + mGlueDbName
          + "; Catalog ID: " + mGlueConfiguration.get(Property.CATALOG_ID)
          + " with validation error: " + e.getMessage(), e);
    } catch (GlueEncryptionException e) {
      throw new IOException("Failed to get table: " + tableName
          + " in Database: " + mGlueDbName
          + "; Catalog ID: " + mGlueConfiguration.get(Property.CATALOG_ID)
          + " error: " + e.getMessage(), e);
    }
  }

  private List<Partition> batchGetPartitions(AWSGlueAsync glueClient, String tableName)
      throws IOException {
    // TODO(shouwei): make getPartition multi-thread to accelerate the large table fetching
    List<Partition> partitions = new ArrayList<>();
    String nextToken = null;
    try {
      do {
        GetPartitionsRequest getPartitionsRequest =
            new GetPartitionsRequest()
                .withCatalogId(mGlueConfiguration.get(Property.CATALOG_ID))
                .withDatabaseName(mGlueDbName)
                .withTableName(tableName)
                .withMaxResults(mGlueConfiguration.getInt(Property.MAX_GLUE_FETCH_PARTITIONS))
                .withNextToken(nextToken);
        GetPartitionsResult getPartitionsResult = glueClient.getPartitions(getPartitionsRequest);
        partitions.addAll(getPartitionsResult.getPartitions());
        nextToken = getPartitionsResult.getNextToken();
        LOG.debug("Glue table {}.{} adding {} batch partitions with total {} partitions.",
            mGlueDbName, tableName, getPartitionsResult.getPartitions().size(), partitions.size());
      } while (nextToken != null);

      if (partitions != null) {
        LOG.info("Glue table {}.{} has {} partitions.",
            mGlueDbName, tableName, partitions.size());
        if (LOG.isDebugEnabled()) {
          partitions.stream().forEach(partition ->
              LOG.debug("Glue table {}.{} with partition: {}.",
                  partition.getDatabaseName(), tableName, partition));
        }
      }
      return partitions;
    } catch (AWSGlueException e) {
      throw new IOException("Cannot get partition information for table: " + tableName
          + " in Database: " + mGlueDbName
          + "; Catalog ID: " + mGlueConfiguration.get(Property.CATALOG_ID)
          + ". error: " + e.getMessage(), e);
    }
  }

  /**
   * Get Glue Client.
   *
   * @return async glue client
   */
  public AWSGlueAsync getClient() {
    return mGlueClient;
  }
}
