improve sql parser

src/main/java/com/alibaba/druid/pool/DruidPooledConnection.java

@@ -1051,6 +1051,10 @@ public class DruidPooledConnection extends PoolableWrapper implements javax.sql.
 
         holder.getStatementEventListeners().remove(listener);
     }
+    
+    public Throwable getDisableError() {
+        return disableError;
+    }
 
     public void checkState() throws SQLException {
         if (holder == null) {

src/main/java/com/alibaba/druid/pool/DruidPooledStatement.java

@@ -71,7 +71,16 @@ public class DruidPooledStatement extends PoolableWrapper implements Statement {
 
     protected void checkOpen() throws SQLException {
         if (closed) {
-            throw new SQLException("statement is closed");
+            Throwable disableError = null;
+            if (this.conn != null) {
+                disableError = this.conn.getDisableError();
+            }
+
+            if (disableError != null) {
+                throw new SQLException("statement is closed", disableError);
+            } else {
+                throw new SQLException("statement is closed");
+            }
         }
     }
 
@@ -90,20 +99,20 @@ public class DruidPooledStatement extends PoolableWrapper implements Statement {
 
     public void incrementExecuteCount() {
         DruidPooledConnection conn = this.getPoolableConnection();
-        
+
         if (conn == null) {
             return;
         }
-        
+
         DruidConnectionHolder holder = conn.getConnectionHolder();
         if (holder == null) {
             return;
         }
-        
+
         if (holder.getDataSource() == null) {
             return;
         }
-        
+
         holder.getDataSource().incrementExecuteCount();
     }
 
@@ -121,7 +130,7 @@ public class DruidPooledStatement extends PoolableWrapper implements Statement {
         conn.beforeExecute();
         try {
             ResultSet rs = stmt.executeQuery(sql);
-            
+
             if (rs == null) {
                 return rs;
             }
@@ -266,14 +275,14 @@ public class DruidPooledStatement extends PoolableWrapper implements Statement {
             throw checkException(t);
         }
     }
-    
+
     @Override
     public void close() throws SQLException {
         if (!this.closed) {
             clearResultSet();
             stmt.close();
             this.closed = true;
-            
+
             if (conn.getConnectionHolder() != null) {
                 conn.getConnectionHolder().removeTrace(this);
             }
@@ -526,7 +535,9 @@ public class DruidPooledStatement extends PoolableWrapper implements Statement {
 
     @Override
     public final void clearBatch() throws SQLException {
-        checkOpen();
+        if (closed) {
+            return;
+        }
 
         try {
             stmt.clearBatch();

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

@@ -0,0 +1,42 @@
+package com.alibaba.druid.sql.ast.statement;
+
+import com.alibaba.druid.sql.ast.SQLExpr;
+import com.alibaba.druid.sql.ast.SQLHint;
+import com.alibaba.druid.sql.ast.SQLObjectImpl;
+import com.alibaba.druid.sql.visitor.SQLASTVisitor;
+
+
+@SuppressWarnings("serial")
+public class SQLExprHint extends SQLObjectImpl implements SQLHint {
+    private SQLExpr expr;
+    
+    public SQLExprHint() {
+        
+    }
+    
+    public SQLExprHint(SQLExpr expr) {
+        this.setExpr(expr);
+    }
+
+    public SQLExpr getExpr() {
+        return expr;
+    }
+    
+    public void setExpr(SQLExpr expr) {
+        if (expr != null) {
+            expr.setParent(this);
+        }
+        
+        this.expr = expr;
+    }
+
+    @Override
+    protected void accept0(SQLASTVisitor visitor) {
+        if (visitor.visit(this)) {
+            acceptChild(visitor, expr);
+        }
+        visitor.endVisit(this);
+    }
+    
+    
+}

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

@@ -15,7 +15,9 @@
  */
 package com.alibaba.druid.sql.dialect.sqlserver.parser;
 
+import com.alibaba.druid.sql.ast.SQLExpr;
 import com.alibaba.druid.sql.ast.SQLSetQuantifier;
+import com.alibaba.druid.sql.ast.statement.SQLExprHint;
 import com.alibaba.druid.sql.ast.statement.SQLExprTableSource;
 import com.alibaba.druid.sql.ast.statement.SQLSelect;
 import com.alibaba.druid.sql.ast.statement.SQLSelectQuery;
@@ -106,4 +108,28 @@ public class SQLServerSelectParser extends SQLSelectParser {
     protected SQLServerExprParser createExprParser() {
         return new SQLServerExprParser(lexer);
     }
+    
+    protected SQLTableSource parseTableSourceRest(SQLTableSource tableSource) {
+        if (lexer.token() == Token.WITH) {
+            lexer.nextToken();
+            accept(Token.LPAREN);
+            
+            for (;;) {
+                SQLExpr expr = this.expr();
+                SQLExprHint hint = new SQLExprHint(expr);
+                hint.setParent(tableSource);
+                tableSource.getHints().add(hint);
+                if (lexer.token() == Token.COMMA) {
+                    lexer.nextToken();
+                    continue;
+                } else {
+                    break;
+                }
+            }
+            
+            accept(Token.RPAREN);
+        }
+        
+        return super.parseTableSourceRest(tableSource);
+    }
 }

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

@@ -16,6 +16,7 @@
 package com.alibaba.druid.sql.dialect.sqlserver.visitor;
 
 import com.alibaba.druid.sql.ast.SQLSetQuantifier;
+import com.alibaba.druid.sql.ast.statement.SQLExprTableSource;
 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;
@@ -210,4 +211,21 @@ public class SQLServerOutputVisitor extends SQLASTOutputVisitor implements SQLSe
     public void endVisit(SQLServerUpdateStatement x) {
         
     }
+    
+    public boolean visit(SQLExprTableSource x) {
+        x.getExpr().accept(this);
+        
+        if (x.getHints() != null && x.getHints().size() > 0) {
+            print(" WITH (");
+            printAndAccept(x.getHints(), ", ");
+            print(")");
+        }
+
+        if (x.getAlias() != null) {
+            print(' ');
+            print(x.getAlias());
+        }
+
+        return false;
+    }
 }

src/main/java/com/alibaba/druid/sql/visitor/SQLASTOutputVisitor.java

@@ -95,6 +95,7 @@ import com.alibaba.druid.sql.ast.statement.SQLDeleteStatement;
 import com.alibaba.druid.sql.ast.statement.SQLDropIndexStatement;
 import com.alibaba.druid.sql.ast.statement.SQLDropTableStatement;
 import com.alibaba.druid.sql.ast.statement.SQLDropViewStatement;
+import com.alibaba.druid.sql.ast.statement.SQLExprHint;
 import com.alibaba.druid.sql.ast.statement.SQLExprTableSource;
 import com.alibaba.druid.sql.ast.statement.SQLInsertStatement;
 import com.alibaba.druid.sql.ast.statement.SQLInsertStatement.ValuesClause;
@@ -1464,4 +1465,10 @@ public class SQLASTOutputVisitor extends SQLASTVisitorAdapter {
         decrementIndent();
         return false;
     }
+    
+    @Override
+    public boolean visit(SQLExprHint x) {
+        x.getExpr().accept(this);
+        return false;
+    }
 }

src/main/java/com/alibaba/druid/sql/visitor/SQLASTVisitor.java

@@ -82,6 +82,7 @@ import com.alibaba.druid.sql.ast.statement.SQLDeleteStatement;
 import com.alibaba.druid.sql.ast.statement.SQLDropIndexStatement;
 import com.alibaba.druid.sql.ast.statement.SQLDropTableStatement;
 import com.alibaba.druid.sql.ast.statement.SQLDropViewStatement;
+import com.alibaba.druid.sql.ast.statement.SQLExprHint;
 import com.alibaba.druid.sql.ast.statement.SQLExprTableSource;
 import com.alibaba.druid.sql.ast.statement.SQLInsertStatement;
 import com.alibaba.druid.sql.ast.statement.SQLJoinTableSource;
@@ -434,44 +435,48 @@ public interface SQLASTVisitor {
     void endVisit(SQLAlterTableAlterColumn x);
 
     boolean visit(SQLAlterTableAlterColumn x);
-    
+
     boolean visit(SQLCheck x);
 
     void endVisit(SQLCheck x);
-    
+
     boolean visit(SQLAlterTableDropForeinKey x);
 
     void endVisit(SQLAlterTableDropForeinKey x);
-    
+
     boolean visit(SQLAlterTableDropPrimaryKey x);
-    
+
     void endVisit(SQLAlterTableDropPrimaryKey x);
-    
+
     boolean visit(SQLAlterTableDisableKeys x);
-    
+
     void endVisit(SQLAlterTableDisableKeys x);
-    
+
     boolean visit(SQLAlterTableEnableKeys x);
-    
+
     void endVisit(SQLAlterTableEnableKeys x);
-    
+
     boolean visit(SQLAlterTableStatement x);
-    
+
     void endVisit(SQLAlterTableStatement x);
-    
+
     boolean visit(SQLAlterTableAddForeignKey x);
-    
+
     void endVisit(SQLAlterTableAddForeignKey x);
 
     boolean visit(SQLAlterTableDisableConstraint x);
-    
+
     void endVisit(SQLAlterTableDisableConstraint x);
-    
+
     boolean visit(SQLAlterTableEnableConstraint x);
-    
+
     void endVisit(SQLAlterTableEnableConstraint x);
 
     boolean visit(SQLColumnCheck x);
-    
+
     void endVisit(SQLColumnCheck x);
+
+    boolean visit(SQLExprHint x);
+
+    void endVisit(SQLExprHint x);
 }

src/main/java/com/alibaba/druid/sql/visitor/SQLASTVisitorAdapter.java

@@ -82,6 +82,7 @@ import com.alibaba.druid.sql.ast.statement.SQLDeleteStatement;
 import com.alibaba.druid.sql.ast.statement.SQLDropIndexStatement;
 import com.alibaba.druid.sql.ast.statement.SQLDropTableStatement;
 import com.alibaba.druid.sql.ast.statement.SQLDropViewStatement;
+import com.alibaba.druid.sql.ast.statement.SQLExprHint;
 import com.alibaba.druid.sql.ast.statement.SQLExprTableSource;
 import com.alibaba.druid.sql.ast.statement.SQLInsertStatement;
 import com.alibaba.druid.sql.ast.statement.SQLInsertStatement.ValuesClause;
@@ -885,4 +886,14 @@ public class SQLASTVisitorAdapter implements SQLASTVisitor {
     public void endVisit(SQLColumnCheck x) {
         
     }
+
+    @Override
+    public boolean visit(SQLExprHint x) {
+        return true;
+    }
+
+    @Override
+    public void endVisit(SQLExprHint x) {
+        
+    }
 }

src/test/java/com/alibaba/druid/bvt/sql/sqlserver/SQLServerSelectTest16.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.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 SQLServerSelectTest16 extends TestCase {
+
+    public void test_simple() throws Exception {
+        String sql = "select GEN_VAL " + //
+                     "from ID_GENERATOR with (updlock, rowlock) " + //
+                     "where GEN_NAME = 'T_USERS'"; //
+
+        String expect = "SELECT GEN_VAL" + //
+                        "\nFROM ID_GENERATOR WITH (updlock, rowlock)" + //
+                        "\nWHERE GEN_NAME = 'T_USERS'";
+
+        SQLServerStatementParser parser = new SQLServerStatementParser(sql);
+        SQLStatement stmt = parser.parseStatementList().get(0);
+
+        String text = TestUtils.outputSqlServer(stmt);
+
+        Assert.assertEquals(expect, text);
+
+        System.out.println(text);
+    }
+}