From f6795c3f8e0645974009741a1de63440a01b8fb0 Mon Sep 17 00:00:00 2001
From: justbk <249396768@qq.com>
Date: Sat, 23 Jul 2022 16:32:21 +0800
Subject: [PATCH] support non standard create package grammer

---
 .../main/java/org/postgresql/core/Parser.java | 33 ++++++++++++++++++-
 .../java/org/postgresql/core/ParserTest.java  | 29 ++++++++++++++++
 2 files changed, 61 insertions(+), 1 deletion(-)

diff --git a/pgjdbc/src/main/java/org/postgresql/core/Parser.java b/pgjdbc/src/main/java/org/postgresql/core/Parser.java
index c18a059..fcc1cbc 100644
--- a/pgjdbc/src/main/java/org/postgresql/core/Parser.java
+++ b/pgjdbc/src/main/java/org/postgresql/core/Parser.java
@@ -607,6 +607,31 @@ public class Parser {
 
     return query.length;
   }
+    
+    /**
+     * Judge if keyword contains END str. like END or END; or *;END;*
+     * @param keyword upperCase string
+     * @return true if contains END str.
+     */
+    private static boolean isPackageEnd(String keyword) {
+      final String search = "END";
+      final char compareChar = ';';
+      if (search.equals(keyword)) {
+          return true;
+      }
+      int position = keyword.indexOf(search);
+      if (position == -1) {
+          return false;
+      }
+      if (position != 0 && keyword.charAt(position - 1) != compareChar) {
+          return false;
+      }
+      int lastPos = position + search.length();
+      if (lastPos != keyword.length() && keyword.charAt(lastPos) != compareChar) {
+          return false;
+      }
+      return true;
+    }
 
     /**
      * Judge whether the statement contains the keywords procedure, function, create, package and declare
@@ -624,7 +649,8 @@ public class Parser {
             if (queryArr[i] == null) {
                 continue;
             }
-            switch (queryArr[i].toUpperCase(Locale.ENGLISH)) {
+            String upperQuery = queryArr[i].toUpperCase(Locale.ENGLISH);
+            switch (upperQuery) {
                 case "PROCEDURE":
                 case "FUNCTION":
                 case "DECLARE":
@@ -644,6 +670,11 @@ public class Parser {
                     }
                     break;
                 default:
+                    if (haveCreate && havePackage) {
+                        if (isPackageEnd(upperQuery)) {
+                            return true;
+                        }
+                    }
                     break;
             }
         }
diff --git a/pgjdbc/src/test/java/org/postgresql/core/ParserTest.java b/pgjdbc/src/test/java/org/postgresql/core/ParserTest.java
index 4274fcd..97c3dba 100644
--- a/pgjdbc/src/test/java/org/postgresql/core/ParserTest.java
+++ b/pgjdbc/src/test/java/org/postgresql/core/ParserTest.java
@@ -216,4 +216,33 @@ public class ParserTest {
     Assert.assertEquals(34, command.getBatchRewriteValuesBraceOpenPosition());
     Assert.assertEquals(56, command.getBatchRewriteValuesBraceClosePosition());
   }
+  
+  @Test
+  public void testPackage() throws SQLException {
+    String sql = "create or replace package test_pkg_end as\n" +
+            "aa int:=1;\n" +
+            "end;";
+    String sql1 = "create or replace package test_pkg_end as\n" +
+            "aa int:=1;\n" +
+            "end ;";
+    String sql2 = "create or replace package test_pkg_end as\n" +
+            "aa int:=1;\n" +
+            "end;\n/";
+    String sql3 = "create or replace package test_pkg_end as\n" +
+            "aa int:=1;" +
+            "end;\n/";
+    String sql4 = "create or replace package test_pkg_end as\n" +
+            "aa int:=1;" +
+            "end;";
+    String sql5 = "create or replace package test_pkg_end as\n" +
+            "aa int:=1;" +
+            "end ;";
+    String sqlTests[] = new String[] {sql, sql1, sql2, sql3, sql4, sql5};
+    for (String mySql: sqlTests) {
+      List<NativeQuery> queries = Parser.parseJdbcSql(mySql,
+              false,false,
+              true, false, new String[0]);
+      assertEquals(1, queries.size());
+    }
+  }
 }