alluxio-master-2.x/integration/tools/hms/src/main/java/alluxio/cli/hms/TableValidationTask.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.cli.hms;

import alluxio.cli.ValidationTaskResult;
import alluxio.cli.ValidationUtils;
import alluxio.collections.Pair;
import alluxio.util.ExceptionUtils;

import org.apache.hadoop.hive.metastore.IMetaStoreClient;
import org.apache.hadoop.hive.metastore.api.FieldSchema;
import org.apache.hadoop.hive.metastore.api.InvalidOperationException;
import org.apache.hadoop.hive.metastore.api.UnknownDBException;
import org.apache.hadoop.hive.metastore.api.UnknownTableException;

import java.util.Arrays;
import java.util.List;
import java.util.stream.Collectors;

/**
 * A task which validates table information and schemas can be obtained from the Hive Metastore.
 */
public class TableValidationTask extends MetastoreValidationTask<IMetaStoreClient, String> {

  // The maximum number of table objects that this test will get.
  // Used to avoid issuing too many calls to the hive metastore
  // which may need a long time based on network conditions
  private static final int GET_TABLE_OBJECT_THRESHOLD = 5;

  private final String mTables;
  private final String mDatabase;

  /**
   * Create a new metastore validation task.
   *
   * @param database the database containing the desired tables
   * @param tableNames a single table name, or comma-delimited list of table names to check for;
   *                   If null or empty, retrieves all tables in the database
   * @param prereq the pre-requisite task supplying the metastore client
   */
  public TableValidationTask(String database, String tableNames,
      MetastoreValidationTask<?, IMetaStoreClient> prereq) {
    super(prereq);
    mTables = tableNames;
    mDatabase = database;
  }

  @Override
  Pair<ValidationTaskResult, String> getValidationWithResult() throws InterruptedException {
    if (mInputTask == null) {
      return new Pair<>(new ValidationTaskResult(ValidationUtils.State.FAILED, getName(),
          "Did not have a pre-requisite task to obtain HMS client for table check", ""), null);
    }
    Pair<ValidationTaskResult, IMetaStoreClient> res = mInputTask.getValidationWithResult();
    if (res.getFirst().getState() != ValidationUtils.State.OK) {
      return new Pair<>(res.getFirst(), null);
    }
    IMetaStoreClient client = res.getSecond();
    try {
      if (mTables != null && !mTables.isEmpty()) {
        return new Pair<>(getTableSchemaTest(client, Arrays.asList(mTables.split(","))), null);
      } else {
        return new Pair<>(getAllTableInfoTest(client), null);
      }
    } finally {
      client.close();
    }
  }

  private ValidationTaskResult getAllTableInfoTest(IMetaStoreClient client) {
    List<String> tables;
    String testName = "GetAllTables";
    try {
      tables = client.getAllTables(mDatabase);
    } catch (Throwable t) {
      return addThrowableWarning(testName, t, "Database: " + mDatabase);
    }
    if (tables.size() == 0) {
      return new ValidationTaskResult(ValidationUtils.State.OK, testName,
          String.format("Database (name: %s, tables: %s)", mDatabase, tables.toString()), "");
    }
    if (tables.size() > GET_TABLE_OBJECT_THRESHOLD) {
      tables = tables.subList(0, GET_TABLE_OBJECT_THRESHOLD);
    }
    return getTableSchemaTest(client, tables);
  }

  private ValidationTaskResult getTableSchemaTest(IMetaStoreClient client,
      List<String> tableNames) {
    try {
      StringBuilder tableFieldsOutput = new StringBuilder();
      for (String table : tableNames) {
        String tableFields = client.getSchema(mDatabase, table).stream()
            .map(FieldSchema::getName).collect(Collectors.joining(","));
        tableFieldsOutput.append(String
            .format("Table (name: %s, fields: %s)%n", table, tableFields));
      }
      return new ValidationTaskResult(ValidationUtils.State.OK, getName(),
          tableFieldsOutput.toString(), "");
    } catch (Throwable t) {
      return addThrowableWarning(getName(), t, "tables: " + String.join(",", tableNames));
    }
  }

  private ValidationTaskResult addThrowableWarning(String opName, Throwable t, String opTarget) {
    ValidationTaskResult taskResult = new ValidationTaskResult()
        .setState(ValidationUtils.State.WARNING).setName(opName)
        .setOutput(ExceptionUtils.asPlainText(t));
    if (t instanceof InvalidOperationException) {
      taskResult.setAdvice(opName + " is invalid");
    } else if (t instanceof UnknownDBException) {
      taskResult.setAdvice("Please make sure the given database name is valid "
          + "and existing in the target hive metastore");
    } else if (t instanceof UnknownTableException) {
      taskResult.setAdvice("Please make sure the given table names are valid "
          + "and existing in the target hive metastore");
    } else {
      taskResult.setAdvice(String.format("Failed to run %s (%s)", opName, opTarget));
    }
    return taskResult;
  }

  @Override
  public String getName() {
    return "TableValidationCheckTask";
  }
}
