improve sqlserver parser

rishards · Sep 24, 2014 · 5049285 · 5049285
1 parent e7829a7
commit 5049285
Show file tree

Hide file tree

Showing 10 changed files with 409 additions and 2 deletions.
diff --git a/src/main/java/com/alibaba/druid/sql/dialect/sqlserver/ast/stmt/SQLServerBlockStatement.java b/src/main/java/com/alibaba/druid/sql/dialect/sqlserver/ast/stmt/SQLServerBlockStatement.java
@@ -0,0 +1,46 @@
+/*
+ * 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 java.util.ArrayList;
+import java.util.List;
+
+import com.alibaba.druid.sql.ast.SQLStatement;
+import com.alibaba.druid.sql.dialect.sqlserver.ast.SQLServerObjectImpl;
+import com.alibaba.druid.sql.dialect.sqlserver.ast.SQLServerStatement;
+import com.alibaba.druid.sql.dialect.sqlserver.visitor.SQLServerASTVisitor;
+
+public class SQLServerBlockStatement  extends SQLServerObjectImpl implements SQLServerStatement  {
+
+    private List<SQLStatement>    statementList = new ArrayList<SQLStatement>();
+
+    public List<SQLStatement> getStatementList() {
+        return statementList;
+    }
+
+    public void setStatementList(List<SQLStatement> statementList) {
+        this.statementList = statementList;
+    }
+
+    @Override
+    public void accept0(SQLServerASTVisitor visitor) {
+        if (visitor.visit(this)) {
+            acceptChild(visitor, statementList);
+        }
+        visitor.endVisit(this);
+    }
+
+}
diff --git a/src/main/java/com/alibaba/druid/sql/dialect/sqlserver/ast/stmt/SQLServerIfStatement.java b/src/main/java/com/alibaba/druid/sql/dialect/sqlserver/ast/stmt/SQLServerIfStatement.java
@@ -0,0 +1,89 @@
+/*
+ * 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 java.util.ArrayList;
+import java.util.List;
+
+import com.alibaba.druid.sql.ast.SQLExpr;
+import com.alibaba.druid.sql.ast.SQLStatement;
+import com.alibaba.druid.sql.dialect.sqlserver.ast.SQLServerObjectImpl;
+import com.alibaba.druid.sql.dialect.sqlserver.ast.SQLServerStatement;
+import com.alibaba.druid.sql.dialect.sqlserver.visitor.SQLServerASTVisitor;
+
+public class SQLServerIfStatement extends SQLServerObjectImpl implements SQLServerStatement {
+
+    private SQLExpr            condition;
+    private List<SQLStatement> statements = new ArrayList<SQLStatement>();
+    private Else               elseItem;
+
+    @Override
+    public void accept0(SQLServerASTVisitor visitor) {
+        if (visitor.visit(this)) {
+            acceptChild(visitor, condition);
+            acceptChild(visitor, statements);
+            acceptChild(visitor, elseItem);
+        }
+        visitor.endVisit(this);
+    }
+
+    public SQLExpr getCondition() {
+        return condition;
+    }
+
+    public void setCondition(SQLExpr condition) {
+        this.condition = condition;
+    }
+
+    public List<SQLStatement> getStatements() {
+        return statements;
+    }
+
+    public void setStatements(List<SQLStatement> statements) {
+        this.statements = statements;
+    }
+
+    public Else getElseItem() {
+        return elseItem;
+    }
+
+    public void setElseItem(Else elseItem) {
+        this.elseItem = elseItem;
+    }
+
+    public static class Else extends SQLServerObjectImpl {
+
+        private List<SQLStatement> statements = new ArrayList<SQLStatement>();
+
+        @Override
+        public void accept0(SQLServerASTVisitor visitor) {
+            if (visitor.visit(this)) {
+                acceptChild(visitor, statements);
+            }
+            visitor.endVisit(this);
+        }
+
+        public List<SQLStatement> getStatements() {
+            return statements;
+        }
+
+        public void setStatements(List<SQLStatement> statements) {
+            this.statements = statements;
+        }
+
+    }
+
+}
diff --git a/src/main/java/com/alibaba/druid/sql/dialect/sqlserver/parser/SQLServerLexer.java b/src/main/java/com/alibaba/druid/sql/dialect/sqlserver/parser/SQLServerLexer.java
@@ -42,6 +42,10 @@ public class SQLServerLexer extends Lexer {
         map.put("PERCENT", Token.PERCENT);
         map.put("IDENTITY", Token.IDENTITY);
         map.put("DECLARE", Token.DECLARE);
+        map.put("IF", Token.IF);
+        map.put("ELSE", Token.ELSE);
+        map.put("BEGIN", Token.BEGIN);
+        map.put("END", Token.END);
 
         DEFAULT_SQL_SERVER_KEYWORDS = new Keywords(map);
     }

diff --git a/src/main/java/com/alibaba/druid/sql/dialect/sqlserver/parser/SQLServerStatementParser.java b/src/main/java/com/alibaba/druid/sql/dialect/sqlserver/parser/SQLServerStatementParser.java
@@ -32,8 +32,10 @@
 import com.alibaba.druid.sql.dialect.sqlserver.ast.SQLServerDeclareItem;
 import com.alibaba.druid.sql.dialect.sqlserver.ast.SQLServerOutput;
 import com.alibaba.druid.sql.dialect.sqlserver.ast.SQLServerTop;
+import com.alibaba.druid.sql.dialect.sqlserver.ast.stmt.SQLServerBlockStatement;
 import com.alibaba.druid.sql.dialect.sqlserver.ast.stmt.SQLServerDeclareStatement;
 import com.alibaba.druid.sql.dialect.sqlserver.ast.stmt.SQLServerExecStatement;
+import com.alibaba.druid.sql.dialect.sqlserver.ast.stmt.SQLServerIfStatement;
 import com.alibaba.druid.sql.dialect.sqlserver.ast.stmt.SQLServerInsertStatement;
 import com.alibaba.druid.sql.dialect.sqlserver.ast.stmt.SQLServerSetStatement;
 import com.alibaba.druid.sql.dialect.sqlserver.ast.stmt.SQLServerSetTransactionIsolationLevelStatement;
@@ -87,6 +89,16 @@ public boolean parseStatementListDialect(List<SQLStatement> statementList) {
             statementList.add(this.parseDeclare());
             return true;
         }
+
+        if (lexer.token() == Token.IF) {
+            statementList.add(this.parseIf());
+            return true;
+        }
+
+        if (lexer.token() == Token.BEGIN) {
+            statementList.add(this.parseBlock());
+            return true;
+        }
 
         return false;
     }
@@ -374,4 +386,40 @@ public SQLStatement parseSet() {
             return stmt;
         }
     }
+
+    public SQLServerIfStatement parseIf() {
+        accept(Token.IF);
+
+        SQLServerIfStatement stmt = new SQLServerIfStatement();
+
+        stmt.setCondition(this.exprParser.expr());
+
+        this.parseStatementList(stmt.getStatements(), 1);
+
+        if(lexer.token() == Token.SEMI) {
+            lexer.nextToken();
+        }
+
+        if (lexer.token() == Token.ELSE) {
+            lexer.nextToken();
+
+            SQLServerIfStatement.Else elseItem = new SQLServerIfStatement.Else();
+            this.parseStatementList(elseItem.getStatements(), 1);
+            stmt.setElseItem(elseItem);
+        }
+
+        return stmt;
+    }
+
+    public SQLServerBlockStatement parseBlock() {
+        SQLServerBlockStatement block = new SQLServerBlockStatement();
+
+        accept(Token.BEGIN);
+
+        parseStatementList(block.getStatementList());
+
+        accept(Token.END);
+
+        return block;
+    }
 }
diff --git a/src/main/java/com/alibaba/druid/sql/dialect/sqlserver/visitor/SQLServerASTVisitor.java b/src/main/java/com/alibaba/druid/sql/dialect/sqlserver/visitor/SQLServerASTVisitor.java
@@ -21,8 +21,11 @@
 import com.alibaba.druid.sql.dialect.sqlserver.ast.SQLServerSelectQueryBlock;
 import com.alibaba.druid.sql.dialect.sqlserver.ast.SQLServerTop;
 import com.alibaba.druid.sql.dialect.sqlserver.ast.expr.SQLServerObjectReferenceExpr;
+import com.alibaba.druid.sql.dialect.sqlserver.ast.stmt.SQLServerBlockStatement;
 import com.alibaba.druid.sql.dialect.sqlserver.ast.stmt.SQLServerDeclareStatement;
 import com.alibaba.druid.sql.dialect.sqlserver.ast.stmt.SQLServerExecStatement;
+import com.alibaba.druid.sql.dialect.sqlserver.ast.stmt.SQLServerIfStatement;
+import com.alibaba.druid.sql.dialect.sqlserver.ast.stmt.SQLServerIfStatement.Else;
 import com.alibaba.druid.sql.dialect.sqlserver.ast.stmt.SQLServerInsertStatement;
 import com.alibaba.druid.sql.dialect.sqlserver.ast.stmt.SQLServerSetStatement;
 import com.alibaba.druid.sql.dialect.sqlserver.ast.stmt.SQLServerSetTransactionIsolationLevelStatement;
@@ -82,5 +85,17 @@ public interface SQLServerASTVisitor extends SQLASTVisitor {
     boolean visit(SQLServerDeclareStatement x);
 
     void endVisit(SQLServerDeclareStatement x);
+
+    boolean visit(Else x);
+
+    void endVisit(Else x);
+
+    boolean visit(SQLServerIfStatement x);
+
+    void endVisit(SQLServerIfStatement x);
+
+    boolean visit(SQLServerBlockStatement x);
+
+    void endVisit(SQLServerBlockStatement x);
 
 }
diff --git a/...main/java/com/alibaba/druid/sql/dialect/sqlserver/visitor/SQLServerASTVisitorAdapter.java b/...main/java/com/alibaba/druid/sql/dialect/sqlserver/visitor/SQLServerASTVisitorAdapter.java
@@ -21,8 +21,11 @@
 import com.alibaba.druid.sql.dialect.sqlserver.ast.SQLServerSelectQueryBlock;
 import com.alibaba.druid.sql.dialect.sqlserver.ast.SQLServerTop;
 import com.alibaba.druid.sql.dialect.sqlserver.ast.expr.SQLServerObjectReferenceExpr;
+import com.alibaba.druid.sql.dialect.sqlserver.ast.stmt.SQLServerBlockStatement;
 import com.alibaba.druid.sql.dialect.sqlserver.ast.stmt.SQLServerDeclareStatement;
 import com.alibaba.druid.sql.dialect.sqlserver.ast.stmt.SQLServerExecStatement;
+import com.alibaba.druid.sql.dialect.sqlserver.ast.stmt.SQLServerIfStatement;
+import com.alibaba.druid.sql.dialect.sqlserver.ast.stmt.SQLServerIfStatement.Else;
 import com.alibaba.druid.sql.dialect.sqlserver.ast.stmt.SQLServerInsertStatement;
 import com.alibaba.druid.sql.dialect.sqlserver.ast.stmt.SQLServerSetStatement;
 import com.alibaba.druid.sql.dialect.sqlserver.ast.stmt.SQLServerSetTransactionIsolationLevelStatement;
@@ -161,4 +164,34 @@ public void endVisit(SQLServerDeclareStatement x) {
 
     }
 
+    @Override
+    public boolean visit(Else x) {
+        return true;
+    }
+
+    @Override
+    public void endVisit(Else x) {
+
+    }
+
+    @Override
+    public boolean visit(SQLServerIfStatement x) {
+        return true;
+    }
+
+    @Override
+    public void endVisit(SQLServerIfStatement x) {
+
+    }
+
+    @Override
+    public boolean visit(SQLServerBlockStatement x) {
+        return true;
+    }
+
+    @Override
+    public void endVisit(SQLServerBlockStatement x) {
+
+    }
+
 }
diff --git a/src/main/java/com/alibaba/druid/sql/dialect/sqlserver/visitor/SQLServerOutputVisitor.java b/src/main/java/com/alibaba/druid/sql/dialect/sqlserver/visitor/SQLServerOutputVisitor.java
@@ -17,6 +17,7 @@
 
 import com.alibaba.druid.sql.ast.SQLName;
 import com.alibaba.druid.sql.ast.SQLSetQuantifier;
+import com.alibaba.druid.sql.ast.SQLStatement;
 import com.alibaba.druid.sql.ast.statement.SQLAssignItem;
 import com.alibaba.druid.sql.ast.statement.SQLColumnConstraint;
 import com.alibaba.druid.sql.ast.statement.SQLColumnDefinition;
@@ -28,8 +29,11 @@
 import com.alibaba.druid.sql.dialect.sqlserver.ast.SQLServerSelectQueryBlock;
 import com.alibaba.druid.sql.dialect.sqlserver.ast.SQLServerTop;
 import com.alibaba.druid.sql.dialect.sqlserver.ast.expr.SQLServerObjectReferenceExpr;
+import com.alibaba.druid.sql.dialect.sqlserver.ast.stmt.SQLServerBlockStatement;
 import com.alibaba.druid.sql.dialect.sqlserver.ast.stmt.SQLServerDeclareStatement;
 import com.alibaba.druid.sql.dialect.sqlserver.ast.stmt.SQLServerExecStatement;
+import com.alibaba.druid.sql.dialect.sqlserver.ast.stmt.SQLServerIfStatement;
+import com.alibaba.druid.sql.dialect.sqlserver.ast.stmt.SQLServerIfStatement.Else;
 import com.alibaba.druid.sql.dialect.sqlserver.ast.stmt.SQLServerInsertStatement;
 import com.alibaba.druid.sql.dialect.sqlserver.ast.stmt.SQLServerSetStatement;
 import com.alibaba.druid.sql.dialect.sqlserver.ast.stmt.SQLServerSetTransactionIsolationLevelStatement;
@@ -463,6 +467,83 @@ public boolean visit(SQLServerDeclareStatement x) {
 
     @Override
     public void endVisit(SQLServerDeclareStatement x) {
-
+
+    }
+
+    @Override
+    public boolean visit(Else x) {
+        print("ELSE");
+        incrementIndent();
+        println();
+
+        for (int i = 0, size = x.getStatements().size(); i < size; ++i) {
+            if (i != 0) {
+                println();
+            }
+            SQLStatement item = x.getStatements().get(i);
+            item.setParent(x);
+            item.accept(this);
+        }
+
+        decrementIndent();
+        return false;
+    }
+
+    @Override
+    public void endVisit(Else x) {
+
+    }
+
+    @Override
+    public boolean visit(SQLServerIfStatement x) {
+        print("IF ");
+        x.getCondition().accept(this);
+        incrementIndent();
+        println();
+        for (int i = 0, size = x.getStatements().size(); i < size; ++i) {
+            SQLStatement item = x.getStatements().get(i);
+            item.setParent(x);
+            item.accept(this);
+            if (i != size - 1) {
+                println();
+            }
+        }
+        decrementIndent();
+
+        if (x.getElseItem() != null) {
+            println();
+            x.getElseItem().accept(this);
+        }
+        return false;
+    }
+
+    @Override
+    public void endVisit(SQLServerIfStatement x) {
+
+    }
+
+    @Override
+    public boolean visit(SQLServerBlockStatement x) {
+        print("BEGIN");
+        incrementIndent();
+        println();
+        for (int i = 0, size = x.getStatementList().size(); i < size; ++i) {
+            if (i != 0) {
+                println();
+            }
+            SQLStatement stmt = x.getStatementList().get(i);
+            stmt.setParent(x);
+            stmt.accept(this);
+            print(";");
+        }
+        decrementIndent();
+        println();
+        print("END");
+        return false;
+    }
+
+    @Override
+    public void endVisit(SQLServerBlockStatement x) {
+
     }
 }