vitess.io/vitess@v0.16.2/go/vt/vttablet/tabletserver/planbuilder/plan_test.go

/*
Copyright 2019 The Vitess Authors.

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 planbuilder

import (
	"bufio"
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"os"
	"path/filepath"
	"strings"
	"testing"

	"vitess.io/vitess/go/vt/vtgate/evalengine"

	"github.com/stretchr/testify/require"

	"vitess.io/vitess/go/vt/sqlparser"
	"vitess.io/vitess/go/vt/tableacl"
	"vitess.io/vitess/go/vt/vttablet/tabletserver/schema"
)

// MarshalJSON returns a JSON of the given Plan.
// This is only for testing.
func (p *Plan) MarshalJSON() ([]byte, error) {
	mplan := struct {
		PlanID            PlanType
		TableName         sqlparser.IdentifierCS `json:",omitempty"`
		Permissions       []Permission           `json:",omitempty"`
		FieldQuery        *sqlparser.ParsedQuery `json:",omitempty"`
		FullQuery         *sqlparser.ParsedQuery `json:",omitempty"`
		NextCount         string                 `json:",omitempty"`
		WhereClause       *sqlparser.ParsedQuery `json:",omitempty"`
		NeedsReservedConn bool                   `json:",omitempty"`
	}{
		PlanID:      p.PlanID,
		TableName:   p.TableName(),
		Permissions: p.Permissions,
		FullQuery:   p.FullQuery,
		WhereClause: p.WhereClause,
	}
	if p.NextCount != nil {
		mplan.NextCount = evalengine.FormatExpr(p.NextCount)
	}
	if p.NeedsReservedConn {
		mplan.NeedsReservedConn = true
	}
	return json.Marshal(&mplan)
}

func TestPlan(t *testing.T) {
	testPlan(t, "exec_cases.txt")
}

func TestDDLPlan(t *testing.T) {
	testPlan(t, "ddl_cases.txt")
}

func testPlan(t *testing.T, fileName string) {
	t.Helper()
	testSchema := loadSchema("schema_test.json")
	for tcase := range iterateExecFile(fileName) {
		t.Run(tcase.input, func(t *testing.T) {
			if strings.Contains(tcase.options, "PassthroughDMLs") {
				PassthroughDMLs = true
			}
			var plan *Plan
			var err error
			statement, err := sqlparser.Parse(tcase.input)
			if err == nil {
				plan, err = Build(statement, testSchema, "dbName", false)
			}
			PassthroughDMLs = false

			var out string
			if err != nil {
				out = err.Error()
			} else {
				bout, err := json.Marshal(plan)
				require.NoError(t, err, "Error marshalling %v: %v", plan, err)
				out = string(bout)
			}
			if out != tcase.output {
				t.Errorf("Line:%v\ngot  = %s\nwant = %s", tcase.lineno, out, tcase.output)
				if err != nil {
					out = fmt.Sprintf("\"%s\"", out)
				} else {
					bout, _ := json.MarshalIndent(plan, "", "  ")
					out = string(bout)
				}
				fmt.Printf("\"in> %s\"\nout>%s\nexpected: %s\n\n", tcase.input, out, tcase.output)
			}
		})
	}
}

func TestPlanInReservedConn(t *testing.T) {
	testSchema := loadSchema("schema_test.json")
	for tcase := range iterateExecFile("exec_cases.txt") {
		t.Run(tcase.input, func(t *testing.T) {
			if strings.Contains(tcase.options, "PassthroughDMLs") {
				PassthroughDMLs = true
			}
			var plan *Plan
			var err error
			statement, err := sqlparser.Parse(tcase.input)
			if err == nil {
				plan, err = Build(statement, testSchema, "dbName", false)
			}
			PassthroughDMLs = false

			var out string
			if err != nil {
				out = err.Error()
			} else {
				bout, err := json.Marshal(plan)
				if err != nil {
					t.Fatalf("Error marshalling %v: %v", plan, err)
				}
				out = string(bout)
			}
			if out != tcase.output {
				t.Errorf("Line:%v\ngot  = %s\nwant = %s", tcase.lineno, out, tcase.output)
				if err != nil {
					out = fmt.Sprintf("\"%s\"", out)
				} else {
					bout, _ := json.MarshalIndent(plan, "", "  ")
					out = string(bout)
				}
				fmt.Printf("\"%s\"\n%s\n\n", tcase.input, out)
			}
		})
	}
}

func TestCustom(t *testing.T) {
	testSchemas, _ := filepath.Glob("testdata/*_schema.json")
	if len(testSchemas) == 0 {
		t.Log("No schemas to test")
		return
	}
	for _, schemFile := range testSchemas {
		schem := loadSchema(schemFile)
		t.Logf("Testing schema %s", schemFile)
		files, err := filepath.Glob(strings.Replace(schemFile, "schema.json", "*.txt", -1))
		if err != nil {
			log.Fatal(err)
		}
		if len(files) == 0 {
			t.Fatalf("No test files for %s", schemFile)
		}
		for _, file := range files {
			t.Logf("Testing file %s", file)
			for tcase := range iterateExecFile(file) {
				statement, err := sqlparser.Parse(tcase.input)
				if err != nil {
					t.Fatalf("Got error: %v, parsing sql: %v", err.Error(), tcase.input)
				}
				plan, err := Build(statement, schem, "dbName", false)
				var out string
				if err != nil {
					out = err.Error()
				} else {
					bout, err := json.Marshal(plan)
					if err != nil {
						t.Fatalf("Error marshalling %v: %v", plan, err)
					}
					out = string(bout)
				}
				if out != tcase.output {
					t.Errorf("File: %s: Line:%v\ngot  = %s\nwant = %s", file, tcase.lineno, out, tcase.output)
				}
			}
		}
	}
}

func TestStreamPlan(t *testing.T) {
	testSchema := loadSchema("schema_test.json")
	for tcase := range iterateExecFile("stream_cases.txt") {
		plan, err := BuildStreaming(tcase.input, testSchema)
		var out string
		if err != nil {
			out = err.Error()
		} else {
			bout, err := json.Marshal(plan)
			if err != nil {
				t.Fatalf("Error marshalling %v: %v", plan, err)
			}
			out = string(bout)
		}
		if out != tcase.output {
			t.Errorf("Line:%v\ngot  = %s\nwant = %s", tcase.lineno, out, tcase.output)
		}
	}
}

func TestMessageStreamingPlan(t *testing.T) {
	testSchema := loadSchema("schema_test.json")
	plan, err := BuildMessageStreaming("msg", testSchema)
	require.NoError(t, err)
	bout, _ := json.Marshal(plan)
	planJSON := string(bout)

	wantPlan := &Plan{
		PlanID: PlanMessageStream,
		Table:  testSchema["msg"],
		Permissions: []Permission{{
			TableName: "msg",
			Role:      tableacl.WRITER,
		}},
	}
	bout, _ = json.Marshal(wantPlan)
	wantJSON := string(bout)

	if planJSON != wantJSON {
		t.Errorf("BuildMessageStreaming: \n%s, want\n%s", planJSON, wantJSON)
	}

	_, err = BuildMessageStreaming("absent", testSchema)
	want := "table absent not found in schema"
	if err == nil || err.Error() != want {
		t.Errorf("BuildMessageStreaming(absent) error: %v, want %s", err, want)
	}

	_, err = BuildMessageStreaming("a", testSchema)
	want = "'a' is not a message table"
	if err == nil || err.Error() != want {
		t.Errorf("BuildMessageStreaming(absent) error: %v, want %s", err, want)
	}
}

func TestLockPlan(t *testing.T) {
	testSchema := loadSchema("schema_test.json")
	for tcase := range iterateExecFile("lock_cases.txt") {
		t.Run(tcase.input, func(t *testing.T) {
			var plan *Plan
			var err error
			statement, err := sqlparser.Parse(tcase.input)
			if err == nil {
				plan, err = Build(statement, testSchema, "dbName", false)
			}

			var out string
			if err != nil {
				out = err.Error()
			} else {
				bout, err := json.Marshal(plan)
				if err != nil {
					t.Fatalf("Error marshalling %v: %v", plan, err)
				}
				out = string(bout)
			}
			if out != tcase.output {
				t.Errorf("Line:%v\ngot  = %s\nwant = %s", tcase.lineno, out, tcase.output)
				if err != nil {
					out = fmt.Sprintf("\"%s\"", out)
				} else {
					bout, _ := json.MarshalIndent(plan, "", "  ")
					out = string(bout)
				}
				fmt.Printf("\"in> %s\"\nout>%s\nexpected: %s\n\n", tcase.input, out, tcase.output)
			}
		})
	}
}

func loadSchema(name string) map[string]*schema.Table {
	b, err := os.ReadFile(locateFile(name))
	if err != nil {
		panic(err)
	}
	tables := make([]*schema.Table, 0, 10)
	err = json.Unmarshal(b, &tables)
	if err != nil {
		panic(err)
	}
	s := make(map[string]*schema.Table)
	for _, t := range tables {
		s[t.Name.String()] = t
	}
	return s
}

type testCase struct {
	file    string
	lineno  int
	options string
	input   string
	output  string
}

func iterateExecFile(name string) (testCaseIterator chan testCase) {
	name = locateFile(name)
	fd, err := os.OpenFile(name, os.O_RDONLY, 0)
	if err != nil {
		panic(fmt.Sprintf("Could not open file %s", name))
	}
	testCaseIterator = make(chan testCase)
	go func() {
		defer close(testCaseIterator)

		r := bufio.NewReader(fd)
		lineno := 0
		options := ""
		for {
			binput, err := r.ReadBytes('\n')
			if err != nil {
				if err != io.EOF {
					fmt.Printf("Line: %d\n", lineno)
					panic(fmt.Errorf("Error reading file %s: %s", name, err.Error()))
				}
				break
			}
			lineno++
			input := string(binput)
			if input == "" || input == "\n" || input[0] == '#' || strings.HasPrefix(input, "Length:") {
				// fmt.Printf("%s\n", input)
				continue
			}

			if strings.HasPrefix(input, "options:") {
				options = input[8:]
				continue
			}
			err = json.Unmarshal(binput, &input)
			if err != nil {
				fmt.Printf("Line: %d, input: %s\n", lineno, binput)
				panic(err)
			}
			input = strings.Trim(input, "\"")
			var output []byte
			for {
				l, err := r.ReadBytes('\n')
				lineno++
				if err != nil {
					fmt.Printf("Line: %d\n", lineno)
					panic(fmt.Errorf("Error reading file %s: %s", name, err.Error()))
				}
				output = append(output, l...)
				if l[0] == '}' {
					output = output[:len(output)-1]
					b := bytes.NewBuffer(make([]byte, 0, 64))
					if err := json.Compact(b, output); err == nil {
						output = b.Bytes()
					}
					break
				}
				if l[0] == '"' {
					output = output[1 : len(output)-2]
					break
				}
			}
			testCaseIterator <- testCase{name, lineno, options, input, string(output)}
			options = ""
		}
	}()
	return testCaseIterator
}

func locateFile(name string) string {
	return "testdata/" + name
}