feat: support SET statements combining with other queries with semico…

…lon in PreparedStatement

Previoulsy, PreparedStatement supported executing several statements delimited with semicolon,
however, it was failing for SELECT and SET combination as the driver was expecting "query results"
from SET.

Now we ignore SET results, so it enables users to save network roundtrip

pgjdbc/src/main/java/org/postgresql/core/Parser.java

@@ -235,6 +235,8 @@ public static List<NativeQuery> parseJdbcSql(String query, boolean standardConfo
             currentCommandType = SqlCommandType.SELECT;
           } else if (wordLength == 4 && parseWithKeyword(aChars, keywordStart)) {
             currentCommandType = SqlCommandType.WITH;
+          } else if (wordLength == 3 && parseSetKeyword(aChars, keywordStart)) {
+            currentCommandType = SqlCommandType.SET;
           } else if (wordLength == 6 && parseInsertKeyword(aChars, keywordStart)) {
             if (!isInsertPresent && (nativeQueries == null || nativeQueries.isEmpty())) {
               // Only allow rewrite for insert command starting with the insert keyword.
@@ -732,6 +734,23 @@ public static boolean parseSelectKeyword(final char[] query, int offset) {
         && (query[offset + 5] | 32) == 't';
   }
 
+  /**
+   * Parse string to check presence of SET keyword regardless of case.
+   *
+   * @param query char[] of the query statement
+   * @param offset position of query to start checking
+   * @return boolean indicates presence of word
+   */
+  public static boolean parseSetKeyword(final char[] query, int offset) {
+    if (query.length < (offset + 3)) {
+      return false;
+    }
+
+    return (query[offset] | 32) == 's'
+        && (query[offset + 1] | 32) == 'e'
+        && (query[offset + 2] | 32) == 't';
+  }
+
   /**
    * Parse string to check presence of CREATE keyword regardless of case.
    *

pgjdbc/src/main/java/org/postgresql/core/SqlCommandType.java

@@ -25,5 +25,24 @@ public enum SqlCommandType {
   SELECT,
   WITH,
   CREATE,
-  ALTER;
+  ALTER,
+  SET;
+
+  private static final SqlCommandType[] SQL_COMMAND_TYPES = SqlCommandType.values();
+
+  /**
+   * Returns the SqlCommandType for the given command status.
+   * Returns BLANK if the no match is found.
+   *
+   * @param commandStatus the command status
+   * @return the SqlCommandType for the given command status
+   */
+  public static SqlCommandType fromCommandStatus(String commandStatus) {
+    for (SqlCommandType type : SQL_COMMAND_TYPES) {
+      if (type.name().equalsIgnoreCase(commandStatus)) {
+        return type;
+      }
+    }
+    return SqlCommandType.BLANK;
+  }
 }

pgjdbc/src/main/java/org/postgresql/jdbc/BatchResultHandler.java

@@ -98,14 +98,19 @@ public void handleCommandStatus(String status, long updateCount, long insertOID)
       this.latestGeneratedRows = null;
     }
 
-    if (resultIndex >= queries.length) {
+    //SET commands can be ignored for the purposes of update counts
+    boolean statusIsFromSet = status.regionMatches(true, 0, "SET", 0, 3);
+
+    if (!statusIsFromSet && resultIndex >= queries.length) {
       handleError(new PSQLException(GT.tr("Too many update results were returned."),
           PSQLState.TOO_MANY_RESULTS));
       return;
     }
     latestGeneratedKeysRs = null;
 
-    longUpdateCounts[resultIndex++] = updateCount;
+    if (!statusIsFromSet) {
+      longUpdateCounts[resultIndex++] = updateCount;
+    }
   }
 
   private boolean isAutoCommit() {

pgjdbc/src/main/java/org/postgresql/jdbc/PgStatement.java

@@ -19,6 +19,7 @@
 import org.postgresql.core.ResultCursor;
 import org.postgresql.core.ResultHandlerBase;
 import org.postgresql.core.SqlCommand;
+import org.postgresql.core.SqlCommandType;
 import org.postgresql.core.Tuple;
 import org.postgresql.util.GT;
 import org.postgresql.util.PSQLException;
@@ -213,7 +214,7 @@ private void append(ResultWrapper newResult) {
       if (results == null) {
         lastResult = results = newResult;
       } else {
-        castNonNull(lastResult).append(newResult);
+        lastResult = results = castNonNull(lastResult).append(newResult);
       }
     }
 
@@ -222,15 +223,23 @@ public void handleResultRows(Query fromQuery, Field[] fields, List<Tuple> tuples
         @Nullable ResultCursor cursor) {
       try {
         ResultSet rs = PgStatement.this.createResultSet(fromQuery, fields, tuples, cursor);
-        append(new ResultWrapper(rs));
+        append(new ResultWrapper(rs, commandTypeOfQuery(fromQuery)));
       } catch (SQLException e) {
         handleError(e);
       }
     }
 
+    private SqlCommandType commandTypeOfQuery(Query query) {
+      SqlCommand sqlCommand = query.getSqlCommand();
+      if (sqlCommand == null) {
+        return SqlCommandType.BLANK;
+      }
+      return sqlCommand.getType();
+    }
+
     @Override
     public void handleCommandStatus(String status, long updateCount, long insertOID) {
-      append(new ResultWrapper(updateCount, insertOID));
+      append(new ResultWrapper(updateCount, insertOID, SqlCommandType.fromCommandStatus(status)));
     }
 
     @Override
@@ -901,7 +910,7 @@ private BatchResultHandler internalExecuteBatch() throws SQLException {
       try (ResourceLock ignore = lock.obtain()) {
         checkClosed();
         if (wantsGeneratedKeysAlways) {
-          generatedKeys = new ResultWrapper(handler.getGeneratedKeys());
+          generatedKeys = new ResultWrapper(handler.getGeneratedKeys(), SqlCommandType.BLANK);
         }
       }
     }

pgjdbc/src/main/java/org/postgresql/jdbc/ResultWrapper.java

@@ -6,6 +6,8 @@
 
 package org.postgresql.jdbc;
 
+import org.postgresql.core.SqlCommandType;
+
 import org.checkerframework.checker.nullness.qual.Nullable;
 import org.checkerframework.dataflow.qual.Pure;
 
@@ -18,13 +20,15 @@
  * @author Oliver Jowett (oliver@opencloud.com)
  */
 public class ResultWrapper {
-  public ResultWrapper(@Nullable ResultSet rs) {
+  public ResultWrapper(@Nullable ResultSet rs, SqlCommandType commandType) {
     this.rs = rs;
+    this.commandType = commandType;
     this.updateCount = -1;
     this.insertOID = -1;
   }
 
-  public ResultWrapper(long updateCount, long insertOID) {
+  public ResultWrapper(long updateCount, long insertOID, SqlCommandType commandType) {
+    this.commandType = commandType;
     this.rs = null;
     this.updateCount = updateCount;
     this.insertOID = insertOID;
@@ -47,17 +51,35 @@ public long getInsertOID() {
     return next;
   }
 
-  public void append(ResultWrapper newResult) {
+  /**
+   * Append a result to its internal chain of results.
+   * It has a special behavior for {@code SET} commands as {@code SET} is discarded if there are
+   * other results in the chain.
+   * If this is a {@code SET} command, the {@code newResult} is returned has the new head of
+   * the chain.
+   * If the newResult is a {@code SET} command, it's not appended and this is returned.
+   *
+   * @param newResult the result to append
+   * @return the head of the chain
+   */
+  public ResultWrapper append(ResultWrapper newResult) {
+    if (commandType == SqlCommandType.SET) {
+      return newResult;
+    }
+    if (newResult.commandType == SqlCommandType.SET) {
+      return this;
+    }
     ResultWrapper tail = this;
     while (tail.next != null) {
       tail = tail.next;
     }
-
     tail.next = newResult;
+    return this;
   }
 
   private final @Nullable ResultSet rs;
   private final long updateCount;
   private final long insertOID;
+  private final SqlCommandType commandType;
   private @Nullable ResultWrapper next;
 }

pgjdbc/src/test/java/org/postgresql/core/ParserTest.java

@@ -13,6 +13,8 @@
 import org.junit.Assert;
 import org.junit.Ignore;
 import org.junit.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.ValueSource;
 
 import java.sql.SQLException;
 import java.util.List;
@@ -114,6 +116,16 @@ public void testSelectCommandParsing() {
     assertTrue("Failed to correctly parse lower case command.", Parser.parseSelectKeyword(command, 0));
   }
 
+  /**
+   * Test SET command parsing.
+   */
+  @ParameterizedTest
+  @ValueSource(strings = {"SET", "set", "sEt", "seT", "Set", "sET", "SeT", "seT"})
+  public void testSetCommandParsing(String set) {
+    char[] command = set.toCharArray();
+    assertTrue("Parser.parseSetKeyword(\"" + set + "\", 0)", Parser.parseSetKeyword(command, 0));
+  }
+
   @Test
   public void testEscapeProcessing() throws Exception {
     assertEquals("DATE '1999-01-09'", Parser.replaceProcessing("{d '1999-01-09'}", true, false));
@@ -230,6 +242,15 @@ public void insertMultiInsert() throws SQLException {
     Assert.assertEquals(56, command.getBatchRewriteValuesBraceClosePosition());
   }
 
+  @Test
+  public void setVariable() throws SQLException {
+    String query =
+        "set search_path to 'public'";
+    List<NativeQuery> qry = Parser.parseJdbcSql(query, true, true, true, true, true);
+    SqlCommand command = qry.get(0).getCommand();
+    Assert.assertEquals("command type of " + query, SqlCommandType.SET, command.getType());
+  }
+
   @Test
   public void valuesTableParse() throws SQLException {
     String query = "insert into values_table (id, name) values (?,?)";