From af51ed1317154fed3e0a5b39cd9e932cf1efbafc Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E9=AB=98=E9=93=81?= <shaojin.wensj@alibaba-inc.com>
Date: Sun, 31 Jan 2021 02:57:50 +0800
Subject: [PATCH] improved postgresql parser, fix #4154

---
 .../postgresql/ast/stmt/PGSelectQueryBlock.java       |  9 +++++++++
 .../sql/dialect/postgresql/parser/PGSelectParser.java |  5 +++++
 .../dialect/postgresql/visitor/PGOutputVisitor.java   |  7 +++++--
 src/test/resources/bvt/parser/postgresql-8.txt        | 11 +++++++++++
 4 files changed, 30 insertions(+), 2 deletions(-)
 create mode 100644 src/test/resources/bvt/parser/postgresql-8.txt

diff --git a/src/main/java/com/alibaba/druid/sql/dialect/postgresql/ast/stmt/PGSelectQueryBlock.java b/src/main/java/com/alibaba/druid/sql/dialect/postgresql/ast/stmt/PGSelectQueryBlock.java
index 93c4f8cb2c..94461ce8b7 100644
--- a/src/main/java/com/alibaba/druid/sql/dialect/postgresql/ast/stmt/PGSelectQueryBlock.java
+++ b/src/main/java/com/alibaba/druid/sql/dialect/postgresql/ast/stmt/PGSelectQueryBlock.java
@@ -145,6 +145,7 @@ public class PGSelectQueryBlock extends SQLSelectQueryBlock implements PGSQLObje
 
         private List<SQLExpr> of = new ArrayList<SQLExpr>(2);
         private boolean       noWait;
+        private boolean       skipLocked;
         private Option        option;
 
         public Option getOption() {
@@ -171,6 +172,14 @@ public class PGSelectQueryBlock extends SQLSelectQueryBlock implements PGSQLObje
             this.noWait = noWait;
         }
 
+        public boolean isSkipLocked() {
+            return skipLocked;
+        }
+
+        public void setSkipLocked(boolean skipLocked) {
+            this.skipLocked = skipLocked;
+        }
+
         @Override
         public void accept0(PGASTVisitor visitor) {
             if (visitor.visit(this)) {
diff --git a/src/main/java/com/alibaba/druid/sql/dialect/postgresql/parser/PGSelectParser.java b/src/main/java/com/alibaba/druid/sql/dialect/postgresql/parser/PGSelectParser.java
index 21afadec14..2899c712c4 100644
--- a/src/main/java/com/alibaba/druid/sql/dialect/postgresql/parser/PGSelectParser.java
+++ b/src/main/java/com/alibaba/druid/sql/dialect/postgresql/parser/PGSelectParser.java
@@ -26,6 +26,7 @@ import com.alibaba.druid.sql.dialect.postgresql.ast.stmt.PGSelectQueryBlock;
 import com.alibaba.druid.sql.dialect.postgresql.ast.stmt.PGSelectQueryBlock.IntoOption;
 import com.alibaba.druid.sql.ast.statement.SQLValuesQuery;
 import com.alibaba.druid.sql.parser.*;
+import com.alibaba.druid.util.FnvHash;
 
 import java.util.List;
 
@@ -231,6 +232,10 @@ public class PGSelectParser extends SQLSelectParser {
             if (lexer.token() == Token.NOWAIT) {
                 lexer.nextToken();
                 forClause.setNoWait(true);
+            } else  if (lexer.identifierEquals(FnvHash.Constants.SKIP)) {
+                lexer.nextToken();
+                acceptIdentifier("LOCKED");
+                forClause.setSkipLocked(true);
             }
 
             queryBlock.setForClause(forClause);
diff --git a/src/main/java/com/alibaba/druid/sql/dialect/postgresql/visitor/PGOutputVisitor.java b/src/main/java/com/alibaba/druid/sql/dialect/postgresql/visitor/PGOutputVisitor.java
index 70d914cd62..4486608f44 100644
--- a/src/main/java/com/alibaba/druid/sql/dialect/postgresql/visitor/PGOutputVisitor.java
+++ b/src/main/java/com/alibaba/druid/sql/dialect/postgresql/visitor/PGOutputVisitor.java
@@ -67,12 +67,13 @@ public class PGOutputVisitor extends SQLASTOutputVisitor implements PGASTVisitor
     public boolean visit(ForClause x) {
         print0(ucase ? "FOR " : "for ");
         if (ForClause.Option.UPDATE.equals(x.getOption())) {
-            print0(ucase ? "UPDATE " : "update ");
+            print0(ucase ? "UPDATE" : "update");
         } else if (ForClause.Option.SHARE.equals(x.getOption())) {
-            print0(ucase ? "SHARE " : "share ");
+            print0(ucase ? "SHARE" : "share");
         }
 
         if (x.getOf().size() > 0) {
+            print(' ');
             for (int i = 0; i < x.getOf().size(); ++i) {
                 if (i != 0) {
                     println(", ");
@@ -83,6 +84,8 @@ public class PGOutputVisitor extends SQLASTOutputVisitor implements PGASTVisitor
 
         if (x.isNoWait()) {
             print0(ucase ? " NOWAIT" : " nowait");
+        } else if (x.isSkipLocked()) {
+            print0(ucase ? " SKIP LOCKED" : " skip locked");
         }
 
         return false;
diff --git a/src/test/resources/bvt/parser/postgresql-8.txt b/src/test/resources/bvt/parser/postgresql-8.txt
new file mode 100644
index 0000000000..5603959f87
--- /dev/null
+++ b/src/test/resources/bvt/parser/postgresql-8.txt
@@ -0,0 +1,11 @@
+select * from msg.msg_pool
+where status=?
+order by create_time
+limit 100 for update skip locked
+---------------------------
+SELECT *
+FROM msg.msg_pool
+WHERE status = ?
+ORDER BY create_time
+LIMIT 100
+FOR UPDATE SKIP LOCKED
\ No newline at end of file