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..927cb97e14 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,6 +13,7 @@ public class SQLUnnestTableSource extends SQLTableSourceImpl private final List items = new ArrayList(); protected List columns = new ArrayList(); private boolean ordinality; + private boolean withOffset; private SQLExpr offset; public SQLUnnestTableSource() { @@ -70,6 +71,14 @@ public void setOffset(SQLExpr x) { this.offset = x; } + public boolean isWithOffset() { + return withOffset; + } + + public void setWithOffset(boolean withOffset) { + this.withOffset = withOffset; + } + public SQLUnnestTableSource clone() { SQLUnnestTableSource x = new SQLUnnestTableSource(); @@ -91,6 +100,8 @@ public SQLUnnestTableSource clone() { x.setOffset(offset); } + x.withOffset = withOffset; + 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..8871888987 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,19 @@ 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.setOffset( + this.exprParser.expr() + ); + } else { + String offsetAlias = this.tableAlias(false); + if (offsetAlias != null) { + unnest.setOffset( + new com.alibaba.druid.sql.ast.expr.SQLIdentifierExpr(offsetAlias) + ); + } + } } 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 ac293549ae..0c29a3a9a0 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 @@ -5152,6 +5152,8 @@ public boolean visit(SQLUnnestTableSource x) { if (x.getOffset() != null) { print0(ucase ? " WITH OFFSET AS " : " with offset as "); x.getOffset().accept(this); + } else if (x.isWithOffset()) { + print0(ucase ? " WITH OFFSET" : " with offset"); } 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..c4a425a00f 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 @@ -8,7 +8,7 @@ import com.alibaba.druid.sql.ast.statement.SQLUnnestTableSource; import org.junit.Test; -import static org.junit.Assert.assertTrue; +import static org.junit.Assert.*; public class UnnestTest { @Test @@ -29,4 +29,51 @@ public void test_0() throws Exception { SQLJoinTableSource left = (SQLJoinTableSource) join.getLeft(); assertTrue(left.getRight() instanceof SQLUnnestTableSource); } + + // https://github.com/alibaba/druid/issues/6547 + @Test + public void test_unnest_with_offset_no_alias() 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; + assertEquals("numbers", unnest.getAlias()); + assertTrue(unnest.isWithOffset()); + assertNull(unnest.getOffset()); + + String output = SQLUtils.toSQLString(stmt, DbType.bigquery); + assertTrue(output.toUpperCase().contains("WITH OFFSET")); + assertFalse(output.toUpperCase().contains("WITH OFFSET AS")); + } + + // https://github.com/alibaba/druid/issues/6547 + @Test + public void test_unnest_with_offset_as_alias() throws Exception { + String sql = "SELECT * FROM UNNEST ([10,20,30]) as numbers WITH OFFSET AS off"; + SQLSelectStatement stmt = (SQLSelectStatement) SQLUtils.parseSingleStatement(sql, DbType.bigquery); + SQLTableSource from = stmt.getSelect().getQueryBlock().getFrom(); + assertTrue(from instanceof SQLUnnestTableSource); + SQLUnnestTableSource unnest = (SQLUnnestTableSource) from; + assertEquals("numbers", unnest.getAlias()); + assertTrue(unnest.isWithOffset()); + assertNotNull(unnest.getOffset()); + + String output = SQLUtils.toSQLString(stmt, DbType.bigquery); + assertTrue(output.toUpperCase().contains("WITH OFFSET AS")); + } + + // https://github.com/alibaba/druid/issues/6547 + @Test + public void test_unnest_with_offset_implicit_alias() throws Exception { + // BigQuery allows: WITH OFFSET offset_alias (without AS keyword) + String sql = "SELECT * FROM UNNEST ([10,20,30]) as numbers WITH OFFSET off"; + SQLSelectStatement stmt = (SQLSelectStatement) SQLUtils.parseSingleStatement(sql, DbType.bigquery); + SQLTableSource from = stmt.getSelect().getQueryBlock().getFrom(); + assertTrue(from instanceof SQLUnnestTableSource); + SQLUnnestTableSource unnest = (SQLUnnestTableSource) from; + assertEquals("numbers", unnest.getAlias()); + assertTrue(unnest.isWithOffset()); + assertNotNull(unnest.getOffset()); + } }