package com.continuuity.hive.datasets;

import com.continuuity.api.data.batch.RecordScannable;
import com.continuuity.api.dataset.Dataset;
import com.continuuity.common.conf.Constants;
import com.continuuity.data2.dataset2.DatasetFramework;
import com.continuuity.data2.dataset2.DatasetManagementException;
import com.continuuity.data2.transaction.Transaction;
import com.continuuity.data2.transaction.TransactionAware;
import com.continuuity.hive.context.ConfigurationUtil;
import com.continuuity.hive.context.ContextManager;
import com.continuuity.hive.context.TxnCodec;
import com.google.common.collect.Maps;
import org.apache.hadoop.conf.Configuration;

import java.io.IOException;
import java.lang.reflect.Type;
import java.util.Map;

/**
 * Helps in instantiating a dataset.
 */
public class DatasetAccessor {

  // TODO: this will go away when dataset manager does not return datasets having classloader conflict - REACTOR-276
  private static final Map<String, ClassLoader> DATASET_CLASSLOADERS = Maps.newHashMap();

  public static RecordScannable getRecordScannable(Configuration conf) throws IOException {
    RecordScannable recordScannable = instantiate(conf);

    if (recordScannable instanceof TransactionAware) {
      Transaction tx = ConfigurationUtil.get(conf, Constants.Explore.TX_QUERY_KEY, TxnCodec.INSTANCE);
        ((TransactionAware) recordScannable).startTx(tx);
    }

    return recordScannable;
  }

  public static Type getRecordScannableType(Configuration conf) throws IOException {
    return instantiate(conf).getRecordType();
  }

  private static RecordScannable instantiate(Configuration conf) throws IOException {
    String datasetName = conf.get(Constants.Explore.DATASET_NAME);
    if (datasetName == null) {
      throw new IOException(String.format("Dataset name property %s not defined.", Constants.Explore.DATASET_NAME));
    }

    DatasetFramework framework = ContextManager.getDatasetManager(conf);

    try {
      ClassLoader classLoader = DATASET_CLASSLOADERS.get(datasetName);
      if (classLoader == null) {
        classLoader = conf.getClassLoader();
      }

      Dataset dataset = framework.getDataset(datasetName, classLoader);
      if (dataset != null && !DATASET_CLASSLOADERS.containsKey(datasetName)) {
        DATASET_CLASSLOADERS.put(datasetName, dataset.getClass().getClassLoader());
      }

      if (!(dataset instanceof RecordScannable)) {
        throw new IOException(
          String.format("Dataset %s does not implement RecordScannable, and hence cannot be queried in Hive.",
                        datasetName));
      }

      return (RecordScannable) dataset;
    } catch (DatasetManagementException e) {
      throw new IOException(e);
    }
  }
}

@@ -1010,12 +1010,18 @@
 .new
-Hash
+Concurrent
 Map(
@@ -1944,24 +1944,45 @@
 ger(conf);%0A%0A
+    Dataset dataset;%0A
     try %7B%0A  
@@ -2138,231 +2138,136 @@
     
-%7D%0A%0A      Dataset dataset = framework.getDataset(datasetName, classLoader);%0A      if (dataset != null && !DATASET_CLASSLOADERS.containsKey(datasetName)) %7B%0A        DATASET_CLASSLOADERS.put(datasetName, dataset.getClass().getC
+  dataset = firstLoad(framework, datasetName, classLoader);%0A      %7D else %7B%0A        dataset = framework.getDataset(datasetName, c
 lass
@@ -2272,18 +2272,16 @@
 ssLoader
-()
 );%0A     
@@ -2638,24 +2638,24 @@
 ception(e);%0A
-
     %7D%0A  %7D%0A%7D%0A
@@ -2651,11 +2651,762 @@
    %7D%0A  %7D
+%0A%0A  private static synchronized Dataset firstLoad(DatasetFramework framework, String datasetName, ClassLoader classLoader)%0A    throws DatasetManagementException, IOException %7B%0A    ClassLoader datasetClassLoader = DATASET_CLASSLOADERS.get(datasetName);%0A    if (datasetClassLoader != null) %7B%0A      // Some other call in parallel may have already loaded it, so use the same classlaoder%0A      return framework.getDataset(datasetName, datasetClassLoader);%0A    %7D%0A%0A    // No classloader for dataset exists, load the dataset and save the classloader.%0A    Dataset dataset = framework.getDataset(datasetName, classLoader);%0A    if (dataset != null) %7B%0A      DATASET_CLASSLOADERS.put(datasetName, dataset.getClass().getClassLoader());%0A    %7D%0A    return dataset;%0A  %7D
 %0A%7D%0A

Allow concurrent access to dataset classloader map.