diff --git llap-ext-client/src/java/org/apache/hadoop/hive/llap/LlapBaseInputFormat.java llap-ext-client/src/java/org/apache/hadoop/hive/llap/LlapBaseInputFormat.java index ef03be660e..30f372003f 100644 --- llap-ext-client/src/java/org/apache/hadoop/hive/llap/LlapBaseInputFormat.java +++ llap-ext-client/src/java/org/apache/hadoop/hive/llap/LlapBaseInputFormat.java @@ -37,6 +37,7 @@ import java.util.Set; import java.util.UUID; import java.util.concurrent.LinkedBlockingQueue; +import java.util.regex.Pattern; import org.apache.commons.collections4.ListUtils; import org.apache.hadoop.hive.conf.HiveConf; @@ -114,6 +115,8 @@ public static final String PWD_KEY = "llap.if.pwd"; public static final String HANDLE_ID = "llap.if.handleid"; public static final String DB_KEY = "llap.if.database"; + public static final String SESSION_QUERIES_FOR_GET_NUM_SPLITS = "llap.session.queries.for.get.num.splits"; + public static final Pattern SET_QUERY_PATTERN = Pattern.compile("^\\s*set\\s+.*=.+$", Pattern.CASE_INSENSITIVE); public final String SPLIT_QUERY = "select get_splits(\"%s\",%d)"; public static final LlapServiceInstance[] serviceInstanceArray = new LlapServiceInstance[0]; @@ -259,6 +262,20 @@ public LlapBaseInputFormat() { if (database != null && !database.isEmpty()) { stmt.execute("USE " + database); } + String sessionQueries = job.get(SESSION_QUERIES_FOR_GET_NUM_SPLITS); + if (sessionQueries != null && !sessionQueries.trim().isEmpty()) { + String[] queries = sessionQueries.trim().split(","); + for (String q : queries) { + //allow only set queries + if (SET_QUERY_PATTERN.matcher(q).matches()) { + LOG.debug("Executing session query: {}", q); + stmt.execute(q); + } else { + LOG.warn("Only SET queries are allowed, not executing this query: {}", q); + } + } + } + ResultSet res = stmt.executeQuery(sql); while (res.next()) { // deserialize split