344 lines
10 KiB
Go
344 lines
10 KiB
Go
// Copyright 2024 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 executor
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"net"
|
|
"net/http"
|
|
"net/url"
|
|
"strconv"
|
|
"sync"
|
|
"testing"
|
|
|
|
"github.com/pingcap/tidb/pkg/domain/infosync"
|
|
"github.com/pingcap/tidb/pkg/executor/internal/exec"
|
|
"github.com/pingcap/tidb/pkg/infoschema"
|
|
"github.com/pingcap/tidb/pkg/meta/model"
|
|
"github.com/pingcap/tidb/pkg/parser"
|
|
"github.com/pingcap/tidb/pkg/parser/mysql"
|
|
plannercore "github.com/pingcap/tidb/pkg/planner/core"
|
|
"github.com/pingcap/tidb/pkg/planner/core/resolve"
|
|
"github.com/pingcap/tidb/pkg/types"
|
|
"github.com/pingcap/tidb/pkg/util/chunk"
|
|
"github.com/pingcap/tidb/pkg/util/hint"
|
|
"github.com/pingcap/tidb/pkg/util/mock"
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
func TestTrafficForm(t *testing.T) {
|
|
tests := []struct {
|
|
sql string
|
|
method string
|
|
path string
|
|
form url.Values
|
|
hasStartTime bool
|
|
}{
|
|
{
|
|
sql: "traffic capture to '/tmp' duration='1s' encryption_method='aes' compress=false",
|
|
method: http.MethodPost,
|
|
path: "/api/traffic/capture",
|
|
form: url.Values{
|
|
"output": []string{"/tmp"},
|
|
"duration": []string{"1s"},
|
|
"encrypt-method": []string{"aes"},
|
|
"compress": []string{"false"},
|
|
},
|
|
hasStartTime: true,
|
|
},
|
|
{
|
|
sql: "traffic capture to '/tmp' duration='1s'",
|
|
method: http.MethodPost,
|
|
path: "/api/traffic/capture",
|
|
form: url.Values{
|
|
"output": []string{"/tmp"},
|
|
"duration": []string{"1s"},
|
|
},
|
|
hasStartTime: true,
|
|
},
|
|
{
|
|
sql: "traffic replay from '/tmp' user='root' password='123456' speed=1.0 read_only=true",
|
|
method: http.MethodPost,
|
|
path: "/api/traffic/replay",
|
|
form: url.Values{
|
|
"input": []string{"/tmp"},
|
|
"username": []string{"root"},
|
|
"password": []string{"123456"},
|
|
"speed": []string{"1.0"},
|
|
"readonly": []string{"true"},
|
|
},
|
|
hasStartTime: true,
|
|
},
|
|
{
|
|
sql: "traffic replay from '/tmp' user='root'",
|
|
method: http.MethodPost,
|
|
path: "/api/traffic/replay",
|
|
form: url.Values{
|
|
"input": []string{"/tmp"},
|
|
"username": []string{"root"},
|
|
},
|
|
hasStartTime: true,
|
|
},
|
|
{
|
|
sql: "cancel traffic jobs",
|
|
method: http.MethodPost,
|
|
path: "/api/traffic/cancel",
|
|
form: url.Values{},
|
|
},
|
|
{
|
|
sql: "show traffic jobs",
|
|
method: http.MethodGet,
|
|
path: "/api/traffic/show",
|
|
form: url.Values{},
|
|
},
|
|
}
|
|
|
|
suite := newTrafficTestSuite(t, 10)
|
|
ctx := context.TODO()
|
|
httpHandler := &mockHTTPHandler{t: t, httpOK: true}
|
|
server, port := runServer(t, httpHandler)
|
|
defer server.Close()
|
|
ctx = fillCtxWithTiProxyAddr(ctx, []int{port})
|
|
for i, test := range tests {
|
|
executor := suite.build(ctx, test.sql)
|
|
require.NoError(t, executor.Open(ctx))
|
|
chk := exec.NewFirstChunk(executor)
|
|
require.NoError(t, executor.Next(ctx, chk), "case %d", i)
|
|
require.Equal(t, test.method, httpHandler.getMethod(), "case %d", i)
|
|
require.Equal(t, test.path, httpHandler.getPath(), "case %d", i)
|
|
actualForm := httpHandler.getForm()
|
|
if test.hasStartTime {
|
|
require.NotEmpty(t, actualForm.Get("start-time"), "case %d", i)
|
|
test.form.Add("start-time", actualForm.Get("start-time"))
|
|
}
|
|
require.Equal(t, test.form, actualForm, "case %d", i)
|
|
}
|
|
}
|
|
|
|
func TestTrafficError(t *testing.T) {
|
|
suite := newTrafficTestSuite(t, 10)
|
|
ctx := context.TODO()
|
|
exec := suite.build(ctx, "cancel traffic jobs")
|
|
|
|
// no tiproxy
|
|
m := make(map[string]*infosync.TiProxyServerInfo)
|
|
tempCtx := context.WithValue(ctx, tiproxyAddrKey, m)
|
|
require.ErrorContains(t, exec.Next(tempCtx, nil), "no tiproxy server found")
|
|
|
|
// tiproxy no response
|
|
m["127.0.0.1:0"] = &infosync.TiProxyServerInfo{IP: "127.0.0.1", StatusPort: "0"}
|
|
require.ErrorContains(t, exec.Next(tempCtx, nil), "dial tcp")
|
|
|
|
// tiproxy responds with error
|
|
httpHandler := &mockHTTPHandler{t: t, httpOK: false}
|
|
server, port := runServer(t, httpHandler)
|
|
defer server.Close()
|
|
tempCtx = fillCtxWithTiProxyAddr(ctx, []int{port})
|
|
require.ErrorContains(t, exec.Next(tempCtx, nil), "500 Internal Server Error")
|
|
}
|
|
|
|
func TestTrafficShow(t *testing.T) {
|
|
suite := newTrafficTestSuite(t, 2)
|
|
ctx := context.TODO()
|
|
fields := trafficJobFields()
|
|
|
|
handlers := make([]*mockHTTPHandler, 0, 2)
|
|
servers := make([]*http.Server, 0, 2)
|
|
ports := make([]int, 0, 2)
|
|
for i := 0; i < 2; i++ {
|
|
httpHandler := &mockHTTPHandler{t: t, httpOK: true}
|
|
handlers = append(handlers, httpHandler)
|
|
server, port := runServer(t, httpHandler)
|
|
servers = append(servers, server)
|
|
ports = append(ports, port)
|
|
}
|
|
defer func() {
|
|
for _, server := range servers {
|
|
server.Close()
|
|
}
|
|
}()
|
|
if strconv.Itoa(ports[0]) > strconv.Itoa(ports[1]) {
|
|
ports[0], ports[1] = ports[1], ports[0]
|
|
handlers[0], handlers[1] = handlers[1], handlers[0]
|
|
}
|
|
ctx = fillCtxWithTiProxyAddr(ctx, ports)
|
|
|
|
marshaledTime1, marshaledTime2 := "2020-01-01T00:00:00Z", "2020-01-01T01:00:00Z"
|
|
marshaledJob := `{
|
|
"type": "capture",
|
|
"status": "canceled",
|
|
"start_time": "%s",
|
|
"end_time": "2020-01-01T02:01:01Z",
|
|
"progress": "50%%",
|
|
"error": "mock error"
|
|
}`
|
|
showTime1, showTime2 := "2020-01-01 00:00:00.000000", "2020-01-01 01:00:00.000000"
|
|
showResult := "%s, 2020-01-01 02:01:01.000000, 127.0.0.1:%d, capture, 50%%, canceled, mock error\n"
|
|
tests := []struct {
|
|
resp []string
|
|
chks []string
|
|
}{
|
|
{
|
|
resp: []string{"[]", "[]"},
|
|
chks: []string{},
|
|
},
|
|
{
|
|
resp: []string{fmt.Sprintf("[%s]", fmt.Sprintf(marshaledJob, marshaledTime1)), "[]"},
|
|
chks: []string{fmt.Sprintf(showResult, showTime1, ports[0])},
|
|
},
|
|
{
|
|
resp: []string{fmt.Sprintf("[%s]", fmt.Sprintf(marshaledJob, marshaledTime1)), fmt.Sprintf("[%s]", fmt.Sprintf(marshaledJob, marshaledTime1))},
|
|
chks: []string{fmt.Sprintf("%s%s", fmt.Sprintf(showResult, showTime1, ports[0]), fmt.Sprintf(showResult, showTime1, ports[1]))},
|
|
},
|
|
{
|
|
resp: []string{fmt.Sprintf("[%s,%s]", fmt.Sprintf(marshaledJob, marshaledTime1), fmt.Sprintf(marshaledJob, marshaledTime2)),
|
|
fmt.Sprintf("[%s,%s]", fmt.Sprintf(marshaledJob, marshaledTime1), fmt.Sprintf(marshaledJob, marshaledTime2))},
|
|
chks: []string{fmt.Sprintf("%s%s", fmt.Sprintf(showResult, showTime2, ports[0]), fmt.Sprintf(showResult, showTime2, ports[1])),
|
|
fmt.Sprintf("%s%s", fmt.Sprintf(showResult, showTime1, ports[0]), fmt.Sprintf(showResult, showTime1, ports[1]))},
|
|
},
|
|
}
|
|
|
|
for _, test := range tests {
|
|
for j := range test.resp {
|
|
handlers[j].setResponse(test.resp[j])
|
|
}
|
|
executor := suite.build(ctx, "show traffic jobs")
|
|
require.NoError(t, executor.Open(ctx))
|
|
chk := chunk.New(fields, 2, 2)
|
|
for j := 0; j < len(test.chks); j++ {
|
|
require.NoError(t, executor.Next(ctx, chk))
|
|
require.Equal(t, test.chks[j], chk.ToString(fields))
|
|
}
|
|
require.NoError(t, executor.Next(ctx, chk))
|
|
require.Equal(t, 0, chk.NumRows())
|
|
}
|
|
}
|
|
|
|
type trafficTestSuite struct {
|
|
t *testing.T
|
|
parser *parser.Parser
|
|
planBuilder *plannercore.PlanBuilder
|
|
execBuilder *MockExecutorBuilder
|
|
}
|
|
|
|
func newTrafficTestSuite(t *testing.T, chunkSize int) *trafficTestSuite {
|
|
parser := parser.New()
|
|
sctx := mock.NewContext()
|
|
sctx.GetSessionVars().MaxChunkSize = chunkSize
|
|
is := infoschema.MockInfoSchema([]*model.TableInfo{plannercore.MockSignedTable(), plannercore.MockUnsignedTable()})
|
|
planBuilder, _ := plannercore.NewPlanBuilder().Init(sctx, nil, hint.NewQBHintHandler(nil))
|
|
execBuilder := NewMockExecutorBuilderForTest(sctx, is)
|
|
return &trafficTestSuite{
|
|
t: t,
|
|
parser: parser,
|
|
planBuilder: planBuilder,
|
|
execBuilder: execBuilder,
|
|
}
|
|
}
|
|
|
|
func (suite *trafficTestSuite) build(ctx context.Context, sql string) exec.Executor {
|
|
stmt, err := suite.parser.ParseOneStmt(sql, "", "")
|
|
require.NoError(suite.t, err)
|
|
p, err := suite.planBuilder.Build(ctx, resolve.NewNodeW(stmt))
|
|
require.NoError(suite.t, err)
|
|
executor := suite.execBuilder.build(p)
|
|
require.NotEmpty(suite.t, executor)
|
|
return executor
|
|
}
|
|
|
|
type mockHTTPHandler struct {
|
|
t *testing.T
|
|
sync.Mutex
|
|
form url.Values
|
|
method string
|
|
path string
|
|
resp string
|
|
httpOK bool
|
|
}
|
|
|
|
func (handler *mockHTTPHandler) setResponse(resp string) {
|
|
handler.Lock()
|
|
defer handler.Unlock()
|
|
handler.resp = resp
|
|
}
|
|
|
|
func (handler *mockHTTPHandler) getForm() url.Values {
|
|
handler.Lock()
|
|
defer handler.Unlock()
|
|
return handler.form
|
|
}
|
|
|
|
func (handler *mockHTTPHandler) getMethod() string {
|
|
handler.Lock()
|
|
defer handler.Unlock()
|
|
return handler.method
|
|
}
|
|
|
|
func (handler *mockHTTPHandler) getPath() string {
|
|
handler.Lock()
|
|
defer handler.Unlock()
|
|
return handler.path
|
|
}
|
|
|
|
func (handler *mockHTTPHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
|
|
handler.Lock()
|
|
defer handler.Unlock()
|
|
handler.method = r.Method
|
|
handler.path = r.URL.Path
|
|
require.NoError(handler.t, r.ParseForm())
|
|
handler.form = r.PostForm
|
|
if handler.httpOK {
|
|
w.WriteHeader(http.StatusOK)
|
|
resp := handler.resp
|
|
if len(resp) == 0 && r.Method == http.MethodGet {
|
|
resp = "[]"
|
|
}
|
|
_, err := w.Write([]byte(resp))
|
|
require.NoError(handler.t, err)
|
|
} else {
|
|
w.WriteHeader(http.StatusInternalServerError)
|
|
}
|
|
}
|
|
|
|
func runServer(t *testing.T, handler http.Handler) (*http.Server, int) {
|
|
listener, err := net.Listen("tcp", ":0")
|
|
require.NoError(t, err)
|
|
server := &http.Server{Addr: listener.Addr().String(), Handler: handler}
|
|
go server.Serve(listener)
|
|
return server, listener.Addr().(*net.TCPAddr).Port
|
|
}
|
|
|
|
func fillCtxWithTiProxyAddr(ctx context.Context, ports []int) context.Context {
|
|
m := make(map[string]*infosync.TiProxyServerInfo)
|
|
for _, port := range ports {
|
|
addr := fmt.Sprintf("127.0.0.1:%d", port)
|
|
m[addr] = &infosync.TiProxyServerInfo{IP: "127.0.0.1", StatusPort: strconv.Itoa(port)}
|
|
}
|
|
return context.WithValue(ctx, tiproxyAddrKey, m)
|
|
}
|
|
|
|
func trafficJobFields() []*types.FieldType {
|
|
return []*types.FieldType{
|
|
types.NewFieldType(mysql.TypeDate),
|
|
types.NewFieldType(mysql.TypeDate),
|
|
types.NewFieldType(mysql.TypeString),
|
|
types.NewFieldType(mysql.TypeString),
|
|
types.NewFieldType(mysql.TypeString),
|
|
types.NewFieldType(mysql.TypeString),
|
|
types.NewFieldType(mysql.TypeString),
|
|
}
|
|
}
|