Add schema sub-command

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

go/cmd/root.go

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

go/cmd/schema.go

@@ -0,0 +1,52 @@
+/*
+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"
+	"github.com/vitessio/vt/go/schema"
+	"vitess.io/vitess/go/mysql"
+)
+
+var vtParams mysql.ConnParams
+
+func schemaCmd() *cobra.Command {
+
+	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, args []string) error {
+			cfg := schema.Config{
+				VTParams: vtParams,
+			}
+
+			return schema.Run(cfg)
+		},
+	}
+	registerFlags(cmd)
+	return cmd
+}
+
+func registerFlags(cmd *cobra.Command) {
+	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")
+}

go/schema/schema.go

@@ -0,0 +1,103 @@
+/*
+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
+	}
+	out.Write(b)
+	return nil
+}
+
+type TableInfo struct {
+	Name string
+	Rows int
+}
+
+type SchemaInfo struct {
+	Tables []TableInfo
+}
+
+func Get(cfg Config) (*SchemaInfo, 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 := &SchemaInfo{
+		Tables: tables,
+	}
+	return schemaInfo, nil
+}
+
+func Load(fileName string) (*SchemaInfo, error) {
+	b, err := os.ReadFile(fileName)
+	if err != nil {
+		return nil, err
+	}
+	var si SchemaInfo
+	err = json.Unmarshal(b, &si)
+	if err != nil {
+		return nil, err
+	}
+	return &si, nil
+}

go/schema/schema_test.go

@@ -0,0 +1,31 @@
+package schema
+
+import (
+	"github.com/stretchr/testify/require"
+	"testing"
+)
+
+func TestSchema(t *testing.T) {
+	si, err := Load("../../t/sakila/sakila-schema-info.json")
+	require.NoError(t, err)
+	require.NotNil(t, si)
+	require.NotEmpty(t, si.Tables)
+	require.Equal(t, 16, len(si.Tables))
+	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.Greater(t, table.Rows, 0, "table %s has no rows", table.Name)
+		}
+	}
+}

t/sakila/sakila-schema-info.json

@@ -0,0 +1,68 @@
+{
+  "Tables": [
+    {
+      "Name": "actor",
+      "Rows": 200
+    },
+    {
+      "Name": "address",
+      "Rows": 603
+    },
+    {
+      "Name": "category",
+      "Rows": 16
+    },
+    {
+      "Name": "city",
+      "Rows": 600
+    },
+    {
+      "Name": "country",
+      "Rows": 109
+    },
+    {
+      "Name": "customer",
+      "Rows": 599
+    },
+    {
+      "Name": "film",
+      "Rows": 1000
+    },
+    {
+      "Name": "film_actor",
+      "Rows": 5462
+    },
+    {
+      "Name": "film_category",
+      "Rows": 1000
+    },
+    {
+      "Name": "film_text",
+      "Rows": 1000
+    },
+    {
+      "Name": "inventory",
+      "Rows": 4581
+    },
+    {
+      "Name": "language",
+      "Rows": 6
+    },
+    {
+      "Name": "payment",
+      "Rows": 16086
+    },
+    {
+      "Name": "rental",
+      "Rows": 15425
+    },
+    {
+      "Name": "staff",
+      "Rows": 2
+    },
+    {
+      "Name": "store",
+      "Rows": 2
+    }
+  ]
+}