improved odps support

alibaba · Sep 5, 2021 · 0be11c3 · 0be11c3
1 parent 8ef33c3
commit 0be11c3
Show file tree

Hide file tree

Showing 5 changed files with 290 additions and 5 deletions.
diff --git a/src/main/java/com/alibaba/druid/sql/parser/Lexer.java b/src/main/java/com/alibaba/druid/sql/parser/Lexer.java
@@ -2198,18 +2198,33 @@ public void scanVariable() {
             pos++;
             bufPos++;
 
+            boolean ident = false;
             for (;;) {
                 ch = charAt(++pos);
+                if (isEOF()) {
+                    pos--;
+                    bufPos--;
+                    break;
+                }
 
-                if (ch == '}') {
+                if (ch == '}' && !ident) {
+                    if (isIdentifierChar(charAt(pos + 1))) {
+                        bufPos++;
+                        ident = true;
+                        continue;
+                    }
+                    break;
+                }
+
+                if (ident && isWhitespace(ch)) {
                     break;
                 }
 
                 bufPos++;
                 continue;
             }
 
-            if (ch != '}') {
+            if (ch != '}' && !ident) {
                 throw new ParserException("syntax error. " + info());
             }
             ++pos;
@@ -2226,7 +2241,7 @@ public void scanVariable() {
             }
 
             stringVal = addSymbol();
-            token = Token.VARIANT;
+            token = ident ? IDENTIFIER : Token.VARIANT;
             return;
         } else if (c1 == '$' && charAt(pos + 2) == '{') {
             pos += 2;
@@ -2789,12 +2804,30 @@ public void scanNumber() {
                     && !(ch == 'b' && bufPos == 1 && charAt(pos - 1) == '0' && dbType != DbType.odps)
             ) {
                 bufPos++;
+                boolean brace = false;
                 for (;;) {
                     char c0 = ch;
                     ch = charAt(++pos);
 
+                    if (isEOF()) {
+                        break;
+                    }
+
                     if (!isIdentifierChar(ch)) {
-                        if (ch == '（'  || ch == '）' && c0 > 256) {
+                        if (ch == '{' && charAt(pos - 1) == '$' && !brace) {
+                            bufPos++;
+                            brace = true;
+                            continue;
+                        }
+
+                        if (ch == '}' && brace) {
+                            bufPos++;
+                            brace = false;
+                            continue;
+                        }
+
+                        if ((ch == '（'  || ch == '）')
+                                && c0 > 256) {
                             bufPos++;
                             continue;
                         }

diff --git a/src/main/java/com/alibaba/druid/sql/visitor/SQLASTOutputVisitor.java b/src/main/java/com/alibaba/druid/sql/visitor/SQLASTOutputVisitor.java
@@ -6511,7 +6511,7 @@ public boolean visit(SQLAlterViewRenameStatement x) {
         }
 
         SQLName changeOwnerTo = x.getChangeOwnerTo();
-        if (to != null) {
+        if (changeOwnerTo != null) {
             print0(ucase ? " CHANGEOWNER TO " : " changeowner to ");
             printExpr(changeOwnerTo);
         }

diff --git a/src/main/java/com/alibaba/druid/support/opds/udf/SqlTypeUDF.java b/src/main/java/com/alibaba/druid/support/opds/udf/SqlTypeUDF.java
@@ -6,6 +6,7 @@
 import com.alibaba.druid.sql.ast.statement.*;
 import com.alibaba.druid.sql.dialect.hive.ast.HiveInsertStatement;
 import com.alibaba.druid.sql.dialect.hive.ast.HiveMultiInsertStatement;
+import com.alibaba.druid.sql.dialect.hive.stmt.HiveLoadDataStatement;
 import com.alibaba.druid.sql.dialect.odps.ast.OdpsDeclareVariableStatement;
 import com.alibaba.druid.sql.dialect.odps.ast.OdpsQueryAliasStatement;
 import com.alibaba.druid.sql.parser.*;
@@ -97,6 +98,8 @@ static SQLType getStmtSqlType(SQLStatement stmt, SQLType sqlType) {
                         ? SQLType.INSERT_OVERWRITE_VALUES
                         : SQLType.INSERT_INTO_VALUES;
             }
+        } else if (stmt instanceof HiveLoadDataStatement) {
+            sqlType = SQLType.LOAD;
         } else if (stmt instanceof SQLUpdateStatement) {
             sqlType = SQLType.UPDATE;
         } else if (stmt instanceof SQLDeleteStatement) {

diff --git a/src/test/java/com/alibaba/druid/bvt/sql/odps/OdpsLexerTest.java b/src/test/java/com/alibaba/druid/bvt/sql/odps/OdpsLexerTest.java
@@ -0,0 +1,60 @@
+package com.alibaba.druid.bvt.sql.odps;
+
+import com.alibaba.druid.sql.dialect.odps.parser.OdpsLexer;
+import com.alibaba.druid.sql.parser.Token;
+import junit.framework.TestCase;
+
+public class OdpsLexerTest extends TestCase {
+    public void test_0() throws Exception {
+        String str = "DESC；";
+        OdpsLexer lexer = new OdpsLexer(str);
+        lexer.nextToken();
+        assertEquals(Token.DESC, lexer.token());
+        lexer.nextToken();
+        assertEquals(Token.SEMI, lexer.token());
+    }
+
+    public void test_1() throws Exception {
+        String str = "——\n" +
+                "\n" +
+                "drop table if exists tdl_idle_mem_portrait_pred_feats_20210831;";
+        OdpsLexer lexer = new OdpsLexer(str);
+        lexer.nextToken();
+        assertEquals(Token.DROP, lexer.token());
+    }
+
+    public void test_2() throws Exception {
+        String str = "drop table graph_embedding_dev.04_s1_${bizdate}_0108;";
+        OdpsLexer lexer = new OdpsLexer(str);
+        lexer.nextToken();
+        assertEquals(Token.DROP, lexer.token());
+        lexer.nextToken();
+        assertEquals(Token.TABLE, lexer.token());
+        lexer.nextToken();
+        assertEquals(Token.IDENTIFIER, lexer.token());
+        lexer.nextToken();
+        assertEquals(Token.DOT, lexer.token());
+        lexer.nextToken();
+        assertEquals(Token.IDENTIFIER, lexer.token());
+        assertEquals("04_s1_${bizdate}_0108", lexer.stringVal());
+    }
+
+    public void test_3() throws Exception {
+        String str = "dm_claim_unwww_unsettled_dt@@{yyyyMMdd}_${rundt}";
+        OdpsLexer lexer = new OdpsLexer(str);
+        lexer.nextToken();
+        assertEquals(Token.IDENTIFIER, lexer.token());
+        lexer.nextToken();
+        assertEquals(Token.IDENTIFIER, lexer.token());
+        assertEquals("@{yyyyMMdd}_${rundt}", lexer.stringVal());
+    }
+
+    public void test_4() throws Exception {
+        String str = "${PN}_events";
+        OdpsLexer lexer = new OdpsLexer(str);
+        lexer.nextToken();
+        assertEquals(Token.IDENTIFIER, lexer.token());
+        assertEquals("${PN}_events", lexer.stringVal());
+    }
+
+}
diff --git a/src/test/java/com/alibaba/druid/bvt/sql/schemaStat/SchemaStatTest24.java b/src/test/java/com/alibaba/druid/bvt/sql/schemaStat/SchemaStatTest24.java
@@ -0,0 +1,189 @@
+package com.alibaba.druid.bvt.sql.schemaStat;
+
+import com.alibaba.druid.sql.SQLUtils;
+import com.alibaba.druid.sql.ast.SQLStatement;
+import com.alibaba.druid.sql.parser.SQLParserUtils;
+import com.alibaba.druid.sql.parser.SQLStatementParser;
+import com.alibaba.druid.sql.repository.SchemaRepository;
+import com.alibaba.druid.sql.visitor.SchemaStatVisitor;
+import com.alibaba.druid.util.JdbcConstants;
+import junit.framework.TestCase;
+
+public class SchemaStatTest24 extends TestCase {
+    SchemaRepository repository;
+
+    protected void setUp() throws Exception {
+        repository = new SchemaRepository(JdbcConstants.ODPS);
+    }
+
+    public void test_0() throws Exception {
+        String sql = "clone table t1 to t2;";
+
+        SQLStatementParser parser = SQLParserUtils.createSQLStatementParser(sql, JdbcConstants.ODPS);
+        SQLStatement stmt = parser.parseStatementList().get(0);
+
+        SchemaStatVisitor statVisitor = SQLUtils.createSchemaStatVisitor(repository);
+        stmt.accept(statVisitor);
+
+        System.out.println("Tables : " + statVisitor.getTables());
+        System.out.println(statVisitor.getColumns());
+//        System.out.println(statVisitor.getGroupByColumns()); // group by
+        System.out.println("relationships : " + statVisitor.getRelationships()); // group by
+        System.out.println(statVisitor.getConditions());
+
+        assertEquals(0, statVisitor.getColumns().size());
+        assertEquals(0, statVisitor.getConditions().size());
+        assertEquals(0, statVisitor.getFunctions().size());
+
+        assertTrue(statVisitor.containsTable("t1"));
+        assertTrue(statVisitor.containsTable("t2"));
+    }
+
+    public void test_1() throws Exception {
+        String sql = "MERGE INTO t1 AS t USING (select * from ods_dayima_topic_delta where pt='20210831') AS s \n" +
+                "ON s.id = t.id \n" +
+                "WHEN matched THEN UPDATE SET t.id=s.id,t.topic_category_id=s.topic_category_id,t.uid=s.uid,t.title=s.title,t.content=s.content,t.STATUS=s.STATUS,t.dateline=s.dateline,t.ip=s.ip,t.replynum=s.replynum,t.lastreplytime=s.lastreplytime,t.settop=s.settop,t.lastoptime=s.lastoptime,t.LOCK=s.LOCK,t.setbottom=s.setbottom,t.flag=s.flag,t.digest=s.digest,t.createtime=s.createtime,t.lastpost=s.lastpost,t.favnum=s.favnum,t.attach_pictures=s.attach_pictures,t.viewnum=s.viewnum,t.settoptime=s.settoptime,t.hidden=s.hidden,t.TYPE=s.TYPE,t.hashtag_code=s.hashtag_code,t.like_num=s.like_num,t.join_num=s.join_num,t.is_anonymous=s.is_anonymous,t.has_goods=s.has_goods,t.extra_info=s.extra_info,t.created=s.created\n" +
+                "WHEN NOT matched THEN INSERT VALUES(s.id,s.topic_category_id,s.uid,s.title,s.content,s.STATUS,s.dateline,s.ip,s.replynum,s.lastreplytime,s.settop,s.lastoptime,s.LOCK,s.setbottom,s.flag,s.digest,s.createtime,s.lastpost,s.favnum,s.attach_pictures,s.viewnum,s.settoptime,s.hidden,s.TYPE,s.hashtag_code,s.like_num,s.join_num,s.is_anonymous,s.has_goods,s.extra_info,s.created);";
+
+        SQLStatementParser parser = SQLParserUtils.createSQLStatementParser(sql, JdbcConstants.ODPS);
+        SQLStatement stmt = parser.parseStatementList().get(0);
+
+        SchemaStatVisitor statVisitor = SQLUtils.createSchemaStatVisitor(repository);
+        stmt.accept(statVisitor);
+
+        System.out.println("Tables : " + statVisitor.getTables());
+        System.out.println(statVisitor.getColumns());
+//        System.out.println(statVisitor.getGroupByColumns()); // group by
+        System.out.println("relationships : " + statVisitor.getRelationships()); // group by
+        System.out.println(statVisitor.getConditions());
+
+        assertEquals(34, statVisitor.getColumns().size());
+//        assertEquals(1, statVisitor.getConditions().size());
+        assertEquals(0, statVisitor.getFunctions().size());
+
+        assertTrue(statVisitor.containsTable("t1"));
+        assertTrue(statVisitor.containsColumn("t1", "id"));
+    }
+
+    public void test_2() throws Exception {
+        String sql = "unload from t  partition  (ds='20210829')\n" +
+                "into\n" +
+                "location 'oss://oss-cn-shanghai-internal.aliyuncs.com/xx/xxx'\n" +
+                "stored by 'com.aliyun.odps.TsvStorageHandler'\n" +
+                "with serdeproperties ('odps.properties.rolearn'='acs:ram::123:role/xxx', 'odps.text.option.gzip.output.enabled'='true');";
+
+        SQLStatementParser parser = SQLParserUtils.createSQLStatementParser(sql, JdbcConstants.ODPS);
+        SQLStatement stmt = parser.parseStatementList().get(0);
+
+        SchemaStatVisitor statVisitor = SQLUtils.createSchemaStatVisitor(repository);
+        stmt.accept(statVisitor);
+
+        System.out.println("Tables : " + statVisitor.getTables());
+        System.out.println(statVisitor.getColumns());
+//        System.out.println(statVisitor.getGroupByColumns()); // group by
+        System.out.println("relationships : " + statVisitor.getRelationships()); // group by
+        System.out.println(statVisitor.getConditions());
+
+        assertEquals(1, statVisitor.getColumns().size());
+        assertEquals(0, statVisitor.getConditions().size());
+        assertEquals(0, statVisitor.getFunctions().size());
+
+        assertTrue(statVisitor.containsTable("t"));
+//        assertTrue(statVisitor.containsColumn("t", "f1"));
+//        assertEquals("t.f1 < 4", statVisitor.getConditions().get(0).toString());
+    }
+
+    public void test_3() throws Exception {
+        String sql = "analyze table t partition(ds='20210829') compute statistics for columns (member_ids);";
+
+        SQLStatementParser parser = SQLParserUtils.createSQLStatementParser(sql, JdbcConstants.ODPS);
+        SQLStatement stmt = parser.parseStatementList().get(0);
+
+        SchemaStatVisitor statVisitor = SQLUtils.createSchemaStatVisitor(repository);
+        stmt.accept(statVisitor);
+
+        System.out.println("Tables : " + statVisitor.getTables());
+        System.out.println(statVisitor.getColumns());
+//        System.out.println(statVisitor.getGroupByColumns()); // group by
+        System.out.println("relationships : " + statVisitor.getRelationships()); // group by
+        System.out.println(statVisitor.getConditions());
+
+        assertEquals(1, statVisitor.getColumns().size());
+        assertEquals(0, statVisitor.getConditions().size());
+        assertEquals(0, statVisitor.getFunctions().size());
+
+        assertTrue(statVisitor.containsTable("t"));
+        assertTrue(statVisitor.containsColumn("t", "ds"));
+    }
+
+    public void test_4() throws Exception {
+        String sql = "count t PARTITION (type='20210830');";
+
+        SQLStatementParser parser = SQLParserUtils.createSQLStatementParser(sql, JdbcConstants.ODPS);
+        SQLStatement stmt = parser.parseStatementList().get(0);
+
+        SchemaStatVisitor statVisitor = SQLUtils.createSchemaStatVisitor(repository);
+        stmt.accept(statVisitor);
+
+        System.out.println("Tables : " + statVisitor.getTables());
+        System.out.println(statVisitor.getColumns());
+//        System.out.println(statVisitor.getGroupByColumns()); // group by
+        System.out.println("relationships : " + statVisitor.getRelationships()); // group by
+        System.out.println(statVisitor.getConditions());
+
+        assertEquals(1, statVisitor.getColumns().size());
+        assertEquals(0, statVisitor.getConditions().size());
+        assertEquals(0, statVisitor.getFunctions().size());
+
+        assertTrue(statVisitor.containsTable("t"));
+        assertTrue(statVisitor.containsColumn("t", "type"));
+    }
+
+    public void test_5() throws Exception {
+        String sql = "EXSTORE t PARTITION(ds='20210901');";
+
+        SQLStatementParser parser = SQLParserUtils.createSQLStatementParser(sql, JdbcConstants.ODPS);
+        SQLStatement stmt = parser.parseStatementList().get(0);
+
+        SchemaStatVisitor statVisitor = SQLUtils.createSchemaStatVisitor(repository);
+        stmt.accept(statVisitor);
+
+        System.out.println("Tables : " + statVisitor.getTables());
+        System.out.println(statVisitor.getColumns());
+//        System.out.println(statVisitor.getGroupByColumns()); // group by
+        System.out.println("relationships : " + statVisitor.getRelationships()); // group by
+        System.out.println(statVisitor.getConditions());
+
+        assertEquals(1, statVisitor.getColumns().size());
+        assertEquals(0, statVisitor.getConditions().size());
+        assertEquals(0, statVisitor.getFunctions().size());
+
+        assertTrue(statVisitor.containsTable("t"));
+        assertTrue(statVisitor.containsColumn("t", "ds"));
+    }
+
+    public void test_6() throws Exception {
+        String sql = "alter table t partition(pt='20210827',pval='0to30')\n" +
+                "rename to partition(pt='20210829',pval='0to30');";
+
+        SQLStatementParser parser = SQLParserUtils.createSQLStatementParser(sql, JdbcConstants.ODPS);
+        SQLStatement stmt = parser.parseStatementList().get(0);
+
+        SchemaStatVisitor statVisitor = SQLUtils.createSchemaStatVisitor(repository);
+        stmt.accept(statVisitor);
+
+        System.out.println("Tables : " + statVisitor.getTables());
+        System.out.println(statVisitor.getColumns());
+//        System.out.println(statVisitor.getGroupByColumns()); // group by
+        System.out.println("relationships : " + statVisitor.getRelationships()); // group by
+        System.out.println(statVisitor.getConditions());
+
+        assertEquals(2, statVisitor.getColumns().size());
+        assertEquals(0, statVisitor.getConditions().size());
+        assertEquals(0, statVisitor.getFunctions().size());
+
+        assertTrue(statVisitor.containsTable("t"));
+        assertTrue(statVisitor.containsColumn("t", "pt"));
+    }
+
+}