improved sql parser. for issue #1979

alibaba · Sep 17, 2017 · b54a987 · b54a987
1 parent ff9364c
commit b54a987
Show file tree

Hide file tree

Showing 5 changed files with 115 additions and 5 deletions.
diff --git a/src/main/java/com/alibaba/druid/sql/dialect/sqlserver/ast/stmt/SQLServerSetStatement.java b/src/main/java/com/alibaba/druid/sql/dialect/sqlserver/ast/stmt/SQLServerSetStatement.java
@@ -23,8 +23,8 @@
 import com.alibaba.druid.sql.dialect.sqlserver.visitor.SQLServerASTVisitor;
 
 public class SQLServerSetStatement extends SQLServerStatementImpl implements SQLServerStatement {
-
-    private SQLAssignItem item = new SQLAssignItem();
+    private Option        option;
+    private SQLAssignItem item    = new SQLAssignItem();
 
     public SQLServerSetStatement(){
     }
@@ -38,6 +38,9 @@ public SQLAssignItem getItem() {
     }
 
     public void setItem(SQLAssignItem item) {
+        if (item != null) {
+            item.setParent(this);
+        }
         this.item = item;
     }
 
@@ -46,11 +49,28 @@ public void output(StringBuffer buf) {
         item.output(buf);
     }
 
+    public Option getOption() {
+        return option;
+    }
+
+    public void setOption(Option option) {
+        this.option = option;
+    }
+
     @Override
     public void accept0(SQLServerASTVisitor visitor) {
         if (visitor.visit(this)) {
             acceptChild(visitor, this.item);
         }
         visitor.endVisit(this);
     }
+
+    public void set(SQLExpr target, SQLExpr value) {
+        SQLAssignItem assignItem = new SQLAssignItem(target, value);
+        setItem(assignItem);
+    }
+
+    public static enum Option {
+        IDENTITY_INSERT
+    }
 }
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
@@ -362,7 +362,7 @@ public SQLServerExprParser getExprParser() {
     public SQLStatement parseSet() {
         accept(Token.SET);
 
-        if (lexer.identifierEquals("TRANSACTION")) {
+        if (lexer.identifierEquals(FnvHash.Constants.TRANSACTION)) {
             lexer.nextToken();
             acceptIdentifier("ISOLATION");
             acceptIdentifier("LEVEL");
@@ -402,7 +402,7 @@ public SQLStatement parseSet() {
             return stmt;
         }
 
-        if (lexer.identifierEquals("STATISTICS")) {
+        if (lexer.identifierEquals(FnvHash.Constants.STATISTICS)) {
             lexer.nextToken();
 
             SQLServerSetStatement stmt = new SQLServerSetStatement();
@@ -415,14 +415,33 @@ public SQLStatement parseSet() {
                 if (lexer.token() == Token.ON) {
                     stmt.getItem().setValue(new SQLIdentifierExpr("ON"));
                     lexer.nextToken();
-                } else if (lexer.identifierEquals("OFF")) {
+                } else if (lexer.identifierEquals(FnvHash.Constants.OFF)) {
                     stmt.getItem().setValue(new SQLIdentifierExpr("OFF"));
                     lexer.nextToken();
                 }
             }
             return stmt;
         }
 
+        if (lexer.identifierEquals(FnvHash.Constants.IDENTITY_INSERT)) {
+            SQLServerSetStatement stmt = new SQLServerSetStatement();
+            stmt.setOption(SQLServerSetStatement.Option.IDENTITY_INSERT);
+
+            lexer.nextToken();
+            SQLName table = this.exprParser.name();
+
+            if (lexer.token() == Token.ON) {
+                SQLExpr value = new SQLIdentifierExpr("ON");
+                stmt.set(table, value);
+                lexer.nextToken();
+            } else if (lexer.identifierEquals(FnvHash.Constants.OFF)) {
+                SQLExpr value = new SQLIdentifierExpr("OFF");
+                stmt.set(table, value);
+                lexer.nextToken();
+            }
+            return stmt;
+        }
+
         if (lexer.token() == Token.VARIANT) {
             SQLSetStatement stmt = new SQLSetStatement(getDbType());
             parseAssignItems(stmt.getItems(), stmt);

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
@@ -350,6 +350,13 @@ public void endVisit(SQLServerSetTransactionIsolationLevelStatement x) {
     @Override
     public boolean visit(SQLServerSetStatement x) {
         print0(ucase ? "SET " : "set ");
+
+        SQLServerSetStatement.Option option = x.getOption();
+        if (option != null) {
+            print(option.name());
+            print(' ');
+        }
+
         SQLAssignItem item = x.getItem();
         item.getTarget().accept(this);
         print(' ');

diff --git a/src/main/java/com/alibaba/druid/util/FnvHash.java b/src/main/java/com/alibaba/druid/util/FnvHash.java
@@ -36,6 +36,25 @@ public static long fnv1a_64(String input) {
         return hash;
     }
 
+    public static long fnv1a_64(String input, int offset, int end) {
+        if (input == null) {
+            return 0;
+        }
+
+        if (input.length() < end) {
+            end = input.length();
+        }
+
+        long hash = BASIC;
+        for (int i = offset; i < end; ++i) {
+            char c = input.charAt(i);
+            hash ^= c;
+            hash *= PRIME;
+        }
+
+        return hash;
+    }
+
     public static long fnv1a_64(char[] chars) {
         if (chars == null) {
             return 0;

diff --git a/src/test/java/com/alibaba/druid/bvt/sql/sqlserver/SQLServerSelectTest27.java b/src/test/java/com/alibaba/druid/bvt/sql/sqlserver/SQLServerSelectTest27.java
@@ -0,0 +1,45 @@
+/*
+ * Copyright 1999-2017 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 com.alibaba.druid.sql.SQLUtils;
+import com.alibaba.druid.sql.ast.SQLStatement;
+import com.alibaba.druid.util.JdbcConstants;
+import junit.framework.TestCase;
+import org.junit.Assert;
+
+import java.util.List;
+
+public class SQLServerSelectTest27 extends TestCase {
+
+    public void test_simple() throws Exception {
+        String sql = //
+                "set identity_insert tb_coupon OFF"; //
+
+        List<SQLStatement> stmtList = SQLUtils.parseStatements(sql, JdbcConstants.SQL_SERVER);
+
+        {
+            String text = SQLUtils.toSQLString(stmtList, JdbcConstants.SQL_SERVER);
+
+            Assert.assertEquals("SET IDENTITY_INSERT tb_coupon OFF", text);
+        }
+        {
+            String text = SQLUtils.toSQLString(stmtList, JdbcConstants.SQL_SERVER, SQLUtils.DEFAULT_LCASE_FORMAT_OPTION);
+
+            Assert.assertEquals("set IDENTITY_INSERT tb_coupon OFF", text);
+        }
+    }
+}