db2 sql parser support select_for_update. fixed issue #1322

src/main/java/com/alibaba/druid/sql/ast/statement/SQLSelectQueryBlock.java

@@ -33,7 +33,10 @@ public class SQLSelectQueryBlock extends SQLObjectImpl implements SQLSelectQuery
     protected SQLExpr                   where;
     protected SQLSelectGroupByClause    groupBy;
     protected SQLOrderBy                orderBy;
-    protected boolean parenthesized = false;
+    protected boolean                   parenthesized = false;
+    protected boolean                   forUpdate     = false;
+    protected boolean                   noWait        = false;
+    protected SQLExpr                   waitTime;
 
     public SQLSelectQueryBlock(){
 
@@ -117,6 +120,30 @@ public class SQLSelectQueryBlock extends SQLObjectImpl implements SQLSelectQuery
 	public void setParenthesized(boolean parenthesized) {
 		this.parenthesized = parenthesized;
 	}
+	
+    public boolean isForUpdate() {
+        return forUpdate;
+    }
+
+    public void setForUpdate(boolean forUpdate) {
+        this.forUpdate = forUpdate;
+    }
+    
+    public boolean isNoWait() {
+        return noWait;
+    }
+
+    public void setNoWait(boolean noWait) {
+        this.noWait = noWait;
+    }
+    
+    public SQLExpr getWaitTime() {
+        return waitTime;
+    }
+    
+    public void setWaitTime(SQLExpr waitTime) {
+        this.waitTime = waitTime;
+    }
 
 	@Override
     protected void accept0(SQLASTVisitor visitor) {

src/main/java/com/alibaba/druid/sql/dialect/db2/parser/DB2SelectParser.java

@@ -116,9 +116,15 @@ public class DB2SelectParser extends SQLSelectParser {
             
             if (lexer.token() == Token.FOR) {
                 lexer.nextToken();
-                acceptIdentifier("READ");
-                accept(Token.ONLY);
-                queryBlock.setForReadOnly(true);
+                
+                if (lexer.token() == Token.UPDATE) {
+                    queryBlock.setForUpdate(true);
+                    lexer.nextToken();
+                } else {
+                    acceptIdentifier("READ");
+                    accept(Token.ONLY);
+                    queryBlock.setForReadOnly(true);
+                }
             }
             
             if (lexer.token() == Token.OPTIMIZE) {

src/main/java/com/alibaba/druid/sql/dialect/db2/visitor/DB2OutputVisitor.java

@@ -58,6 +58,9 @@ public class DB2OutputVisitor extends SQLASTOutputVisitor implements DB2ASTVisit
         if (x.isForReadOnly()) {
             println();
             print0(ucase ? "FOR READ ONLY" : "for read only");
+        } else if (x.isForUpdate()) {
+            println();
+            print0(ucase ? "FOR UPDATE" : "for update");
         }
 
         if (x.getIsolation() != null) {

src/main/java/com/alibaba/druid/sql/dialect/mysql/ast/statement/MySqlSelectQueryBlock.java

@@ -22,7 +22,6 @@ import com.alibaba.druid.sql.ast.SQLCommentHint;
 import com.alibaba.druid.sql.ast.SQLExpr;
 import com.alibaba.druid.sql.ast.SQLName;
 import com.alibaba.druid.sql.ast.SQLObjectImpl;
-import com.alibaba.druid.sql.ast.SQLOrderBy;
 import com.alibaba.druid.sql.ast.statement.SQLSelectQueryBlock;
 import com.alibaba.druid.sql.dialect.mysql.ast.MySqlObject;
 import com.alibaba.druid.sql.dialect.mysql.visitor.MySqlASTVisitor;
@@ -44,9 +43,6 @@ public class MySqlSelectQueryBlock extends SQLSelectQueryBlock implements MySqlO
     private SQLName              procedureName;
     private List<SQLExpr>        procedureArgumentList;
 
-    private boolean              forUpdate       = false;
-    private boolean              noWait          = false;
-    private SQLExpr              waitTime;
     private boolean              lockInShareMode = false;
 
     private List<SQLCommentHint> hints;
@@ -74,30 +70,6 @@ public class MySqlSelectQueryBlock extends SQLSelectQueryBlock implements MySqlO
         this.hints = hints;
     }
 
-    public boolean isForUpdate() {
-        return forUpdate;
-    }
-
-    public void setForUpdate(boolean forUpdate) {
-        this.forUpdate = forUpdate;
-    }
-    
-    public boolean isNoWait() {
-        return noWait;
-    }
-
-    public void setNoWait(boolean noWait) {
-        this.noWait = noWait;
-    }
-    
-    public SQLExpr getWaitTime() {
-        return waitTime;
-    }
-    
-    public void setWaitTime(SQLExpr waitTime) {
-        this.waitTime = waitTime;
-    }
-
     public boolean isLockInShareMode() {
         return lockInShareMode;
     }

src/test/java/com/alibaba/druid/bvt/sql/db2/DB2SelectTest_13.java

@@ -0,0 +1,76 @@
+/*
+ * Copyright 1999-2101 Alibaba Group Holding Ltd.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.alibaba.druid.bvt.sql.db2;
+
+import java.util.List;
+
+import org.junit.Assert;
+
+import com.alibaba.druid.sql.DB2Test;
+import com.alibaba.druid.sql.SQLUtils;
+import com.alibaba.druid.sql.ast.SQLStatement;
+import com.alibaba.druid.sql.dialect.db2.parser.DB2StatementParser;
+import com.alibaba.druid.sql.dialect.db2.visitor.DB2SchemaStatVisitor;
+import com.alibaba.druid.stat.TableStat;
+import com.alibaba.druid.stat.TableStat.Column;
+import com.alibaba.druid.util.JdbcConstants;
+
+public class DB2SelectTest_13 extends DB2Test {
+
+    public void test_0() throws Exception {
+        String sql = "SELECT WORKDEPT, EMPNO, SALARY, BONUS, COMM "//
+                     + " FROM DSN8B10.EMP"//
+                     + " WHERE WORKDEPT IN ('D11','D21')" //
+                     + " FOR UPDATE;";
+
+        DB2StatementParser parser = new DB2StatementParser(sql);
+        List<SQLStatement> statementList = parser.parseStatementList();
+        SQLStatement stmt = statementList.get(0);
+        print(statementList);
+
+        Assert.assertEquals(1, statementList.size());
+
+        DB2SchemaStatVisitor visitor = new DB2SchemaStatVisitor();
+        stmt.accept(visitor);
+
+//        System.out.println("Tables : " + visitor.getTables());
+//        System.out.println("fields : " + visitor.getColumns());
+//        System.out.println("coditions : " + visitor.getConditions());
+//        System.out.println("orderBy : " + visitor.getOrderByColumns());
+
+        Assert.assertEquals(1, visitor.getTables().size());
+        Assert.assertEquals(5, visitor.getColumns().size());
+        Assert.assertEquals(1, visitor.getConditions().size());
+
+        Assert.assertTrue(visitor.getTables().containsKey(new TableStat.Name("DSN8B10.EMP")));
+
+         Assert.assertTrue(visitor.getColumns().contains(new Column("DSN8B10.EMP", "WORKDEPT")));
+        // Assert.assertTrue(visitor.getColumns().contains(new Column("mytable", "first_name")));
+        // Assert.assertTrue(visitor.getColumns().contains(new Column("mytable", "full_name")));
+
+        Assert.assertEquals("SELECT WORKDEPT, EMPNO, SALARY, BONUS, COMM"
+                + "\nFROM DSN8B10.EMP"
+                + "\nWHERE WORKDEPT IN ('D11', 'D21')"
+                + "\nFOR UPDATE", //
+                            SQLUtils.toSQLString(stmt, JdbcConstants.DB2));
+        
+        Assert.assertEquals("select WORKDEPT, EMPNO, SALARY, BONUS, COMM"
+                + "\nfrom DSN8B10.EMP"
+                + "\nwhere WORKDEPT in ('D11', 'D21')"
+                + "\nfor update", //
+                            SQLUtils.toSQLString(stmt, JdbcConstants.DB2, SQLUtils.DEFAULT_LCASE_FORMAT_OPTION));
+    }
+}