diff --git a/sqle/driver/mysql/splitter/block.go b/sqle/driver/mysql/splitter/block.go index 488a93a4a..462f87a42 100644 --- a/sqle/driver/mysql/splitter/block.go +++ b/sqle/driver/mysql/splitter/block.go @@ -11,6 +11,10 @@ type Block interface { MatchEnd(token *parser.Token) bool } +func isThenToken(token *parser.Token) bool { + return strings.ToUpper(token.Ident()) == "THEN" +} + var allBlocks []Block = []Block{ BeginEndBlock{}, IfEndIfBlock{}, diff --git a/sqle/driver/mysql/splitter/splitter.go b/sqle/driver/mysql/splitter/splitter.go index bf9d77fc0..6ada2c85a 100644 --- a/sqle/driver/mysql/splitter/splitter.go +++ b/sqle/driver/mysql/splitter/splitter.go @@ -187,19 +187,13 @@ func (s *splitter) skipBeginEndBlock(token *parser.Token) *parser.Token { } } if matchedBlock != nil { - // 如果匹配到的是 IF 块,需要向前看一个 token 来确定它是一个控制流语句还是一个函数。 + // 如果匹配到的是 IF 块,需要向前看以区分控制流 IF 与 IF() 函数。 + // 此处 scanner 正位于 IF 之后(currentOffset),前向扫描后需复位以便后续 END 检测。 if _, ok := matchedBlock.(IfEndIfBlock); ok { - // 向前看一个 token - nextToken := s.scanner.NextToken() - // **非常重要**:将扫描器位置恢复,以便后续逻辑可以重新处理这个 token - s.scanner.SetCursor(currentOffset) - - if nextToken.Ident() == "(" { - // 如果 IF 后面是 '(',说明是 IF() 函数,而不是块语句,继续处理下一个 token。 - } else { - // 否则,它是一个真正的 IF 块语句,将其压入堆栈。 + if s.ifIsControlFlowStatement() { blockStack = append(blockStack, matchedBlock) } + s.scanner.SetCursor(currentOffset) } else { // 对于其他类型的块,直接压入堆栈 blockStack = append(blockStack, matchedBlock) @@ -235,6 +229,28 @@ func (s *splitter) skipBeginEndBlock(token *parser.Token) *parser.Token { return token } +// ifIsControlFlowStatement 从 IF 之后向前看,判断是控制流 IF 块还是 IF() 函数。 +// 规则:跟踪括号深度,若在深度 0 处遇到 THEN,则为控制流 IF。 +func (s *splitter) ifIsControlFlowStatement() bool { + depth := 0 + for s.scanner.Offset() < len(s.scanner.Text()) { + token := s.scanner.NextToken() + switch token.Ident() { + case "(": + depth++ + case ")": + if depth > 0 { + depth-- + } + default: + if isThenToken(token) && depth == 0 { + return true + } + } + } + return false +} + // isBeginQuit 判断BEGIN语句跟的是否就是结束符 func (s *splitter) isBeginQuit(token *parser.Token, tokenIndex int) bool { if token.Ident() == s.delimiter.DelimiterStr && tokenIndex == 1 { diff --git a/sqle/driver/mysql/splitter/splitter_test.go b/sqle/driver/mysql/splitter/splitter_test.go index ec94ad10d..463b788ae 100644 --- a/sqle/driver/mysql/splitter/splitter_test.go +++ b/sqle/driver/mysql/splitter/splitter_test.go @@ -1426,3 +1426,104 @@ SELECT * FROM users;`, }) } } + +// TestIfControlFlowStatementInProcedure 测试 IF 控制流与 IF() 函数在存储过程中的拆分。 +func TestIfControlFlowStatementInProcedure(t *testing.T) { + parser := NewSplitter() + + testCases := []struct { + name string + sql string + expected int + contains string + }{ + { + name: "IF with parenthesized condition in procedure", + sql: `CREATE PROCEDURE test_proc(IN Arg0 INT) +BEGIN + DECLARE v_lng DECIMAL(10,6); + DECLARE v_lat DECIMAL(10,6); + IF ( v_lng IS NULL OR v_lat IS NULL ) AND Arg0 IS NOT NULL THEN + SET v_lng = 0; + END IF; +END;`, + expected: 1, + contains: "END;", + }, + { + name: "IF without parenthesized condition in procedure", + sql: `CREATE PROCEDURE test_proc() +BEGIN + DECLARE v_variable INT DEFAULT 0; + IF v_variable = 0 THEN + SET v_variable = 1; + END IF; +END;`, + expected: 1, + contains: "END;", + }, + { + name: "IF with ELSEIF in procedure", + sql: `CREATE PROCEDURE test_proc() +BEGIN + DECLARE v_variable INT DEFAULT 0; + IF v_variable = 0 THEN + SET v_variable = 1; + ELSEIF v_variable = 1 THEN + SET v_variable = 2; + ELSE + SET v_variable = 0; + END IF; +END;`, + expected: 1, + contains: "END;", + }, + { + name: "IF function in trigger body", + sql: `CREATE TRIGGER trg BEFORE INSERT ON account +FOR EACH ROW +BEGIN + SET @deposits = @deposits + IF(NEW.amount>0,NEW.amount,0); +END;`, + expected: 1, + contains: "IF(NEW.amount>0,NEW.amount,0)", + }, + { + name: "IF function with nested CASE THEN in arguments", + sql: `CREATE TRIGGER trg BEFORE INSERT ON account +FOR EACH ROW +BEGIN + SET @x = IF(CASE WHEN NEW.amount > 0 THEN 1 ELSE 0 END, NEW.amount, 0); +END;`, + expected: 1, + contains: "IF(CASE WHEN NEW.amount > 0 THEN 1 ELSE 0 END, NEW.amount, 0)", + }, + { + name: "procedure followed by select", + sql: `CREATE PROCEDURE test_proc() +BEGIN + IF ( 1 = 1 ) AND 2 > 0 THEN + SELECT 1; + END IF; +END; +SELECT 1;`, + expected: 2, + contains: "END;", + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + stmts, err := parser.ParseSqlText(tc.sql) + if err != nil { + t.Fatalf("解析失败: %v", err) + } + if len(stmts) != tc.expected { + t.Fatalf("期望分割出 %d 个语句,实际得到 %d 个", tc.expected, len(stmts)) + } + if tc.contains != "" && !strings.Contains(stmts[0].Text(), tc.contains) { + t.Fatalf("第一条语句应包含 %q,实际为: %s", tc.contains, stmts[0].Text()) + } + }) + } +}