fix: Parser does not handle version 14 function syntax Issue pgjdbc#2507

davecramer · Dec 21, 2022 · dda9845 · dda9845
1 parent df3278e
commit dda9845
Show file tree

Hide file tree

Showing 2 changed files with 39 additions and 1 deletion.
diff --git a/pgjdbc/src/main/java/org/postgresql/core/Parser.java b/pgjdbc/src/main/java/org/postgresql/core/Parser.java
@@ -72,6 +72,7 @@ public static List<NativeQuery> parseJdbcSql(String query, boolean standardConfo
     boolean isInsertPresent = false;
     boolean isReturningPresent = false;
     boolean isReturningPresentPrev = false;
+    boolean isAtomicPresent = false;
     SqlCommandType currentCommandType = SqlCommandType.BLANK;
     SqlCommandType prevCommandType = SqlCommandType.BLANK;
     int numberOfStatements = 0;
@@ -139,7 +140,7 @@ public static List<NativeQuery> parseJdbcSql(String query, boolean standardConfo
           break;
 
         case ';':
-          if (inParen == 0) {
+          if (!isAtomicPresent && inParen == 0) {
             if (!whitespaceOnly) {
               numberOfStatements++;
               nativeSql.append(aChars, fragmentStart, i - fragmentStart);
@@ -238,12 +239,17 @@ public static List<NativeQuery> parseJdbcSql(String query, boolean standardConfo
               isCurrentReWriteCompatible = false;
             }
           }
+
         } else if (currentCommandType == SqlCommandType.WITH
             && inParen == 0) {
           SqlCommandType command = parseWithCommandType(aChars, i, keywordStart, wordLength);
           if (command != null) {
             currentCommandType = command;
           }
+        } else if (currentCommandType == SqlCommandType.CREATE) {
+          if ( wordLength == 6 && parseAtomicKeyword(aChars, keywordStart)) {
+            isAtomicPresent = true;
+          }
         }
         if (inParen != 0 || aChar == ')') {
           // RETURNING and VALUES cannot be present in braces
@@ -608,6 +614,24 @@ public static boolean parseInsertKeyword(final char[] query, int offset) {
         && (query[offset + 5] | 32) == 't';
   }
 
+  /**
+   Parse string to check presence of ATOMIC 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 parseAtomicKeyword(final char[] query, int offset) {
+    if (query.length < (offset + 7)) {
+      return false;
+    }
+    return (query[offset] | 32) == 'a'
+        && (query[offset + 1] | 32) == 't'
+        && (query[offset + 2] | 32) == 'o'
+        && (query[offset + 3] | 32) == 'm'
+        && (query[offset + 4] | 32) == 'i'
+        && (query[offset + 5] | 32) == 'c';
+  }
   /**
    * Parse string to check presence of MOVE keyword regardless of case.
    *

diff --git a/pgjdbc/src/test/java/org/postgresql/core/ParserTest.java b/pgjdbc/src/test/java/org/postgresql/core/ParserTest.java
@@ -284,4 +284,18 @@ public void alterTableParseWithOnUpdateClause() throws SQLException {
     Assert.assertFalse("No returning keyword should be present", command.isReturningKeywordPresent());
     Assert.assertEquals(SqlCommandType.ALTER, command.getType());
   }
+
+  @Test
+  public void testParseV14functions() throws SQLException {
+    String[] returningColumns = {"*"};
+    String query = "CREATE OR REPLACE FUNCTION asterisks(n int)\n"
+        + "  RETURNS SETOF text\n"
+        + "  LANGUAGE sql IMMUTABLE STRICT PARALLEL SAFE\n"
+        + "BEGIN ATOMIC\n"
+        + "SELECT repeat('*', g) FROM generate_series (1, n) g; \n"
+        + "END;";
+    List<NativeQuery> qry = Parser.parseJdbcSql(query, true, true, true, true, true, returningColumns);
+    Assert.assertNotNull(qry);
+    Assert.assertEquals("There should only be one query returned here", 1, qry.size());
+  }
 }