diff --git a/core/src/main/java/com/alibaba/druid/sql/ast/statement/SQLUnnestTableSource.java b/core/src/main/java/com/alibaba/druid/sql/ast/statement/SQLUnnestTableSource.java index 4963cdd258..fcf790e4a5 100644 --- a/core/src/main/java/com/alibaba/druid/sql/ast/statement/SQLUnnestTableSource.java +++ b/core/src/main/java/com/alibaba/druid/sql/ast/statement/SQLUnnestTableSource.java @@ -13,7 +13,8 @@ public class SQLUnnestTableSource extends SQLTableSourceImpl private final List items = new ArrayList(); protected List columns = new ArrayList(); private boolean ordinality; - private SQLExpr offset; + private boolean withOffset; + private SQLExpr offsetAs; public SQLUnnestTableSource() { } @@ -23,7 +24,7 @@ protected void accept0(SQLASTVisitor v) { if (v.visit(this)) { acceptChild(v, items); acceptChild(v, columns); - acceptChild(v, offset); + acceptChild(v, offsetAs); super.accept0(v); } v.endVisit(this); @@ -59,15 +60,23 @@ public void setItem(int i, SQLExpr item) { this.items.set(i, item); } - public SQLExpr getOffset() { - return offset; + public boolean isWithOffset() { + return withOffset; } - public void setOffset(SQLExpr x) { + public void setWithOffset(boolean withOffset) { + this.withOffset = withOffset; + } + + public SQLExpr getOffsetAs() { + return offsetAs; + } + + public void setOffsetAs(SQLExpr x) { if (x != null) { x.setParent(this); } - this.offset = x; + this.offsetAs = x; } public SQLUnnestTableSource clone() { @@ -86,9 +95,10 @@ public SQLUnnestTableSource clone() { } x.alias = alias; + x.offsetAs = offsetAs; - if (offset != null) { - x.setOffset(offset); + if (offsetAs != null) { + x.setOffsetAs(offsetAs); } return x; diff --git a/core/src/main/java/com/alibaba/druid/sql/parser/SQLSelectParser.java b/core/src/main/java/com/alibaba/druid/sql/parser/SQLSelectParser.java index a426de86f8..ef21a6c0dd 100644 --- a/core/src/main/java/com/alibaba/druid/sql/parser/SQLSelectParser.java +++ b/core/src/main/java/com/alibaba/druid/sql/parser/SQLSelectParser.java @@ -1400,10 +1400,13 @@ protected SQLTableSource parseUnnestTableSource() { if (lexer.nextIf(Token.WITH)) { acceptIdentifier("OFFSET"); - lexer.nextIf(Token.AS); - unnest.setOffset( - this.exprParser.expr() - ); + unnest.setWithOffset(true); + + if (lexer.nextIf(Token.AS)) { + unnest.setOffsetAs( + this.exprParser.expr() + ); + } } return unnest; } else { 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 index 992793b84a..39a1f5e8d9 100644 --- a/core/src/main/java/com/alibaba/druid/sql/visitor/SQLASTOutputVisitor.java +++ b/core/src/main/java/com/alibaba/druid/sql/visitor/SQLASTOutputVisitor.java @@ -5141,9 +5141,13 @@ public boolean visit(SQLUnnestTableSource x) { print(')'); } - if (x.getOffset() != null) { - print0(ucase ? " WITH OFFSET AS " : " with offset as "); - x.getOffset().accept(this); + if (x.isWithOffset()) { + print0(ucase ? " WITH OFFSET " : " with offset "); + + if (x.getOffsetAs() != null) { + print0(ucase ? "AS " : "as "); + x.getOffsetAs().accept(this); + } } printPivot(x.getPivot()); printUnpivot(x.getUnpivot()); diff --git a/core/src/test/java/com/alibaba/druid/bvt/sql/bigquery/UnnestTest.java b/core/src/test/java/com/alibaba/druid/bvt/sql/bigquery/UnnestTest.java index 28ade585ba..f02836301c 100644 --- a/core/src/test/java/com/alibaba/druid/bvt/sql/bigquery/UnnestTest.java +++ b/core/src/test/java/com/alibaba/druid/bvt/sql/bigquery/UnnestTest.java @@ -2,12 +2,20 @@ import com.alibaba.druid.DbType; import com.alibaba.druid.sql.SQLUtils; +import com.alibaba.druid.sql.ast.SQLStructDataType; +import com.alibaba.druid.sql.ast.expr.SQLArrayExpr; +import com.alibaba.druid.sql.ast.expr.SQLIdentifierExpr; import com.alibaba.druid.sql.ast.statement.SQLJoinTableSource; import com.alibaba.druid.sql.ast.statement.SQLSelectStatement; import com.alibaba.druid.sql.ast.statement.SQLTableSource; import com.alibaba.druid.sql.ast.statement.SQLUnnestTableSource; +import com.alibaba.druid.sql.visitor.SQLASTOutputVisitor; +import com.alibaba.druid.sql.visitor.VisitorFeature; import org.junit.Test; +import java.util.Objects; + +import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertTrue; public class UnnestTest { @@ -29,4 +37,76 @@ public void test_0() throws Exception { SQLJoinTableSource left = (SQLJoinTableSource) join.getLeft(); assertTrue(left.getRight() instanceof SQLUnnestTableSource); } + + @Test + public void test_1() throws Exception { + String sql = "SELECT *\n" + + "FROM UNNEST(\n" + + " ARRAY<\n" + + " STRUCT<\n" + + " x INT64,\n" + + " y STRING,\n" + + " z STRUCT>>[\n" + + " (1, 'foo', (10, 11)),\n" + + " (3, 'bar', (20, 21))]);"; + + SQLSelectStatement stmt = (SQLSelectStatement) SQLUtils.parseSingleStatement(sql, DbType.bigquery); + + SQLTableSource from = stmt.getSelect().getQueryBlock().getFrom(); + assertTrue(from instanceof SQLUnnestTableSource); + + SQLUnnestTableSource unnest = (SQLUnnestTableSource) from; + assertEquals(1, unnest.getItems().size()); + + assertTrue(unnest.getItems().get(0) instanceof SQLArrayExpr); + SQLArrayExpr array = (SQLArrayExpr) unnest.getItems().get(0); + + assertEquals(2, array.getValues().size()); + assertTrue(array.getDataType() instanceof SQLStructDataType); + } + + @Test + public void test_2() throws Exception { + String sql = "SELECT * FROM UNNEST ([10,20,30]) as numbers WITH OFFSET;"; + + SQLSelectStatement stmt = (SQLSelectStatement) SQLUtils.parseSingleStatement(sql, DbType.bigquery); + + SQLTableSource from = stmt.getSelect().getQueryBlock().getFrom(); + assertTrue(from instanceof SQLUnnestTableSource); + + SQLUnnestTableSource unnest = (SQLUnnestTableSource) from; + assertTrue(unnest.isWithOffset()); + + assertTrue(Objects.isNull(unnest.getOffsetAs())); + } + + @Test + public void test_3() throws Exception { + String sql = "SELECT * FROM UNNEST ([10,20,30]) as numbers WITH OFFSET AS ROW_INDEX;"; + + SQLSelectStatement stmt = (SQLSelectStatement) SQLUtils.parseSingleStatement(sql, DbType.bigquery); + + SQLTableSource from = stmt.getSelect().getQueryBlock().getFrom(); + assertTrue(from instanceof SQLUnnestTableSource); + + SQLUnnestTableSource unnest = (SQLUnnestTableSource) from; + assertTrue(unnest.isWithOffset()); + assertTrue(Objects.nonNull(unnest.getOffsetAs())); + assertTrue(unnest.getOffsetAs() instanceof SQLIdentifierExpr); + assertEquals("ROW_INDEX", unnest.getOffsetAs().toString()); + } + + @Test + public void test_4() throws Exception { + String sql = "SELECT *\n" + + "FROM UNNEST([10, 20, 30]) AS numbers WITH OFFSET "; + + SQLSelectStatement stmt = (SQLSelectStatement) SQLUtils.parseSingleStatement(sql, DbType.bigquery); + StringBuilder out = new StringBuilder(); + SQLASTOutputVisitor visitor = SQLUtils.createOutputVisitor(out, DbType.bigquery); + visitor.config(VisitorFeature.OutputUCase); + stmt.accept(visitor); + + assertEquals(sql, out.toString()); + } }