Implement reference sub-command. Add sakila files for testing

Signed-off-by: Rohit Nayak <[email protected]>

go/keys/keys.go

@@ -22,7 +22,7 @@ import (
 	"io"
 	"os"
 	"sort"
-
+	"strconv"
 	querypb "vitess.io/vitess/go/vt/proto/query"
 	"vitess.io/vitess/go/vt/sqlparser"
 	"vitess.io/vitess/go/vt/vtgate/planbuilder/operators"
@@ -41,7 +41,22 @@ func Run(cfg Config) error {
 	return run(os.Stdout, cfg)
 }
 
-func run(out io.Writer, cfg Config) error {
+func GetKeysInfo(cfg Config) (*Output, error) {
+	keysInfo, err := getKeysInfo(cfg)
+	if err != nil {
+		return nil, err
+	}
+	var output Output
+	for _, query := range keysInfo.queries {
+		output.Queries = append(output.Queries, *query)
+	}
+	for _, failed := range keysInfo.failed {
+		output.Failed = append(output.Failed, *failed)
+	}
+	return &output, nil
+}
+
+func getKeysInfo(cfg Config) (*queryList, error) {
 	si := &schemaInfo{
 		tables: make(map[string]columns),
 	}
@@ -62,7 +77,7 @@ func run(out io.Writer, cfg Config) error {
 		case data.Skip, data.Error, data.VExplain:
 			skip = true
 		case data.Unknown:
-			return fmt.Errorf("unknown command type: %s", query.Type)
+			return nil, fmt.Errorf("unknown command type: %s", query.Type)
 		case data.Comment, data.CommentWithCommand, data.EmptyLine, data.WaitForAuthoritative, data.SkipIfBelowVersion:
 			// no-op for keys
 		case data.QueryT:
@@ -71,13 +86,22 @@ func run(out io.Writer, cfg Config) error {
 				continue
 			}
 			process(query, si, ql)
+		default:
+			return nil, fmt.Errorf("unknown command type: %s", query.Type)
 		}
 	}
 
 	if err := loader.Close(); err != nil {
-		return err
+		return nil, err
 	}
+	return ql, nil
+}
 
+func run(out io.Writer, cfg Config) error {
+	ql, err := getKeysInfo(cfg)
+	if err != nil {
+		return err
+	}
 	return ql.writeJSONTo(out)
 }
 
@@ -132,11 +156,12 @@ func (ql *queryList) processQuery(si *schemaInfo, ast sqlparser.Statement, q dat
 		ReservedVars: reservedVars,
 		SemTable:     st,
 	}
+	usageCount := getUsageCount(ast)
 
 	structure := sqlparser.CanonicalString(ast)
 	r, found := ql.queries[structure]
 	if found {
-		r.UsageCount++
+		r.UsageCount += usageCount
 		r.LineNumbers = append(r.LineNumbers, q.Line)
 		return
 	}
@@ -154,7 +179,7 @@ func (ql *queryList) processQuery(si *schemaInfo, ast sqlparser.Statement, q dat
 	ql.queries[structure] = &QueryAnalysisResult{
 		QueryStructure:  structure,
 		StatementType:   result.StatementType,
-		UsageCount:      1,
+		UsageCount:      usageCount,
 		LineNumbers:     []int{q.Line},
 		TableNames:      tableNames,
 		GroupingColumns: result.GroupingColumns,
@@ -163,6 +188,27 @@ func (ql *queryList) processQuery(si *schemaInfo, ast sqlparser.Statement, q dat
 	}
 }
 
+func getUsageCount(ast sqlparser.Statement) int {
+	var comments *sqlparser.ParsedComments
+	switch stmt := ast.(type) {
+	case *sqlparser.Select:
+		comments = stmt.Comments
+	case *sqlparser.Insert:
+		comments = stmt.Comments
+	case *sqlparser.Update:
+		comments = stmt.Comments
+	case *sqlparser.Delete:
+		comments = stmt.Comments
+	}
+	directives := comments.Directives()
+	usageCountStr, _ := directives.GetString("VT_USAGE_COUNT", "1")
+	usageCount, _ := strconv.Atoi(usageCountStr)
+	if usageCount == 0 {
+		usageCount = 1
+	}
+	return usageCount
+}
+
 func (ql *queryList) addFailedQuery(q data.Query, err error) {
 	key := q.Query + err.Error()
 	if v, exists := ql.failed[key]; exists {

go/keys/keys_test.go

@@ -20,6 +20,7 @@ import (
 	"os"
 	"strings"
 	"testing"
+	"vitess.io/vitess/go/vt/sqlparser"
 
 	"github.com/stretchr/testify/require"
 
@@ -78,3 +79,36 @@ func TestKeysNonAuthoritativeTable(t *testing.T) {
 		require.NotEmpty(t, result.FilterColumns)
 	}
 }
+
+func TestGetUsageCount(t *testing.T) {
+	type testCase struct {
+		query      string
+		usageCount int
+	}
+
+	cases := []testCase{
+		{
+			query:      "select /*vt+ VT_USAGE_COUNT=11 */ * from actor where first_name = 'Scarlett'",
+			usageCount: 11,
+		},
+		{
+			query:      "select /*vt+ VT_USAGE_COUNT=11.5 */ * from actor where first_name = 'Scarlett'",
+			usageCount: 1,
+		},
+		{
+			query:      "select /*vt+ VT_USAGE_COUNT=0 */ * from actor where first_name = 'Scarlett'",
+			usageCount: 1,
+		},
+		{
+			query:      "select  * from actor where first_name = 'Scarlett' and last_name = 'Johansson'",
+			usageCount: 1,
+		},
+	}
+
+	for _, tcase := range cases {
+		ast, err := sqlparser.NewTestParser().Parse(tcase.query)
+		require.NoError(t, err)
+		usageCount := getUsageCount(ast)
+		require.Equal(t, tcase.usageCount, usageCount)
+	}
+}

go/reference/reference.go

@@ -17,9 +17,12 @@ limitations under the License.
 package reference
 
 import (
+	"fmt"
 	"github.com/vitessio/vt/go/data"
+	"github.com/vitessio/vt/go/keys"
 	"io"
 	"os"
+	"strings"
 )
 
 type Config struct {
@@ -33,9 +36,109 @@ func Run(cfg Config) error {
 	return run(os.Stdout, cfg)
 }
 
+func Find(cfg Config) (*ReferenceInfo, error) {
+	ri, err := GetReferenceInfo(cfg)
+	if err != nil {
+		return nil, err
+	}
+	totalJoins := 0
+	for _, ts := range ri.TableSummaries {
+		totalJoins += ts.JoinCount
+	}
+
+	thresholdJoins := totalJoins / 20 // 5%
+	writePercentage := func(ts *TableSummary) float64 {
+		return float64(ts.NumWrites) / float64(ts.NumWrites+ts.NumReads)
+	}
+	writePercentageThreshold := 1 / 100.0 // 1%
+	for _, ts := range ri.TableSummaries {
+		if ts.JoinCount > thresholdJoins && writePercentage(ts) < writePercentageThreshold {
+			ri.ChosenTables = append(ri.ChosenTables, ts.TableName)
+		}
+	}
+	return ri, nil
+}
+
 func run(out io.Writer, cfg Config) error {
-	if _, err := out.Write([]byte("Placeholder for reference output\n")); err != nil {
+	ri, err := Find(cfg)
+	if err != nil {
 		return err
 	}
+	for _, table := range ri.ChosenTables {
+		fmt.Fprintf(out, "%s:: %+v\n", table, ri.TableSummaries[table])
+	}
 	return nil
 }
+
+type TableSummary struct {
+	TableName string
+	NumReads  int
+	NumWrites int
+	JoinCount int
+}
+
+func (ts TableSummary) String() string {
+	return fmt.Sprintf("Table: %s, Reads: %d, Writes: %d, Joins: %d", ts.TableName, ts.NumReads, ts.NumWrites, ts.JoinCount)
+}
+
+type ReferenceInfo struct {
+	TableSummaries map[string]*TableSummary
+	ChosenTables   []string
+}
+
+func GetReferenceInfo(cfg Config) (*ReferenceInfo, error) {
+	ri := &ReferenceInfo{
+		TableSummaries: make(map[string]*TableSummary),
+	}
+	keysConfig := keys.Config{
+		FileName: cfg.FileName,
+		Loader:   cfg.Loader,
+	}
+	keysOutput, err := keys.GetKeysInfo(keysConfig)
+	if err != nil {
+		return nil, err
+	}
+	getRit := func(table string) *TableSummary {
+		summary, ok := ri.TableSummaries[table]
+		if !ok {
+			summary = &TableSummary{
+				TableName: table,
+			}
+			ri.TableSummaries[table] = summary
+		}
+		return summary
+	}
+
+	for _, query := range keysOutput.Queries {
+		usageCount := query.UsageCount
+		isWrite := false
+		isRead := false
+		st := strings.ToLower(query.StatementType)
+		switch st {
+		case "select":
+			isRead = true
+		case "insert", "update", "delete":
+			isWrite = true
+		default:
+			continue
+		}
+
+		for _, table := range query.TableNames {
+			rit := getRit(table)
+			if isRead {
+				rit.NumReads += usageCount
+			}
+			if isWrite {
+				rit.NumWrites += usageCount
+			}
+		}
+
+		for _, pred := range query.JoinPredicates {
+			rit1 := getRit(pred.LHS.Table)
+			rit2 := getRit(pred.RHS.Table)
+			rit1.JoinCount += usageCount
+			rit2.JoinCount += usageCount
+		}
+	}
+	return ri, nil
+}

go/reference/reference_test.go

@@ -0,0 +1,23 @@
+package reference
+
+import (
+	"github.com/stretchr/testify/require"
+	"github.com/vitessio/vt/go/data"
+	"testing"
+)
+
+func TestReference(t *testing.T) {
+	cfg := Config{
+		FileName: "../../t/sakila/sakila.test",
+		Loader:   data.SQLScriptLoader{},
+	}
+
+	ri, err := Find(cfg)
+	require.NoError(t, err)
+	require.NotNil(t, ri)
+	require.NotEmpty(t, ri.TableSummaries)
+	validReferenceTables := []string{"actor", "address", "category", "city", "country", "film", "language", "staff"}
+	for _, table := range ri.ChosenTables {
+		require.Containsf(t, validReferenceTables, table, "table %s is not a valid reference table", table)
+	}
+}

go/summarize/summarize.go

@@ -57,6 +57,13 @@ func Run(args []string) {
 	}
 }
 
+func SummarizeKeysFile(fileName string) ([]TableSummary, error) {
+	trace := readTraceFile(fileName)
+	tableSummaries, failuresSummaries := summarizeKeysQueries(trace.AnalysedQueries)
+	_ = failuresSummaries
+	return tableSummaries, nil
+}
+
 func exit(msg string) {
 	fmt.Println(msg)
 	os.Exit(1)