Merge pull request #64 from vitessio/rohit/reference2

Add table row count to summary output

go/cmd/keys.go

@@ -17,7 +17,7 @@ limitations under the License.
 package cmd
 
 import (
-	"errors"
+	"fmt"
 
 	"github.com/spf13/cobra"
 
@@ -53,9 +53,11 @@ func keysCmd() *cobra.Command {
 	return cmd
 }
 
+const allowedInputTypes = "'sql', 'mysql-log' or 'vtgate-log'"
+
 func addInputTypeFlag(cmd *cobra.Command, s *string) {
 	*s = "sql"
-	cmd.Flags().StringVar(s, "input-type", "sql", "Specifies the type of input file: 'sql' or 'mysql-log'")
+	cmd.Flags().StringVar(s, "input-type", "sql", fmt.Sprintf("Specifies the type of input file: %s", allowedInputTypes))
 }
 
 func configureLoader(inputType string, needsBindVars bool) (data.Loader, error) {
@@ -67,6 +69,6 @@ func configureLoader(inputType string, needsBindVars bool) (data.Loader, error)
 	case "vtgate-log":
 		return data.VtGateLogLoader{NeedsBindVars: needsBindVars}, nil
 	default:
-		return nil, errors.New("invalid input type: must be 'sql' or 'mysql-log'")
+		return nil, fmt.Errorf("invalid input type: must be %s", allowedInputTypes)
 	}
 }

go/cmd/root.go

@@ -37,6 +37,7 @@ func Execute() {
 	root.AddCommand(testerCmd())
 	root.AddCommand(tracerCmd())
 	root.AddCommand(keysCmd())
+	root.AddCommand(schemaCmd())
 
 	err := root.Execute()
 	if err != nil {

go/cmd/schema.go

@@ -0,0 +1,50 @@
+/*
+Copyright 2024 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 cmd
+
+import (
+	"github.com/spf13/cobra"
+	"vitess.io/vitess/go/mysql"
+
+	"github.com/vitessio/vt/go/schema"
+)
+
+func schemaCmd() *cobra.Command {
+	var vtParams mysql.ConnParams
+
+	cmd := &cobra.Command{
+		Use:     "schema ",
+		Short:   "Loads info from the database including row counts",
+		Example: "vt schema",
+		Args:    cobra.ExactArgs(0),
+		RunE: func(_ *cobra.Command, _ []string) error {
+			cfg := schema.Config{
+				VTParams: vtParams,
+			}
+
+			return schema.Run(cfg)
+		},
+	}
+
+	cmd.Flags().StringVarP(&vtParams.Host, "host", "", "127.0.0.1", "Database host")
+	cmd.Flags().IntVarP(&vtParams.Port, "port", "", 3306, "Database port")
+	cmd.Flags().StringVarP(&vtParams.Uname, "user", "", "root", "Database user")
+	cmd.Flags().StringVarP(&vtParams.Pass, "password", "", "", "Database password")
+	cmd.Flags().StringVarP(&vtParams.DbName, "database", "", "", "Database name")
+
+	return cmd
+}

go/data/query.go

@@ -18,6 +18,8 @@ package data
 
 import (
 	"fmt"
+	"strconv"
+	"strings"
 
 	log "github.com/sirupsen/logrus"
 )
@@ -36,10 +38,11 @@ type (
 	}
 
 	Query struct {
-		FirstWord string
-		Query     string
-		Line      int
-		Type      CmdType
+		FirstWord  string
+		Query      string
+		Line       int
+		Type       CmdType
+		UsageCount int
 
 		// These fields are only set if the log file is a slow query log
 		QueryTime, LockTime    float64
@@ -56,13 +59,20 @@ type (
 // If the query log contains directives, they will be read and queries will be skipped as necessary.
 func ForeachSQLQuery(loader IteratorLoader, f func(Query) error) error {
 	skip := false
+	usageCount := 1
 	for {
 		query, kontinue := loader.Next()
 		if !kontinue {
 			break
 		}
 
 		switch query.Type {
+		case UsageCount:
+			var err error
+			usageCount, err = strconv.Atoi(strings.TrimSpace(query.Query))
+			if err != nil {
+				return fmt.Errorf("usage Count is incorrectly specified: %s", query.Query)
+			}
 		case Skip, Error, VExplain:
 			skip = true
 		case Unknown:
@@ -74,9 +84,11 @@ func ForeachSQLQuery(loader IteratorLoader, f func(Query) error) error {
 				skip = false
 				continue
 			}
+			query.UsageCount = usageCount
 			if err := f(query); err != nil {
 				return err
 			}
+			usageCount = 1
 		}
 	}
 

go/data/typ.go

@@ -34,6 +34,7 @@ const (
 	VitessOnly
 	MysqlOnly
 	Reference
+	UsageCount
 )
 
 var commandMap = map[string]CmdType{ //nolint:gochecknoglobals // this is instead of a const
@@ -46,6 +47,7 @@ var commandMap = map[string]CmdType{ //nolint:gochecknoglobals // this is instea
 	"vitess_only":           VitessOnly,
 	"mysql_only":            MysqlOnly,
 	"reference":             Reference,
+	"usage_count":           UsageCount,
 }
 
 func (cmd CmdType) String() string {

go/keys/keys.go

@@ -140,8 +140,12 @@ func (ql *queryList) processQuery(si *schemaInfo, ast sqlparser.Statement, q dat
 
 	structure := sqlparser.CanonicalString(ast)
 	r, found := ql.queries[structure]
+	usageCount := q.UsageCount
+	if usageCount == 0 {
+		usageCount = 1
+	}
 	if found {
-		r.UsageCount++
+		r.UsageCount += usageCount
 		r.LineNumbers = append(r.LineNumbers, q.Line)
 		r.QueryTime += q.QueryTime
 		r.LockTime += q.LockTime
@@ -163,7 +167,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,

go/schema/schema.go

@@ -0,0 +1,105 @@
+/*
+Copyright 2024 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 schema
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"io"
+	"os"
+
+	"vitess.io/vitess/go/mysql"
+)
+
+type Config struct {
+	VTParams mysql.ConnParams
+}
+
+func Run(cfg Config) error {
+	return run(os.Stdout, cfg)
+}
+
+func run(out io.Writer, cfg Config) error {
+	si, err := Get(cfg)
+	if err != nil {
+		return err
+	}
+	b, err := json.MarshalIndent(si, "", "  ")
+	if err != nil {
+		return err
+	}
+	_, err = out.Write(b)
+	return err
+}
+
+type TableInfo struct {
+	Name string `json:"name"`
+	Rows int    `json:"rows"`
+}
+
+type Info struct {
+	FileType string      `json:"fileType"`
+	Tables   []TableInfo `json:"tables"`
+}
+
+func Get(cfg Config) (*Info, error) {
+	vtParams := &mysql.ConnParams{
+		Host:   cfg.VTParams.Host,
+		Port:   cfg.VTParams.Port,
+		Uname:  cfg.VTParams.Uname,
+		Pass:   cfg.VTParams.Pass,
+		DbName: cfg.VTParams.DbName,
+	}
+
+	vtConn, err := mysql.Connect(context.Background(), vtParams)
+	if err != nil {
+		return nil, err
+	}
+	defer vtConn.Close()
+	queryTableSizes := "SELECT table_name, table_rows FROM information_schema.tables WHERE table_schema = '%s' and table_type = 'BASE TABLE'"
+	qr, err := vtConn.ExecuteFetch(fmt.Sprintf(queryTableSizes, cfg.VTParams.DbName), -1, false)
+	if err != nil {
+		return nil, err
+	}
+	var tables []TableInfo
+	for _, row := range qr.Rows {
+		tableName := row[0].ToString()
+		tableRows, _ := row[1].ToInt64()
+		tables = append(tables, TableInfo{
+			Name: tableName,
+			Rows: int(tableRows),
+		})
+	}
+	schemaInfo := &Info{
+		Tables: tables,
+	}
+	return schemaInfo, nil
+}
+
+func Load(fileName string) (*Info, error) {
+	b, err := os.ReadFile(fileName)
+	if err != nil {
+		return nil, err
+	}
+	var si Info
+	err = json.Unmarshal(b, &si)
+	if err != nil {
+		return nil, err
+	}
+	return &si, nil
+}

go/schema/schema_test.go

@@ -0,0 +1,48 @@
+/*
+Copyright 2024 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 schema
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestSchema(t *testing.T) {
+	si, err := Load("../testdata/sakila-schema-info.json")
+	require.NoError(t, err)
+	require.NotNil(t, si)
+	require.NotEmpty(t, si.Tables)
+	require.Len(t, si.Tables, 16)
+	var tables []string
+	for _, table := range si.Tables {
+		tables = append(tables, table.Name)
+	}
+	require.Contains(t, tables, "actor")
+	require.NotContains(t, tables, "foo")
+	for _, table := range si.Tables {
+		require.NotEmpty(t, table.Name)
+		switch table.Name {
+		case "language":
+			require.Equal(t, 6, table.Rows)
+		case "film":
+			require.Equal(t, 1000, table.Rows)
+		default:
+			require.Positive(t, table.Rows, "table %s has no rows", table.Name)
+		}
+	}
+}