Function TestCursorFetchErrorInFetch

pkg/server/tests/cursor/cursor_test.go:40–93

(t *testing.T)

)

func TestCursorFetchErrorInFetch(t *testing.T) {
	tmpStoragePath := t.TempDir()
	restore := config.RestoreFunc()
	defer restore()
	config.UpdateGlobal(func(conf *config.Config) {
		conf.TempStoragePath = tmpStoragePath
	})

	store, dom := testkit.CreateMockStoreAndDomain(t)
	srv := server2.CreateMockServer(t, store)
	srv.SetDomain(dom)
	defer srv.Close()

	appendUint32 := binary.LittleEndian.AppendUint32
	ctx := context.Background()
	c := server2.CreateMockConn(t, srv)

	tk := testkit.NewTestKitWithSession(t, store, c.Context().Session)
	tk.MustExec("use test")
	tk.MustExec("drop table if exists t")
	tk.MustExec("create table t(id int, payload BLOB)")
	payload := make([]byte, 512)
	for i := 0; i < 2048; i++ {
		rand.Read(payload)
		tk.MustExec("insert into t values (?, ?)", i, payload)
	}

	tk.MustExec("set global tidb_enable_tmp_storage_on_oom = ON")
	tk.MustExec("set global tidb_mem_oom_action = 'CANCEL'")
	defer tk.MustExec("set global tidb_mem_oom_action= DEFAULT")
	require.Len(t, tk.Session().GetSessionVars().MemTracker.GetChildrenForTest(), 1)

	// execute a normal statement, it'll spill to disk
	stmt, _, _, err := c.Context().Prepare("select * from t")
	require.NoError(t, err)

	tk.MustExec(fmt.Sprintf("set tidb_mem_quota_query=%d", 1))

	require.NoError(t, failpoint.Enable("github.com/pingcap/tidb/pkg/util/chunk/get-chunk-error", "return(true)"))
	defer func() {
		require.NoError(t, failpoint.Disable("github.com/pingcap/tidb/pkg/util/chunk/get-chunk-error"))
	}()

	require.NoError(t, c.Dispatch(ctx, append(
		appendUint32([]byte{tmysql.ComStmtExecute}, uint32(stmt.ID())),
		tmysql.CursorTypeReadOnly, 0x1, 0x0, 0x0, 0x0,
	)))

	require.ErrorContains(t, c.Dispatch(ctx, appendUint32(appendUint32([]byte{tmysql.ComStmtFetch}, uint32(stmt.ID())), 1024)), "fail to get chunk for test")
	// after getting a failed FETCH, the cursor should have been reseted
	require.False(t, stmt.GetCursorActive())
	require.Len(t, tk.Session().GetSessionVars().MemTracker.GetChildrenForTest(), 0)
	require.Len(t, tk.Session().GetSessionVars().DiskTracker.GetChildrenForTest(), 0)
}

func TestCursorFetchShouldSpill(t *testing.T) {
	restore := config.RestoreFunc()