Fix bq elseif.

alibaba · Dec 5, 2024 · ed630ea · ed630ea
1 parent 99a5a99
commit ed630ea
Show file tree

Hide file tree

Showing 6 changed files with 109 additions and 3 deletions.
diff --git a/core/src/main/java/com/alibaba/druid/sql/ast/statement/SQLIfStatement.java b/core/src/main/java/com/alibaba/druid/sql/ast/statement/SQLIfStatement.java
@@ -121,6 +121,15 @@ public boolean replace(SQLExpr expr, SQLExpr target) {
     public static class ElseIf extends SQLObjectImpl implements SQLReplaceable {
         private SQLExpr condition;
         private List<SQLStatement> statements = new ArrayList<SQLStatement>();
+        private boolean isConcatenated;
+
+        public boolean isConcatenated() {
+            return isConcatenated;
+        }
+
+        public void setConcatenated(boolean isConcatenated) {
+            this.isConcatenated = isConcatenated;
+        }
 
         @Override
         public void accept0(SQLASTVisitor visitor) {

diff --git a/core/src/main/java/com/alibaba/druid/sql/dialect/bigquery/parser/BigQueryLexer.java b/core/src/main/java/com/alibaba/druid/sql/dialect/bigquery/parser/BigQueryLexer.java
@@ -109,6 +109,7 @@ protected Keywords loadKeywords() {
         map.put("TABLE", Token.TABLE);
         map.put("EXCEPTION", Token.EXCEPTION);
         map.put("RAISE", Token.RAISE);
+        map.put("ELSEIF", Token.ELSEIF);
 
         return new Keywords(map);
     }

diff --git a/core/src/main/java/com/alibaba/druid/sql/dialect/bigquery/visitor/BigQueryOutputVisitor.java b/core/src/main/java/com/alibaba/druid/sql/dialect/bigquery/visitor/BigQueryOutputVisitor.java
@@ -394,4 +394,35 @@ protected void printCreateViewAs(SQLCreateViewStatement x) {
         printOptions(x.getOptions());
         super.printCreateViewAs(x);
     }
+
+    @Override
+    public boolean visit(SQLIfStatement x) {
+        print0(ucase ? "IF " : "if ");
+        x.getCondition().accept(this);
+        this.indentCount++;
+        println();
+        print0(ucase ? "THEN" : "then");
+        println();
+        for (int i = 0, size = x.getStatements().size(); i < size; ++i) {
+            SQLStatement item = x.getStatements().get(i);
+            item.accept(this);
+            if (i != size - 1) {
+                println();
+            }
+        }
+        this.indentCount--;
+
+        for (SQLIfStatement.ElseIf elseIf : x.getElseIfList()) {
+            println();
+            elseIf.accept(this);
+        }
+
+        if (x.getElseItem() != null) {
+            println();
+            x.getElseItem().accept(this);
+        }
+        println();
+        print0(ucase ? "END IF" : "end if");
+        return false;
+    }
 }
diff --git a/core/src/main/java/com/alibaba/druid/sql/parser/SQLStatementParser.java b/core/src/main/java/com/alibaba/druid/sql/parser/SQLStatementParser.java
@@ -1271,9 +1271,11 @@ public SQLStatement parseIf() {
     }
 
     protected void parseIfElse(SQLIfStatement stmt) {
-        while (lexer.nextIf(Token.ELSE)) {
-            if (lexer.nextIf(Token.IF)) {
+        while (lexer.token == Token.ELSE || lexer.token == Token.ELSEIF) {
+            boolean isConcatenated = lexer.token == Token.ELSEIF;
+            if (lexer.nextIf(Token.ELSEIF) || (lexer.nextIf(Token.ELSE) && lexer.nextIf(Token.IF))) {
                 SQLIfStatement.ElseIf elseIf = new SQLIfStatement.ElseIf();
+                elseIf.setConcatenated(isConcatenated);
                 elseIf.setCondition(this.exprParser.expr());
                 elseIf.setParent(stmt);
                 accept(Token.THEN);

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
@@ -7642,8 +7642,13 @@ public boolean visit(SQLIfStatement.Else x) {
 
     @Override
     public boolean visit(SQLIfStatement.ElseIf x) {
-        print0(ucase ? "ELSE IF" : "else if");
+        if (x.isConcatenated()) {
+            print0(ucase ? "ELSEIF " : "elseif ");
+        } else {
+            print0(ucase ? "ELSE IF " : "else if ");
+        }
         x.getCondition().accept(this);
+        println();
         print0(ucase ? " THEN" : " then");
         this.indentCount++;
         println();

diff --git a/core/src/test/resources/bvt/parser/bigquery/0.txt b/core/src/test/resources/bvt/parser/bigquery/0.txt
@@ -6,6 +6,64 @@ SELECT 1, 2
 UNION ALL
 SELECT 2, 3
 ------------------------------------------------------------------------------------------------------------------------
+IF source_start_date <= latest_old_table_partition
+  THEN
+  SET previous_partition = (
+          SELECT MAX(_partitiontime)
+              from `p-gofin-id-data-raw.data_science_us.dm_customer_cs_cscore`
+          WHERE
+              _partitiontime < TIMESTAMP(source_start_date)
+      );
+
+  SET latest_partition = (
+        SELECT MAX(_partitiontime) as partitiontime
+            from `p-gofin-id-data-raw.data_science_us.dm_customer_cs_cscore`
+      );
+ELSEIF source_start_date = first_new_table_partition
+  THEN
+  SET previous_partition = TIMESTAMP(latest_old_table_partition);
+  SET latest_partition = TIMESTAMP(first_new_table_partition);
+ELSE
+  SET previous_partition = (
+        SELECT MAX(_partitiontime) as partitiontime
+            from `p-gofin-id-data-raw.data_science_us.m_cust_pl_c_scores`
+        WHERE
+            _partitiontime < TIMESTAMP(source_start_date)
+      );
+
+  SET latest_partition = (
+        SELECT MAX(_partitiontime) 
+            from `p-gofin-id-data-raw.data_science_us.m_cust_pl_c_scores`
+      );
+END IF;
+--------------------
+IF source_start_date <= latest_old_table_partition
+	THEN
+	SET previous_partition =(
+		SELECT MAX(_partitiontime)
+		FROM `p-gofin-id-data-raw`.data_science_us.dm_customer_cs_cscore
+		WHERE _partitiontime < TIMESTAMP(source_start_date)
+	);
+	SET latest_partition =(
+		SELECT MAX(_partitiontime) AS partitiontime
+		FROM `p-gofin-id-data-raw`.data_science_us.dm_customer_cs_cscore
+	);
+ELSEIF source_start_date = first_new_table_partition
+ THEN
+	SET previous_partition = TIMESTAMP(latest_old_table_partition);
+	SET latest_partition = TIMESTAMP(first_new_table_partition);
+ELSE
+	SET previous_partition =(
+		SELECT MAX(_partitiontime) AS partitiontime
+		FROM `p-gofin-id-data-raw`.data_science_us.m_cust_pl_c_scores
+		WHERE _partitiontime < TIMESTAMP(source_start_date)
+	);
+	SET latest_partition =(
+		SELECT MAX(_partitiontime)
+		FROM `p-gofin-id-data-raw`.data_science_us.m_cust_pl_c_scores
+	);
+END IF;
+------------------------------------------------------------------------------------------------------------------------
 select a, b, c
 from t1
 group by all