增加set hivevar:name=value的解析支持#5270

先前的代码重置了，重新提交。mvn vlidate 和mvn clean install验证ok
alibaba · Jul 9, 2023 · 16d9677 · 16d9677
1 parent 3694c6d
commit 16d9677
Show file tree

Hide file tree

Showing 5 changed files with 135 additions and 7 deletions.
diff --git a/core/src/main/java/com/alibaba/druid/sql/ast/expr/SQLPropertyExpr.java b/core/src/main/java/com/alibaba/druid/sql/ast/expr/SQLPropertyExpr.java
@@ -25,9 +25,11 @@
 import java.io.IOException;
 import java.util.Collections;
 import java.util.List;
+import java.util.Objects;
 
 public final class SQLPropertyExpr extends SQLExprImpl implements SQLName, SQLReplaceable, Comparable<SQLPropertyExpr> {
     private SQLExpr owner;
+    private String splitString;
     private String name;
 
     protected long nameHashCod64;
@@ -66,6 +68,14 @@ public SQLExpr getOwner() {
         return this.owner;
     }
 
+    public String getSplitString() {
+        return splitString;
+    }
+
+    public void setSplitString(String splitString) {
+        this.splitString = splitString;
+    }
+
     @Deprecated
     public String getOwnernName() {
         if (owner instanceof SQLName) {
@@ -112,6 +122,9 @@ protected void computeHashCode64() {
             hash ^= '.';
             hash *= FnvHash.PRIME;
         }
+        if (splitString != null) {
+            hash = FnvHash.hashCode64(hash, splitString);
+        }
         hash = FnvHash.hashCode64(hash, name);
         hashCode64 = hash;
     }
@@ -216,6 +229,9 @@ public boolean equals(Object obj) {
             return false;
         }
 
+        if (!Objects.equals(this.splitString, other.splitString)) {
+            return false;
+        }
         if (owner == null) {
             if (other.owner != null) {
                 return false;
@@ -396,7 +412,11 @@ public String toString() {
             return this.name;
         }
 
-        return owner.toString() + '.' + name;
+        if (splitString == null) {
+            return owner.toString() + '.' + name;
+        } else {
+            return owner.toString() + splitString + name;
+        }
     }
 
     @Override

diff --git a/core/src/main/java/com/alibaba/druid/sql/dialect/hive/parser/HiveExprParser.java b/core/src/main/java/com/alibaba/druid/sql/dialect/hive/parser/HiveExprParser.java
@@ -60,11 +60,15 @@ public HiveExprParser(Lexer lexer) {
     }
 
     public SQLExpr primaryRest(SQLExpr expr) {
-//        if(lexer.token() == Token.COLON) {
-//            lexer.nextToken();
-//            expr = dotRest(expr);
-//            return expr;
-//        }
+        if (lexer.token() == Token.COLON) {
+            lexer.nextToken();
+            expr = dotRest(expr);
+            if (expr instanceof SQLPropertyExpr) {
+                SQLPropertyExpr spe = (SQLPropertyExpr) expr;
+                spe.setSplitString(":");
+            }
+            return expr;
+        }
 
         switch (lexer.token()) {
             case LBRACKET:

diff --git a/core/src/main/java/com/alibaba/druid/sql/dialect/odps/parser/OdpsExprParser.java b/core/src/main/java/com/alibaba/druid/sql/dialect/odps/parser/OdpsExprParser.java
@@ -193,6 +193,10 @@ public SQLExpr primaryRest(SQLExpr expr) {
                 return propertyExpr;
             }
             expr = dotRest(expr);
+            if (expr instanceof SQLPropertyExpr) {
+                SQLPropertyExpr spe = (SQLPropertyExpr) expr;
+                spe.setSplitString(":");
+            }
             return expr;
         }
 

diff --git a/core/src/main/java/com/alibaba/druid/sql/visitor/SQLASTOutputVisitor.java b/core/src/main/java/com/alibaba/druid/sql/visitor/SQLASTOutputVisitor.java
@@ -2282,7 +2282,11 @@ public boolean visit(SQLPropertyExpr x) {
                 printExpr(owner, parameterized);
             }
         }
-        print('.');
+        if (x.getSplitString() != null) {
+            print(x.getSplitString());
+        } else {
+            print('.');
+        }
         String name = x.getName();
         if ("*".equals(name)) {
             print0(name);

diff --git a/core/src/test/java/com/alibaba/druid/bvt/sql/hive/HiveSetTest_1.java b/core/src/test/java/com/alibaba/druid/bvt/sql/hive/HiveSetTest_1.java
@@ -0,0 +1,96 @@
+package com.alibaba.druid.bvt.sql.hive;
+
+import java.util.List;
+
+import com.alibaba.druid.DbType;
+import com.alibaba.druid.sql.SQLUtils;
+import com.alibaba.druid.sql.ast.SQLStatement;
+import com.alibaba.druid.sql.visitor.SchemaStatVisitor;
+import com.alibaba.druid.util.JdbcConstants;
+
+import junit.framework.TestCase;
+
+/**
+ * @author lizongbo
+ */
+public class HiveSetTest_1 extends TestCase {
+
+    static String sql1 = "set hivevar:exp_id = 8001\n" +
+        ";";
+
+    static String sql2 = "set hivevar:exp_id = \"8001\"\n" +
+        ";";
+
+    static String sql3 = "set hivevar.exp_id = 8001\n" +
+        ";";
+
+    static String sql4 = "set hivevar.exp_id = \"8001\"\n" +
+        ";";
+    public void test_setHiveVar1() throws Exception {
+        {
+            DbType dbType = JdbcConstants.ODPS;
+            List<SQLStatement> statementList = SQLUtils.parseStatements(sql1, dbType);
+            SQLStatement stmt = statementList.get(0);
+            //System.out.println(dbType + "==" + statementList);
+            assertEquals(1, statementList.size());
+            //这里会多一个换行符， SQLExprParser的第 5259行针对ODPS类型，调用的是nextTokenForSet，而不是走nextToken，不敢随便改，先这样吧
+            assertEquals("SET hivevar:exp_id = 8001\n;", stmt.toString());
+        }
+        {
+            DbType dbType = JdbcConstants.HIVE;
+            List<SQLStatement> statementList = SQLUtils.parseStatements(sql1, dbType);
+            SQLStatement stmt = statementList.get(0);
+            //System.out.println(dbType + "==" + statementList);
+            assertEquals(1, statementList.size());
+            assertEquals("SET hivevar:exp_id = 8001;", stmt.toString());
+        }
+
+    }
+
+    public void test_setHiveVar2() throws Exception {
+        for (DbType dbType : new DbType[]{JdbcConstants.ODPS, JdbcConstants.HIVE}) {
+            List<SQLStatement> statementList = SQLUtils.parseStatements(sql2, dbType);
+            SQLStatement stmt = statementList.get(0);
+            //System.out.println(dbType + "==" + statementList);
+            assertEquals(1, statementList.size());
+            assertEquals("SET hivevar:exp_id = '8001';", stmt.toString());
+
+        }
+
+    }
+
+
+    public void test_setHiveVar3() throws Exception {
+        {
+            DbType dbType = JdbcConstants.ODPS;
+            List<SQLStatement> statementList = SQLUtils.parseStatements(sql3, dbType);
+            SQLStatement stmt = statementList.get(0);
+            //System.out.println(dbType + "==" + statementList);
+            assertEquals(1, statementList.size());
+            //这里会多一个换行符， SQLExprParser的第 5259行针对ODPS类型，调用的是nextTokenForSet，而不是走nextToken，不敢随便改，先这样吧
+            assertEquals("SET hivevar.exp_id = 8001\n;", stmt.toString());
+        }
+        {
+            DbType dbType = JdbcConstants.HIVE;
+            List<SQLStatement> statementList = SQLUtils.parseStatements(sql3, dbType);
+            SQLStatement stmt = statementList.get(0);
+            //System.out.println(dbType + "==" + statementList);
+            assertEquals(1, statementList.size());
+            assertEquals("SET hivevar.exp_id = 8001;", stmt.toString());
+        }
+
+    }
+
+
+    public void test_setHiveVar4() throws Exception {
+        for (DbType dbType : new DbType[]{JdbcConstants.ODPS, JdbcConstants.HIVE}) {
+            List<SQLStatement> statementList = SQLUtils.parseStatements(sql4, dbType);
+            SQLStatement stmt = statementList.get(0);
+            //System.out.println(dbType + "==" + statementList);
+            assertEquals(1, statementList.size());
+            assertEquals("SET hivevar.exp_id = '8001';", stmt.toString());
+
+        }
+
+    }
+}