346 lines
9.7 KiB
Go
346 lines
9.7 KiB
Go
// Copyright 2023 PingCAP, Inc.
|
|
//
|
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
|
// you may not use this file except in compliance with the License.
|
|
// You may obtain a copy of the License at
|
|
//
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
|
//
|
|
// Unless required by applicable law or agreed to in writing, software
|
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
// See the License for the specific language governing permissions and
|
|
// limitations under the License.
|
|
|
|
package asyncloaddata_test
|
|
|
|
import (
|
|
"context"
|
|
"testing"
|
|
"time"
|
|
|
|
. "github.com/pingcap/tidb/executor/asyncloaddata"
|
|
"github.com/pingcap/tidb/testkit"
|
|
"github.com/pingcap/tidb/util/sqlexec"
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
func checkEqualIgnoreTimes(t *testing.T, expected, got *JobInfo) {
|
|
cloned := *expected
|
|
cloned.CreateTime = got.CreateTime
|
|
cloned.StartTime = got.StartTime
|
|
cloned.EndTime = got.EndTime
|
|
require.Equal(t, &cloned, got)
|
|
}
|
|
|
|
func createJob(t *testing.T, conn sqlexec.SQLExecutor, user string) (*Job, *JobInfo) {
|
|
job, err := CreateLoadDataJob(context.Background(), conn, "/tmp/test.csv", "test", "t", "logical", user)
|
|
require.NoError(t, err)
|
|
info, err := job.GetJobInfo(context.Background())
|
|
require.NoError(t, err)
|
|
expected := &JobInfo{
|
|
JobID: job.ID,
|
|
User: user,
|
|
DataSource: "/tmp/test.csv",
|
|
TableSchema: "test",
|
|
TableName: "t",
|
|
ImportMode: "logical",
|
|
Progress: "",
|
|
Status: JobPending,
|
|
StatusMessage: "",
|
|
}
|
|
checkEqualIgnoreTimes(t, expected, info)
|
|
return job, info
|
|
}
|
|
|
|
func TestHappyPath(t *testing.T) {
|
|
store := testkit.CreateMockStore(t)
|
|
tk := testkit.NewTestKit(t, store)
|
|
ctx := context.Background()
|
|
|
|
// job is created
|
|
|
|
job, expected := createJob(t, tk.Session(), "user")
|
|
|
|
// job is started by a worker
|
|
|
|
backup := OfflineThresholdInSec
|
|
OfflineThresholdInSec = 1000
|
|
t.Cleanup(func() {
|
|
OfflineThresholdInSec = backup
|
|
})
|
|
err := job.StartJob(ctx)
|
|
require.NoError(t, err)
|
|
info, err := job.GetJobInfo(ctx)
|
|
require.NoError(t, err)
|
|
expected.Status = JobRunning
|
|
checkEqualIgnoreTimes(t, expected, info)
|
|
|
|
// job is periodically updated by worker
|
|
|
|
ok, err := job.UpdateJobProgress(ctx, "imported 10%")
|
|
require.NoError(t, err)
|
|
require.True(t, ok)
|
|
info, err = job.GetJobInfo(ctx)
|
|
require.NoError(t, err)
|
|
expected.Progress = "imported 10%"
|
|
checkEqualIgnoreTimes(t, expected, info)
|
|
|
|
// job is paused
|
|
|
|
err = UpdateJobExpectedStatus(ctx, tk.Session(), job.ID, JobExpectedPaused)
|
|
require.NoError(t, err)
|
|
info, err = job.GetJobInfo(ctx)
|
|
require.NoError(t, err)
|
|
expected.Status = JobPaused
|
|
checkEqualIgnoreTimes(t, expected, info)
|
|
|
|
// worker still can update progress, maybe response to pausing is delayed
|
|
|
|
ok, err = job.UpdateJobProgress(ctx, "imported 20%")
|
|
require.NoError(t, err)
|
|
require.True(t, ok)
|
|
info, err = job.GetJobInfo(ctx)
|
|
require.NoError(t, err)
|
|
expected.Progress = "imported 20%"
|
|
checkEqualIgnoreTimes(t, expected, info)
|
|
|
|
// job is resumed
|
|
|
|
err = UpdateJobExpectedStatus(ctx, tk.Session(), job.ID, JobExpectedRunning)
|
|
require.NoError(t, err)
|
|
info, err = job.GetJobInfo(ctx)
|
|
require.NoError(t, err)
|
|
expected.Status = JobRunning
|
|
checkEqualIgnoreTimes(t, expected, info)
|
|
|
|
// job is finished
|
|
|
|
err = job.FinishJob(ctx, "finished message")
|
|
require.NoError(t, err)
|
|
info, err = job.GetJobInfo(ctx)
|
|
require.NoError(t, err)
|
|
expected.Status = JobFinished
|
|
expected.StatusMessage = "finished message"
|
|
checkEqualIgnoreTimes(t, expected, info)
|
|
}
|
|
|
|
func TestKeepAlive(t *testing.T) {
|
|
store := testkit.CreateMockStore(t)
|
|
tk := testkit.NewTestKit(t, store)
|
|
ctx := context.Background()
|
|
|
|
// job is created
|
|
|
|
job, expected := createJob(t, tk.Session(), "user")
|
|
|
|
backup := OfflineThresholdInSec
|
|
OfflineThresholdInSec = 1
|
|
t.Cleanup(func() {
|
|
OfflineThresholdInSec = backup
|
|
})
|
|
|
|
// before job is started, worker don't need to keepalive
|
|
// TODO:👆not correct!
|
|
|
|
time.Sleep(2 * time.Second)
|
|
info, err := job.GetJobInfo(ctx)
|
|
require.NoError(t, err)
|
|
checkEqualIgnoreTimes(t, expected, info)
|
|
|
|
err = job.StartJob(ctx)
|
|
require.NoError(t, err)
|
|
info, err = job.GetJobInfo(ctx)
|
|
require.NoError(t, err)
|
|
expected.Status = JobRunning
|
|
checkEqualIgnoreTimes(t, expected, info)
|
|
|
|
// if worker failed to keepalive, job will fail
|
|
|
|
time.Sleep(2 * time.Second)
|
|
info, err = job.GetJobInfo(ctx)
|
|
require.NoError(t, err)
|
|
expected.Status = JobFailed
|
|
expected.StatusMessage = "job expected running but the node is timeout"
|
|
checkEqualIgnoreTimes(t, expected, info)
|
|
|
|
// after the worker is failed to keepalive, further keepalive will fail
|
|
|
|
ok, err := job.UpdateJobProgress(ctx, "imported 20%")
|
|
require.NoError(t, err)
|
|
require.False(t, ok)
|
|
info, err = job.GetJobInfo(ctx)
|
|
require.NoError(t, err)
|
|
checkEqualIgnoreTimes(t, expected, info)
|
|
|
|
// when worker fails to keepalive, before it calls FailJob, it still can
|
|
// change expected status to some extent.
|
|
|
|
err = UpdateJobExpectedStatus(ctx, tk.Session(), job.ID, JobExpectedPaused)
|
|
require.NoError(t, err)
|
|
info, err = job.GetJobInfo(ctx)
|
|
require.NoError(t, err)
|
|
expected.StatusMessage = "job expected paused but the node is timeout"
|
|
checkEqualIgnoreTimes(t, expected, info)
|
|
err = UpdateJobExpectedStatus(ctx, tk.Session(), job.ID, JobExpectedRunning)
|
|
require.NoError(t, err)
|
|
info, err = job.GetJobInfo(ctx)
|
|
require.NoError(t, err)
|
|
expected.StatusMessage = "job expected running but the node is timeout"
|
|
checkEqualIgnoreTimes(t, expected, info)
|
|
err = UpdateJobExpectedStatus(ctx, tk.Session(), job.ID, JobExpectedCanceled)
|
|
require.NoError(t, err)
|
|
info, err = job.GetJobInfo(ctx)
|
|
require.NoError(t, err)
|
|
expected.Status = JobCanceled
|
|
expected.StatusMessage = ""
|
|
checkEqualIgnoreTimes(t, expected, info)
|
|
|
|
// Now the worker calls FailJob, but the status should still be canceled,
|
|
// that's more friendly.
|
|
|
|
err = job.FailJob(ctx, "failed to keepalive")
|
|
require.NoError(t, err)
|
|
info, err = job.GetJobInfo(ctx)
|
|
require.NoError(t, err)
|
|
expected.Status = JobCanceled
|
|
expected.StatusMessage = "failed to keepalive"
|
|
checkEqualIgnoreTimes(t, expected, info)
|
|
}
|
|
|
|
func TestJobIsFailedAndGetAllJobs(t *testing.T) {
|
|
store := testkit.CreateMockStore(t)
|
|
tk := testkit.NewTestKit(t, store)
|
|
ctx := context.Background()
|
|
|
|
// job is created
|
|
|
|
job, expected := createJob(t, tk.Session(), "user")
|
|
|
|
// job can be failed directly when it's pending
|
|
|
|
err := job.FailJob(ctx, "failed message")
|
|
require.NoError(t, err)
|
|
info, err := job.GetJobInfo(ctx)
|
|
require.NoError(t, err)
|
|
expected.Status = JobFailed
|
|
expected.StatusMessage = "failed message"
|
|
checkEqualIgnoreTimes(t, expected, info)
|
|
|
|
// create another job and fail it
|
|
|
|
job, expected = createJob(t, tk.Session(), "user")
|
|
|
|
err = job.StartJob(ctx)
|
|
require.NoError(t, err)
|
|
info, err = job.GetJobInfo(ctx)
|
|
require.NoError(t, err)
|
|
expected.Status = JobRunning
|
|
checkEqualIgnoreTimes(t, expected, info)
|
|
|
|
err = job.FailJob(ctx, "failed message")
|
|
require.NoError(t, err)
|
|
info, err = job.GetJobInfo(ctx)
|
|
require.NoError(t, err)
|
|
expected.Status = JobFailed
|
|
expected.StatusMessage = "failed message"
|
|
checkEqualIgnoreTimes(t, expected, info)
|
|
|
|
// test change expected status of a failed job.
|
|
|
|
err = UpdateJobExpectedStatus(ctx, tk.Session(), job.ID, JobExpectedPaused)
|
|
require.NoError(t, err)
|
|
info, err = job.GetJobInfo(ctx)
|
|
require.NoError(t, err)
|
|
checkEqualIgnoreTimes(t, expected, info)
|
|
err = UpdateJobExpectedStatus(ctx, tk.Session(), job.ID, JobExpectedRunning)
|
|
require.NoError(t, err)
|
|
info, err = job.GetJobInfo(ctx)
|
|
require.NoError(t, err)
|
|
checkEqualIgnoreTimes(t, expected, info)
|
|
err = job.CancelJob(ctx)
|
|
require.ErrorContains(t, err, "The current job status cannot perform the operation. need status running or paused, but got failed")
|
|
|
|
// add job of another user and test GetAllJobInfo
|
|
|
|
job, _ = createJob(t, tk.Session(), "user2")
|
|
|
|
jobs, err := GetAllJobInfo(ctx, tk.Session(), "user")
|
|
require.NoError(t, err)
|
|
require.Equal(t, 2, len(jobs))
|
|
require.Equal(t, JobFailed, jobs[0].Status)
|
|
require.Equal(t, JobFailed, jobs[1].Status)
|
|
|
|
jobs, err = GetAllJobInfo(ctx, tk.Session(), "user2")
|
|
require.NoError(t, err)
|
|
require.Equal(t, 1, len(jobs))
|
|
require.Equal(t, JobPending, jobs[0].Status)
|
|
require.Equal(t, job.ID, jobs[0].JobID)
|
|
|
|
job.User = "wrong_user"
|
|
_, err = job.GetJobInfo(ctx)
|
|
require.ErrorContains(t, err, "doesn't exist")
|
|
}
|
|
|
|
func TestGetJobStatus(t *testing.T) {
|
|
store := testkit.CreateMockStore(t)
|
|
tk := testkit.NewTestKit(t, store)
|
|
ctx := context.Background()
|
|
|
|
// job is created
|
|
|
|
job, _ := createJob(t, tk.Session(), "user")
|
|
|
|
// job is pending
|
|
|
|
status, msg, err := job.GetJobStatus(ctx)
|
|
require.NoError(t, err)
|
|
require.Equal(t, JobPending, status)
|
|
require.Equal(t, "", msg)
|
|
|
|
// job is running
|
|
|
|
backup := OfflineThresholdInSec
|
|
OfflineThresholdInSec = 1000
|
|
t.Cleanup(func() {
|
|
OfflineThresholdInSec = backup
|
|
})
|
|
err = job.StartJob(ctx)
|
|
require.NoError(t, err)
|
|
status, msg, err = job.GetJobStatus(ctx)
|
|
require.NoError(t, err)
|
|
require.Equal(t, JobRunning, status)
|
|
require.Equal(t, "", msg)
|
|
|
|
// job is finished
|
|
|
|
err = job.FinishJob(ctx, "finished message")
|
|
require.NoError(t, err)
|
|
status, msg, err = job.GetJobStatus(ctx)
|
|
require.NoError(t, err)
|
|
require.Equal(t, JobFinished, status)
|
|
require.Equal(t, "finished message", msg)
|
|
|
|
// wrong ID
|
|
|
|
job.ID += 1
|
|
status, msg, err = job.GetJobStatus(ctx)
|
|
require.NoError(t, err)
|
|
require.Equal(t, JobFailed, status)
|
|
require.Contains(t, msg, "doesn't exist")
|
|
}
|
|
|
|
func TestCreateLoadDataJobRedact(t *testing.T) {
|
|
store := testkit.CreateMockStore(t)
|
|
tk := testkit.NewTestKit(t, store)
|
|
ctx := context.Background()
|
|
|
|
_, err := CreateLoadDataJob(ctx, tk.Session(),
|
|
"s3://bucket/a.csv?access-key=hideme&secret-access-key=hideme",
|
|
"db", "table", "mode", "user")
|
|
require.NoError(t, err)
|
|
result := tk.MustQuery("SELECT * FROM mysql.load_data_jobs;")
|
|
result.CheckContain("a.csv")
|
|
result.CheckNotContain("hideme")
|
|
}
|