Skip to content

Commit

Permalink
improve bigquery support
Browse files Browse the repository at this point in the history
  • Loading branch information
wenshao committed Jun 22, 2024
1 parent 67cee52 commit 8a38402
Show file tree
Hide file tree
Showing 27 changed files with 380 additions and 323 deletions.
3 changes: 3 additions & 0 deletions core/src/main/java/com/alibaba/druid/sql/SQLUtils.java
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@
import com.alibaba.druid.sql.ast.expr.*;
import com.alibaba.druid.sql.ast.statement.*;
import com.alibaba.druid.sql.dialect.ads.visitor.AdsOutputVisitor;
import com.alibaba.druid.sql.dialect.bigquery.visitor.BigQueryOutputVisitor;
import com.alibaba.druid.sql.dialect.blink.vsitor.BlinkOutputVisitor;
import com.alibaba.druid.sql.dialect.clickhouse.visitor.ClickSchemaStatVisitor;
import com.alibaba.druid.sql.dialect.clickhouse.visitor.ClickhouseOutputVisitor;
Expand Down Expand Up @@ -532,6 +533,8 @@ public static SQLASTOutputVisitor createFormatOutputVisitor(StringBuilder out,
return new OscarOutputVisitor(out);
case starrocks:
return new StarRocksOutputVisitor(out);
case bigquery:
return new BigQueryOutputVisitor(out);
default:
return new SQLASTOutputVisitor(out, dbType);
}
Expand Down
12 changes: 6 additions & 6 deletions core/src/main/java/com/alibaba/druid/sql/ast/SQLPartitionBy.java
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,7 @@ public abstract class SQLPartitionBy extends SQLObjectImpl {
protected List<SQLName> storeIn = new ArrayList<SQLName>();
protected List<SQLExpr> columns = new ArrayList<SQLExpr>();

protected SQLIntegerExpr lifecycle;
protected SQLIntegerExpr lifeCycle;

public List<SQLPartition> getPartitions() {
return partitions;
Expand Down Expand Up @@ -110,7 +110,7 @@ public void cloneTo(SQLPartitionBy x) {
x.storeIn.add(name2);
}

x.lifecycle = lifecycle;
x.lifeCycle = lifeCycle;
}

public boolean isPartitionByColumn(long columnNameHashCode64) {
Expand All @@ -128,12 +128,12 @@ public boolean isPartitionByColumn(long columnNameHashCode64) {
return false;
}

public SQLIntegerExpr getLifecycle() {
return lifecycle;
public SQLIntegerExpr getLifeCycle() {
return lifeCycle;
}

public void setLifecycle(SQLIntegerExpr lifecycle) {
this.lifecycle = lifecycle;
public void setLifeCycle(SQLIntegerExpr x) {
this.lifeCycle = x;
}

public abstract SQLPartitionBy clone();
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -212,6 +212,7 @@ public List<SQLObject> getChildren() {
DbType.odps,
DbType.mariadb,
DbType.bigquery,
DbType.spark,
DbType.tidb
);

Expand Down

This file was deleted.

Original file line number Diff line number Diff line change
@@ -0,0 +1,72 @@
package com.alibaba.druid.sql.dialect.bigquery.parser;

import com.alibaba.druid.sql.ast.SQLName;
import com.alibaba.druid.sql.ast.expr.SQLIdentifierExpr;
import com.alibaba.druid.sql.ast.statement.SQLColumnDefinition;
import com.alibaba.druid.sql.ast.statement.SQLCreateTableStatement;
import com.alibaba.druid.sql.ast.statement.SQLSelectOrderByItem;
import com.alibaba.druid.sql.dialect.db2.parser.DB2ExprParser;
import com.alibaba.druid.sql.parser.SQLCreateTableParser;
import com.alibaba.druid.sql.parser.SQLExprParser;
import com.alibaba.druid.sql.parser.Token;

public class BigQueryCreateTableParser extends SQLCreateTableParser {
public BigQueryCreateTableParser(String sql) {
super(new DB2ExprParser(sql));
}

public BigQueryCreateTableParser(SQLExprParser exprParser) {
super(exprParser);
}

protected void parseCreateTableRest(SQLCreateTableStatement stmt) {
for (;;) {
if (lexer.token() == Token.PARTITION) {
lexer.nextToken();
accept(Token.BY);

boolean brace = lexer.nextIf(Token.LPAREN);
for (; ; ) {
SQLName name;
name = exprParser.name();
if (name instanceof SQLIdentifierExpr
&& ((SQLIdentifierExpr) name).getName().equalsIgnoreCase("DATE")
&& lexer.nextIf(Token.LPAREN)
) {
name = exprParser.name();
accept(Token.RPAREN);
name.putAttribute("function", "DATE");
}
stmt.addPartitionColumn(new SQLColumnDefinition(name));
if (lexer.nextIf(Token.COMMA)) {
continue;
}
break;
}
if (brace) {
accept(Token.RPAREN);
}
continue;
}

if (lexer.nextIfIdentifier("CLUSTERED")) {
accept(Token.BY);
SQLSelectOrderByItem item = exprParser.parseSelectOrderByItem();
item.setParent(stmt);
stmt.getClusteredBy().add(item);
}

if (lexer.nextIfIdentifier("OPTIONS")) {
exprParser.parseAssignItem(stmt.getTableOptions(), stmt);
continue;
}

if (lexer.nextIfIdentifier("CLONE")) {
stmt.setLike(exprParser.name());
continue;
}

break;
}
}
}
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
package com.alibaba.druid.sql.dialect.bigquery;
package com.alibaba.druid.sql.dialect.bigquery.parser;

import com.alibaba.druid.DbType;
import com.alibaba.druid.sql.ast.SQLExpr;
Expand Down
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
package com.alibaba.druid.sql.dialect.bigquery;
package com.alibaba.druid.sql.dialect.bigquery.parser;

import com.alibaba.druid.DbType;
import com.alibaba.druid.sql.parser.Keywords;
Expand Down
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
package com.alibaba.druid.sql.dialect.bigquery;
package com.alibaba.druid.sql.dialect.bigquery.parser;

import com.alibaba.druid.DbType;
import com.alibaba.druid.sql.parser.SQLExprParser;
Expand Down
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
package com.alibaba.druid.sql.dialect.bigquery;
package com.alibaba.druid.sql.dialect.bigquery.parser;

import com.alibaba.druid.sql.parser.Lexer;
import com.alibaba.druid.sql.parser.SQLCreateTableParser;
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,74 @@
package com.alibaba.druid.sql.dialect.bigquery.visitor;

import com.alibaba.druid.DbType;
import com.alibaba.druid.sql.ast.statement.SQLColumnDefinition;
import com.alibaba.druid.sql.ast.statement.SQLCreateTableStatement;
import com.alibaba.druid.sql.ast.statement.SQLExprTableSource;
import com.alibaba.druid.sql.visitor.SQLASTOutputVisitor;

import java.util.List;

public class BigQueryOutputVisitor extends SQLASTOutputVisitor {
public BigQueryOutputVisitor(StringBuilder appender) {
super(appender, DbType.db2);
}

public BigQueryOutputVisitor(StringBuilder appender, boolean parameterized) {
super(appender, parameterized);
this.dbType = DbType.bigquery;
}

protected void printPartitionedBy(SQLCreateTableStatement x) {
List<SQLColumnDefinition> partitionColumns = x.getPartitionColumns();
int partitionSize = partitionColumns.size();
if (partitionSize == 0) {
return;
}

println();
print0(ucase ? "PARTITION BY (" : "partition by (");
this.indentCount++;
println();
for (int i = 0; i < partitionSize; ++i) {
SQLColumnDefinition column = partitionColumns.get(i);
printPartitoinedByColumn(column);

if (i != partitionSize - 1) {
print(',');
}
if (this.isPrettyFormat() && column.hasAfterComment()) {
print(' ');
printlnComment(column.getAfterCommentsDirect());
}

if (i != partitionSize - 1) {
println();
}
}
this.indentCount--;
println();
print(')');
}

protected void printPartitoinedByColumn(SQLColumnDefinition column) {
String function = (String) column.getName().getAttribute("function");
if (function != null) {
print0(function);
print('(');
}
column.accept(this);
if (function != null) {
print(')');
}
}

protected void printCreateTableLike(SQLCreateTableStatement x) {
SQLExprTableSource like = x.getLike();
if (like == null) {
return;
}
println();
print0(ucase ? "CLONE " : "clone ");
like.accept(this);
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -19,10 +19,8 @@
import com.alibaba.druid.sql.ast.SQLCommentHint;
import com.alibaba.druid.sql.ast.SQLExpr;
import com.alibaba.druid.sql.ast.SQLName;
import com.alibaba.druid.sql.ast.expr.SQLArrayExpr;
import com.alibaba.druid.sql.ast.expr.SQLBinaryOpExpr;
import com.alibaba.druid.sql.ast.expr.SQLCharExpr;
import com.alibaba.druid.sql.ast.expr.SQLIdentifierExpr;
import com.alibaba.druid.sql.ast.statement.*;
import com.alibaba.druid.sql.dialect.hive.ast.HiveAddJarStatement;
import com.alibaba.druid.sql.dialect.hive.ast.HiveInsert;
Expand Down Expand Up @@ -505,116 +503,4 @@ public boolean visit(HiveAddJarStatement x) {
print0(x.getPath());
return false;
}

protected void printStoredAs(SQLCreateTableStatement x) {
SQLExpr storedAs = x.getStoredAs();
if (storedAs != null) {
print0(ucase ? " STORE AS " : " store as ");
printExpr(storedAs, parameterized);
return;
}

SQLAssignItem format = null;
List<SQLAssignItem> options = x.getTableOptions();
boolean trino = x.getDbType() == DbType.presto || x.getDbType() == DbType.trino;
if (trino) {
for (int i = 0, size = options.size(); i < size; ++i) {
SQLAssignItem option = options.get(i);
if (option.getTarget() instanceof SQLIdentifierExpr) {
SQLIdentifierExpr target = (SQLIdentifierExpr) option.getTarget();
if ("FORMAT".equalsIgnoreCase(target.getName())) {
format = option;
}
}
}
}

if (format != null) {
print0(ucase ? " STORED AS " : " stored as ");
if (format.getValue() instanceof SQLCharExpr) {
String formatValue = ((SQLCharExpr) format.getValue()).getText();
print0(formatValue);
} else {
printExpr(format.getValue(), parameterized);
}
}
}

protected void printCreateTableOptions(SQLCreateTableStatement x) {
List<SQLAssignItem> options = x.getTableOptions();
if (options.isEmpty()) {
return;
}

int sepecial = 0;
SQLAssignItem partitionedBy = null, format = null;
boolean trino = x.getDbType() == DbType.presto || x.getDbType() == DbType.trino;
for (int i = 0, size = options.size(); i < size; ++i) {
SQLAssignItem option = options.get(i);
if (option.getTarget() instanceof SQLIdentifierExpr) {
SQLIdentifierExpr target = (SQLIdentifierExpr) option.getTarget();
if (trino) {
if ("PARTITIONED_BY".equalsIgnoreCase(target.getName())) {
sepecial++;
partitionedBy = option;
} else if ("FORMAT".equalsIgnoreCase(target.getName()) && x.getStoredAs() == null) {
sepecial++;
format = option;
}
}
}
}

if (partitionedBy != null) {
println();
print0(ucase ? "PARTITIONED BY (" : "partitioned by (");
SQLExpr value = partitionedBy.getValue();
if (value instanceof SQLArrayExpr) {
List<SQLExpr> values = ((SQLArrayExpr) value).getValues();
for (int i = 0, size = values.size(); i < size; ++i) {
if (i != 0) {
print0(", ");
}
SQLExpr column = values.get(i);
if (column instanceof SQLCharExpr) {
String formatValue = ((SQLCharExpr) column).getText();
print0(formatValue);
} else {
printExpr(column, parameterized);
}
}
} else {
value.accept(this);
}
print(')');
}

if (x.getLifeCycle() != null) {
println();
print0(ucase ? "LIFECYCLE " : "lifecycle ");
x.getLifeCycle().accept(this);
}

if (options.size() == sepecial) {
return;
}

println();
print0(ucase ? "TBLPROPERTIES (" : "tblproperties (");
for (int i = 0, p = 0, size = options.size(); i < size; ++i) {
SQLAssignItem option = options.get(i);
if (option == partitionedBy
|| option == format
) {
continue;
}

if (p != 0) {
print0(", ");
}
option.accept(this);
p++;
}
print(')');
}
}
Loading

0 comments on commit 8a38402

Please sign in to comment.