diff --git a/common/src/java/org/apache/hadoop/hive/conf/HiveConf.java b/common/src/java/org/apache/hadoop/hive/conf/HiveConf.java index 01cd731..2fc5142 100644 --- a/common/src/java/org/apache/hadoop/hive/conf/HiveConf.java +++ b/common/src/java/org/apache/hadoop/hive/conf/HiveConf.java @@ -3106,6 +3106,7 @@ private static String getSQLStdAuthDefaultWhiteListPattern() { ConfVars.OUTPUT_FILE_EXTENSION.varname, ConfVars.SHOW_JOB_FAIL_DEBUG_INFO.varname, ConfVars.TASKLOG_DEBUG_TIMEOUT.varname, + ConfVars.HIVEQUERYID.varname, }; /** diff --git a/itests/hive-unit/src/test/java/org/apache/hive/service/cli/session/TestHiveSessionImpl.java b/itests/hive-unit/src/test/java/org/apache/hive/service/cli/session/TestHiveSessionImpl.java index 2d9ad03..8a32a07 100644 --- a/itests/hive-unit/src/test/java/org/apache/hive/service/cli/session/TestHiveSessionImpl.java +++ b/itests/hive-unit/src/test/java/org/apache/hive/service/cli/session/TestHiveSessionImpl.java @@ -74,6 +74,7 @@ protected synchronized void release(boolean userAccess) { try { //Running a normal async query with no exceptions,then no need to close opHandle + session.open(new HashMap()); session.executeStatementAsync(hql, confOverlay); Mockito.verify(operationManager, Mockito.times(0)).closeOperation(opHandle); diff --git a/ql/src/java/org/apache/hadoop/hive/ql/Driver.java b/ql/src/java/org/apache/hadoop/hive/ql/Driver.java index 93c7a54..9a5495b 100644 --- a/ql/src/java/org/apache/hadoop/hive/ql/Driver.java +++ b/ql/src/java/org/apache/hadoop/hive/ql/Driver.java @@ -384,9 +384,13 @@ public int compile(String command, boolean resetTaskIds) { } saveSession(queryState); - // generate new query id - String queryId = QueryPlan.makeQueryId(); - conf.setVar(HiveConf.ConfVars.HIVEQUERYID, queryId); + // Generate new query id if it's not set for CLI case. If it's session based, + // query id is passed in from the client or initialized when the session starts. + String queryId = conf.getVar(HiveConf.ConfVars.HIVEQUERYID); + if (queryId == null || queryId.isEmpty()) { + queryId = QueryPlan.makeQueryId(); + conf.setVar(HiveConf.ConfVars.HIVEQUERYID, queryId); + } SessionState.get().setupQueryCurrentTimestamp(); diff --git a/service/src/java/org/apache/hive/service/cli/operation/Operation.java b/service/src/java/org/apache/hive/service/cli/operation/Operation.java index 4ca0561..a851936 100644 --- a/service/src/java/org/apache/hive/service/cli/operation/Operation.java +++ b/service/src/java/org/apache/hive/service/cli/operation/Operation.java @@ -31,6 +31,7 @@ import org.apache.hadoop.hive.conf.HiveConf; import org.apache.hadoop.hive.ql.processors.CommandProcessorResponse; import org.apache.hadoop.hive.ql.session.OperationLog; +import org.apache.hadoop.hive.ql.session.SessionState; import org.apache.hive.service.cli.FetchOrientation; import org.apache.hive.service.cli.HiveSQLException; import org.apache.hive.service.cli.OperationHandle; @@ -41,8 +42,13 @@ import org.apache.hive.service.cli.TableSchema; import org.apache.hive.service.cli.session.HiveSession; import org.apache.hive.service.cli.thrift.TProtocolVersion; +import org.apache.logging.log4j.ThreadContext; public abstract class Operation { + // Constants of the key strings for the log4j ThreadContext. + private static final String QUERYID = "QueryId"; + private static final String SESSIONID = "SessionId"; + protected final HiveSession parentSession; private OperationState state = OperationState.INITIALIZED; private final OperationHandle opHandle; @@ -238,6 +244,22 @@ protected void unregisterOperationLog() { */ protected void beforeRun() { createOperationLog(); + registerLoggingContext(); + } + + /** + * Register logging context so that Log4J can print QueryId and/or SessionId for each message + */ + protected void registerLoggingContext() { + ThreadContext.put(QUERYID, SessionState.get().getQueryId()); + ThreadContext.put(SESSIONID, SessionState.get().getSessionId()); + } + + /** + * Unregister logging context + */ + protected void unregisterLoggingContext() { + ThreadContext.clearAll(); } /** @@ -245,6 +267,7 @@ protected void beforeRun() { * Clean up resources, which was set up in beforeRun(). */ protected void afterRun() { + unregisterLoggingContext(); unregisterOperationLog(); } diff --git a/service/src/java/org/apache/hive/service/cli/operation/SQLOperation.java b/service/src/java/org/apache/hive/service/cli/operation/SQLOperation.java index 1c798eb..8b42265 100644 --- a/service/src/java/org/apache/hive/service/cli/operation/SQLOperation.java +++ b/service/src/java/org/apache/hive/service/cli/operation/SQLOperation.java @@ -225,12 +225,14 @@ public Object run() throws HiveSQLException { SessionState.setCurrentSessionState(parentSessionState); // Set current OperationLog in this async thread for keeping on saving query log. registerCurrentOperationLog(); + registerLoggingContext(); try { runQuery(opConfig); } catch (HiveSQLException e) { setOperationException(e); LOG.error("Error running hive query: ", e); } finally { + unregisterLoggingContext(); unregisterOperationLog(); } return null; diff --git a/service/src/java/org/apache/hive/service/cli/session/HiveSessionBase.java b/service/src/java/org/apache/hive/service/cli/session/HiveSessionBase.java index 449d4c7..1ab5652 100644 --- a/service/src/java/org/apache/hive/service/cli/session/HiveSessionBase.java +++ b/service/src/java/org/apache/hive/service/cli/session/HiveSessionBase.java @@ -18,8 +18,6 @@ package org.apache.hive.service.cli.session; -import java.util.Map; - import org.apache.hadoop.hive.conf.HiveConf; import org.apache.hadoop.hive.ql.session.SessionState; import org.apache.hive.service.cli.SessionHandle; @@ -27,7 +25,6 @@ import org.apache.hive.service.cli.thrift.TProtocolVersion; import java.io.File; -import java.util.Map; /** * Methods that don't need to be executed under a doAs diff --git a/service/src/java/org/apache/hive/service/cli/session/HiveSessionImpl.java b/service/src/java/org/apache/hive/service/cli/session/HiveSessionImpl.java index 27d11df..2d784f0 100644 --- a/service/src/java/org/apache/hive/service/cli/session/HiveSessionImpl.java +++ b/service/src/java/org/apache/hive/service/cli/session/HiveSessionImpl.java @@ -24,6 +24,7 @@ import java.io.IOException; import java.io.InputStreamReader; import java.util.ArrayList; +import java.util.HashMap; import java.util.HashSet; import java.util.List; import java.util.Map; @@ -39,6 +40,7 @@ import org.apache.hadoop.hive.conf.HiveConf.ConfVars; import org.apache.hadoop.hive.metastore.IMetaStoreClient; import org.apache.hadoop.hive.metastore.api.MetaException; +import org.apache.hadoop.hive.ql.QueryPlan; import org.apache.hadoop.hive.ql.exec.FetchFormatter; import org.apache.hadoop.hive.ql.exec.ListSinkOperator; import org.apache.hadoop.hive.ql.exec.Utilities; @@ -432,6 +434,18 @@ private OperationHandle executeStatementInternal(String statement, Map(); + } + String queryId = confOverlay.get(HiveConf.ConfVars.HIVEQUERYID.varname); + if (queryId == null || queryId.isEmpty()) { + queryId = QueryPlan.makeQueryId(); + confOverlay.put(HiveConf.ConfVars.HIVEQUERYID.varname, queryId); + sessionState.getConf().setVar(HiveConf.ConfVars.HIVEQUERYID, queryId); + } + OperationManager operationManager = getOperationManager(); ExecuteStatementOperation operation = operationManager .newExecuteStatementOperation(getSession(), statement, confOverlay, runAsync);