Merge pull request #2224 from actiontech/fis-issue2175

use 1 instead select fields

Author: ColdWaterLW

sqle/driver/mysql/audit_test.go

@@ -6376,6 +6376,62 @@ func TestDMLCheckSelectRows(t *testing.T) {
 		WillReturnRows(sqlmock.NewRows([]string{"COUNT(1)"}).AddRow("100000000"))
 	runSingleRuleInspectCase(rule, t, "", inspect6, "select * from exist_tb_2 where user_id in (select v3 from exist_tb_3)", newTestResult().addResult(rulepkg.DMLCheckSelectRows))
 
+	inspect7 := NewMockInspect(e)
+	handler.ExpectQuery(regexp.QuoteMeta("select id, v1 as id from exist_tb_2 limit 10, 10")).
+		WillReturnRows(sqlmock.NewRows([]string{"type"}).AddRow(executor.ExplainRecordAccessTypeIndex).AddRow("range"))
+	handler.ExpectQuery(regexp.QuoteMeta("select count(*) from (SELECT 1 FROM `exist_tb_2` LIMIT 10,10) as t")).
+		WillReturnRows(sqlmock.NewRows([]string{"count(*)"}).AddRow("100000000"))
+	runSingleRuleInspectCase(rule, t, "", inspect7, "select id, v1 as id from exist_tb_2 limit 10, 10", newTestResult().addResult(rulepkg.DMLCheckSelectRows))
+
+	inspect8 := NewMockInspect(e)
+	handler.ExpectQuery(regexp.QuoteMeta("select id, v1 as id from exist_tb_2 group by id, v1")).
+		WillReturnRows(sqlmock.NewRows([]string{"type"}).AddRow(executor.ExplainRecordAccessTypeIndex).AddRow("range"))
+	handler.ExpectQuery(regexp.QuoteMeta("select count(*) from (SELECT 1 FROM `exist_tb_2` GROUP BY `id`,`v1`) as t")).
+		WillReturnRows(sqlmock.NewRows([]string{"count(*)"}).AddRow("100000000"))
+	runSingleRuleInspectCase(rule, t, "", inspect8, "select id, v1 as id from exist_tb_2 group by id, v1", newTestResult().addResult(rulepkg.DMLCheckSelectRows))
+
+	inspect9 := NewMockInspect(e)
+	handler.ExpectQuery(regexp.QuoteMeta("select id, v1 as id from exist_tb_2 limit 10, 10")).
+		WillReturnRows(sqlmock.NewRows([]string{"type"}).AddRow(executor.ExplainRecordAccessTypeIndex).AddRow("range"))
+	handler.ExpectQuery(regexp.QuoteMeta("select count(*) from (SELECT 1 FROM `exist_tb_2` LIMIT 10,10) as t")).
+		WillReturnRows(sqlmock.NewRows([]string{"count(*)"}).AddRow("10"))
+	runSingleRuleInspectCase(rule, t, "", inspect9, "select id, v1 as id from exist_tb_2 limit 10, 10", newTestResult())
+
+	inspect10 := NewMockInspect(e)
+	handler.ExpectQuery(regexp.QuoteMeta("select id, v1 as id from exist_tb_2 group by id, v1")).
+		WillReturnRows(sqlmock.NewRows([]string{"type"}).AddRow(executor.ExplainRecordAccessTypeIndex).AddRow("range"))
+	handler.ExpectQuery(regexp.QuoteMeta("select count(*) from (SELECT 1 FROM `exist_tb_2` GROUP BY `id`,`v1`) as t")).
+		WillReturnRows(sqlmock.NewRows([]string{"count(*)"}).AddRow("10"))
+	runSingleRuleInspectCase(rule, t, "", inspect10, "select id, v1 as id from exist_tb_2 group by id, v1", newTestResult())
+
+	inspect11 := NewMockInspect(e)
+	handler.ExpectQuery(regexp.QuoteMeta("select max(v1) from exist_tb_2 group by id")).
+		WillReturnRows(sqlmock.NewRows([]string{"type"}).AddRow(executor.ExplainRecordAccessTypeIndex).AddRow("range"))
+	handler.ExpectQuery(regexp.QuoteMeta("select count(*) from (SELECT 1 FROM `exist_tb_2` GROUP BY `id`) as t")).
+		WillReturnRows(sqlmock.NewRows([]string{"count(*)"}).AddRow("10"))
+	runSingleRuleInspectCase(rule, t, "", inspect11, "select max(v1) from exist_tb_2 group by id", newTestResult())
+
+	inspect12 := NewMockInspect(e)
+	handler.ExpectQuery(regexp.QuoteMeta("select max(v1) from exist_tb_2 group by id")).
+		WillReturnRows(sqlmock.NewRows([]string{"type"}).AddRow(executor.ExplainRecordAccessTypeIndex).AddRow("range"))
+	handler.ExpectQuery(regexp.QuoteMeta("select count(*) from (SELECT 1 FROM `exist_tb_2` GROUP BY `id`) as t")).
+		WillReturnRows(sqlmock.NewRows([]string{"count(*)"}).AddRow("10000000"))
+	runSingleRuleInspectCase(rule, t, "", inspect12, "select max(v1) from exist_tb_2 group by id", newTestResult().addResult(rulepkg.DMLCheckSelectRows))
+
+	inspect13 := NewMockInspect(e)
+	handler.ExpectQuery(regexp.QuoteMeta("select max(v1) as id, id from exist_tb_2 group by id")).
+		WillReturnRows(sqlmock.NewRows([]string{"type"}).AddRow(executor.ExplainRecordAccessTypeIndex).AddRow("range"))
+	handler.ExpectQuery(regexp.QuoteMeta("select count(*) from (SELECT 1 FROM `exist_tb_2` GROUP BY `id`) as t")).
+		WillReturnRows(sqlmock.NewRows([]string{"count(*)"}).AddRow("10"))
+	runSingleRuleInspectCase(rule, t, "", inspect13, "select max(v1) as id, id from exist_tb_2 group by id", newTestResult())
+
+	inspect14 := NewMockInspect(e)
+	handler.ExpectQuery(regexp.QuoteMeta("select max(v1) as id, id from exist_tb_2 group by id")).
+		WillReturnRows(sqlmock.NewRows([]string{"type"}).AddRow(executor.ExplainRecordAccessTypeIndex).AddRow("range"))
+	handler.ExpectQuery(regexp.QuoteMeta("select count(*) from (SELECT 1 FROM `exist_tb_2` GROUP BY `id`) as t")).
+		WillReturnRows(sqlmock.NewRows([]string{"count(*)"}).AddRow("10000000"))
+	runSingleRuleInspectCase(rule, t, "", inspect14, "select max(v1) as id, id from exist_tb_2 group by id", newTestResult().addResult(rulepkg.DMLCheckSelectRows))
+
 }
 
 func TestDMLCheckScanRows(t *testing.T) {

sqle/driver/mysql/util/util.go

@@ -68,8 +68,15 @@ func GetAffectedRowNum(ctx context.Context, originSql string, conn *executor.Exe
 	// 2. SELECT COUNT(1) FROM test LIMIT 10,10 类型的SQL结果集为空
 	// 已上两种情况,使用子查询 select count(*) from (输入的sql) as t的方式来获取影响行数
 	if cannotConvert {
+		// 将select语句中的查询字段替换为数字1
+		// https://github.com/actiontech/sqle/issues/2175
+		newSql, err := useIntReplaceSelectFields(node)
+		if err != nil {
+			log.NewEntry().Errorf("replace select fields failed, err: %v", err)
+			newSql = originSql
+		}
 		// 移除后缀分号，避免sql语法错误
-		trimSuffix := strings.TrimRight(originSql, ";")
+		trimSuffix := strings.TrimRight(newSql, ";")
 		affectRowSql = fmt.Sprintf("select count(*) from (%s) as t", trimSuffix)
 	} else {
 		sqlBuilder := new(strings.Builder)
@@ -112,6 +119,25 @@ func GetAffectedRowNum(ctx context.Context, originSql string, conn *executor.Exe
 	return affectCount, nil
 }
 
+func useIntReplaceSelectFields(node ast.StmtNode) (string, error) {
+	stmt, ok := node.(*ast.SelectStmt)
+	if !ok {
+		return "", errors.New("pass parameter is not select node")
+	}
+	newValue := &driver.ValueExpr{}
+	newValue.SetInt64(1)
+	selectFields := &ast.SelectField{Expr: newValue}
+	stmt.Fields.Fields = []*ast.SelectField{selectFields}
+
+	sqlBuilder := new(strings.Builder)
+	err := stmt.Restore(format.NewRestoreCtx(format.DefaultRestoreFlags, sqlBuilder))
+	if err != nil {
+		return "", err
+	}
+	affectRowSql := sqlBuilder.String()
+	return affectRowSql, nil
+}
+
 func getSelectNodeFromDelete(stmt *ast.DeleteStmt) *ast.SelectStmt {
 	newSelect := newSelectWithCount()