/**
* Submits query via JDBC
*
* @param query SQL query string
* @param statement sql statement to execute the query with
* @param outputFilename name of file result set is to be written to
* @param timeout time allowed for query execution
* @throws Exception
*/
public void submitQueryJDBC(
String query, Statement statement, String outputFilename, long timeout) throws Exception {
BufferedWriter writer = null;
if (outputFilename != null) {
writer = new BufferedWriter(new FileWriter(new File(outputFilename)));
}
ResultSet resultSet = null;
try {
RunThread runThread = new RunThread(statement, query);
processThread(runThread, timeout);
resultSet = runThread.getResultSet();
if (resultSet == null) {
throw runThread.getException();
}
if (outputFilename == null) {
return;
}
int columnCount = resultSet.getMetaData().getColumnCount();
columnLabels = new ArrayList<String>();
for (int i = 1; i <= columnCount; i++) {
columnLabels.add(resultSet.getMetaData().getColumnLabel(i));
}
Object[] types = new Object[columnCount];
for (int i = 1; i <= columnCount; i++) {
types[i - 1] = resultSet.getMetaData().getColumnType(i);
}
ColumnList.setTypes(types);
LOG.debug("Result set data types:");
LOG.debug(Utils.getTypesInStrings(ColumnList.getTypes()));
while (resultSet.next()) {
Object[] values = new Object[columnCount];
for (int i = 1; i <= columnCount; i++) {
try {
if (resultSet.getObject(i) == null) {
values[i - 1] = null;
continue;
}
values[i - 1] = new String(resultSet.getBytes(i));
} catch (Exception e) {
if (resultSet.getMetaData().getColumnType(i) == Types.DATE) {
values[i - 1] = resultSet.getDate(i);
} else {
values[i - 1] = resultSet.getObject(i);
}
}
}
ColumnList columnList = new ColumnList(values);
if (writer != null) {
writer.write(columnList + "\n");
}
}
if (writer != null) {
writer.close();
}
} finally {
if (resultSet != null) {
resultSet.close();
}
}
}
/**
* Submit query via the submit_plan tool
*
* @param submitPlanCommand command of submit_plan
* @param queryFileName name of file containing input query
* @param outputFileName name of file containing query results
* @param queryType type of query: sql, logical or physical
* @throws Exception
*/
public void submitQueriesSubmitPlan(
String submitPlanCommand,
String queryFileName,
String outputFileName,
String queryType,
long timeout)
throws Exception {
String command =
submitPlanCommand
+ " -f "
+ queryFileName
+ " --format tsv -t "
+ queryType
+ " -z "
+ Utils.getDrillTestProperties().get("ZOOKEEPERS");
LOG.debug("Executing " + command + ".");
RunThread runThread = new RunThread(command);
processThread(runThread, timeout);
Process process = runThread.getProcess();
Scanner scanner = new Scanner(process.getInputStream());
scanner.useDelimiter("\\A");
String output = scanner.hasNext() ? scanner.next() : "";
PrintWriter writer = new PrintWriter(outputFileName);
writer.write(output);
writer.close();
scanner.close();
}
// this function should only be used while the sql resultset is still valid,
// i.e. the connection related to the resultset is still valid (not closed)
public static String getSqlResult(ResultSet resultSet) throws SQLException {
StringBuffer stringBuffer = new StringBuffer();
List columnLabels = new ArrayList<String>();
try {
int columnCount = resultSet.getMetaData().getColumnCount();
for (int i = 1; i <= columnCount; i++) {
columnLabels.add(resultSet.getMetaData().getColumnLabel(i));
}
List<Integer> types = Lists.newArrayList();
for (int i = 1; i <= columnCount; i++) {
types.add(resultSet.getMetaData().getColumnType(i));
}
LOG.debug("Result set data types:");
LOG.debug(Utils.getTypesInStrings(types));
stringBuffer.append(new ColumnList(types, columnLabels).toString() + "\n");
while (resultSet.next()) {
List<Object> values = Lists.newArrayList();
for (int i = 1; i <= columnCount; i++) {
try {
if (resultSet.getObject(i) == null) {
values.add(null);
continue;
}
if (resultSet.getMetaData().getColumnType(i) == Types.NVARCHAR) {
values.add(new String(resultSet.getBytes(i), "UTF-16"));
} else {
values.add(new String(resultSet.getBytes(i), "UTF-8"));
}
} catch (Exception e) {
if (resultSet.getMetaData().getColumnType(i) == Types.DATE) {
values.add(resultSet.getDate(i));
} else {
values.add(resultSet.getObject(i));
}
}
}
stringBuffer.append(new ColumnList(types, values).toString() + "\n");
}
} catch (IllegalArgumentException | IllegalAccessException e1) {
// TODO Auto-generated catch block
e1.printStackTrace();
} finally {
if (resultSet != null) {
resultSet.close();
}
}
return stringBuffer.toString();
}
/**
* Constructs an iteration of test case definitions from various test data sources, obtained from
* the mvn command line option. See README.md for more details.
*
* @return an iteration of object arrays that defines the set of tests to be executed.
* @throws Exception
*/
public static List<DrillTestCase> getDrillTestCases() throws IOException {
String[] testDefSources = null;
try {
testDefSources = TestDriver.OPTIONS.sources.split(",");
} catch (Exception e) {
testDefSources = new String[] {""}; // Look at the default location for test definition files
}
String[] testGroups = null;
try {
testGroups = TestDriver.OPTIONS.groups.split(",");
} catch (Exception e) {
LOG.info("Test groups not specified. Will run all collected tests.");
}
List<DrillTestCase> drillTestCases = new ArrayList<>();
for (String testDefSource : testDefSources) {
testDefSource = Utils.getAbsolutePath(testDefSource, "DRILL_TEST_DATA_DIR");
File testDefSourceFile = new File(testDefSource);
List<File> testDefFiles = searchFiles(testDefSourceFile, ".*.json");
for (File testDefFile : testDefFiles) {
// try {
TestCaseModeler modeler;
try {
modeler = getTestCaseModeler(testDefFile.getAbsolutePath());
} catch (JsonParseException e) {
LOG.warn(
"Caught exception parsing " + testDefFile + ". This test will not be executed.", e);
continue;
}
List<String> categories = modeler.categories;
boolean foundTests = false;
for (String testGroup : testGroups) {
if (categories != null && !categories.contains(testGroup)) {
continue;
} else {
foundTests = true;
break;
}
}
if (!foundTests) {
continue;
}
String queryFileExtension = modeler.matrices.get(0).inputFile;
String expectedFileExtension = modeler.matrices.get(0).expectedFile;
boolean skipSuite = false;
if (modeler.dependencies != null) {
for (String dependency : modeler.dependencies) {
if (TestDriver.OPTIONS.excludeDependenciesAsList().contains(dependency)) {
skipSuite = true;
}
}
}
if (skipSuite) {
continue;
}
List<File> testQueryFiles = searchFiles(testDefFile.getParentFile(), queryFileExtension);
for (File testQueryFile : testQueryFiles) {
String expectedFileName =
getExpectedFile(
testQueryFile.getAbsolutePath(), queryFileExtension, expectedFileExtension);
drillTestCases.add(
new DrillTestCase(modeler, testQueryFile.getAbsolutePath(), expectedFileName));
}
}
}
if (drillTestCases.size() == 0) {
LOG.warn("Warning: No test cases have been collected.");
}
return drillTestCases;
}
/**
* The submitter of queries to drill. There are a variety of interfaces via which to submit queries.
*
* @author Zhiyong Liu
*/
public class QuerySubmitter {
protected static final Logger LOG = Logger.getLogger(Utils.getInvokingClassName());
private String schema;
private static Map<String, String> drillProperties = Utils.getDrillTestProperties();
public static final long TIMEOUT_SECONDS =
Integer.parseInt(drillProperties.get("TIME_OUT_SECONDS"));
private List<String> columnLabels = null;
/**
* Constructor with the schema name
*
* @param schema name of schema to use
*/
public QuerySubmitter(String schema) {
this.schema = schema;
}
/**
* Submits query(ies) via CLI
*
* @param queryFileName name of file to use for query submission
* @throws IOException
* @throws InterruptedException
*/
public void submitQueriesSqlline(String sqllineCommand, String queryFileName, long timeout)
throws IOException, InterruptedException {
// TODO: need to parameterize the command line; need to finalize
// treatment of schemas
String command =
sqllineCommand
+ " -n admin -p admin -u jdbc:drill:schema="
+ schema
+ " -f "
+ queryFileName;
LOG.debug("Executing " + command + ".");
RunThread runThread = new RunThread(command);
processThread(runThread, timeout);
}
/**
* Executes a JDBC Query and iterates through the resultSet
*
* @param query
* @param queryFileName
* @param statement
* @return
*/
public boolean executeQueryJDBC(String query, String queryFileName, Statement statement)
throws Exception {
boolean status = true;
ResultSet resultSet = null;
long startTime = 0l;
long connTime = Long.MIN_VALUE;
long executeTime = Long.MIN_VALUE;
long firstRowFetchTime = Long.MIN_VALUE;
long endTime = Long.MIN_VALUE;
long lastRowFetchTime = Long.MIN_VALUE;
long rowCount = 0l;
int columnCount = 0;
try {
LOG.info("Extracted Query from : " + queryFileName);
String basicFileName = (new File(queryFileName)).getName();
String queryLabel = basicFileName.subSequence(0, basicFileName.lastIndexOf(".q")).toString();
LOG.info("Executing Query : " + queryLabel);
startTime = System.currentTimeMillis();
// Time to Connect
connTime = System.currentTimeMillis();
LOG.info("Connect Time: " + ((connTime - startTime) / 1000f) + " sec");
RunThread runThread = new RunThread(statement, query);
boolean success = processThread(runThread);
if (!success) {
return false;
}
resultSet = runThread.getResultSet();
// Time to Execute
executeTime = System.currentTimeMillis();
LOG.info("Execute Time: " + ((executeTime - connTime) / 1000f) + " sec");
columnCount = resultSet.getMetaData().getColumnCount();
while (resultSet.next()) {
if (rowCount == 0) firstRowFetchTime = System.currentTimeMillis();
rowCount++;
lastRowFetchTime = System.currentTimeMillis();
}
} catch (SQLException e) {
e.printStackTrace();
LOG.error(e.getMessage());
status = false;
} finally {
endTime = System.currentTimeMillis();
try {
LOG.info("Closing connections");
if (resultSet != null) {
resultSet.close();
}
} catch (SQLException e) {
LOG.error("[ERROR] During close: " + e.getMessage());
status = false;
}
if (!status)
LOG.error(
"Last row was fetched at "
+ new Date(lastRowFetchTime)
+ " ["
+ ((endTime - lastRowFetchTime) / 1000f)
+ " secs ago]");
LOG.info("Time to fetch 1st row : " + ((firstRowFetchTime - executeTime) / 1000f) + " sec");
LOG.info(
"Fetched "
+ rowCount
+ " rows with "
+ columnCount
+ " columns in "
+ ((endTime - executeTime) / 1000f)
+ " sec");
LOG.info("Fetch Rate: " + (rowCount * 1000f / (endTime - executeTime)) + " rows/sec ");
}
LOG.info("Total Time: " + (System.currentTimeMillis() - startTime) / 1000f + " sec ");
return status;
}
/**
* Submit query via the submit_plan tool
*
* @param submitPlanCommand command of submit_plan
* @param queryFileName name of file containing input query
* @param outputFileName name of file containing query results
* @param queryType type of query: sql, logical or physical
* @throws Exception
*/
public void submitQueriesSubmitPlan(
String submitPlanCommand,
String queryFileName,
String outputFileName,
String queryType,
long timeout)
throws Exception {
String command =
submitPlanCommand
+ " -f "
+ queryFileName
+ " --format tsv -t "
+ queryType
+ " -z "
+ Utils.getDrillTestProperties().get("ZOOKEEPERS");
LOG.debug("Executing " + command + ".");
RunThread runThread = new RunThread(command);
processThread(runThread, timeout);
Process process = runThread.getProcess();
Scanner scanner = new Scanner(process.getInputStream());
scanner.useDelimiter("\\A");
String output = scanner.hasNext() ? scanner.next() : "";
PrintWriter writer = new PrintWriter(outputFileName);
writer.write(output);
writer.close();
scanner.close();
}
/**
* Generates physical or logical plan files.
*
* @param statement statement used
* @param planFileName name of plan file
* @param queryString query for which plan is being generated
* @param type physical or logical
* @param timeout time allowed for generation of plan file
* @throws Exception
*/
public void generatePlan(
Statement statement, String planFileName, String queryString, String type, long timeout)
throws Exception {
String query = "explain plan ";
if (type.equals("logical")) {
query += "without implementation ";
}
query += "for " + queryString;
ResultSet resultSet = null;
String planString = "";
try {
LOG.debug("Submitting query:\n" + query.trim());
RunThread runThread = new RunThread(statement, query);
processThread(runThread, timeout);
resultSet = runThread.getResultSet();
resultSet.next();
resultSet.next();
planString = new String(resultSet.getBytes(2));
} finally {
if (resultSet != null) {
resultSet.close();
}
}
BufferedWriter writer = new BufferedWriter(new FileWriter(new File(planFileName)));
writer.write(planString);
writer.close();
}
/**
* Submits query via JDBC
*
* @param query SQL query string
* @param statement sql statement to execute the query with
* @param timeout time allowed for query execution
* @throws Exception
*/
public void submitQueryJDBC(String query, Statement statement, long timeout) throws Exception {
submitQueryJDBC(query, statement, null, timeout);
}
/**
* Submits query via JDBC
*
* @param query SQL query string
* @param statement sql statement to execute the query with
* @param outputFilename name of file result set is to be written to
* @param timeout time allowed for query execution
* @throws Exception
*/
public void submitQueryJDBC(
String query, Statement statement, String outputFilename, long timeout) throws Exception {
BufferedWriter writer = null;
if (outputFilename != null) {
writer = new BufferedWriter(new FileWriter(new File(outputFilename)));
}
ResultSet resultSet = null;
try {
RunThread runThread = new RunThread(statement, query);
processThread(runThread, timeout);
resultSet = runThread.getResultSet();
if (resultSet == null) {
throw runThread.getException();
}
if (outputFilename == null) {
return;
}
int columnCount = resultSet.getMetaData().getColumnCount();
columnLabels = new ArrayList<String>();
for (int i = 1; i <= columnCount; i++) {
columnLabels.add(resultSet.getMetaData().getColumnLabel(i));
}
Object[] types = new Object[columnCount];
for (int i = 1; i <= columnCount; i++) {
types[i - 1] = resultSet.getMetaData().getColumnType(i);
}
ColumnList.setTypes(types);
LOG.debug("Result set data types:");
LOG.debug(Utils.getTypesInStrings(ColumnList.getTypes()));
while (resultSet.next()) {
Object[] values = new Object[columnCount];
for (int i = 1; i <= columnCount; i++) {
try {
if (resultSet.getObject(i) == null) {
values[i - 1] = null;
continue;
}
values[i - 1] = new String(resultSet.getBytes(i));
} catch (Exception e) {
if (resultSet.getMetaData().getColumnType(i) == Types.DATE) {
values[i - 1] = resultSet.getDate(i);
} else {
values[i - 1] = resultSet.getObject(i);
}
}
}
ColumnList columnList = new ColumnList(values);
if (writer != null) {
writer.write(columnList + "\n");
}
}
if (writer != null) {
writer.close();
}
} finally {
if (resultSet != null) {
resultSet.close();
}
}
}
private boolean processThread(RunThread runThread) throws InterruptedException {
return processThread(runThread, TIMEOUT_SECONDS);
}
private boolean processThread(RunThread runThread, long timeout) throws InterruptedException {
runThread.start();
runThread.join(timeout * 1000);
if (runThread.isAlive()) {
TestVerifier.testStatus = TestVerifier.TEST_STATUS.TIMEOUT;
runThread.interrupt();
return false;
}
return true;
}
private class RunThread extends Thread {
private Statement statement;
private String query;
private String command;
private ResultSet resultSet;
private Process process;
private Exception exception;
public RunThread(Statement statement, String query) {
this.statement = statement;
this.query = query;
}
public RunThread(String command) {
this.command = command;
}
@Override
public void run() {
try {
if (statement != null) {
resultSet = statement.executeQuery(query);
} else {
process = Runtime.getRuntime().exec(command);
process.waitFor();
}
} catch (Exception e) {
TestVerifier.testStatus = TestVerifier.TEST_STATUS.EXECUTION_FAILURE;
LOG.debug("Fatal: execution of query failed. Result set size: 0.", e);
this.exception = e;
}
}
public ResultSet getResultSet() {
return resultSet;
}
public Process getProcess() {
return process;
}
public Exception getException() {
return exception;
}
}
public List<String> getColumnLabels() {
return columnLabels;
}
}
