sqlreplay, api: add privilege argument to the cancel traffic api (#741)

Author: djshow832

pkg/server/api/traffic.go

@@ -11,6 +11,7 @@ import (
 
 	"github.com/gin-gonic/gin"
 	"github.com/pingcap/tiproxy/pkg/sqlreplay/capture"
+	"github.com/pingcap/tiproxy/pkg/sqlreplay/manager"
 	"github.com/pingcap/tiproxy/pkg/sqlreplay/replay"
 	"go.uber.org/zap"
 )
@@ -115,7 +116,17 @@ func (h *Server) TrafficCancel(c *gin.Context) {
 		return
 	}
 
-	result := h.mgr.ReplayJobMgr.Stop()
+	cfg := manager.CancelConfig{}
+	cfg.Type = manager.Capture | manager.Replay
+	if tp := c.PostForm("type"); tp != "" {
+		switch strings.ToLower(tp) {
+		case "capture":
+			cfg.Type = manager.Capture
+		case "replay":
+			cfg.Type = manager.Replay
+		}
+	}
+	result := h.mgr.ReplayJobMgr.Stop(cfg)
 	c.String(http.StatusOK, result)
 }
 

pkg/server/api/traffic_test.go

@@ -201,7 +201,7 @@ func (m *mockReplayJobManager) StartReplay(replayCfg replay.ReplayConfig) error
 	return nil
 }
 
-func (m *mockReplayJobManager) Stop() string {
+func (m *mockReplayJobManager) Stop(manager.CancelConfig) string {
 	m.curJob = ""
 	return "stopped"
 }

pkg/sqlreplay/manager/job.go

@@ -13,15 +13,15 @@ import (
 	"github.com/siddontang/go/hack"
 )
 
-type jobType int
+type JobType int
 
 const (
-	Capture jobType = iota
+	Capture JobType = 1 << iota
 	Replay
 )
 
 type Job interface {
-	Type() jobType
+	Type() JobType
 	String() string
 	MarshalJSON() ([]byte, error)
 	SetProgress(progress float64, endTime time.Time, done bool, err error)
@@ -89,7 +89,7 @@ type captureJob struct {
 	cfg capture.CaptureConfig
 }
 
-func (job *captureJob) Type() jobType {
+func (job *captureJob) Type() JobType {
 	return Capture
 }
 
@@ -116,7 +116,7 @@ type replayJob struct {
 	cfg replay.ReplayConfig
 }
 
-func (job *replayJob) Type() jobType {
+func (job *replayJob) Type() JobType {
 	return Replay
 }
 

pkg/sqlreplay/manager/job_test.go

@@ -16,7 +16,7 @@ import (
 func TestIsRunning(t *testing.T) {
 	tests := []struct {
 		job     Job
-		tp      jobType
+		tp      JobType
 		running bool
 	}{
 		{

pkg/sqlreplay/manager/manager.go

@@ -21,6 +21,10 @@ const (
 	maxJobHistoryCount = 10
 )
 
+type CancelConfig struct {
+	Type JobType
+}
+
 type CertManager interface {
 	SQLTLS() *tls.Config
 }
@@ -29,7 +33,7 @@ type JobManager interface {
 	StartCapture(capture.CaptureConfig) error
 	StartReplay(replay.ReplayConfig) error
 	GetCapture() capture.Capture
-	Stop() string
+	Stop(CancelConfig) string
 	Jobs() string
 	Close()
 }
@@ -158,11 +162,14 @@ func (jm *jobManager) Jobs() string {
 	return hack.String(b)
 }
 
-func (jm *jobManager) Stop() string {
+func (jm *jobManager) Stop(cfg CancelConfig) string {
 	job := jm.runningJob()
 	if job == nil {
 		return "no job running"
 	}
+	if job.Type()&cfg.Type == 0 {
+		return "no privilege to stop the job"
+	}
 	switch job.Type() {
 	case Capture:
 		jm.capture.Stop(errors.Errorf("manually stopped"))

pkg/sqlreplay/manager/manager_test.go

@@ -22,24 +22,26 @@ func TestStartAndStop(t *testing.T) {
 	cpt, rep := &mockCapture{}, &mockReplay{}
 	mgr.capture, mgr.replay = cpt, rep
 
-	require.Contains(t, mgr.Stop(), "no job running")
+	require.Contains(t, mgr.Stop(CancelConfig{Type: Capture | Replay}), "no job running")
 	require.NotNil(t, mgr.GetCapture())
 
 	require.NoError(t, mgr.StartCapture(capture.CaptureConfig{}))
 	require.Error(t, mgr.StartCapture(capture.CaptureConfig{}))
 	require.Error(t, mgr.StartReplay(replay.ReplayConfig{}))
 	require.Len(t, mgr.jobHistory, 1)
 	require.NotEmpty(t, mgr.Jobs())
-	require.Contains(t, mgr.Stop(), "stopped")
-	require.Contains(t, mgr.Stop(), "no job running")
+	require.Contains(t, mgr.Stop(CancelConfig{Type: Replay}), "no privilege to stop the job")
+	require.Contains(t, mgr.Stop(CancelConfig{Type: Capture}), "stopped")
+	require.Contains(t, mgr.Stop(CancelConfig{Type: Capture}), "no job running")
 	require.Len(t, mgr.jobHistory, 1)
 
 	require.NoError(t, mgr.StartReplay(replay.ReplayConfig{}))
 	require.Error(t, mgr.StartCapture(capture.CaptureConfig{}))
 	require.Error(t, mgr.StartReplay(replay.ReplayConfig{}))
 	require.Len(t, mgr.jobHistory, 2)
-	require.Contains(t, mgr.Stop(), "stopped")
-	require.Contains(t, mgr.Stop(), "no job running")
+	require.Contains(t, mgr.Stop(CancelConfig{Type: Capture}), "no privilege to stop the job")
+	require.Contains(t, mgr.Stop(CancelConfig{Type: Replay}), "stopped")
+	require.Contains(t, mgr.Stop(CancelConfig{Type: Replay}), "no job running")
 	require.Len(t, mgr.jobHistory, 2)
 
 	// Test that Jobs() also update progress.
@@ -144,7 +146,7 @@ func TestHistoryLen(t *testing.T) {
 
 	for i := 0; i < maxJobHistoryCount+1; i++ {
 		require.NoError(t, mgr.StartCapture(capture.CaptureConfig{}))
-		require.Contains(t, mgr.Stop(), "stopped")
+		require.Contains(t, mgr.Stop(CancelConfig{Type: Capture}), "stopped")
 		expectedLen := i + 1
 		if expectedLen > maxJobHistoryCount {
 			expectedLen = maxJobHistoryCount