/*
    * Licensed to the Apache Software Foundation (ASF) under one
    * or more contributor license agreements.  See the NOTICE file
    * distributed with this work for additional information
    * regarding copyright ownership.  The ASF licenses this file
    * to you 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
    *
   *     http://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 org.apache.hcatalog.mapreduce;
  
  import org.apache.giraph.io.hcatalog.GiraphHCatInputFormat;
  import org.apache.hadoop.conf.Configuration;
  import org.apache.hadoop.hive.conf.HiveConf;
  import org.apache.hadoop.hive.metastore.HiveMetaStoreClient;
  import org.apache.hadoop.hive.metastore.api.MetaException;
  import org.apache.hadoop.hive.metastore.api.NoSuchObjectException;
  import org.apache.hadoop.hive.metastore.api.Partition;
  import org.apache.hadoop.hive.metastore.api.StorageDescriptor;
  import org.apache.hadoop.hive.ql.metadata.HiveException;
  import org.apache.hadoop.hive.ql.metadata.Table;
  import org.apache.hadoop.mapreduce.InputSplit;
  import org.apache.hadoop.mapreduce.RecordReader;
  import org.apache.hcatalog.common.ErrorType;
  import org.apache.hcatalog.common.HCatException;
  import org.apache.hcatalog.common.HCatUtil;
  import org.apache.hcatalog.data.schema.HCatSchema;
  import org.apache.thrift.TException;
  
  import java.io.IOException;
  import java.util.ArrayList;
  import java.util.HashMap;
  import java.util.List;
  import java.util.Map;
  import java.util.Properties;
  
  /**
   * Utility methods copied from HCatalog because of visibility restrictions.
   */
  public class HCatUtils {
    /**
     * Don't instantiate.
     */
    private HCatUtils() { }
  
    /**
     * Returns the given InputJobInfo after populating with data queried from the
     * metadata service.
     *
     * @param conf Configuration
     * @param inputJobInfo Input job info
     * @return Populated input job info
     * @throws IOException
     */
    public static InputJobInfo getInputJobInfo(
        Configuration conf, InputJobInfo inputJobInfo)
      throws IOException {
      HiveMetaStoreClient client = null;
      HiveConf hiveConf;
      try {
        if (conf != null) {
          hiveConf = HCatUtil.getHiveConf(conf);
        } else {
          hiveConf = new HiveConf(GiraphHCatInputFormat.class);
        }
        client = HCatUtil.getHiveClient(hiveConf);
        Table table = HCatUtil.getTable(client, inputJobInfo.getDatabaseName(),
            inputJobInfo.getTableName());
  
        List<PartInfo> partInfoList = new ArrayList<PartInfo>();
  
        inputJobInfo.setTableInfo(HCatTableInfo.valueOf(table.getTTable()));
        if (table.getPartitionKeys().size() != 0) {
          // Partitioned table
          List<Partition> parts = client.listPartitionsByFilter(
              inputJobInfo.getDatabaseName(),
              inputJobInfo.getTableName(),
              inputJobInfo.getFilter(),
              (short) -1);
  
          if (parts != null) {
            // Default to 100,000 partitions if hive.metastore.maxpartition is not
            // defined
            int maxPart = hiveConf.getInt("hcat.metastore.maxpartitions", 100000);
            if (parts.size() > maxPart) {
              throw new HCatException(ErrorType.ERROR_EXCEED_MAXPART,
                  "total number of partitions is " + parts.size());
            }
  
            // Populate partition info
           for (Partition ptn : parts) {
             HCatSchema schema = HCatUtil.extractSchema(
                 new org.apache.hadoop.hive.ql.metadata.Partition(table, ptn));
             PartInfo partInfo = extractPartInfo(schema, ptn.getSd(),
                 ptn.getParameters(), conf, inputJobInfo);
             partInfo.setPartitionValues(InternalUtil.createPtnKeyValueMap(table,
                 ptn));
             partInfoList.add(partInfo);
           }
         }
       } else {
         // Non partitioned table
         HCatSchema schema = HCatUtil.extractSchema(table);
         PartInfo partInfo = extractPartInfo(schema, table.getTTable().getSd(),
             table.getParameters(), conf, inputJobInfo);
         partInfo.setPartitionValues(new HashMap<String, String>());
         partInfoList.add(partInfo);
       }
       inputJobInfo.setPartitions(partInfoList);
     } catch (MetaException e) {
       throw new IOException("Got MetaException", e);
     } catch (NoSuchObjectException e) {
       throw new IOException("Got NoSuchObjectException", e);
     } catch (TException e) {
       throw new IOException("Got TException", e);
     } catch (HiveException e) {
       throw new IOException("Got HiveException", e);
     } finally {
       HCatUtil.closeHiveClientQuietly(client);
     }
     return inputJobInfo;
   }
 
   /**
    * Extract partition info.
    *
    * @param schema Table schema
    * @param sd Storage descriptor
    * @param parameters Parameters
    * @param conf Configuration
    * @param inputJobInfo Input job info
    * @return Partition info
    * @throws IOException
    */
   private static PartInfo extractPartInfo(
       HCatSchema schema, StorageDescriptor sd, Map<String, String> parameters,
       Configuration conf, InputJobInfo inputJobInfo) throws IOException {
     StorerInfo storerInfo = InternalUtil.extractStorerInfo(sd, parameters);
 
     Properties hcatProperties = new Properties();
     HCatStorageHandler storageHandler = HCatUtil.getStorageHandler(conf,
         storerInfo);
 
     // Copy the properties from storageHandler to jobProperties
     Map<String, String> jobProperties =
         HCatUtil.getInputJobProperties(storageHandler, inputJobInfo);
 
     for (Map.Entry<String, String> param : parameters.entrySet()) {
       hcatProperties.put(param.getKey(), param.getValue());
     }
 
     return new PartInfo(schema, storageHandler, sd.getLocation(),
         hcatProperties, jobProperties, inputJobInfo.getTableInfo());
   }
 
   /**
    * Create a new {@link HCatRecordReader}.
    *
    * @param storageHandler Storage handler
    * @param valuesNotInDataCols Values not in data columns
    * @return Record reader
    */
   public static RecordReader newHCatReader(
       HCatStorageHandler storageHandler,
       Map<String, String> valuesNotInDataCols) {
     return new HCatRecordReader(storageHandler, valuesNotInDataCols);
   }
 
   /**
    * Cast an {@link InputSplit} to {@link HCatSplit}.
    *
    * @param split Input split
    * @return {@link HCatSplit}
    * @throws IOException
    */
   public static HCatSplit castToHCatSplit(InputSplit split)
     throws IOException {
     return InternalUtil.castToHCatSplit(split);
   }
 }