DRILL-6834: Added possibility to run queries that don't return result…

… set. (#521)

An option "exec.return_result_set_for_ddl" was introduced in Drill, when it is false, some queries don't return result set, that cause an SQL exception in Drill Test Framework.
To fix that, Statement execute() method now used instead of executeQuery(). Also added required null checks.

framework/src/main/java/org/apache/drill/test/framework/DrillTestJdbc.java

@@ -189,7 +189,7 @@ private void executeSetupQuery(String query) throws SQLException {
       statement = connection.createStatement();
     }
     try {
-      resultSet = statement.executeQuery(query);
+      resultSet = statement.execute(query) ? statement.getResultSet() : null;
     } finally {
       if (resultSet != null) {
         // TODO(DRILL-2560) : Once resolved, we can use statement#executeUpdate instead of exhausting the result set
@@ -207,7 +207,7 @@ private void executeQuery(String query) throws IOException, SQLException {
     CancelQuery c = null;
     try {
       statement = connection.createStatement();
-      resultSet = statement.executeQuery(query);      
+      resultSet = statement.execute(query) ? statement.getResultSet() : null;
       if (cancelQuery) {
     	c = new CancelQuery(statement);
     	c.start();
@@ -238,7 +238,7 @@ private void executeQuery(String query) throws IOException, SQLException {
       columnNullabilities = Lists.newArrayList();
       columnSizes = Lists.newArrayList();
       columnPrecisions = Lists.newArrayList();
-      int columnCount = resultSet.getMetaData().getColumnCount();
+      int columnCount = resultSet != null ? resultSet.getMetaData().getColumnCount() : 0;
       for (int i = 1; i <= columnCount; i++) {
         columnLabels.add(resultSet.getMetaData().getColumnLabel(i));
         columnTypes.add(resultSet.getMetaData().getColumnType(i));
@@ -250,12 +250,13 @@ private void executeQuery(String query) throws IOException, SQLException {
       LOG.debug("Result set data types:");
       LOG.debug(Utils.getTypesInStrings(columnTypes));
 
-      while (resultSet.next()) {
-        List<Object> values = Utils.getRowValues(resultSet);
-        ColumnList columnList = new ColumnList(columnTypes, values);
-        writer.write(columnList + "\n");
+      if (resultSet != null) {
+        while (resultSet.next()) {
+          List<Object> values = Utils.getRowValues(resultSet);
+          ColumnList columnList = new ColumnList(columnTypes, values);
+          writer.write(columnList + "\n");
+        }
       }
-
     } catch (IllegalArgumentException | IllegalAccessException | IOException e1) {
 		LOG.warn(e1);
 	} finally {
@@ -280,45 +281,45 @@ private void executeLimitZeroQuery(String query) throws IOException, SQLExceptio
 
 	BufferedWriter writer = new BufferedWriter(new FileWriter(new File(outputFilename),true));	
     statement = connection.createStatement();
-    ResultSet resultSet = statement.executeQuery(query);
+    ResultSet resultSet = statement.execute(query) ? statement.getResultSet() : null;
     List columnLabels = new ArrayList<String>();
     List<Integer> columnTypes = Lists.newArrayList();
     List<Integer> columnNullabilities = Lists.newArrayList();
     List<Integer> columnSizes = Lists.newArrayList();
     List<Integer> columnPrecisions = Lists.newArrayList();
 
     try {
-      int columnCount = resultSet.getMetaData().getColumnCount();
+      int columnCount = resultSet != null ? resultSet.getMetaData().getColumnCount() : 0;
       for (int i = 1; i <= columnCount; i++) {
         columnLabels.add(resultSet.getMetaData().getColumnLabel(i));
         columnTypes.add(resultSet.getMetaData().getColumnType(i));
         columnNullabilities.add(resultSet.getMetaData().isNullable(i));
         columnSizes.add(resultSet.getMetaData().getColumnDisplaySize(i));
         columnPrecisions.add(resultSet.getMetaData().getPrecision(i));
       }
-      
-      String msg = "\nlimit 0: " + query + "\n" 
-    		  + "limit 0: " + columnLabels + "\n" 
-    		  + "regular: " + this.columnLabels + "\n"
-    		  + "\nlimit 0: " + Utils.getTypesInStrings(columnTypes) + "\n"
-    		  + "regular: " + Utils.getTypesInStrings(this.columnTypes) + "\n"
-    		  + "\nlimit 0: " + Utils.getNullabilitiesInStrings(columnNullabilities) + "\n"
-    		  + "regular: " + Utils.getNullabilitiesInStrings(this.columnNullabilities) + "\n"
-    		  + "\nlimit 0: " + columnSizes + "\n"
-    		  + "regular: " + this.columnSizes + "\n"
-    		  + "\nlimit 0: " + columnPrecisions + "\n"
-    		  + "regular: " + this.columnPrecisions + "\n";
+
+      String msg = "\nlimit 0: " + query + "\n"
+              + "limit 0: " + columnLabels + "\n"
+              + "regular: " + this.columnLabels + "\n"
+              + "\nlimit 0: " + Utils.getTypesInStrings(columnTypes) + "\n"
+              + "regular: " + Utils.getTypesInStrings(this.columnTypes) + "\n"
+              + "\nlimit 0: " + Utils.getNullabilitiesInStrings(columnNullabilities) + "\n"
+              + "regular: " + Utils.getNullabilitiesInStrings(this.columnNullabilities) + "\n"
+              + "\nlimit 0: " + columnSizes + "\n"
+              + "regular: " + this.columnSizes + "\n"
+              + "\nlimit 0: " + columnPrecisions + "\n"
+              + "regular: " + this.columnPrecisions + "\n";
       writer.append(msg);
-      
-      if (!columnLabels.equals(this.columnLabels) 
-    	  || !columnTypes.equals(this.columnTypes)
+
+      if (!columnLabels.equals(this.columnLabels)
+          || !columnTypes.equals(this.columnTypes)
           || !isNullabilityCompatible(columnNullabilities, this.columnNullabilities)
-          || !columnSizes.equals(this.columnSizes) 
+          || !columnSizes.equals(this.columnSizes)
           || !columnPrecisions.equals(this.columnPrecisions))  {
         LOG.info(msg);
         setTestStatus(TestStatus.DATA_VERIFICATION_FAILURE);
         exception = exception == null? new VerificationException(msg)
-        	: new VerificationException(exception + "\n" + msg);
+            : new VerificationException(exception + "\n" + msg);
       }
     } catch (IllegalArgumentException | IllegalAccessException e1) {
       LOG.warn(e1);

framework/src/main/java/org/apache/drill/test/framework/TestDriver.java

@@ -587,13 +587,16 @@ public void setup() throws IOException, InterruptedException, URISyntaxException
       connection = connectionPool.getOrCreateConnection();
       for (String query : setupQueries) {
         LOG.info(">> Query: " + query + ";");
-        LOG.info(Utils.getSqlResult(Utils.execSQL(query, connection)));
+        ResultSet resultSet = Utils.execSQL(query, connection);
+        if (resultSet != null) {
+          LOG.info(Utils.getSqlResult(resultSet));
+        }
       }
 
       // Initializing variables for reporting
       String getCommitId = "SELECT version, commit_id from sys.version";
       ResultSet resultSet = Utils.execSQL(getCommitId, connection);
-      while(resultSet.next()) {
+      while (resultSet.next()) {
         commitId = resultSet.getString("commit_id");
         version = resultSet.getString("version");
       }
@@ -631,7 +634,10 @@ private void teardown() {
 	  String[] teardownQueries = Utils.getSqlStatements(afterRunQueryFilename);
       for (String query : teardownQueries) {
         LOG.info(">> Query: " + query + ";");
-        LOG.info(Utils.getSqlResult(Utils.execSQL(query, connection)));
+        ResultSet resultSet = Utils.execSQL(query, connection);
+        if(resultSet != null) {
+          LOG.info(Utils.getSqlResult(resultSet));
+        }
       }
       LOG.info(DrillTestDefaults.LINE_BREAK);
     } catch (IOException e) {
@@ -800,7 +806,7 @@ private void queryMemoryUsage() throws IOException, SQLException {
       resultSet = Utils.execSQL(query, connection);
 
       List columnLabels = new ArrayList<String>();
-      int columnCount = resultSet.getMetaData().getColumnCount();
+      int columnCount = resultSet != null ? resultSet.getMetaData().getColumnCount() : 0;
       for (int i = 1; i <= columnCount; i++) {
         columnLabels.add(resultSet.getMetaData().getColumnLabel(i));
       }
@@ -812,13 +818,15 @@ private void queryMemoryUsage() throws IOException, SQLException {
       LOG.debug("Result set data types:");
       LOG.debug(Utils.getTypesInStrings(types));
 
-      while (resultSet.next()) {
-        List<Object> values = Utils.getRowValues(resultSet);
-        ColumnList columnList = new ColumnList(types, values);
-        writer.write(columnList + "\n");
-        for (int i = 0; i < 3; i++) {
-          memUsage[0][i] = memUsage[1][i];
-          memUsage[1][i] = (Long) values.get(i)/1024/1024;
+      if (resultSet != null) {
+        while (resultSet.next()) {
+          List<Object> values = Utils.getRowValues(resultSet);
+          ColumnList columnList = new ColumnList(types, values);
+          writer.write(columnList + "\n");
+          for (int i = 0; i < 3; i++) {
+            memUsage[0][i] = memUsage[1][i];
+            memUsage[1][i] = (Long) values.get(i) / 1024 / 1024;
+          }
         }
       }
       connectionPool.releaseConnection(connection);

framework/src/main/java/org/apache/drill/test/framework/Utils.java

@@ -770,14 +770,15 @@ public static CmdConsOut execCmd(String cmd){
   public static ResultSet execSQL(String sql, Connection connection) throws SQLException {
 	try {
 	  Statement statement = connection.createStatement();
-	  return statement.executeQuery(sql);
+      return statement.execute(sql) ? statement.getResultSet() : null;
 	} catch (SQLException e) {
 	  LOG.error(e.getMessage());
 	  e.printStackTrace();
       try {
 		connection.close();
 	  } catch (SQLException e1) {
-		LOG.error(e.getMessage());e1.printStackTrace();
+		LOG.error(e.getMessage());
+		e1.printStackTrace();
 	  }
       throw e;
 	}