sql parser improve sql server support

src/main/java/com/alibaba/druid/sql/dialect/mysql/parser/MySqlStatementParser.java

@@ -23,6 +23,7 @@ import com.alibaba.druid.sql.ast.SQLName;
 import com.alibaba.druid.sql.ast.SQLOrderBy;
 import com.alibaba.druid.sql.ast.SQLStatement;
 import com.alibaba.druid.sql.ast.expr.SQLCharExpr;
+import com.alibaba.druid.sql.ast.expr.SQLListExpr;
 import com.alibaba.druid.sql.ast.expr.SQLLiteralExpr;
 import com.alibaba.druid.sql.ast.expr.SQLQueryExpr;
 import com.alibaba.druid.sql.ast.expr.SQLVariantRefExpr;
@@ -199,27 +200,7 @@ public class MySqlStatementParser extends SQLStatementParser {
             stmt.setTableSource(tableSource);
         }
 
-        accept(Token.SET);
-
-        for (;;) {
-            SQLUpdateSetItem item = new SQLUpdateSetItem();
-            item.setColumn(this.exprParser.name());
-            if (lexer.token() == Token.EQ) {
-                lexer.nextToken();
-            } else {
-                accept(Token.COLONEQ);
-            }
-            item.setValue(this.exprParser.expr());
-
-            stmt.getItems().add(item);
-
-            if (lexer.token() == (Token.COMMA)) {
-                lexer.nextToken();
-                continue;
-            }
-
-            break;
-        }
+        parseUpdateSet(stmt);
 
         if (lexer.token() == (Token.WHERE)) {
             lexer.nextToken();
@@ -2361,4 +2342,35 @@ public class MySqlStatementParser extends SQLStatementParser {
 
         return stmt;
     }
+    
+    protected void parseUpdateSet(SQLUpdateStatement update) {
+        accept(Token.SET);
+
+        for (;;) {
+            SQLUpdateSetItem item = new SQLUpdateSetItem();
+
+            if (lexer.token() == (Token.LPAREN)) {
+                lexer.nextToken();
+                SQLListExpr list = new SQLListExpr();
+                this.exprParser.exprList(list.getItems());
+                accept(Token.RPAREN);
+                item.setColumn(list);
+            } else {
+                item.setColumn(this.exprParser.primary());
+            }
+            if (lexer.token() == Token.COLONEQ) {
+                lexer.nextToken();
+            } else {
+                accept(Token.EQ);
+            }
+            item.setValue(this.exprParser.expr());
+            update.getItems().add(item);
+
+            if (lexer.token() != Token.COMMA) {
+                break;
+            }
+
+            lexer.nextToken();
+        }
+    }
 }

src/main/java/com/alibaba/druid/sql/dialect/oracle/parser/OracleUpdateParser.java

@@ -55,7 +55,7 @@ public class OracleUpdateParser extends SQLStatementParser {
             }
         }
 
-        parseSet(update);
+        parseUpdateSet(update);
 
         parseWhere(update);
 
@@ -117,30 +117,4 @@ public class OracleUpdateParser extends SQLStatementParser {
         }
     }
 
-    private void parseSet(OracleUpdateStatement update) {
-        accept(Token.SET);
-
-        for (;;) {
-            SQLUpdateSetItem item = new SQLUpdateSetItem();
-
-            if (lexer.token() == (Token.LPAREN)) {
-                lexer.nextToken();
-                SQLListExpr list = new SQLListExpr();
-                this.exprParser.exprList(list.getItems());
-                accept(Token.RPAREN);
-                item.setColumn(list);
-            } else {
-                item.setColumn(this.exprParser.primary());
-            }
-            accept(Token.EQ);
-            item.setValue(this.exprParser.expr());
-            update.getItems().add(item);
-
-            if (lexer.token() != Token.COMMA) {
-                break;
-            }
-
-            lexer.nextToken();
-        }
-    }
 }

src/main/java/com/alibaba/druid/sql/dialect/postgresql/parser/PGSQLStatementParser.java

@@ -60,23 +60,7 @@ public class PGSQLStatementParser extends SQLStatementParser {
         SQLTableSource tableSource = this.exprParser.createSelectParser().parseTableSource();
         udpateStatement.setTableSource(tableSource);
 
-        accept(Token.SET);
-
-        for (;;) {
-            SQLUpdateSetItem item = new SQLUpdateSetItem();
-            item.setColumn(this.exprParser.name());
-            accept(Token.EQ);
-            item.setValue(this.exprParser.expr());
-
-            udpateStatement.getItems().add(item);
-
-            if (lexer.token() == (Token.COMMA)) {
-                lexer.nextToken();
-                continue;
-            }
-
-            break;
-        }
+        parseUpdateSet(udpateStatement);
 
         if (lexer.token() == (Token.WHERE)) {
             lexer.nextToken();

src/main/java/com/alibaba/druid/sql/dialect/sqlserver/ast/SQLServerStatement.java

@@ -17,6 +17,6 @@ package com.alibaba.druid.sql.dialect.sqlserver.ast;
 
 import com.alibaba.druid.sql.ast.SQLStatement;
 
-public interface SQLServerStatement extends SQLStatement {
+public interface SQLServerStatement extends SQLStatement, SQLServerObject {
 
 }

src/main/java/com/alibaba/druid/sql/dialect/sqlserver/ast/stmt/SQLServerUpdateStatement.java

@@ -0,0 +1,55 @@
+/*
+ * Copyright 1999-2011 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.sql.dialect.sqlserver.ast.stmt;
+
+import com.alibaba.druid.sql.ast.statement.SQLTableSource;
+import com.alibaba.druid.sql.ast.statement.SQLUpdateStatement;
+import com.alibaba.druid.sql.dialect.sqlserver.ast.SQLServerStatement;
+import com.alibaba.druid.sql.dialect.sqlserver.visitor.SQLServerASTVisitor;
+import com.alibaba.druid.sql.visitor.SQLASTVisitor;
+
+public class SQLServerUpdateStatement extends SQLUpdateStatement implements SQLServerStatement {
+    private static final long serialVersionUID = 1L;
+    
+    private SQLTableSource from;
+
+    
+    public SQLTableSource getFrom() {
+        return from;
+    }
+
+    
+    public void setFrom(SQLTableSource from) {
+        this.from = from;
+    }
+    
+    @Override
+    protected void accept0(SQLASTVisitor visitor) {
+        this.accept0((SQLServerASTVisitor) visitor);
+    }
+
+    @Override
+    public void accept0(SQLServerASTVisitor visitor) {
+        if (visitor.visit(this)) {
+            acceptChild(visitor, tableSource);
+            acceptChild(visitor, items);
+            acceptChild(visitor, from);
+            acceptChild(visitor, where);
+        }
+        visitor.endVisit(this);
+    }
+    
+}

src/main/java/com/alibaba/druid/sql/dialect/sqlserver/parser/SQLServerStatementParser.java

@@ -22,7 +22,11 @@ import com.alibaba.druid.sql.ast.SQLStatement;
 import com.alibaba.druid.sql.ast.expr.SQLQueryExpr;
 import com.alibaba.druid.sql.ast.statement.SQLInsertInto;
 import com.alibaba.druid.sql.ast.statement.SQLInsertStatement;
+import com.alibaba.druid.sql.ast.statement.SQLTableSource;
+import com.alibaba.druid.sql.ast.statement.SQLUpdateSetItem;
+import com.alibaba.druid.sql.ast.statement.SQLUpdateStatement;
 import com.alibaba.druid.sql.dialect.sqlserver.ast.stmt.SQLServerInsertStatement;
+import com.alibaba.druid.sql.dialect.sqlserver.ast.stmt.SQLServerUpdateStatement;
 import com.alibaba.druid.sql.parser.Lexer;
 import com.alibaba.druid.sql.parser.SQLSelectParser;
 import com.alibaba.druid.sql.parser.SQLStatementParser;
@@ -112,4 +116,34 @@ public class SQLServerStatementParser extends SQLStatementParser {
             insertStatement.setQuery(queryExpr.getSubQuery());
         }
     }
+    
+    protected SQLServerUpdateStatement createUpdateStatement() {
+        return new SQLServerUpdateStatement();
+    }
+    
+    public SQLUpdateStatement parseUpdateStatement() {
+        SQLServerUpdateStatement udpateStatement = createUpdateStatement();
+
+        if (lexer.token() == Token.UPDATE) {
+            lexer.nextToken();
+
+            SQLTableSource tableSource = this.exprParser.createSelectParser().parseTableSource();
+            udpateStatement.setTableSource(tableSource);
+        }
+
+        parseUpdateSet(udpateStatement);
+        
+        if (lexer.token() == Token.FROM) {
+            lexer.nextToken();
+            SQLTableSource from = this.exprParser.createSelectParser().parseTableSource();
+            udpateStatement.setFrom(from);
+        }
+        
+        if (lexer.token() == (Token.WHERE)) {
+            lexer.nextToken();
+            udpateStatement.setWhere(this.exprParser.expr());
+        }
+
+        return udpateStatement;
+    }
 }

src/main/java/com/alibaba/druid/sql/dialect/sqlserver/visitor/SQLServerASTVisitor.java

@@ -19,6 +19,7 @@ import com.alibaba.druid.sql.dialect.sqlserver.ast.SQLServerSelectQueryBlock;
 import com.alibaba.druid.sql.dialect.sqlserver.ast.Top;
 import com.alibaba.druid.sql.dialect.sqlserver.ast.expr.SQLServerObjectReferenceExpr;
 import com.alibaba.druid.sql.dialect.sqlserver.ast.stmt.SQLServerInsertStatement;
+import com.alibaba.druid.sql.dialect.sqlserver.ast.stmt.SQLServerUpdateStatement;
 import com.alibaba.druid.sql.visitor.SQLASTVisitor;
 
 public interface SQLServerASTVisitor extends SQLASTVisitor {
@@ -38,4 +39,8 @@ public interface SQLServerASTVisitor extends SQLASTVisitor {
     boolean visit(SQLServerInsertStatement x);
     
     void endVisit(SQLServerInsertStatement x);
+
+    boolean visit(SQLServerUpdateStatement x);
+    
+    void endVisit(SQLServerUpdateStatement x);
 }

src/main/java/com/alibaba/druid/sql/dialect/sqlserver/visitor/SQLServerASTVisitorAdapter.java

@@ -19,6 +19,7 @@ import com.alibaba.druid.sql.dialect.sqlserver.ast.SQLServerSelectQueryBlock;
 import com.alibaba.druid.sql.dialect.sqlserver.ast.Top;
 import com.alibaba.druid.sql.dialect.sqlserver.ast.expr.SQLServerObjectReferenceExpr;
 import com.alibaba.druid.sql.dialect.sqlserver.ast.stmt.SQLServerInsertStatement;
+import com.alibaba.druid.sql.dialect.sqlserver.ast.stmt.SQLServerUpdateStatement;
 import com.alibaba.druid.sql.visitor.SQLASTVisitorAdapter;
 
 public class SQLServerASTVisitorAdapter extends SQLASTVisitorAdapter implements SQLServerASTVisitor {
@@ -63,4 +64,14 @@ public class SQLServerASTVisitorAdapter extends SQLASTVisitorAdapter implements
         
     }
 
+    @Override
+    public boolean visit(SQLServerUpdateStatement x) {
+        return true;
+    }
+
+    @Override
+    public void endVisit(SQLServerUpdateStatement x) {
+        
+    }
+
 }

src/main/java/com/alibaba/druid/sql/dialect/sqlserver/visitor/SQLServerOutputVisitor.java

@@ -20,6 +20,7 @@ import com.alibaba.druid.sql.dialect.sqlserver.ast.SQLServerSelectQueryBlock;
 import com.alibaba.druid.sql.dialect.sqlserver.ast.Top;
 import com.alibaba.druid.sql.dialect.sqlserver.ast.expr.SQLServerObjectReferenceExpr;
 import com.alibaba.druid.sql.dialect.sqlserver.ast.stmt.SQLServerInsertStatement;
+import com.alibaba.druid.sql.dialect.sqlserver.ast.stmt.SQLServerUpdateStatement;
 import com.alibaba.druid.sql.visitor.SQLASTOutputVisitor;
 
 public class SQLServerOutputVisitor extends SQLASTOutputVisitor implements SQLServerASTVisitor {
@@ -144,4 +145,40 @@ public class SQLServerOutputVisitor extends SQLASTOutputVisitor implements SQLSe
 
     }
 
+    @Override
+    public boolean visit(SQLServerUpdateStatement x) {
+        print("UPDATE ");
+
+        x.getTableSource().accept(this);
+
+        println();
+        print("SET ");
+        for (int i = 0, size = x.getItems().size(); i < size; ++i) {
+            if (i != 0) {
+                print(", ");
+            }
+            x.getItems().get(i).accept(this);
+        }
+        
+        if (x.getFrom() != null) {
+            println();
+            print("FROM ");
+            x.getFrom().setParent(x);
+            x.getFrom().accept(this);
+        }
+
+        if (x.getWhere() != null) {
+            println();
+            print("WHERE ");
+            x.getWhere().setParent(x);
+            x.getWhere().accept(this);
+        }
+
+        return false;
+    }
+
+    @Override
+    public void endVisit(SQLServerUpdateStatement x) {
+        
+    }
 }

src/main/java/com/alibaba/druid/sql/dialect/sqlserver/visitor/SQLServerSchemaStatVisitor.java

@@ -15,13 +15,17 @@
  */
 package com.alibaba.druid.sql.dialect.sqlserver.visitor;
 
+import java.util.Map;
+
 import com.alibaba.druid.sql.ast.statement.SQLInsertStatement;
 import com.alibaba.druid.sql.ast.statement.SQLSelectQueryBlock;
 import com.alibaba.druid.sql.dialect.sqlserver.ast.SQLServerSelectQueryBlock;
 import com.alibaba.druid.sql.dialect.sqlserver.ast.Top;
 import com.alibaba.druid.sql.dialect.sqlserver.ast.expr.SQLServerObjectReferenceExpr;
 import com.alibaba.druid.sql.dialect.sqlserver.ast.stmt.SQLServerInsertStatement;
+import com.alibaba.druid.sql.dialect.sqlserver.ast.stmt.SQLServerUpdateStatement;
 import com.alibaba.druid.sql.visitor.SchemaStatVisitor;
+import com.alibaba.druid.stat.TableStat;
 import com.alibaba.druid.util.JdbcUtils;
 
 
@@ -73,5 +77,30 @@ public class SQLServerSchemaStatVisitor extends SchemaStatVisitor implements SQL
         this.endVisit((SQLInsertStatement) x);        
     }
 
+    @Override
+    public boolean visit(SQLServerUpdateStatement x) {
+        setAliasMap();
+
+        String ident = x.getTableName().toString();
+        setCurrentTable(ident);
+
+        TableStat stat = getTableStat(ident);
+        stat.incrementUpdateCount();
+
+        Map<String, String> aliasMap = getAliasMap();
+        aliasMap.put(ident, ident);
+
+        accept(x.getItems());
+        accept(x.getFrom());
+        accept(x.getWhere());
+
+        return false;
+    }
+
+    @Override
+    public void endVisit(SQLServerUpdateStatement x) {
+        
+    }
+
 
 }

src/main/java/com/alibaba/druid/sql/parser/SQLStatementParser.java

@@ -21,6 +21,7 @@ import java.util.List;
 import com.alibaba.druid.sql.ast.SQLExpr;
 import com.alibaba.druid.sql.ast.SQLName;
 import com.alibaba.druid.sql.ast.SQLStatement;
+import com.alibaba.druid.sql.ast.expr.SQLListExpr;
 import com.alibaba.druid.sql.ast.expr.SQLQueryExpr;
 import com.alibaba.druid.sql.ast.statement.SQLAssignItem;
 import com.alibaba.druid.sql.ast.statement.SQLCallStatement;
@@ -565,23 +566,7 @@ public class SQLStatementParser extends SQLParser {
             udpateStatement.setTableSource(tableSource);
         }
 
-        accept(Token.SET);
-
-        for (;;) {
-            SQLUpdateSetItem item = new SQLUpdateSetItem();
-            item.setColumn(this.exprParser.name());
-            accept(Token.EQ);
-            item.setValue(this.exprParser.expr());
-
-            udpateStatement.getItems().add(item);
-
-            if (lexer.token() == (Token.COMMA)) {
-                lexer.nextToken();
-                continue;
-            }
-
-            break;
-        }
+        parseUpdateSet(udpateStatement);
 
         if (lexer.token() == (Token.WHERE)) {
             lexer.nextToken();
@@ -590,6 +575,33 @@ public class SQLStatementParser extends SQLParser {
 
         return udpateStatement;
     }
+    
+    protected void parseUpdateSet(SQLUpdateStatement update) {
+        accept(Token.SET);
+
+        for (;;) {
+            SQLUpdateSetItem item = new SQLUpdateSetItem();
+
+            if (lexer.token() == (Token.LPAREN)) {
+                lexer.nextToken();
+                SQLListExpr list = new SQLListExpr();
+                this.exprParser.exprList(list.getItems());
+                accept(Token.RPAREN);
+                item.setColumn(list);
+            } else {
+                item.setColumn(this.exprParser.primary());
+            }
+            accept(Token.EQ);
+            item.setValue(this.exprParser.expr());
+            update.getItems().add(item);
+
+            if (lexer.token() != Token.COMMA) {
+                break;
+            }
+
+            lexer.nextToken();
+        }
+    }
 
     protected SQLUpdateStatement createUpdateStatement() {
         return new SQLUpdateStatement();

src/test/java/com/alibaba/druid/bvt/sql/sqlserver/SQLServerUpdateTest.java

@@ -0,0 +1,62 @@
+/*
+ * Copyright 1999-2011 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.sqlserver;
+
+import junit.framework.TestCase;
+
+import org.junit.Assert;
+
+import com.alibaba.druid.sql.ast.SQLStatement;
+import com.alibaba.druid.sql.dialect.sqlserver.parser.SQLServerStatementParser;
+import com.alibaba.druid.sql.test.TestUtils;
+
+public class SQLServerUpdateTest extends TestCase {
+
+    public void test_isEmpty() throws Exception {
+        String sql = "update reg_student_charge_item " + //
+                     "set FAmountReceived = b.amount   " + //
+                     "from reg_student_charge_item a" + //
+                     "    ,(" + //
+                     "          select a.FId,      " + //
+                     "                   case when sum(b.FChargeAmount) is null then 0 " + //
+                     "                        else sum(b.FChargeAmount)" + //
+                     "                   end as amount " + //
+                     "           from reg_student_charge_item a " + //
+                     "           left join reg_student_charge_daybook b on a.FId = b.FChargeItemId" + //
+                     "           where a.FId=?    group by a.FId" + //
+                     "     ) b " + //
+                     "where a.FId = b.FId and a.FId = ?";
+
+        String expect = "UPDATE reg_student_charge_item" +
+        		"\nSET FAmountReceived = b.amount" +
+        		"\nFROM reg_student_charge_item a, (SELECT a.FId, CASE WHEN SUM(b.FChargeAmount) IS NULL THEN 0 ELSE SUM(b.FChargeAmount) END AS amount" +
+        		"\n\t\tFROM reg_student_charge_item a" +
+        		"\n\t\t\tLEFT JOIN reg_student_charge_daybook b ON a.FId = b.FChargeItemId" +
+        		"\n\t\tWHERE a.FId = ? GROUP BY a.FId" +
+        		"\n\t\t) b" +
+        		"\nWHERE a.FId = b.FId" +
+        		"\nAND a.FId = ?";
+
+        SQLServerStatementParser parser = new SQLServerStatementParser(sql);
+        SQLStatement stmt = parser.parseStatementList().get(0);
+
+        String text = TestUtils.outputSqlServer(stmt);
+
+        Assert.assertEquals(expect, text);
+
+        System.out.println(text);
+    }
+}