dxf, importinto: compare subtask task IDs as strings

.agents/skills/tidb-test-guidelines/references/dxf-case-map.md

@@ -99,7 +99,7 @@
 - `pkg/dxf/importinto/collect_conflicts_test.go` - dxf/importinto: Tests collect conflicts step executor.
 - `pkg/dxf/importinto/conflict_resolution_test.go` - dxf/importinto: Tests conflict resolution step executor.
 - `pkg/dxf/importinto/encode_and_sort_operator_test.go` - dxf/importinto: Tests encode and sort operator.
-- `pkg/dxf/importinto/job_testkit_test.go` - dxf/importinto: Tests submit task next-gen.
+- `pkg/dxf/importinto/job_testkit_test.go` - dxf/importinto: Tests task submission and exact GetJobLastUpdateTime lookups across active/history subtasks.
 - `pkg/dxf/importinto/metrics_test.go` - dxf/importinto: Tests metric manager.
 - `pkg/dxf/importinto/planner_test.go` - dxf/importinto: Tests logical plan.
 - `pkg/dxf/importinto/proto_test.go` - dxf/importinto: Tests KV conflict info aggregation.
@@ -109,6 +109,11 @@
 - `pkg/dxf/importinto/task_executor_testkit_test.go` - dxf/importinto: Tests post process step executor.
 - `pkg/dxf/importinto/wrapper_test.go` - dxf/importinto: Tests chunk convert.
 
+## pkg/dxf/importinto/jobhistory
+
+### Tests
+- `pkg/dxf/importinto/jobhistory/history_test.go` - dxf/importinto/jobhistory: Tests target-only history aggregation for adjacent task IDs around 2^53.
+
 ## pkg/dxf/importinto/conflictedkv
 
 ### Tests

pkg/dxf/framework/integrationtests/framework_err_handling_test.go

@@ -83,7 +83,7 @@ func TestOnTaskError(t *testing.T) {
 		taskKey := "key3-1"
 		taskID := prepareForAwaitingResolutionTestFn(t, taskKey)
 		tk := testkit.NewTestKit(t, c.Store)
-		tk.MustExec(fmt.Sprintf("update mysql.tidb_background_subtask set state='pending' where state='failed' and task_key= %d", taskID))
+		tk.MustExec(fmt.Sprintf("update mysql.tidb_background_subtask set state='pending' where state='failed' and task_key= '%d'", taskID))
 		tk.MustExec(fmt.Sprintf("update mysql.tidb_global_task set state='running' where id = %d", taskID))
 		task := testutil.WaitTaskDone(c.Ctx, t, taskKey)
 		require.Equal(t, proto.TaskStateSucceed, task.State)

pkg/dxf/framework/integrationtests/modify_test.go

@@ -24,6 +24,7 @@ import (
 	"github.com/pingcap/tidb/pkg/dxf/framework/handle"
 	mockexecute "github.com/pingcap/tidb/pkg/dxf/framework/mock/execute"
 	"github.com/pingcap/tidb/pkg/dxf/framework/proto"
+	"github.com/pingcap/tidb/pkg/dxf/framework/storage"
 	"github.com/pingcap/tidb/pkg/dxf/framework/taskexecutor/execute"
 	"github.com/pingcap/tidb/pkg/dxf/framework/testutil"
 	"github.com/pingcap/tidb/pkg/sessionctx"
@@ -513,7 +514,7 @@ func TestModifyTaskMaxNodeCountForSubtaskBalance(t *testing.T) {
 			return runtimeInfo.activeSubtaskCount.Load() == 1
 		}, 10*time.Second, 100*time.Millisecond)
 		require.NoError(t, c.TaskMgr.WithNewSession(func(se sessionctx.Context) error {
-			rows, err2 := sqlexec.ExecSQL(c.Ctx, se.GetSQLExecutor(), "select count(distinct exec_id) from mysql.tidb_background_subtask where task_key = %?", task.ID)
+			rows, err2 := sqlexec.ExecSQL(c.Ctx, se.GetSQLExecutor(), "select count(distinct exec_id) from mysql.tidb_background_subtask where task_key = %?", storage.TaskIDToKey(task.ID))
 			require.NoError(t, err2)
 			require.Equal(t, 1, len(rows))
 			require.EqualValues(t, 1, rows[0].GetInt64(0))
@@ -533,7 +534,7 @@ func TestModifyTaskMaxNodeCountForSubtaskBalance(t *testing.T) {
 			return runtimeInfo.activeSubtaskCount.Load() == 2
 		}, 10*time.Second, 100*time.Millisecond)
 		require.NoError(t, c.TaskMgr.WithNewSession(func(se sessionctx.Context) error {
-			rows, err2 := sqlexec.ExecSQL(c.Ctx, se.GetSQLExecutor(), "select count(distinct exec_id) from mysql.tidb_background_subtask where task_key = %?", task.ID)
+			rows, err2 := sqlexec.ExecSQL(c.Ctx, se.GetSQLExecutor(), "select count(distinct exec_id) from mysql.tidb_background_subtask where task_key = %?", storage.TaskIDToKey(task.ID))
 			require.NoError(t, err2)
 			require.Equal(t, 1, len(rows))
 			require.EqualValues(t, 2, rows[0].GetInt64(0))

pkg/dxf/framework/storage/converter.go

@@ -90,6 +90,14 @@ func Row2Task(r chunk.Row) *proto.Task {
 	return task
 }
 
+// TaskIDToKey returns the canonical decimal key stored in background-subtask tables.
+// due to history reason, the task_key column inside subtask tables are defined
+// as varchar, but it actually stores the ID of the task which is integer. to
+// make sure TiDB query can correct handle this, we need to convert it to string.
+func TaskIDToKey(taskID int64) string {
+	return strconv.FormatInt(taskID, 10)
+}
+
 // row2BasicSubTask converts a row to a subtask with basic info
 func row2BasicSubTask(r chunk.Row) *proto.SubtaskBase {
 	taskIDStr := r.GetString(2)

pkg/dxf/framework/storage/history.go

@@ -41,12 +41,12 @@ const (
 // TransferSubtasks2HistoryWithSession transfer the selected subtasks into tidb_background_subtask_history table by taskID.
 func (*TaskManager) TransferSubtasks2HistoryWithSession(ctx context.Context, se sessionctx.Context, taskID int64) error {
 	exec := se.GetSQLExecutor()
-	_, err := sqlexec.ExecSQL(ctx, exec, `insert into mysql.tidb_background_subtask_history select * from mysql.tidb_background_subtask where task_key = %?`, taskID)
+	_, err := sqlexec.ExecSQL(ctx, exec, `insert into mysql.tidb_background_subtask_history select * from mysql.tidb_background_subtask where task_key = %?`, TaskIDToKey(taskID))
 	if err != nil {
 		return err
 	}
 	// delete taskID subtask
-	_, err = sqlexec.ExecSQL(ctx, exec, "delete from mysql.tidb_background_subtask where task_key = %?", taskID)
+	_, err = sqlexec.ExecSQL(ctx, exec, "delete from mysql.tidb_background_subtask where task_key = %?", TaskIDToKey(taskID))
 	return err
 }
 

pkg/dxf/framework/storage/subtask_state.go

@@ -80,7 +80,7 @@ func (mgr *TaskManager) FailSubtask(ctx context.Context, execID string, taskID i
 		proto.SubtaskStateFailed,
 		serializeErr(err),
 		execID,
-		taskID,
+		TaskIDToKey(taskID),
 		proto.SubtaskStatePending,
 		proto.SubtaskStateRunning)
 	return err1
@@ -102,7 +102,7 @@ func (mgr *TaskManager) CancelSubtask(ctx context.Context, execID string, taskID
 		state in (%?, %?);`,
 		proto.SubtaskStateCanceled,
 		execID,
-		taskID,
+		TaskIDToKey(taskID),
 		proto.SubtaskStatePending,
 		proto.SubtaskStateRunning)
 	return err1
@@ -114,7 +114,7 @@ func (mgr *TaskManager) PauseSubtasks(ctx context.Context, execID string, taskID
 		return err
 	}
 	_, err := mgr.ExecuteSQLWithNewSession(ctx,
-		`update mysql.tidb_background_subtask set state = "paused" where task_key = %? and state in ("running", "pending") and exec_id = %?`, taskID, execID)
+		`update mysql.tidb_background_subtask set state = "paused" where task_key = %? and state in ("running", "pending") and exec_id = %?`, TaskIDToKey(taskID), execID)
 	return err
 }
 
@@ -124,7 +124,7 @@ func (mgr *TaskManager) ResumeSubtasks(ctx context.Context, taskID int64) error
 		return err
 	}
 	_, err := mgr.ExecuteSQLWithNewSession(ctx,
-		`update mysql.tidb_background_subtask set state = "pending", error = null where task_key = %? and state = "paused"`, taskID)
+		`update mysql.tidb_background_subtask set state = "pending", error = null where task_key = %? and state = "paused"`, TaskIDToKey(taskID))
 	return err
 }
 

pkg/dxf/framework/storage/table_test.go

@@ -329,7 +329,7 @@ func TestSwitchTaskStep(t *testing.T) {
 		require.Equal(t, proto.ExtraParams{}, persistedTask.ExtraParams)
 		require.Zero(t, persistedTask.StartTime)
 		require.GreaterOrEqual(t, persistedTask.StateUpdateTime, switchTime)
-		tk.MustQuery(fmt.Sprintf("select count(1) from mysql.tidb_background_subtask where task_key = %d", prepareTaskID)).
+		tk.MustQuery(fmt.Sprintf("select count(1) from mysql.tidb_background_subtask where task_key = '%d'", prepareTaskID)).
 			Check(testkit.Rows("0"))
 
 		prepareTask.Meta = []byte(`{"prepare":"stale-owner"}`)
@@ -376,7 +376,7 @@ func TestGetSubtaskSummaries(t *testing.T) {
 	bytes, err := json.Marshal(summary)
 	require.NoError(t, err)
 
-	tk.MustExec(fmt.Sprintf("update mysql.tidb_background_subtask set summary = '%s' where task_key = %d", string(bytes), task.ID))
+	tk.MustExec(fmt.Sprintf("update mysql.tidb_background_subtask set summary = '%s' where task_key = '%d'", string(bytes), task.ID))
 	summaries, err := tm.GetAllSubtaskSummaryByStep(ctx, subtasks[0].TaskID, proto.StepOne)
 	require.NoError(t, err)
 	require.Len(t, summaries, len(subtasks))
@@ -388,7 +388,7 @@ func TestGetSubtaskSummaries(t *testing.T) {
 	// If the JSON value is wrong, we still get an empty summary.
 	// This can only happen if the summary field is manually updated.
 	// It's acceptable, as the correct summary will be set by the executor later.
-	tk.MustExec(fmt.Sprintf(`update mysql.tidb_background_subtask set summary = '{"wrong_key": 123}' where task_key = %d`, task.ID))
+	tk.MustExec(fmt.Sprintf(`update mysql.tidb_background_subtask set summary = '{"wrong_key": 123}' where task_key = '%d'`, task.ID))
 	summaries, err = tm.GetAllSubtaskSummaryByStep(ctx, subtasks[0].TaskID, proto.StepOne)
 	require.NoError(t, err)
 	for _, summary := range summaries {
@@ -613,6 +613,26 @@ func TestGetActiveSubtasks(t *testing.T) {
 	require.Equal(t, proto.SubtaskStateRunning, activeSubtasks[0].State)
 	require.Equal(t, int64(3), activeSubtasks[1].ID)
 	require.Equal(t, proto.SubtaskStatePending, activeSubtasks[1].State)
+
+	const (
+		adjacentTaskID = int64(9007199254740992)
+		targetTaskID   = int64(9007199254740993)
+		largeIDExec    = "large-id-exec"
+	)
+	testutil.InsertSubtask(t, tm, adjacentTaskID, proto.StepOne, largeIDExec, nil, proto.SubtaskStatePending, proto.TaskTypeExample, 8)
+	testutil.InsertSubtask(t, tm, targetTaskID, proto.StepOne, largeIDExec, nil, proto.SubtaskStatePending, proto.TaskTypeExample, 8)
+
+	activeSubtasks, err = tm.GetActiveSubtasks(ctx, targetTaskID)
+	require.NoError(t, err)
+	require.Len(t, activeSubtasks, 1)
+	require.Equal(t, targetTaskID, activeSubtasks[0].TaskID)
+
+	filteredSubtasks, err := tm.GetSubtasksByExecIDAndStepAndStates(
+		ctx, largeIDExec, targetTaskID, proto.StepOne, proto.SubtaskStatePending,
+	)
+	require.NoError(t, err)
+	require.Len(t, filteredSubtasks, 1)
+	require.Equal(t, targetTaskID, filteredSubtasks[0].TaskID)
 }
 
 func TestSubTaskTable(t *testing.T) {

pkg/dxf/framework/storage/task_state.go

@@ -118,7 +118,7 @@ func (mgr *TaskManager) PauseTaskOnError(ctx context.Context, taskID int64, task
 				state_update_time = unix_timestamp(),
 				end_time = null
 			where task_key = %? and step = %? and state = %?`,
-			proto.SubtaskStatePaused, taskID, step, proto.SubtaskStateFailed,
+			proto.SubtaskStatePaused, TaskIDToKey(taskID), step, proto.SubtaskStateFailed,
 		)
 		return err
 	})
@@ -302,7 +302,7 @@ func (mgr *TaskManager) ModifiedTask(ctx context.Context, task *proto.Task) erro
 			update mysql.tidb_background_subtask
 			set concurrency = %?, state_update_time = unix_timestamp()
 			where task_key = %? and state in (%?, %?, %?)`,
-			task.RequiredSlots, task.ID,
+			task.RequiredSlots, TaskIDToKey(task.ID),
 			proto.SubtaskStatePending, proto.SubtaskStateRunning, proto.SubtaskStatePaused)
 		return err
 	})

pkg/dxf/framework/storage/task_table.go

@@ -593,7 +593,7 @@ func (mgr *TaskManager) GetSubtasksByExecIDAndStepAndStates(ctx context.Context,
 	if err := injectfailpoint.DXFRandomErrorWithOnePercent(); err != nil {
 		return nil, err
 	}
-	args := []any{execID, taskID, step}
+	args := []any{execID, TaskIDToKey(taskID), step}
 	for _, state := range states {
 		args = append(args, state)
 	}
@@ -616,7 +616,7 @@ func (mgr *TaskManager) GetFirstSubtaskInStates(ctx context.Context, tidbID stri
 	if err := injectfailpoint.DXFRandomErrorWithOnePercent(); err != nil {
 		return nil, err
 	}
-	args := []any{tidbID, taskID, step}
+	args := []any{tidbID, TaskIDToKey(taskID), step}
 	for _, state := range states {
 		args = append(args, state)
 	}
@@ -641,7 +641,7 @@ func (mgr *TaskManager) GetActiveSubtasks(ctx context.Context, taskID int64) ([]
 	rs, err := mgr.ExecuteSQLWithNewSession(ctx, `
 		select `+basicSubtaskColumns+` from mysql.tidb_background_subtask
 		where task_key = %? and state in (%?, %?)`,
-		taskID, proto.SubtaskStatePending, proto.SubtaskStateRunning)
+		TaskIDToKey(taskID), proto.SubtaskStatePending, proto.SubtaskStateRunning)
 	if err != nil {
 		return nil, err
 	}
@@ -659,7 +659,7 @@ func (mgr *TaskManager) GetAllSubtasksByStepAndState(ctx context.Context, taskID
 	}
 	rs, err := mgr.ExecuteSQLWithNewSession(ctx, `select `+SubtaskColumns+` from mysql.tidb_background_subtask
 		where task_key = %? and state = %? and step = %?`,
-		taskID, state, step)
+		TaskIDToKey(taskID), state, step)
 	if err != nil {
 		return nil, err
 	}
@@ -684,7 +684,7 @@ func (mgr *TaskManager) GetAllSubtaskSummaryByStep(
 	rs, err := mgr.ExecuteSQLWithNewSession(ctx,
 		`select summary from mysql.tidb_background_subtask
 		where task_key = %? and step = %?`,
-		taskID, step)
+		TaskIDToKey(taskID), step)
 	if err != nil {
 		return nil, err
 	}
@@ -714,7 +714,7 @@ func (mgr *TaskManager) GetSubtaskRowCount(ctx context.Context, taskID int64, st
 			union all
 			select summary from mysql.tidb_background_subtask_history where task_key = %? and step = %?
 		) as combined`,
-		taskID, step, taskID, step)
+		TaskIDToKey(taskID), step, TaskIDToKey(taskID), step)
 	if err != nil {
 		return 0, err
 	}
@@ -750,7 +750,7 @@ func (mgr *TaskManager) GetSubtaskCntGroupByStates(ctx context.Context, taskID i
 		from mysql.tidb_background_subtask
 		where task_key = %? and step = %?
 		group by state`,
-		taskID, step)
+		TaskIDToKey(taskID), step)
 	if err != nil {
 		return nil, err
 	}
@@ -773,7 +773,7 @@ func (mgr *TaskManager) GetSubtaskStateCntAndErrorsByStep(ctx context.Context, t
 		`select state, error
 			from mysql.tidb_background_subtask
 			where task_key = %? and step = %?`,
-		taskID, step)
+		TaskIDToKey(taskID), step)
 	if err != nil {
 		return nil, nil, err
 	}
@@ -806,7 +806,7 @@ func (mgr *TaskManager) GetSubtaskErrors(ctx context.Context, taskID int64) ([]e
 	}
 	rs, err := mgr.ExecuteSQLWithNewSession(ctx,
 		`select error from mysql.tidb_background_subtask
-             where task_key = %? AND state in (%?, %?)`, taskID, proto.SubtaskStateFailed, proto.SubtaskStateCanceled)
+             where task_key = %? AND state in (%?, %?)`, TaskIDToKey(taskID), proto.SubtaskStateFailed, proto.SubtaskStateCanceled)
 	if err != nil {
 		return nil, err
 	}
@@ -991,7 +991,7 @@ func (mgr *TaskManager) SwitchTaskStepInBatch(
 		// some subtasks may be inserted by other schedulers, we can skip them.
 		rs, err := sqlexec.ExecSQL(ctx, se.GetSQLExecutor(), `
 			select count(1) from mysql.tidb_background_subtask
-			where task_key = %? and step = %?`, task.ID, nextStep)
+			where task_key = %? and step = %?`, TaskIDToKey(task.ID), nextStep)
 		if err != nil {
 			return err
 		}
@@ -1063,7 +1063,7 @@ func (mgr *TaskManager) GetSubtasksWithHistory(ctx context.Context, taskID int64
 	err = mgr.WithNewTxn(ctx, func(se sessionctx.Context) error {
 		rs, err = sqlexec.ExecSQL(ctx, se.GetSQLExecutor(),
 			`select `+SubtaskColumns+` from mysql.tidb_background_subtask where task_key = %? and step = %?`,
-			taskID, step,
+			TaskIDToKey(taskID), step,
 		)
 		if err != nil {
 			return err
@@ -1073,7 +1073,7 @@ func (mgr *TaskManager) GetSubtasksWithHistory(ctx context.Context, taskID int64
 		// when the user show import jobs, we need to check the history table.
 		rsFromHistory, err := sqlexec.ExecSQL(ctx, se.GetSQLExecutor(),
 			`select `+SubtaskColumns+` from mysql.tidb_background_subtask_history where task_key = %? and step = %?`,
-			taskID, step,
+			TaskIDToKey(taskID), step,
 		)
 		if err != nil {
 			return err

pkg/dxf/framework/testutil/table_util.go

@@ -99,7 +99,7 @@ func GetSubtasksFromHistory(ctx context.Context, mgr *storage.TaskManager) (int,
 // GetSubtasksFromHistoryByTaskID gets subtasks by taskID from history table for test.
 func GetSubtasksFromHistoryByTaskID(ctx context.Context, mgr *storage.TaskManager, taskID int64) (int, error) {
 	rs, err := mgr.ExecuteSQLWithNewSession(ctx,
-		`select `+storage.SubtaskColumns+` from mysql.tidb_background_subtask_history where task_key = %?`, taskID)
+		`select `+storage.SubtaskColumns+` from mysql.tidb_background_subtask_history where task_key = %?`, storage.TaskIDToKey(taskID))
 	if err != nil {
 		return 0, err
 	}
@@ -109,7 +109,7 @@ func GetSubtasksFromHistoryByTaskID(ctx context.Context, mgr *storage.TaskManage
 // GetSubtasksByTaskID gets subtasks by taskID for test.
 func GetSubtasksByTaskID(ctx context.Context, mgr *storage.TaskManager, taskID int64) ([]*proto.Subtask, error) {
 	rs, err := mgr.ExecuteSQLWithNewSession(ctx,
-		`select `+storage.SubtaskColumns+` from mysql.tidb_background_subtask where task_key = %?`, taskID)
+		`select `+storage.SubtaskColumns+` from mysql.tidb_background_subtask where task_key = %?`, storage.TaskIDToKey(taskID))
 	if err != nil {
 		return nil, err
 	}
@@ -170,7 +170,7 @@ func GetSubtaskNodes(ctx context.Context, mgr *storage.TaskManager, taskID int64
 	rs, err := mgr.ExecuteSQLWithNewSession(ctx, `
 		select distinct(exec_id) from mysql.tidb_background_subtask where task_key=%?
 		union
-		select distinct(exec_id) from mysql.tidb_background_subtask_history where task_key=%?`, taskID, taskID)
+		select distinct(exec_id) from mysql.tidb_background_subtask_history where task_key=%?`, storage.TaskIDToKey(taskID), storage.TaskIDToKey(taskID))
 	if err != nil {
 		return nil, err
 	}
@@ -220,7 +220,7 @@ func GetTasksFromHistoryInStates(ctx context.Context, mgr *storage.TaskManager,
 // DeleteSubtasksByTaskID deletes the subtask of the given task ID.
 func DeleteSubtasksByTaskID(ctx context.Context, mgr *storage.TaskManager, taskID int64) error {
 	_, err := mgr.ExecuteSQLWithNewSession(ctx, `delete from mysql.tidb_background_subtask
-		where task_key = %?`, taskID)
+		where task_key = %?`, storage.TaskIDToKey(taskID))
 	if err != nil {
 		return err
 	}
@@ -244,9 +244,9 @@ func IsTaskCancelling(ctx context.Context, mgr *storage.TaskManager, taskID int6
 // PrintSubtaskInfo log the subtask info by taskKey for test.
 func PrintSubtaskInfo(ctx context.Context, mgr *storage.TaskManager, taskID int64) {
 	rs, _ := mgr.ExecuteSQLWithNewSession(ctx,
-		`select `+storage.SubtaskColumns+` from mysql.tidb_background_subtask_history where task_key = %?`, taskID)
+		`select `+storage.SubtaskColumns+` from mysql.tidb_background_subtask_history where task_key = %?`, storage.TaskIDToKey(taskID))
 	rs2, _ := mgr.ExecuteSQLWithNewSession(ctx,
-		`select `+storage.SubtaskColumns+` from mysql.tidb_background_subtask where task_key = %?`, taskID)
+		`select `+storage.SubtaskColumns+` from mysql.tidb_background_subtask where task_key = %?`, storage.TaskIDToKey(taskID))
 	rs = append(rs, rs2...)
 
 	for _, r := range rs {

pkg/dxf/importinto/BUILD.bazel

@@ -112,7 +112,7 @@ go_test(
     ],
     embed = [":importinto"],
     flaky = True,
-    shard_count = 29,
+    shard_count = 30,
     deps = [
         "//pkg/config",
         "//pkg/config/kerneltype",

pkg/dxf/importinto/job.go

@@ -404,7 +404,7 @@ func GetJobLastUpdateTime(ctx context.Context, jobID int64) (types.Time, error)
 					union
 				select state_update_time from mysql.tidb_background_subtask_history where task_key = %?
 			) t`,
-			task.ID, task.ID,
+			storage.TaskIDToKey(task.ID), storage.TaskIDToKey(task.ID),
 		)
 		return err
 	})