improved presto/trino support

This commit is contained in:
高铁
2021-10-03 12:21:06 +08:00
parent d85142ca66
commit eca685c228
18 changed files with 256 additions and 19 deletions
@@ -56,6 +56,8 @@ public enum DbType {
*/
gaussdb (1 << 38),
trino (1 << 39),
ingres (0),
cloudscape (0),
@@ -206,6 +206,7 @@ public class SQLDateExpr extends SQLExprImpl implements SQLLiteralExpr, SQLValua
public static long supportDbTypes = DbType.of(DbType.mysql
, DbType.oracle
, DbType.presto
, DbType.trino
, DbType.postgresql
, DbType.ads
, DbType.hive
@@ -144,7 +144,12 @@ public class SQLDateTimeExpr extends SQLExprImpl implements SQLLiteralExpr, SQLV
return Collections.emptyList();
}
public static long supportDbTypes = DbType.of(DbType.mysql, DbType.oracle, DbType.presto, DbType.postgresql);
public static long supportDbTypes = DbType.of(DbType.mysql
, DbType.oracle
, DbType.presto
, DbType.trino
, DbType.postgresql
);
public static boolean isSupport(DbType dbType) {
return (dbType.mask & supportDbTypes) != 0;
@@ -131,7 +131,7 @@ public class SQLExprUtils {
if (quote == '`') {
expr = new SQLIdentifierExpr(str);
} else if (quote == '"') {
if (dbType == DbType.oracle || dbType == DbType.presto) {
if (dbType == DbType.oracle || dbType == DbType.presto || dbType == DbType.trino) {
expr = new SQLIdentifierExpr(str);
}
else {
@@ -145,7 +145,12 @@ public class SQLTimeExpr extends SQLExprImpl implements SQLLiteralExpr, SQLValua
return Collections.emptyList();
}
public static long supportDbTypes = DbType.of(DbType.mysql, DbType.oracle, DbType.presto, DbType.postgresql);
public static long supportDbTypes = DbType.of(DbType.mysql
, DbType.oracle
, DbType.presto
, DbType.trino
, DbType.postgresql
);
public static boolean isSupport(DbType dbType) {
return (dbType.mask & supportDbTypes) != 0;
@@ -62,6 +62,7 @@ public class SQLFunctionBuilder {
switch (dbType) {
case ads:
case presto:
case trino:
case odps:
return new SQLMethodInvokeExpr("coalesce", null, expr1, expr2);
case oracle:
@@ -2447,6 +2447,7 @@ public class MySqlStatementParser extends SQLStatementParser {
case mysql:
case ads:
case presto:
case trino:
Lexer.SavePoint mark = lexer.mark();
if (lexer.token() == Token.LPAREN) {
@@ -29,6 +29,7 @@ public class PrestoLexer extends Lexer {
map.put("USING", Token.USING);
map.put("MATCHED", Token.MATCHED);
map.put("UPSERT", Token.UPSERT);
map.put("ARRAY", Token.ARRAY);
DEFAULT_PHOENIX_KEYWORDS = new Keywords(map);
}
@@ -1296,6 +1296,17 @@ public class SQLExprParser extends SQLParser {
break;
}
throw new ParserException("ERROR. " + lexer.info());
case ARRAY: {
SQLArrayExpr array = new SQLArrayExpr();
array.setExpr(new SQLIdentifierExpr("ARRAY"));
lexer.nextToken();
accept(Token.LBRACKET);
this.exprList(array.getValues(), array);
accept(Token.RBRACKET);
sqlExpr = array;
break;
}
default:
throw new ParserException("ERROR. " + lexer.info());
}
@@ -130,6 +130,7 @@ public class SQLParserUtils {
case hive:
return new HiveStatementParser(sql, features);
case presto:
case trino:
return new PrestoStatementParser(sql);
case ads:
return new AdsStatementParser(sql);
@@ -173,6 +174,7 @@ public class SQLParserUtils {
case phoenix:
return new PhoenixExprParser(sql, features);
case presto:
case trino:
return new PrestoExprParser(sql, features);
case hive:
return new HiveExprParser(sql, features);
@@ -213,6 +215,7 @@ public class SQLParserUtils {
case phoenix:
return new PhoenixLexer(sql);
case presto:
case trino:
return new PrestoLexer(sql);
case antspark:
return new AntsparkLexer(sql);
@@ -4793,7 +4793,7 @@ public class SQLStatementParser extends SQLParser {
}
}
if (DbType.mysql == dbType || DbType.ads == dbType || DbType.presto == dbType) {
if (DbType.mysql == dbType || DbType.ads == dbType || DbType.presto == dbType || DbType.trino == dbType) {
if (lexer.token == Token.LPAREN) {
lexer.nextToken();
@@ -2314,6 +2314,16 @@ public class SQLASTOutputVisitor extends SQLASTVisitorAdapter implements Paramet
this.indentCount--;
} else if (parent instanceof SQLOpenStatement) {
visit(subQuery);
} else if (parent instanceof SQLMethodInvokeExpr
&& ((SQLMethodInvokeExpr) parent).getArguments().size() == 1
&& (((SQLMethodInvokeExpr) parent).methodNameHashCode64() == FnvHash.Constants.LATERAL
|| ((SQLMethodInvokeExpr) parent).methodNameHashCode64() == FnvHash.Constants.ARRAY)
) {
this.indentCount++;
println();
visit(subQuery);
this.indentCount--;
println();
} else {
print('(');
this.indentCount++;
@@ -4752,7 +4762,11 @@ public class SQLASTOutputVisitor extends SQLASTVisitorAdapter implements Paramet
}
if (x.getDbProperties().size() > 0) {
if (dbType == DbType.mysql || dbType == DbType.presto || dbType == DbType.ads || dbType == DbType.mariadb) {
if (dbType == DbType.mysql
|| dbType == DbType.presto
|| dbType == DbType.trino
|| dbType == DbType.ads
|| dbType == DbType.mariadb) {
println();
print0(ucase ? "WITH (" : "with (");
} else if (dbType == DbType.hive) {
@@ -626,6 +626,8 @@ public final class JdbcUtils implements JdbcConstants {
return DbType.clickhouse;
} else if (rawUrl.startsWith("jdbc:presto:")) {
return DbType.presto;
} else if (rawUrl.startsWith("jdbc:trino:")) {
return DbType.trino;
} else if (rawUrl.startsWith("jdbc:inspur:")) {
return DbType.kdb;
} else if (rawUrl.startsWith("jdbc:polardb")) {
@@ -136,6 +136,7 @@ public class WallFilter extends FilterAdapter implements WallFilterMBean {
case mariadb:
case h2:
case presto:
case trino:
if (config == null) {
config = new WallConfig(MySqlWallProvider.DEFAULT_CONFIG_DIR);
}
@@ -31,9 +31,9 @@ public class MySqlSelectTest_292_lateral extends MysqlTest {
assertEquals("SELECT *\n" +
"FROM t\n" +
"\tCROSS JOIN LATERAL((\n" +
"\tCROSS JOIN LATERAL(\n" +
"\t\tVALUES (1)\n" +
"\t))", stmt.toString());
"\t)", stmt.toString());
System.out.println(stmt.toString());
}
@@ -46,9 +46,9 @@ public class MySqlSelectTest_292_lateral extends MysqlTest {
assertEquals("SELECT *\n" +
"FROM t\n" +
"\tFULL JOIN LATERAL((\n" +
"\tFULL JOIN LATERAL(\n" +
"\t\tVALUES (1)\n" +
"\t)) ON true", stmt.toString());
"\t) ON true", stmt.toString());
System.out.println(stmt.toString());
}
@@ -60,9 +60,9 @@ public class MySqlSelectTest_292_lateral extends MysqlTest {
.parseSingleStatement(sql, DbType.mysql);
assertEquals("SELECT *\n" +
"FROM t, LATERAL((\n" +
"FROM t, LATERAL(\n" +
"\t\tVALUES (1)\n" +
"\t)) AS a (x)", stmt.toString());
"\t) AS a (x)", stmt.toString());
System.out.println(stmt.toString());
}
@@ -33,13 +33,13 @@ public class PGSelectTest62_json extends TestCase {
List<SQLStatement> stmtList = SQLUtils.parseStatements(sql, dbType);
SQLStatement stmt = stmtList.get(0);
assertEquals("SELECT array_to_string(array((\n" +
assertEquals("SELECT array_to_string(array(\n" +
"\t\tSELECT jsonb_array_elements('[{\"aa\":\"123\",\"bb\":\"123\"},{\"aa\":\"12356\",\"bb\":\"12356\"}]') ->> 'aa'\n" +
"\t)), ',') AS aa", SQLUtils.toPGString(stmt));
"\t), ',') AS aa", SQLUtils.toPGString(stmt));
assertEquals("select array_to_string(array((\n" +
assertEquals("select array_to_string(array(\n" +
"\t\tselect jsonb_array_elements('[{\"aa\":\"123\",\"bb\":\"123\"},{\"aa\":\"12356\",\"bb\":\"12356\"}]') ->> 'aa'\n" +
"\t)), ',') as aa", SQLUtils.toPGString(stmt, SQLUtils.DEFAULT_LCASE_FORMAT_OPTION));
"\t), ',') as aa", SQLUtils.toPGString(stmt, SQLUtils.DEFAULT_LCASE_FORMAT_OPTION));
assertEquals(1, stmtList.size());
@@ -53,13 +53,13 @@ public class PGSelectTest68 extends TestCase {
"\t\tELSE -1\n" +
"\tEND\n" +
"FROM (\n" +
"\t(SELECT ARRAY_TO_STRING(ARRAY((\n" +
"\t(SELECT ARRAY_TO_STRING(ARRAY(\n" +
"\t\t\tSELECT DISTINCT area_info\n" +
"\t\t\tFROM md_mesh\n" +
"\t\t\tWHERE mesh IN (\n" +
"\t\t\t\tSELECT regexp_split_to_table(?, ',')\n" +
"\t\t\t)\n" +
"\t\t)), ',') AS area_string)\n" +
"\t\t), ',') AS area_string)\n" +
") t", SQLUtils.toPGString(stmt));
assertEquals("select case \n" +
@@ -69,13 +69,13 @@ public class PGSelectTest68 extends TestCase {
"\t\telse -1\n" +
"\tend\n" +
"from (\n" +
"\t(select ARRAY_TO_STRING(ARRAY((\n" +
"\t(select ARRAY_TO_STRING(ARRAY(\n" +
"\t\t\tselect distinct area_info\n" +
"\t\t\tfrom md_mesh\n" +
"\t\t\twhere mesh in (\n" +
"\t\t\t\tselect regexp_split_to_table(?, ',')\n" +
"\t\t\t)\n" +
"\t\t)), ',') as area_string)\n" +
"\t\t), ',') as area_string)\n" +
") t", SQLUtils.toPGString(stmt, SQLUtils.DEFAULT_LCASE_FORMAT_OPTION));
assertEquals(1, stmtList.size());
@@ -0,0 +1,190 @@
package com.alibaba.druid.bvt.sql.presto;
import com.alibaba.druid.DbType;
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.visitor.SchemaStatVisitor;
import com.alibaba.druid.stat.TableStat;
import org.junit.Test;
import java.util.Map;
import static org.junit.Assert.*;
public class PrestoSelect_0 {
protected final DbType dbType = DbType.presto;
@Test
public void test_with() throws Exception {
String sql = "WITH RECURSIVE t(n) AS (\n" +
" VALUES (1)\n" +
" UNION ALL\n" +
" SELECT n + 1 FROM t WHERE n < 4\n" +
")\n" +
"SELECT sum(n) FROM t;";
SQLStatementParser parser = SQLParserUtils.createSQLStatementParser(sql, dbType);
SQLStatement stmt = parser.parseStatement();
SchemaStatVisitor visitor = SQLUtils.createSchemaStatVisitor(dbType);
stmt.accept(visitor);
Map<TableStat.Name, TableStat> tableMap = visitor.getTables();
assertTrue(tableMap.isEmpty());
assertEquals("WITH RECURSIVE t (n) AS (\n" +
"\t\tVALUES (1)\n" +
"\t\tUNION ALL\n" +
"\t\tSELECT n + 1\n" +
"\t\tFROM t\n" +
"\t\tWHERE n < 4\n" +
"\t)\n" +
"SELECT sum(n)\n" +
"FROM t", stmt.toString());
}
@Test
public void test_unnest_0() throws Exception {
String sql = "SELECT student, score\n" +
"FROM tests\n" +
"CROSS JOIN UNNEST(scores) AS t (score);";
SQLStatementParser parser = SQLParserUtils.createSQLStatementParser(sql, dbType);
SQLStatement stmt = parser.parseStatement();
SchemaStatVisitor visitor = SQLUtils.createSchemaStatVisitor(dbType);
stmt.accept(visitor);
Map<TableStat.Name, TableStat> tableMap = visitor.getTables();
assertFalse(tableMap.isEmpty());
assertEquals("SELECT student, score\n" +
"FROM tests\n" +
"\tCROSS JOIN UNNEST(scores) AS t (score)", stmt.toString());
}
@Test
public void test_unnest_1() throws Exception {
String sql = "SELECT numbers, animals, n, a\n" +
"FROM (\n" +
" VALUES\n" +
" (ARRAY[2, 5], ARRAY['dog', 'cat', 'bird']),\n" +
" (ARRAY[7, 8, 9], ARRAY['cow', 'pig'])\n" +
") AS x (numbers, animals)\n" +
"CROSS JOIN UNNEST(numbers, animals) AS t (n, a);";
SQLStatementParser parser = SQLParserUtils.createSQLStatementParser(sql, dbType);
SQLStatement stmt = parser.parseStatement();
SchemaStatVisitor visitor = SQLUtils.createSchemaStatVisitor(dbType);
stmt.accept(visitor);
Map<TableStat.Name, TableStat> tableMap = visitor.getTables();
assertTrue(tableMap.isEmpty());
assertEquals("SELECT numbers, animals, n, a\n" +
"FROM (VALUES (ARRAY[2, 5], ARRAY['dog', 'cat', 'bird']), (ARRAY[7, 8, 9], ARRAY['cow', 'pig'])) AS x (numbers, animals)\n" +
"\tCROSS JOIN UNNEST(numbers, animals) AS t (n, a)", stmt.toString());
}
@Test
public void test_unnest_2() throws Exception {
String sql = "SELECT numbers, n, a\n" +
"FROM (\n" +
" VALUES\n" +
" (ARRAY[2, 5]),\n" +
" (ARRAY[7, 8, 9])\n" +
") AS x (numbers)\n" +
"CROSS JOIN UNNEST(numbers) WITH ORDINALITY AS t (n, a);";
SQLStatementParser parser = SQLParserUtils.createSQLStatementParser(sql, dbType);
SQLStatement stmt = parser.parseStatement();
SchemaStatVisitor visitor = SQLUtils.createSchemaStatVisitor(dbType);
stmt.accept(visitor);
Map<TableStat.Name, TableStat> tableMap = visitor.getTables();
assertTrue(tableMap.isEmpty());
assertEquals("SELECT numbers, n, a\n" +
"FROM (VALUES (ARRAY[2, 5]), (ARRAY[7, 8, 9])) AS x (numbers)\n" +
"\tCROSS JOIN UNNEST(numbers) WITH ORDINALITY AS t (n, a)", stmt.toString());
}
@Test
public void test_lateral() throws Exception {
String sql = "SELECT name, x, y\n" +
"FROM nation\n" +
"CROSS JOIN LATERAL (SELECT name || ' :-' AS x)\n" +
"CROSS JOIN LATERAL (SELECT x || ')' AS y);";
SQLStatementParser parser = SQLParserUtils.createSQLStatementParser(sql, dbType);
SQLStatement stmt = parser.parseStatement();
SchemaStatVisitor visitor = SQLUtils.createSchemaStatVisitor(dbType);
stmt.accept(visitor);
Map<TableStat.Name, TableStat> tableMap = visitor.getTables();
assertFalse(tableMap.isEmpty());
assertEquals("SELECT name, x, y\n" +
"FROM nation\n" +
"\tCROSS JOIN LATERAL(\n" +
"\t\tSELECT name || ' :-' AS x\n" +
"\t)\n" +
"\tCROSS JOIN LATERAL(\n" +
"\t\tSELECT x || ')' AS y\n" +
"\t)", stmt.toString());
}
@Test
public void test_exists() throws Exception {
String sql = "SELECT name\n" +
"FROM nation\n" +
"WHERE EXISTS (\n" +
" SELECT *\n" +
" FROM region\n" +
" WHERE region.regionkey = nation.regionkey\n" +
");";
SQLStatementParser parser = SQLParserUtils.createSQLStatementParser(sql, dbType);
SQLStatement stmt = parser.parseStatement();
SchemaStatVisitor visitor = SQLUtils.createSchemaStatVisitor(dbType);
stmt.accept(visitor);
Map<TableStat.Name, TableStat> tableMap = visitor.getTables();
assertFalse(tableMap.isEmpty());
assertEquals("SELECT name\n" +
"FROM nation\n" +
"WHERE EXISTS (\n" +
"\tSELECT *\n" +
"\tFROM region\n" +
"\tWHERE region.regionkey = nation.regionkey\n" +
")", stmt.toString());
}
@Test
public void test_in() throws Exception {
String sql = "SELECT name\n" +
"FROM nation\n" +
"WHERE regionkey IN (\n" +
" SELECT regionkey\n" +
" FROM region\n" +
" WHERE name = 'AMERICA' OR name = 'AFRICA'\n" +
");";
SQLStatementParser parser = SQLParserUtils.createSQLStatementParser(sql, dbType);
SQLStatement stmt = parser.parseStatement();
SchemaStatVisitor visitor = SQLUtils.createSchemaStatVisitor(dbType);
stmt.accept(visitor);
Map<TableStat.Name, TableStat> tableMap = visitor.getTables();
assertFalse(tableMap.isEmpty());
assertEquals("SELECT name\n" +
"FROM nation\n" +
"WHERE regionkey IN (\n" +
"\tSELECT regionkey\n" +
"\tFROM region\n" +
"\tWHERE name = 'AMERICA'\n" +
"\t\tOR name = 'AFRICA'\n" +
")", stmt.toString());
}
@Test
public void test_scalar() throws Exception {
String sql = "SELECT name\n" +
"FROM nation\n" +
"WHERE regionkey = (SELECT max(regionkey) FROM region);";
SQLStatementParser parser = SQLParserUtils.createSQLStatementParser(sql, dbType);
SQLStatement stmt = parser.parseStatement();
SchemaStatVisitor visitor = SQLUtils.createSchemaStatVisitor(dbType);
stmt.accept(visitor);
Map<TableStat.Name, TableStat> tableMap = visitor.getTables();
assertFalse(tableMap.isEmpty());
assertEquals("SELECT name\n" +
"FROM nation\n" +
"WHERE regionkey = (\n" +
"\tSELECT max(regionkey)\n" +
"\tFROM region\n" +
")", stmt.toString());
}
}