Parser support from generated source

alibaba · Nov 1, 2024 · f060c27 · f060c27
1 parent 1613a76
commit f060c27
Show file tree

Hide file tree

Showing 7 changed files with 180 additions and 4 deletions.
diff --git a/core/src/main/java/com/alibaba/druid/sql/ast/statement/SQLGeneratedTableSource.java b/core/src/main/java/com/alibaba/druid/sql/ast/statement/SQLGeneratedTableSource.java
@@ -0,0 +1,72 @@
+package com.alibaba.druid.sql.ast.statement;
+
+import com.alibaba.druid.sql.ast.SQLExpr;
+import com.alibaba.druid.sql.ast.SQLName;
+import com.alibaba.druid.sql.ast.SQLReplaceable;
+import com.alibaba.druid.sql.ast.expr.SQLIdentifierExpr;
+import com.alibaba.druid.sql.visitor.SQLASTVisitor;
+
+import java.util.ArrayList;
+import java.util.List;
+
+public class SQLGeneratedTableSource extends SQLTableSourceImpl
+        implements SQLReplaceable {
+    private final List<SQLExpr> items = new ArrayList<SQLExpr>();
+    protected List<SQLName> columns = new ArrayList<SQLName>();
+    private SQLIdentifierExpr methodName;
+
+    public List<SQLExpr> getItems() {
+        return items;
+    }
+
+    public List<SQLName> getColumns() {
+        return columns;
+    }
+
+    @Override
+    protected void accept0(SQLASTVisitor v) {
+        if (v.visit(this)) {
+            acceptChild(v, methodName);
+            acceptChild(v, columns);
+            acceptChild(v, items);
+        }
+        v.endVisit(this);
+    }
+
+    @Override
+    public boolean replace(SQLExpr expr, SQLExpr target) {
+        for (int i = 0; i < items.size(); i++) {
+            if (items.get(i) == expr) {
+                target.setParent(this);
+                items.set(i, target);
+                return true;
+            }
+        }
+        if (target instanceof SQLName) {
+            SQLName targetName = (SQLName) target;
+            for (int i = 0; i < columns.size(); i++) {
+                if (columns.get(i) == expr) {
+                    target.setParent(this);
+                    columns.set(i, targetName);
+                    return true;
+                }
+            }
+        }
+        if (target instanceof SQLIdentifierExpr) {
+            if (methodName == expr) {
+                target.setParent(this);
+                methodName = (SQLIdentifierExpr) target;
+                return true;
+            }
+        }
+        return false;
+    }
+
+    public SQLIdentifierExpr getMethodName() {
+        return methodName;
+    }
+
+    public void setMethodName(SQLIdentifierExpr methodName) {
+        this.methodName = methodName;
+    }
+}
diff --git a/core/src/main/java/com/alibaba/druid/sql/dialect/gaussdb/parser/GaussDbStatementParser.java b/core/src/main/java/com/alibaba/druid/sql/dialect/gaussdb/parser/GaussDbStatementParser.java
@@ -5,9 +5,9 @@
 import com.alibaba.druid.sql.ast.expr.SQLBinaryOpExpr;
 import com.alibaba.druid.sql.ast.expr.SQLBinaryOperator;
 import com.alibaba.druid.sql.ast.expr.SQLListExpr;
-import com.alibaba.druid.sql.ast.expr.SQLQueryExpr;
 import com.alibaba.druid.sql.ast.statement.SQLCreateTableStatement;
 import com.alibaba.druid.sql.ast.statement.SQLInsertStatement;
+import com.alibaba.druid.sql.ast.statement.SQLSelect;
 import com.alibaba.druid.sql.ast.statement.SQLUpdateSetItem;
 import com.alibaba.druid.sql.dialect.gaussdb.ast.stmt.GaussDbInsertStatement;
 import com.alibaba.druid.sql.dialect.postgresql.ast.stmt.PGInsertStatement;
@@ -90,8 +90,8 @@ public PGInsertStatement parseInsert() {
                 break;
             }
         } else if (lexer.token() == (Token.SELECT)) {
-            SQLQueryExpr queryExpr = (SQLQueryExpr) this.exprParser.expr();
-            stmt.setQuery(queryExpr.getSubQuery());
+            SQLSelect select = this.createSQLSelectParser().select();
+            stmt.setQuery(select);
         }
 
         if (lexer.nextIf(Token.ON)) {

diff --git a/core/src/main/java/com/alibaba/druid/sql/dialect/postgresql/parser/PGSelectParser.java b/core/src/main/java/com/alibaba/druid/sql/dialect/postgresql/parser/PGSelectParser.java
@@ -24,6 +24,7 @@
 import com.alibaba.druid.sql.dialect.postgresql.ast.stmt.PGSelectQueryBlock.IntoOption;
 import com.alibaba.druid.sql.parser.*;
 import com.alibaba.druid.util.FnvHash;
+import com.google.common.collect.Lists;
 
 import java.util.List;
 
@@ -389,4 +390,9 @@ private void parserParameters(List<SQLParameter> parameters) {
             break;
         }
     }
+
+    @Override
+    protected List<String> getReturningFunctions() {
+        return Lists.newArrayList("GENERATE_SERIES", "GENERATE_SUBSCRIPTS");
+    }
 }
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
@@ -22,6 +22,7 @@
 import com.alibaba.druid.sql.dialect.mysql.ast.expr.MySqlOrderingExpr;
 import com.alibaba.druid.util.FnvHash;
 import com.alibaba.druid.util.StringUtils;
+import com.google.common.collect.Lists;
 
 import java.util.List;
 
@@ -1284,6 +1285,11 @@ public SQLTableSource parseTableSource(boolean forFrom) {
             return parseTableSourceRest(unnestTableSource);
         }
 
+        SQLTableSource generatedTableSource = parseGeneratedTableSource();
+        if (generatedTableSource != null) {
+            return parseTableSourceRest(generatedTableSource);
+        }
+
         SQLExprTableSource tableReference = getTableSource();
 
         parseTableSourceQueryTableExpr(tableReference);
@@ -1368,6 +1374,39 @@ protected SQLTableSource parseUnnestTableSource() {
         return null;
     }
 
+    protected SQLTableSource parseGeneratedTableSource() {
+        for (String returningFunction : getReturningFunctions()) {
+            if (lexer.identifierEquals(returningFunction)) {
+                Lexer.SavePoint mark = lexer.mark();
+                SQLIdentifierExpr methodName = new SQLIdentifierExpr(returningFunction);
+                lexer.nextToken();
+
+                if (lexer.nextIf(Token.LPAREN)) {
+                    SQLGeneratedTableSource generated = new SQLGeneratedTableSource();
+                    generated.setMethodName(methodName);
+                    this.exprParser.exprList(generated.getItems(), generated);
+                    accept(Token.RPAREN);
+
+                    String alias = this.tableAlias();
+                    generated.setAlias(alias);
+
+                    if (lexer.nextIf(Token.LPAREN)) {
+                        this.exprParser.names(generated.getColumns(), generated);
+                        accept(Token.RPAREN);
+                    }
+                    return generated;
+                } else {
+                    lexer.reset(mark);
+                }
+            }
+        }
+        return null;
+    }
+
+    protected List<String> getReturningFunctions() {
+        return Lists.newArrayList("GENERATE_SERIES");
+    }
+
     protected SQLTableSource primaryTableSourceRest(SQLTableSource tableSource) {
         return tableSource;
     }

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
@@ -4862,6 +4862,38 @@ public boolean visit(SQLUnnestTableSource x) {
         return false;
     }
 
+    @Override
+    public boolean visit(SQLGeneratedTableSource x) {
+        printExpr(x.getMethodName());
+        print('(');
+        List<SQLExpr> items = x.getItems();
+        printAndAccept(items, ", ");
+        print(')');
+
+        final List<SQLName> columns = x.getColumns();
+        final String alias = x.getAlias();
+        if (alias != null) {
+            if (columns.size() > 0) {
+                print0(ucase ? " AS " : " as ");
+            } else {
+                print(' ');
+            }
+            print0(alias);
+        }
+
+        if (columns.size() > 0) {
+            print0(" (");
+            for (int i = 0; i < columns.size(); i++) {
+                if (i != 0) {
+                    print0(", ");
+                }
+                printExpr(columns.get(i));
+            }
+            print(')');
+        }
+        return false;
+    }
+
     @Override
     public boolean visit(SQLTruncateStatement x) {
         List<SQLCommentHint> headHints = x.getHeadHintsDirect();

diff --git a/core/src/main/java/com/alibaba/druid/sql/visitor/SQLASTVisitor.java b/core/src/main/java/com/alibaba/druid/sql/visitor/SQLASTVisitor.java
@@ -2262,6 +2262,13 @@ default boolean visit(SQLUnnestTableSource x) {
     default void endVisit(SQLUnnestTableSource x) {
     }
 
+    default boolean visit(SQLGeneratedTableSource x) {
+        return true;
+    }
+
+    default void endVisit(SQLGeneratedTableSource x) {
+    }
+
     default boolean visit(SQLCopyFromStatement x) {
         return true;
     }

diff --git a/core/src/test/resources/bvt/parser/gaussdb/2.txt b/core/src/test/resources/bvt/parser/gaussdb/2.txt
@@ -17,4 +17,24 @@ DISTRIBUTE BY HASH ("dlr_code");
 select group_concat(col1 order by submit_time desc separator '|') user_education from t1
 --------------------
 SELECT group_concat(col1 ORDER BY submit_time DESC SEPARATOR '|') AS user_education
-FROM t1
+FROM t1
+------------------------------------------------------------------------------------------------------------------------
+INSERT INTO public.dim_iptv_datebase_m
+select date_format(daytime::date,'%Y-%m') base_date,
+    retrodict_month,
+    date_format(daytime::date,'%Y-%m') start_date,
+    date_format(now()::date,'%Y-%m-%d') end_date
+from generate_series('2022-07-01', now()::date, '1 month') s(daytime)
+LEFT JOIN public.dim_retrodict_month_d on type='iptv';
+--------------------
+INSERT INTO public.dim_iptv_datebase_m
+SELECT date_format(daytime::date, '%Y-%m') AS base_date, retrodict_month
+	, date_format(daytime::date, '%Y-%m') AS start_date
+	, date_format(now()::date, '%Y-%m-%d') AS end_date
+FROM GENERATE_SERIES('2022-07-01', now()::date, '1 month') AS s (daytime)
+	LEFT JOIN public.dim_retrodict_month_d ON type = 'iptv';
+------------------------------------------------------------------------------------------------------------------------
+select s.a from generate_subscripts('{NULL,1,NULL,2}'::int[], 1) as s(a)
+--------------------
+SELECT s.a
+FROM GENERATE_SUBSCRIPTS('{NULL,1,NULL,2}'::int[], 1) AS s (a)