From 95dd37995e5438e58fab09cbfae973dee1d10eff Mon Sep 17 00:00:00 2001
From: BugsGuru <lwqlwq888@gmail.com>
Date: Fri, 24 Jan 2025 20:47:40 +0800
Subject: [PATCH] fix ai rules mysql

---
 sqle/driver/mysql/mysql.go | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/sqle/driver/mysql/mysql.go b/sqle/driver/mysql/mysql.go
index de26a3350..cdc69ec15 100644
--- a/sqle/driver/mysql/mysql.go
+++ b/sqle/driver/mysql/mysql.go
@@ -377,7 +377,10 @@ func (i *MysqlDriverImpl) audit(ctx context.Context, sql string) (*driverV2.Audi
 
 		handler, ok := rulepkg.RuleHandlerMap[rule.Name]
 		if !ok || handler.Func == nil {
-			continue
+			handler, ok = rulepkg.AIRuleHandlerMap[rule.Name]
+			if !ok || handler.Func == nil {
+				continue
+			}
 		}
 		if i.IsOfflineAudit() && !handler.IsAllowOfflineRule(nodes[0]) {
 			continue