lens-driver-es/src/main/java/org/apache/lens/driver/es/ESDriver.java

/**
 * 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
 * <p>
 * http://www.apache.org/licenses/LICENSE-2.0
 * <p>
 * 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.lens.driver.es;

import java.io.IOException;
import java.io.ObjectInput;
import java.io.ObjectOutput;
import java.lang.reflect.Constructor;
import java.lang.reflect.InvocationTargetException;
import java.util.Map;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicInteger;

import org.apache.lens.api.query.QueryHandle;
import org.apache.lens.api.query.QueryPrepareHandle;
import org.apache.lens.cube.metadata.CubeMetastoreClient;
import org.apache.lens.cube.parse.HQLParser;
import org.apache.lens.driver.es.client.ESClient;
import org.apache.lens.driver.es.client.ESResultSet;
import org.apache.lens.driver.es.client.jest.JestClientImpl;
import org.apache.lens.driver.es.translator.ESVisitor;
import org.apache.lens.server.api.LensConfConstants;
import org.apache.lens.server.api.driver.*;
import org.apache.lens.server.api.error.LensException;
import org.apache.lens.server.api.events.LensEventListener;
import org.apache.lens.server.api.query.AbstractQueryContext;
import org.apache.lens.server.api.query.PreparedQueryContext;
import org.apache.lens.server.api.query.QueryContext;
import org.apache.lens.server.api.query.cost.FactPartitionBasedQueryCost;
import org.apache.lens.server.api.query.cost.QueryCost;

import org.apache.commons.lang.Validate;

import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.hive.conf.HiveConf;
import org.apache.hadoop.hive.ql.metadata.HiveException;
import org.apache.hadoop.hive.ql.metadata.Table;
import org.apache.hadoop.hive.ql.parse.ASTNode;
import org.apache.hadoop.hive.ql.parse.HiveParser;
import org.apache.hadoop.hive.ql.session.SessionState;

import org.antlr.runtime.CommonToken;
import org.antlr.runtime.tree.Tree;

import com.google.common.collect.Maps;

import lombok.extern.slf4j.Slf4j;

/**
 * Driver for elastic search
 */
@Slf4j
public class ESDriver extends AbstractLensDriver {

  private static final AtomicInteger THID = new AtomicInteger();
  private static final double STREAMING_PARTITION_COST = 0;
  private static final QueryCost ES_DRIVER_COST = new FactPartitionBasedQueryCost(STREAMING_PARTITION_COST);

  private ESClient esClient;
  private ExecutorService asyncQueryPool;
  private ESDriverConfig config;

  /**
   * States
   */
  private final Map<String, ESQuery> rewrittenQueriesCache = Maps.newConcurrentMap();
  private final Map<QueryHandle, Future<LensResultSet>> resultSetMap = Maps.newConcurrentMap();
  private final Map<QueryHandle, QueryContext> handleContextMap = Maps.newConcurrentMap();

  @Override
  public QueryCost estimate(AbstractQueryContext qctx) {
    return ES_DRIVER_COST;
  }

  @Override
  public DriverQueryPlan explain(final AbstractQueryContext context) throws LensException {
    final ESQuery esQuery = rewrite(context);
    final String jsonExplanation = esClient.explain(esQuery);
    if (jsonExplanation == null) {
      throw new LensException("Explanation failed, empty json was returned");
    }
    return new DriverQueryPlan() {
      @Override
      public String getPlan() {
        return jsonExplanation;
      }

      @Override
      public QueryCost getCost() {
        return ES_DRIVER_COST;
      }
    };
  }

  @Override
  public void prepare(PreparedQueryContext pContext) throws LensException {
    rewrite(pContext);
  }

  @Override
  public DriverQueryPlan explainAndPrepare(PreparedQueryContext pContext) throws LensException {
    prepare(pContext);
    return explain(pContext);
  }

  @Override
  public void closePreparedQuery(QueryPrepareHandle handle) {
    /**
     * Elastic search does not have a concept of prepared query.
     */
  }

  @Override
  public LensResultSet execute(QueryContext context) throws LensException {
    handleContextMap.put(context.getQueryHandle(), context);
    final ESQuery esQuery = rewrite(context);
    final ESResultSet resultSet = esClient.execute(esQuery);
    context.setDriverStatus(DriverQueryStatus.DriverQueryState.SUCCESSFUL);
    handleContextMap.remove(context.getQueryHandle());
    return resultSet;
  }

  @Override
  public void executeAsync(final QueryContext context) {
    handleContextMap.put(context.getQueryHandle(), context);
    final Future<LensResultSet> futureResult
      = asyncQueryPool.submit(new ESQueryExecuteCallable(context, SessionState.get()));
    resultSetMap.put(context.getQueryHandle(), futureResult);
  }

  @Override
  public void updateStatus(QueryContext context) {
    final QueryHandle queryHandle = context.getQueryHandle();
    final Future<LensResultSet> lensResultSetFuture = resultSetMap.get(queryHandle);
    if (lensResultSetFuture == null) {
      context.getDriverStatus().setState(DriverQueryStatus.DriverQueryState.CLOSED);
      context.getDriverStatus().setStatusMessage(queryHandle + " closed");
      context.getDriverStatus().setResultSetAvailable(false);
    } else if (lensResultSetFuture.isDone()) {
      context.getDriverStatus().setState(DriverQueryStatus.DriverQueryState.SUCCESSFUL);
      context.getDriverStatus().setStatusMessage(queryHandle + " successful");
      context.getDriverStatus().setResultSetAvailable(true);
    } else if (lensResultSetFuture.isCancelled()) {
      context.getDriverStatus().setState(DriverQueryStatus.DriverQueryState.CANCELED);
      context.getDriverStatus().setStatusMessage(queryHandle + " cancelled");
      context.getDriverStatus().setResultSetAvailable(false);
    }
  }

  @Override
  public LensResultSet fetchResultSet(QueryContext context) throws LensException {
    try {
      /**
       * removing the result set as soon as the fetch is done
       */
      return resultSetMap.remove(context.getQueryHandle()).get();
    } catch (NullPointerException e) {
      throw new LensException("The results for the query "
        + context.getQueryHandleString()
        + "has already been fetched");
    } catch (InterruptedException | ExecutionException e) {
      throw new LensException("Error fetching result set!", e);
    }
  }

  @Override
  public void closeResultSet(QueryHandle handle) throws LensException {
    if (resultSetMap.containsKey(handle)) {
      resultSetMap.remove(handle);
    }
  }

  @Override
  public boolean cancelQuery(QueryHandle handle) throws LensException {
    try {
      boolean cancelled = resultSetMap.get(handle).cancel(true);
      if (cancelled) {
        handleContextMap.get(handle).setDriverStatus(DriverQueryStatus.DriverQueryState.CANCELED);
      }
      return cancelled;
    } catch (NullPointerException e) {
      throw new LensException("The query does not exist or was already purged", e);
    }
  }

  @Override
  public void closeQuery(QueryHandle handle) throws LensException {
    cancelQuery(handle);
    closeResultSet(handle);
    handleContextMap.remove(handle);
  }

  @Override
  public void close() throws LensException {
    for(QueryHandle handle : resultSetMap.keySet()) {
      try {
        closeQuery(handle);
      } catch (LensException e) {
        log.error("Error while closing query {}", handle.getHandleIdString(), e);
      }
    }
  }

  @Override
  public void registerDriverEventListener(LensEventListener<DriverEvent> driverEventListener) {

  }

  private ESQuery rewrite(AbstractQueryContext context) throws LensException {
    final String key = keyFor(context);
    if (rewrittenQueriesCache.containsKey(key)) {
      return rewrittenQueriesCache.get(key);
    } else {
      final ASTNode rootQueryNode = HQLParser.parseHQL(context.getDriverQuery(this), new HiveConf());
      setIndexAndTypeIfNotPresent(context, rootQueryNode);
      final ESQuery esQuery = ESVisitor.rewrite(config, rootQueryNode);
      rewrittenQueriesCache.put(key, esQuery);
      return esQuery;
    }
  }

  private void setIndexAndTypeIfNotPresent(AbstractQueryContext context, ASTNode rootQueryNode) throws LensException {
    final ASTNode dbSchemaTable = HQLParser.findNodeByPath(
      rootQueryNode,
      HiveParser.TOK_FROM,
      HiveParser.TOK_TABREF,
      HiveParser.TOK_TABNAME);
    try {
      Validate.notNull(dbSchemaTable);
      if (dbSchemaTable.getChildren().size() == 2) {
        /**
         * Index and type is already set here
         */
        return;
      }
      /**
       * Get the table name, check metastore and set index and actual table name
       */
      final Tree firstChild = dbSchemaTable.getChild(0);
      final String lensTable = firstChild.getText();
      final Table tbl = CubeMetastoreClient.getInstance(context.getHiveConf()).getHiveTable(lensTable);
      final String index = tbl.getProperty(LensConfConstants.ES_INDEX_NAME);
      final String type = tbl.getProperty(LensConfConstants.ES_TYPE_NAME);
      Validate.notNull(index, LensConfConstants.ES_INDEX_NAME + " property missing in table definition");
      Validate.notNull(type, LensConfConstants.ES_TYPE_NAME + " property missing in table definition");
      ((ASTNode) firstChild).getToken().setText(type);
      final ASTNode indexIdentifier = new ASTNode(new CommonToken(HiveParser.Identifier, index));
      indexIdentifier.setParent(dbSchemaTable);
      dbSchemaTable.insertChild(0, indexIdentifier);
    } catch (HiveException e) {
      throw new LensException("Error occured when trying to communicate with metastore");
    }
  }

  private String keyFor(AbstractQueryContext context) {
    return String.valueOf(context.getFinalDriverQuery(this)!=null) + ":" + context.getDriverQuery(this);
  }

  ESClient getESClient() {
    return esClient;
  }

  @Override
  public void configure(Configuration conf, String driverType, String driverName) throws LensException {
    super.configure(conf, driverType, driverName);
    config = new ESDriverConfig(getConf());
    Class klass;
    try {
      klass = Class.forName(getConf().get(ESDriverConfig.CLIENT_CLASS_KEY));
      if (klass != null) {
        log.debug("Picked up class {}", klass);
        if (ESClient.class.isAssignableFrom(klass)) {
          final Constructor constructor = klass.getConstructor(ESDriverConfig.class, Configuration.class);
          esClient = (ESClient) constructor.newInstance(config, getConf());
          log.debug("Successfully instantiated es client of type {}", klass);
        }
      } else {
        log.debug("Client class not provided, falling back to the default Jest client");
        esClient = new JestClientImpl(config, conf);
      }
    } catch (ClassNotFoundException
      | NoSuchMethodException
      | InstantiationException
      | IllegalAccessException
      | InvocationTargetException e) {
      log.error("ES driver {} cannot start!", getFullyQualifiedName(), e);
      throw new LensException("Cannot start es driver", e);
    }
    log.info("ES Driver {} configured", getFullyQualifiedName());
    asyncQueryPool = Executors.newCachedThreadPool(new ThreadFactory() {
      @Override
      public Thread newThread(Runnable runnable) {
        Thread th = new Thread(runnable);
        th.setName("lens-driver-es-" + THID.incrementAndGet());
        return th;
      }
    });
  }

  @Override
  public void writeExternal(ObjectOutput out) throws IOException {
    /**
     * This flow could be abstracted out at the driver level
     */
  }

  @Override
  public void readExternal(ObjectInput in) throws IOException, ClassNotFoundException {
    /**
     * This flow could be abstracted out at the driver level
     */
  }

  protected class ESQueryExecuteCallable implements Callable<LensResultSet> {

    private final QueryContext queryContext;
    private final SessionState sessionState;

    public ESQueryExecuteCallable(QueryContext queryContext, SessionState sessionState) {
      this.queryContext = queryContext;
      this.sessionState = sessionState;
    }

    @Override
    public LensResultSet call() throws Exception {
      SessionState.setCurrentSessionState(sessionState);
      return execute(queryContext);
    }
  }

}