You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@lens.apache.org by am...@apache.org on 2016/04/28 06:07:02 UTC

lens git commit: LENS-1013 : Add postDriverSelection to DriverQueryHook

Repository: lens
Updated Branches:
  refs/heads/master 6995962fd -> 9b7541bcb


LENS-1013 : Add postDriverSelection to DriverQueryHook


Project: http://git-wip-us.apache.org/repos/asf/lens/repo
Commit: http://git-wip-us.apache.org/repos/asf/lens/commit/9b7541bc
Tree: http://git-wip-us.apache.org/repos/asf/lens/tree/9b7541bc
Diff: http://git-wip-us.apache.org/repos/asf/lens/diff/9b7541bc

Branch: refs/heads/master
Commit: 9b7541bcbcc33faf3cbde71452f334783795ea27
Parents: 6995962
Author: Puneet Gupta <pu...@gmail.com>
Authored: Thu Apr 28 09:36:38 2016 +0530
Committer: Amareshwari Sriramadasu <am...@apache.org>
Committed: Thu Apr 28 09:36:38 2016 +0530

----------------------------------------------------------------------
 .../org/apache/lens/driver/hive/HiveDriver.java |  6 ++-
 .../apache/lens/driver/hive/TestHiveDriver.java |  3 --
 .../lens/driver/hive/TestRemoteHiveDriver.java  |  3 --
 .../org/apache/lens/driver/jdbc/JDBCDriver.java |  6 ++-
 .../apache/lens/driver/jdbc/TestJdbcDriver.java |  3 --
 .../server/api/driver/AbstractLensDriver.java   |  7 +++
 .../lens/server/api/driver/DriverQueryHook.java | 39 ++++++++------
 .../lens/server/api/driver/LensDriver.java      |  6 +++
 .../server/api/driver/NoOpDriverQueryHook.java  | 14 +++--
 .../server/api/user/MockDriverQueryHook.java    | 34 +++++++++---
 .../server/query/QueryExecutionServiceImpl.java | 18 ++++---
 .../org/apache/lens/server/LensJerseyTest.java  | 11 +++-
 .../apache/lens/server/TestServerRestart.java   | 55 +++++++++++++++++---
 .../lens/server/query/TestQueryService.java     |  7 ++-
 .../drivers/hive/hive1/hivedriver-site.xml      |  7 +++
 .../drivers/hive/hive2/hivedriver-site.xml      |  7 +++
 16 files changed, 176 insertions(+), 50 deletions(-)
----------------------------------------------------------------------


http://git-wip-us.apache.org/repos/asf/lens/blob/9b7541bc/lens-driver-hive/src/main/java/org/apache/lens/driver/hive/HiveDriver.java
----------------------------------------------------------------------
diff --git a/lens-driver-hive/src/main/java/org/apache/lens/driver/hive/HiveDriver.java b/lens-driver-hive/src/main/java/org/apache/lens/driver/hive/HiveDriver.java
index 6eec3f7..19a010e 100644
--- a/lens-driver-hive/src/main/java/org/apache/lens/driver/hive/HiveDriver.java
+++ b/lens-driver-hive/src/main/java/org/apache/lens/driver/hive/HiveDriver.java
@@ -559,7 +559,6 @@ public class HiveDriver extends AbstractLensDriver {
       Configuration qdconf = ctx.getDriverConf(this);
       qdconf.set("mapred.job.name", ctx.getQueryHandle().toString());
       decidePriority(ctx);
-      queryHook.preLaunch(ctx);
       SessionHandle sessionHandle = getSession(ctx);
       OperationHandle op = getClient().executeStatementAsync(sessionHandle, ctx.getSelectedDriverQuery(),
         qdconf.getValByRegex(".*"));
@@ -1370,4 +1369,9 @@ public class HiveDriver extends AbstractLensDriver {
   public boolean hasLensSession(LensSessionHandle session) {
     return lensToHiveSession.containsKey(session.getPublicId().toString());
   }
+
+  @Override
+  public DriverQueryHook getQueryHook() {
+    return queryHook;
+  }
 }

http://git-wip-us.apache.org/repos/asf/lens/blob/9b7541bc/lens-driver-hive/src/test/java/org/apache/lens/driver/hive/TestHiveDriver.java
----------------------------------------------------------------------
diff --git a/lens-driver-hive/src/test/java/org/apache/lens/driver/hive/TestHiveDriver.java b/lens-driver-hive/src/test/java/org/apache/lens/driver/hive/TestHiveDriver.java
index 3f77b0c..cf51909 100644
--- a/lens-driver-hive/src/test/java/org/apache/lens/driver/hive/TestHiveDriver.java
+++ b/lens-driver-hive/src/test/java/org/apache/lens/driver/hive/TestHiveDriver.java
@@ -39,7 +39,6 @@ import org.apache.lens.server.api.query.QueryContext;
 import org.apache.lens.server.api.query.cost.QueryCost;
 import org.apache.lens.server.api.query.priority.CostRangePriorityDecider;
 import org.apache.lens.server.api.query.priority.CostToPriorityRangeConf;
-import org.apache.lens.server.api.user.MockDriverQueryHook;
 import org.apache.lens.server.api.util.LensUtil;
 
 import org.apache.hadoop.conf.Configuration;
@@ -121,7 +120,6 @@ public class TestHiveDriver {
     conf = new HiveConf();
     conf.addResource("drivers/hive/hive1/hivedriver-site.xml");
     conf.setClass(HiveDriver.HIVE_CONNECTION_CLASS, EmbeddedThriftConnection.class, ThriftConnection.class);
-    conf.setClass(HiveDriver.HIVE_QUERY_HOOK_CLASS, MockDriverQueryHook.class, DriverQueryHook.class);
     conf.set("hive.lock.manager", "org.apache.hadoop.hive.ql.lockmgr.EmbeddedLockManager");
     conf.setBoolean(HiveDriver.HS2_CALCULATE_PRIORITY, true);
     driver = new HiveDriver();
@@ -484,7 +482,6 @@ public class TestHiveDriver {
    */
   protected void validateExecuteAsync(QueryContext ctx, DriverQueryState finalState, boolean isPersistent,
     boolean formatNulls) throws Exception {
-    assertEquals(ctx.getSelectedDriverConf().get(MockDriverQueryHook.KEY), MockDriverQueryHook.VALUE);
     validateExecuteAsync(ctx, finalState, isPersistent, formatNulls, driver);
   }
 

http://git-wip-us.apache.org/repos/asf/lens/blob/9b7541bc/lens-driver-hive/src/test/java/org/apache/lens/driver/hive/TestRemoteHiveDriver.java
----------------------------------------------------------------------
diff --git a/lens-driver-hive/src/test/java/org/apache/lens/driver/hive/TestRemoteHiveDriver.java b/lens-driver-hive/src/test/java/org/apache/lens/driver/hive/TestRemoteHiveDriver.java
index 4f18c24..765bb45 100644
--- a/lens-driver-hive/src/test/java/org/apache/lens/driver/hive/TestRemoteHiveDriver.java
+++ b/lens-driver-hive/src/test/java/org/apache/lens/driver/hive/TestRemoteHiveDriver.java
@@ -28,13 +28,11 @@ import java.util.concurrent.atomic.AtomicInteger;
 
 import org.apache.lens.api.query.QueryHandle;
 import org.apache.lens.server.api.LensConfConstants;
-import org.apache.lens.server.api.driver.DriverQueryHook;
 import org.apache.lens.server.api.driver.DriverQueryPlan;
 import org.apache.lens.server.api.driver.DriverQueryStatus.DriverQueryState;
 import org.apache.lens.server.api.driver.LensDriver;
 import org.apache.lens.server.api.error.LensException;
 import org.apache.lens.server.api.query.QueryContext;
-import org.apache.lens.server.api.user.MockDriverQueryHook;
 
 import org.apache.commons.io.FileUtils;
 import org.apache.hadoop.hive.conf.HiveConf;
@@ -136,7 +134,6 @@ public class TestRemoteHiveDriver extends TestHiveDriver {
     conf.addResource("drivers/hive/hive1/hivedriver-site.xml");
     driver = new HiveDriver();
     conf.setBoolean(HiveDriver.HS2_CALCULATE_PRIORITY, true);
-    conf.setClass(HiveDriver.HIVE_QUERY_HOOK_CLASS, MockDriverQueryHook.class, DriverQueryHook.class);
     driver.configure(conf, "hive", "hive1");
     drivers = Lists.<LensDriver>newArrayList(driver);
     System.out.println("TestRemoteHiveDriver created");

http://git-wip-us.apache.org/repos/asf/lens/blob/9b7541bc/lens-driver-jdbc/src/main/java/org/apache/lens/driver/jdbc/JDBCDriver.java
----------------------------------------------------------------------
diff --git a/lens-driver-jdbc/src/main/java/org/apache/lens/driver/jdbc/JDBCDriver.java b/lens-driver-jdbc/src/main/java/org/apache/lens/driver/jdbc/JDBCDriver.java
index eef4464..044a19c 100644
--- a/lens-driver-jdbc/src/main/java/org/apache/lens/driver/jdbc/JDBCDriver.java
+++ b/lens-driver-jdbc/src/main/java/org/apache/lens/driver/jdbc/JDBCDriver.java
@@ -935,7 +935,6 @@ public class JDBCDriver extends AbstractLensDriver {
     String rewrittenQuery = rewriteQuery(context);
     JdbcQueryContext jdbcCtx = new JdbcQueryContext(context, logSegregationContext);
     jdbcCtx.setRewrittenQuery(rewrittenQuery);
-    queryHook.preLaunch(context);
     try {
       Future<QueryResult> future = asyncQueryPool.submit(new QueryCallable(jdbcCtx, logSegregationContext));
       jdbcCtx.setResultFuture(future);
@@ -1145,4 +1144,9 @@ public class JDBCDriver extends AbstractLensDriver {
   public void writeExternal(ObjectOutput arg0) throws IOException {
     // TODO Auto-generated method stub
   }
+
+  @Override
+  public DriverQueryHook getQueryHook() {
+    return queryHook;
+  }
 }

http://git-wip-us.apache.org/repos/asf/lens/blob/9b7541bc/lens-driver-jdbc/src/test/java/org/apache/lens/driver/jdbc/TestJdbcDriver.java
----------------------------------------------------------------------
diff --git a/lens-driver-jdbc/src/test/java/org/apache/lens/driver/jdbc/TestJdbcDriver.java b/lens-driver-jdbc/src/test/java/org/apache/lens/driver/jdbc/TestJdbcDriver.java
index 81a9552..67f6c1f 100644
--- a/lens-driver-jdbc/src/test/java/org/apache/lens/driver/jdbc/TestJdbcDriver.java
+++ b/lens-driver-jdbc/src/test/java/org/apache/lens/driver/jdbc/TestJdbcDriver.java
@@ -40,7 +40,6 @@ import org.apache.lens.server.api.query.ExplainQueryContext;
 import org.apache.lens.server.api.query.PreparedQueryContext;
 import org.apache.lens.server.api.query.QueryContext;
 import org.apache.lens.server.api.query.cost.QueryCost;
-import org.apache.lens.server.api.user.MockDriverQueryHook;
 import org.apache.lens.server.api.util.LensUtil;
 
 import org.apache.hadoop.conf.Configuration;
@@ -86,7 +85,6 @@ public class TestJdbcDriver {
     baseConf.set(JDBC_USER, "SA");
     baseConf.set(JDBC_PASSWORD, "");
     baseConf.set(JDBC_EXPLAIN_KEYWORD_PARAM, "explain plan for ");
-    baseConf.setClass(JDBC_QUERY_HOOK_CLASS, MockDriverQueryHook.class, DriverQueryHook.class);
     hConf = new HiveConf(baseConf, this.getClass());
 
     driver = new JDBCDriver();
@@ -746,7 +744,6 @@ public class TestJdbcDriver {
 
   private void executeAsync(QueryContext ctx) throws LensException {
     driver.executeAsync(ctx);
-    assertEquals(ctx.getSelectedDriverConf().get(MockDriverQueryHook.KEY), MockDriverQueryHook.VALUE);
   }
 
   /**

http://git-wip-us.apache.org/repos/asf/lens/blob/9b7541bc/lens-server-api/src/main/java/org/apache/lens/server/api/driver/AbstractLensDriver.java
----------------------------------------------------------------------
diff --git a/lens-server-api/src/main/java/org/apache/lens/server/api/driver/AbstractLensDriver.java b/lens-server-api/src/main/java/org/apache/lens/server/api/driver/AbstractLensDriver.java
index 883ad9d..959a5b2 100644
--- a/lens-server-api/src/main/java/org/apache/lens/server/api/driver/AbstractLensDriver.java
+++ b/lens-server-api/src/main/java/org/apache/lens/server/api/driver/AbstractLensDriver.java
@@ -47,6 +47,8 @@ public abstract class AbstractLensDriver implements LensDriver {
   @Getter
   private String fullyQualifiedName = null;
 
+  private DriverQueryHook noOpDriverQueryHook = new NoOpDriverQueryHook();
+
   @Override
   public void configure(Configuration conf, String driverType, String driverName) throws LensException {
     if (StringUtils.isBlank(driverType) || StringUtils.isBlank(driverName)) {
@@ -103,6 +105,11 @@ public abstract class AbstractLensDriver implements LensDriver {
   }
 
   @Override
+  public DriverQueryHook getQueryHook() {
+    return noOpDriverQueryHook;
+  }
+
+  @Override
   public String toString() {
     return getFullyQualifiedName();
   }

http://git-wip-us.apache.org/repos/asf/lens/blob/9b7541bc/lens-server-api/src/main/java/org/apache/lens/server/api/driver/DriverQueryHook.java
----------------------------------------------------------------------
diff --git a/lens-server-api/src/main/java/org/apache/lens/server/api/driver/DriverQueryHook.java b/lens-server-api/src/main/java/org/apache/lens/server/api/driver/DriverQueryHook.java
index ecac6be..da125bd 100644
--- a/lens-server-api/src/main/java/org/apache/lens/server/api/driver/DriverQueryHook.java
+++ b/lens-server-api/src/main/java/org/apache/lens/server/api/driver/DriverQueryHook.java
@@ -23,30 +23,39 @@ package org.apache.lens.server.api.driver;
 
 import org.apache.lens.server.api.error.LensException;
 import org.apache.lens.server.api.query.AbstractQueryContext;
+import org.apache.lens.server.api.query.QueryContext;
 
 /**
- * Drivers can choose to initialize a DriverQueryHook object in their
- * initialization and use that wherever they want.
+ * Drivers should initialize a DriverQueryHook object in their  initialization and expose it
+ * via {@link LensDriver#getQueryHook}. Lens Server will invoke the driver hook at relevant points during
+ * query execution. By default each driver exposes a {@link NoOpDriverQueryHook} which does nothing when invoked.
  *
- * The only use case I see right now is to provide a hook just before query is
- * launched on the driver.
+ * The only use case I see right now is to provide a hook just after driver has been selected for a query and
+ * before query is launched on the driver. One example usage for hive driver would be to add dynamic configuration or
+ * stall execution of a query by looking at the final translated query itself (based on table involved, filters
+ * involved, etc in the query).
  *
- * This interface is meant to unify drivers' needs of having hooks. Each driver
- * can use the methods in their own way. Each driver can pose its own restrictions
- * or guidelines on methods for its hooks.
- * e.g. some driver may choose to not allow any hooks
- * another driver may allow hooks but for restricted usage
- * Some drivers may want their hooks to be initialized with some constructor params
- * Currently, Hivedriver and Jdbcdriver only require their hook implementations to have a default constructor
- *
- * This interface is expected to evolve for some time as more needs of hooks are discovered
+ * This interface is expected to evolve for some time as more needs for hook are discovered
  *
+
  */
 public interface DriverQueryHook {
   /**
-   * Should be Called before launch on the driver
+   * Called just before launching the query on the selected driver.
+   * @param ctx
+   * @throws LensException
+   */
+  void preLaunch(QueryContext ctx) throws LensException;
+
+  /**
+   * Called just after driver has been selected to execute a query.
+   *
+   * Note: Note if this method updates any configuration, same should be reflected in QueryContext
+   * via {@link AbstractQueryContext#updateConf(Map)} to ensure the modified configration is persisted and is available
+   * on server restarts and other bookkeeping needs.
+   *
    * @param ctx
    * @throws LensException
    */
-  void preLaunch(AbstractQueryContext ctx) throws LensException;
+  void postDriverSelection(AbstractQueryContext ctx) throws LensException;
 }

http://git-wip-us.apache.org/repos/asf/lens/blob/9b7541bc/lens-server-api/src/main/java/org/apache/lens/server/api/driver/LensDriver.java
----------------------------------------------------------------------
diff --git a/lens-server-api/src/main/java/org/apache/lens/server/api/driver/LensDriver.java b/lens-server-api/src/main/java/org/apache/lens/server/api/driver/LensDriver.java
index 69295d9..95ea360 100644
--- a/lens-server-api/src/main/java/org/apache/lens/server/api/driver/LensDriver.java
+++ b/lens-server-api/src/main/java/org/apache/lens/server/api/driver/LensDriver.java
@@ -219,4 +219,10 @@ public interface LensDriver extends Externalizable {
    * @param queryContext
    */
   Priority decidePriority(AbstractQueryContext queryContext);
+
+  /**
+   * @return the DriverQueryHook implementation for the driver.
+   * @see DriverQueryHook for more details.
+   */
+  DriverQueryHook getQueryHook();
 }

http://git-wip-us.apache.org/repos/asf/lens/blob/9b7541bc/lens-server-api/src/main/java/org/apache/lens/server/api/driver/NoOpDriverQueryHook.java
----------------------------------------------------------------------
diff --git a/lens-server-api/src/main/java/org/apache/lens/server/api/driver/NoOpDriverQueryHook.java b/lens-server-api/src/main/java/org/apache/lens/server/api/driver/NoOpDriverQueryHook.java
index a5475e4..8cd03cb 100644
--- a/lens-server-api/src/main/java/org/apache/lens/server/api/driver/NoOpDriverQueryHook.java
+++ b/lens-server-api/src/main/java/org/apache/lens/server/api/driver/NoOpDriverQueryHook.java
@@ -21,15 +21,23 @@
  */
 package org.apache.lens.server.api.driver;
 
+import org.apache.lens.server.api.error.LensException;
 import org.apache.lens.server.api.query.AbstractQueryContext;
+import org.apache.lens.server.api.query.QueryContext;
 
 import lombok.extern.slf4j.Slf4j;
 
 @Slf4j
 public class NoOpDriverQueryHook implements DriverQueryHook {
   @Override
-  public void preLaunch(AbstractQueryContext ctx) {
-    log.info("Pre launch for {}, user query: {}, driver query: {}", ctx.getSubmittedUser(), ctx.getUserQuery(),
-      ctx.getSelectedDriverQuery());
+  public void preLaunch(QueryContext ctx) {
+    log.debug("Pre launch for {}, user query: {}, driver {}, driver query: {}", ctx.getSubmittedUser(),
+      ctx.getUserQuery(), ctx.getSelectedDriver().getFullyQualifiedName(), ctx.getSelectedDriverQuery());
+  }
+
+  @Override
+  public void postDriverSelection(AbstractQueryContext ctx) throws LensException {
+    log.debug("Post driver selection for {}, user query: {}, driver {}, driver query: {}", ctx.getSubmittedUser(),
+      ctx.getUserQuery(), ctx.getSelectedDriver().getFullyQualifiedName(), ctx.getSelectedDriverQuery());
   }
 }

http://git-wip-us.apache.org/repos/asf/lens/blob/9b7541bc/lens-server-api/src/test/java/org/apache/lens/server/api/user/MockDriverQueryHook.java
----------------------------------------------------------------------
diff --git a/lens-server-api/src/test/java/org/apache/lens/server/api/user/MockDriverQueryHook.java b/lens-server-api/src/test/java/org/apache/lens/server/api/user/MockDriverQueryHook.java
index 2505ef4..0c4a3f0 100644
--- a/lens-server-api/src/test/java/org/apache/lens/server/api/user/MockDriverQueryHook.java
+++ b/lens-server-api/src/test/java/org/apache/lens/server/api/user/MockDriverQueryHook.java
@@ -18,15 +18,37 @@
  */
 package org.apache.lens.server.api.user;
 
-import org.apache.lens.server.api.driver.DriverQueryHook;
+import java.util.HashMap;
+
+import org.apache.lens.server.api.driver.NoOpDriverQueryHook;
+import org.apache.lens.server.api.error.LensException;
 import org.apache.lens.server.api.query.AbstractQueryContext;
+import org.apache.lens.server.api.query.QueryContext;
+
+public class MockDriverQueryHook extends NoOpDriverQueryHook {
+  public static final String KEY_PRE_LAUNCH = "TEST_KEY_PRE_LAUNCH";
+  public static final String VALUE_PRE_LAUNCH = "TEST_VALUE_PRE_LAUNCH";
 
-public class MockDriverQueryHook implements DriverQueryHook {
-  public static final String KEY = "TEST_KEY";
-  public static final String VALUE = "TEST_VALUE";
+  public static final String KEY_POST_SELECT = "TEST_KEY_POST_SELECT";
+  public static final String VALUE_POST_SELECT = "TEST_VALUE_POST_SELECT";
+  public static final String UNSAVED_KEY_POST_SELECT = "TEST_UNSAVED__KEY_POST_SELECT";
+  public static final String UNSAVED_VALUE_POST_SELECT = "TEST_UNSAVED_VALUE_POST_SELECT";
 
   @Override
-  public void preLaunch(AbstractQueryContext ctx) {
-    ctx.getSelectedDriverConf().set(KEY, VALUE);
+  public void preLaunch(QueryContext ctx) {
+    super.preLaunch(ctx);
+    ctx.getSelectedDriverConf().set(KEY_PRE_LAUNCH, VALUE_PRE_LAUNCH);
+  }
+
+  @Override
+  public void postDriverSelection(AbstractQueryContext ctx) throws LensException {
+    super.postDriverSelection(ctx);
+
+    //Updated both in driver config and LensConf(which gets persisted)
+    ctx.getSelectedDriverConf().set(KEY_POST_SELECT, VALUE_POST_SELECT);
+    ctx.updateConf(new HashMap<String, String>(1){{put(KEY_POST_SELECT, VALUE_POST_SELECT); }});
+
+    //Updated only in driver conf.
+    ctx.getSelectedDriverConf().set(UNSAVED_KEY_POST_SELECT, UNSAVED_VALUE_POST_SELECT);
   }
 }

http://git-wip-us.apache.org/repos/asf/lens/blob/9b7541bc/lens-server/src/main/java/org/apache/lens/server/query/QueryExecutionServiceImpl.java
----------------------------------------------------------------------
diff --git a/lens-server/src/main/java/org/apache/lens/server/query/QueryExecutionServiceImpl.java b/lens-server/src/main/java/org/apache/lens/server/query/QueryExecutionServiceImpl.java
index c72a890..f5552dc 100644
--- a/lens-server/src/main/java/org/apache/lens/server/query/QueryExecutionServiceImpl.java
+++ b/lens-server/src/main/java/org/apache/lens/server/query/QueryExecutionServiceImpl.java
@@ -658,8 +658,12 @@ public class QueryExecutionServiceImpl extends BaseLensService implements QueryE
     @Override
     public void run() {
       log.info("Starting QuerySubmitter thread");
-      while (!pausedForTest && !stopped && !querySubmitter.isInterrupted()) {
+      while (!stopped && !querySubmitter.isInterrupted()) {
         try {
+          if (pausedForTest) {
+            Thread.sleep(100);
+            continue;
+          }
           QueryContext query = queuedQueries.take();
           synchronized (query) {
 
@@ -725,6 +729,7 @@ public class QueryExecutionServiceImpl extends BaseLensService implements QueryE
     private void launchQuery(final QueryContext query) throws LensException {
 
       checkEstimatedQueriesState(query);
+      query.getSelectedDriver().getQueryHook().preLaunch(query);
       QueryStatus oldStatus = query.getStatus();
       QueryStatus newStatus = new QueryStatus(query.getStatus().getProgress(), null,
         QueryStatus.Status.LAUNCHED, "Query is launched on driver", false, null, null, null);
@@ -757,13 +762,13 @@ public class QueryExecutionServiceImpl extends BaseLensService implements QueryE
     }
   }
 
-  // used in tests
 
   /**
    * Pause query submitter.
+   * note : used in tests only
    */
-  public void pauseQuerySubmitter() {
-    querySubmitterRunnable.pausedForTest = true;
+  public void pauseQuerySubmitter(boolean pause) {
+    querySubmitterRunnable.pausedForTest = pause;
   }
 
   /**
@@ -1422,6 +1427,7 @@ public class QueryExecutionServiceImpl extends BaseLensService implements QueryE
       ctx.setSelectedDriverQueryCost(selectedDriverQueryCost);
       Priority priority = driver.decidePriority(ctx);
       ctx.setPriority(priority == null ? Priority.NORMAL : priority);
+      driver.getQueryHook().postDriverSelection(ctx);
       selectGauge.markSuccess();
     } finally {
       parallelCallGauge.markSuccess();
@@ -1927,11 +1933,11 @@ public class QueryExecutionServiceImpl extends BaseLensService implements QueryE
 
   /**
    * Gets the query context.
-   *
+   * note: this method is made public to expose it to test cases
    * @param queryHandle the query handle
    * @return the query context
    */
-  QueryContext getQueryContext(QueryHandle queryHandle) {
+  public QueryContext getQueryContext(QueryHandle queryHandle) {
     return allQueries.get(queryHandle);
   }
 

http://git-wip-us.apache.org/repos/asf/lens/blob/9b7541bc/lens-server/src/test/java/org/apache/lens/server/LensJerseyTest.java
----------------------------------------------------------------------
diff --git a/lens-server/src/test/java/org/apache/lens/server/LensJerseyTest.java b/lens-server/src/test/java/org/apache/lens/server/LensJerseyTest.java
index fb757e4..b5d5482 100644
--- a/lens-server/src/test/java/org/apache/lens/server/LensJerseyTest.java
+++ b/lens-server/src/test/java/org/apache/lens/server/LensJerseyTest.java
@@ -40,6 +40,7 @@ import org.apache.lens.driver.hive.TestRemoteHiveDriver;
 import org.apache.lens.server.api.LensConfConstants;
 import org.apache.lens.server.api.metrics.LensMetricsUtil;
 import org.apache.lens.server.api.metrics.MetricsService;
+import org.apache.lens.server.api.query.QueryExecutionService;
 import org.apache.lens.server.model.LogSegregationContext;
 import org.apache.lens.server.model.MappedDiagnosticLogSegregationContext;
 import org.apache.lens.server.query.QueryExecutionServiceImpl;
@@ -213,20 +214,26 @@ public abstract class LensJerseyTest extends JerseyTest {
    */
   public void restartLensServer() {
     HiveConf h = getServerConf();
-    restartLensServer(h);
+    restartLensServer(h, false);
   }
 
   /**
    * Restart lens server.
    *
    * @param conf the conf
+   * @param pauseQuerySubmitter whether to pause query submitter while starting lens server
    */
-  public void restartLensServer(HiveConf conf) {
+  public void restartLensServer(HiveConf conf, boolean pauseQuerySubmitter) {
     LensServices.get().stop();
     LensMetricsUtil.clearRegistry();
     System.out.println("Lens services stopped!");
     LensServices.setInstance(new LensServices(LensServices.LENS_SERVICES_NAME, this.logSegregationContext));
     LensServices.get().init(conf);
+    if (pauseQuerySubmitter) {
+      QueryExecutionServiceImpl queryService = LensServices.get().getService(QueryExecutionService.NAME);
+      queryService.pauseQuerySubmitter(true);
+      System.out.println("Paused Query Submitter");
+    }
     LensServices.get().start();
     System.out.println("Lens services restarted!");
   }

http://git-wip-us.apache.org/repos/asf/lens/blob/9b7541bc/lens-server/src/test/java/org/apache/lens/server/TestServerRestart.java
----------------------------------------------------------------------
diff --git a/lens-server/src/test/java/org/apache/lens/server/TestServerRestart.java b/lens-server/src/test/java/org/apache/lens/server/TestServerRestart.java
index 0f55d9e..1fa61ef 100644
--- a/lens-server/src/test/java/org/apache/lens/server/TestServerRestart.java
+++ b/lens-server/src/test/java/org/apache/lens/server/TestServerRestart.java
@@ -20,9 +20,10 @@ package org.apache.lens.server;
 
 import static org.apache.lens.server.LensServerTestUtil.createTable;
 import static org.apache.lens.server.LensServerTestUtil.loadData;
+import static org.apache.lens.server.api.user.MockDriverQueryHook.*;
 import static org.apache.lens.server.common.RestAPITestUtil.execute;
 
-import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.*;
 
 import java.io.*;
 import java.util.*;
@@ -39,6 +40,7 @@ import org.apache.lens.api.query.*;
 import org.apache.lens.api.result.LensAPIResult;
 import org.apache.lens.driver.hive.TestRemoteHiveDriver;
 import org.apache.lens.server.api.error.LensException;
+import org.apache.lens.server.api.query.QueryContext;
 import org.apache.lens.server.api.query.QueryExecutionService;
 import org.apache.lens.server.api.session.SessionService;
 import org.apache.lens.server.common.TestResourceFile;
@@ -47,6 +49,7 @@ import org.apache.lens.server.query.TestQueryService;
 import org.apache.lens.server.session.HiveSessionService;
 import org.apache.lens.server.session.LensSessionImpl;
 
+import org.apache.hadoop.conf.Configuration;
 import org.apache.hive.service.Service;
 
 import org.glassfish.jersey.media.multipart.FormDataBodyPart;
@@ -152,18 +155,19 @@ public class TestServerRestart extends LensAllApplicationJerseyTest {
     List<QueryHandle> launchedQueries = new ArrayList<>();
     final int NUM_QUERIES = 10;
 
-    boolean killed = false;
+    boolean isQuerySubmitterPaused = false;
+    QueryHandle handleForMockDriverQueryHookTest = null;
     for (int i = 0; i < NUM_QUERIES; i++) {
-      if (!killed && i > NUM_QUERIES / 3) {
+      if (!isQuerySubmitterPaused && i > NUM_QUERIES / 3) {
         // Kill the query submitter thread to make sure some queries stay in accepted queue
         try {
-          queryService.pauseQuerySubmitter();
+          queryService.pauseQuerySubmitter(true);
           log.info("Stopped query submitter");
           Assert.assertFalse(queryService.getHealthStatus().isHealthy());
         } catch (Exception exc) {
           log.error("Could not kill query submitter", exc);
         }
-        killed = true;
+        isQuerySubmitterPaused = true;
       }
 
       final FormDataMultiPart mp = new FormDataMultiPart();
@@ -182,13 +186,23 @@ public class TestServerRestart extends LensAllApplicationJerseyTest {
         .get(LensQuery.class);
       log.info("{} submitted query {} state: {}", i, handle, ctx.getStatus().getStatus());
       launchedQueries.add(handle);
+      if (i == (NUM_QUERIES-1)) {
+        //checking this only for one of the queued queries. A queued query has all the config information available in
+        // server memory. (Some of the information is lost after query is purged)
+        testMockDriverQueryHookPostDriverSelection(queryService, handle, false);
+        handleForMockDriverQueryHookTest = handle;
+        log.info("Testing query {} for MockDriverQueryHook", handleForMockDriverQueryHookTest);
+      }
     }
 
     // Restart the server
     log.info("Restarting lens server!");
-    restartLensServer();
+    restartLensServer(getServerConf(), true);
     log.info("Restarted lens server!");
     queryService = LensServices.get().getService(QueryExecutionService.NAME);
+    Assert.assertFalse(queryService.getHealthStatus().isHealthy());
+    testMockDriverQueryHookPostDriverSelection(queryService, handleForMockDriverQueryHookTest, true);
+    queryService.pauseQuerySubmitter(false);
     Assert.assertTrue(queryService.getHealthStatus().isHealthy());
 
     // All queries should complete after server restart
@@ -223,6 +237,35 @@ public class TestServerRestart extends LensAllApplicationJerseyTest {
   }
 
   /**
+   * Tests whether the driver configuration updated by mock query driver hook is
+   * 1. updated in LensConf wherever applicable and
+   * 2. is persisted and available even after server startup.
+   *
+   * @param queryService
+   * @param handle
+   * @param afterRestart
+   */
+  private void testMockDriverQueryHookPostDriverSelection(QueryExecutionServiceImpl queryService, QueryHandle handle,
+    boolean afterRestart){
+    QueryContext ctx = queryService.getQueryContext(handle);
+    assertNotNull(ctx, "Make sure that the query has not  been purged");
+    assertTrue(ctx.getStatus().queued(), "Make sure query is still in QUEUED state");
+    LensConf lensQueryConf = queryService.getQueryContext(handle).getLensConf();
+    Configuration driverConf = queryService.getQueryContext(handle).getSelectedDriverConf();
+
+    assertEquals(driverConf.get(KEY_POST_SELECT), VALUE_POST_SELECT);
+    assertEquals(lensQueryConf.getProperty(KEY_POST_SELECT), VALUE_POST_SELECT);
+
+    if (afterRestart) {
+      //This will be unavailable since if was not updated in LensConf by MockDriverQueryHook
+      assertNull(driverConf.get(UNSAVED_KEY_POST_SELECT));
+    } else {
+      assertEquals(driverConf.get(UNSAVED_KEY_POST_SELECT), UNSAVED_VALUE_POST_SELECT);
+    }
+    assertNull(lensQueryConf.getProperty(UNSAVED_KEY_POST_SELECT));
+  }
+
+  /**
    * Test hive server restart.
    *
    * @throws Exception the exception

http://git-wip-us.apache.org/repos/asf/lens/blob/9b7541bc/lens-server/src/test/java/org/apache/lens/server/query/TestQueryService.java
----------------------------------------------------------------------
diff --git a/lens-server/src/test/java/org/apache/lens/server/query/TestQueryService.java b/lens-server/src/test/java/org/apache/lens/server/query/TestQueryService.java
index 6bf077d..fd6bae3 100644
--- a/lens-server/src/test/java/org/apache/lens/server/query/TestQueryService.java
+++ b/lens-server/src/test/java/org/apache/lens/server/query/TestQueryService.java
@@ -23,6 +23,7 @@ import static javax.ws.rs.core.Response.Status.*;
 import static org.apache.lens.server.LensServerTestUtil.DB_WITH_JARS;
 import static org.apache.lens.server.LensServerTestUtil.DB_WITH_JARS_2;
 import static org.apache.lens.server.api.LensServerAPITestUtil.getLensConf;
+import static org.apache.lens.server.api.user.MockDriverQueryHook.*;
 import static org.apache.lens.server.common.RestAPITestUtil.*;
 
 import static org.testng.Assert.*;
@@ -668,6 +669,8 @@ public class TestQueryService extends LensJerseyTest {
       new GenericType<LensAPIResult<QueryHandle>>() {}).getData();
 
     assertNotNull(handle);
+    QueryContext ctx = queryService.getUpdatedQueryContext(lensSessionId, handle);
+    assertEquals(ctx.getSelectedDriverConf().get(KEY_POST_SELECT), VALUE_POST_SELECT);
 
     // Get query
     LensQuery lensQuery = target.path(handle.toString()).queryParam("sessionid", lensSessionId).request(mt)
@@ -696,12 +699,14 @@ public class TestQueryService extends LensJerseyTest {
       }*/
       Thread.sleep(1000);
     }
+    assertEquals(ctx.getSelectedDriverConf().get(KEY_PRE_LAUNCH), VALUE_PRE_LAUNCH);
     assertTrue(lensQuery.getSubmissionTime() > 0);
     assertTrue(lensQuery.getLaunchTime() > 0);
     assertTrue(lensQuery.getDriverStartTime() > 0);
     assertTrue(lensQuery.getDriverFinishTime() > 0);
     assertTrue(lensQuery.getFinishTime() > 0);
-    QueryContext ctx = queryService.getUpdatedQueryContext(lensSessionId, lensQuery.getQueryHandle());
+    ctx = queryService.getUpdatedQueryContext(lensSessionId, lensQuery.getQueryHandle());
+
     assertNotNull(ctx.getPhase1RewrittenQuery());
     assertEquals(ctx.getPhase1RewrittenQuery(), ctx.getUserQuery()); //Since there is no rewriter in this test
     assertEquals(lensQuery.getStatus().getStatus(), QueryStatus.Status.SUCCESSFUL);

http://git-wip-us.apache.org/repos/asf/lens/blob/9b7541bc/lens-server/src/test/resources/drivers/hive/hive1/hivedriver-site.xml
----------------------------------------------------------------------
diff --git a/lens-server/src/test/resources/drivers/hive/hive1/hivedriver-site.xml b/lens-server/src/test/resources/drivers/hive/hive1/hivedriver-site.xml
index 6362473..4d2bc9a 100644
--- a/lens-server/src/test/resources/drivers/hive/hive1/hivedriver-site.xml
+++ b/lens-server/src/test/resources/drivers/hive/hive1/hivedriver-site.xml
@@ -82,4 +82,11 @@
     <value>hive/hive1</value>
   </property>
 
+  <property>
+    <name>lens.driver.hive.query.hook.class</name>
+    <value>org.apache.lens.server.api.user.MockDriverQueryHook</value>
+    <description>The query hook class for hive driver.
+    </description>
+  </property>
+
 </configuration>

http://git-wip-us.apache.org/repos/asf/lens/blob/9b7541bc/lens-server/src/test/resources/drivers/hive/hive2/hivedriver-site.xml
----------------------------------------------------------------------
diff --git a/lens-server/src/test/resources/drivers/hive/hive2/hivedriver-site.xml b/lens-server/src/test/resources/drivers/hive/hive2/hivedriver-site.xml
index d5be96e..f9f345e 100644
--- a/lens-server/src/test/resources/drivers/hive/hive2/hivedriver-site.xml
+++ b/lens-server/src/test/resources/drivers/hive/hive2/hivedriver-site.xml
@@ -82,4 +82,11 @@
     <value>hive/hive2</value>
   </property>
 
+  <property>
+    <name>lens.driver.hive.query.hook.class</name>
+    <value>org.apache.lens.server.api.user.MockDriverQueryHook</value>
+    <description>The query hook class for hive driver.
+    </description>
+  </property>
+
 </configuration>