Merge pull request #2945 from actiontech/2928

fix: panic when analysis insert into select stmt

Author: littleniannian

sqle/driver/mysql/analysis.go

@@ -310,7 +310,12 @@ func (i *MysqlDriverImpl) ExtractSchemaTableList(sql string) ([]SchemaTable, err
 	case *ast.InsertStmt:
 		getMultiTables(stmt.Table.TableRefs)
 		if stmt.Select != nil {
-			getMultiTables(stmt.Select.(*ast.SelectStmt).From.TableRefs)
+			// TODO：INSERT INTO SQLE00115_t1_tmp_employee (id, cname, sex, age, salary) SELECT 4000002, '小张', 0, 25, (SELECT AVG(salary) FROM SQLE00115_t1_employee) 对于这条SQL，解析器无法解析子查询为一个Select语句，而是认为是一个文本
+			if selectStmt, ok := stmt.Select.(*ast.SelectStmt); ok {
+				if selectStmt.From != nil{
+					getMultiTables(selectStmt.From.TableRefs)
+				}
+			}
 		}
 	case *ast.DeleteStmt:
 		getMultiTables(stmt.TableRefs.TableRefs)

sqle/driver/mysql/util/util.go

@@ -50,7 +50,17 @@ func GetAffectedRowNum(ctx context.Context, originSql string, conn *executor.Exe
 		// 包含子查询的insert语句，insert into t1 (name) select name from t2
 		isSelectInsert := stmt.Select != nil && stmt.Lists == nil
 		if isSelectInsert {
-			newNode = getSelectNodeFromSelect(stmt.Select.(*ast.SelectStmt))
+			if selectStmt, ok := stmt.Select.(*ast.SelectStmt); ok {
+				newNode = getSelectNodeFromSelect(selectStmt)
+			}
+			// union语句，无法转换为select count语句
+			if unionStmt, ok := stmt.Select.(*ast.UnionStmt); ok {
+				cannotConvert = true
+				originSql, err = restoreToSqlWithFlag(format.DefaultRestoreFlags, unionStmt)
+				if err != nil {
+					return 0, err
+				}
+			}
 		} else if isCommonInsert {
 			return int64(len(stmt.Lists)), nil
 		} else {
@@ -73,6 +83,10 @@ func GetAffectedRowNum(ctx context.Context, originSql string, conn *executor.Exe
 		trimSuffix := strings.TrimRight(originSql, ";")
 		affectedRowSql = fmt.Sprintf("select count(*) from (%s) as t", trimSuffix)
 	} else {
+		if newNode == nil {
+			log.NewEntry().Errorf("in GetAffectedRowNum, when getting select node from %v failed", originSql)
+			return 0, fmt.Errorf("get select node from %v failed", originSql)
+		}
 		sqlBuilder := new(strings.Builder)
 		err = newNode.Restore(format.NewRestoreCtx(format.DefaultRestoreFlags, sqlBuilder))
 		if err != nil {