sql parser improve sql server support

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) {
+        
+    }
+
 
 }