Method MustQuery

pkg/testkit/testkit.go:175–218

MustQuery query the statements and returns result rows. If expected result is set it asserts the query result equals expected result.

(sql string, args ...any)

// MustQuery query the statements and returns result rows.
// If expected result is set it asserts the query result equals expected result.
func (tk *TestKit) MustQuery(sql string, args ...any) *Result {
	defer func() {
		if tk.alloc != nil {
			tk.alloc.Reset()
		}
	}()
	rs1 := tk.MustQueryWithContext(context.Background(), sql, args...)
	if !strings.Contains(sql, "information_schema") ||
		strings.Contains(sql, "trace") ||
		strings.Contains(sql, "statements_summary") ||
		strings.Contains(sql, "slow_query") ||
		strings.Contains(sql, "cluster_config") ||
		strings.Contains(sql, "CLUSTER_") ||
		strings.Contains(sql, "STATEMENTS_SUMMARY_EVICTED") ||
		strings.Contains(sql, "TIDB_TRX") {
		return rs1
	}
	err := failpoint.Enable("github.com/pingcap/tidb/pkg/planner/core/skipExtractor", "return(true)")
	if err != nil {
		panic(err)
	}
	rs2 := tk.MustQueryWithContext(context.Background(), sql, args...)
	err = failpoint.Disable("github.com/pingcap/tidb/pkg/planner/core/skipExtractor")
	if err != nil {
		panic(err)
	}
	rs1Row := make([][]string, 0, len(rs1.rows))
	for _, row := range rs1.rows {
		rs1SubRow := make([]string, 0, len(row))
		for _, col := range row {
			rs1SubRow = append(rs1SubRow, strings.Clone(col))
		}
		rs1Row = append(rs1Row, rs1SubRow)
	}
	slices.SortFunc(rs1.rows, func(a, b []string) int {
		return slices.Compare(a, b)
	})
	slices.SortFunc(rs2.rows, func(a, b []string) int {
		return slices.Compare(a, b)
	})
	rs2.Check(rs1.Rows())
	rs1.rows = rs1Row
	return rs1
}

// EventuallyMustQueryAndCheck query the statements and assert that
// result rows.lt will equal the expected results in waitFor time, periodically checking equality each tick.