dbinfo: add info for primary keys, indexes, foreign keys

go/dbinfo/dbinfo.go

@@ -20,6 +20,7 @@ import (
 	"encoding/json"
 	"io"
 	"os"
+	"sort"
 
 	"vitess.io/vitess/go/mysql"
 )
@@ -53,10 +54,30 @@ type TableColumn struct {
 	Extra      string `json:"extra,omitempty"`
 }
 
+type PrimaryKey struct {
+	Columns []string `json:"columns"`
+}
+
+type Index struct {
+	Name      string
+	Columns   []string `json:"columns"`
+	NonUnique bool     `json:"nonUnique,omitempty"`
+}
+
+type ForeignKey struct {
+	ColumnName           string `json:"columnName"`
+	ConstraintName       string `json:"constraintName"`
+	ReferencedTableName  string `json:"referencedTableName"`
+	ReferencedColumnName string `json:"referencedColumnName"`
+}
+
 type TableInfo struct {
-	Name    string         `json:"name"`
-	Rows    int            `json:"rows"`
-	Columns []*TableColumn `json:"columns"`
+	Name        string         `json:"name"`
+	Rows        int            `json:"rows"`
+	Columns     []*TableColumn `json:"columns"`
+	PrimaryKey  *PrimaryKey    `json:"primaryKey,omitempty"`
+	Indexes     []*Index       `json:"indexes,omitempty"`
+	ForeignKeys []*ForeignKey  `json:"foreignKeys,omitempty"`
 }
 
 type Info struct {
@@ -65,6 +86,80 @@ type Info struct {
 	GlobalVariables map[string]string `json:"globalVariables"`
 }
 
+func getTableSizes(dbh *DBHelper, tableMap map[string]*TableInfo) error {
+	ts, err := dbh.getTableSizes()
+	if err != nil {
+		return err
+	}
+
+	for tableName, tableRows := range ts {
+		ti, ok := tableMap[tableName]
+		if !ok {
+			ti = &TableInfo{
+				Name: tableName,
+			}
+			tableMap[tableName] = ti
+		}
+		ti.Rows = tableRows
+	}
+	return nil
+}
+
+func getPrimaryKeys(dbh *DBHelper, tableMap map[string]*TableInfo) error {
+	pks, err := dbh.getPrimaryKeys()
+	if err != nil {
+		return err
+	}
+
+	for tableName, pk := range pks {
+		ti, ok := tableMap[tableName]
+		if !ok {
+			ti = &TableInfo{}
+			tableMap[tableName] = ti
+		}
+		ti.PrimaryKey = &PrimaryKey{
+			Columns: pk.columns,
+		}
+	}
+	return nil
+}
+
+func getIndexes(dbh *DBHelper, tableMap map[string]*TableInfo) error {
+	idxs, err := dbh.getIndexes()
+	if err != nil {
+		return err
+	}
+
+	for tableName, tidx := range idxs {
+		ti, ok := tableMap[tableName]
+		if !ok {
+			ti = &TableInfo{}
+			tableMap[tableName] = ti
+		}
+		for _, idx := range tidx.indexes {
+			ti.Indexes = append(ti.Indexes, idx)
+		}
+	}
+	return nil
+}
+
+func getForeignKeys(dbh *DBHelper, tableMap map[string]*TableInfo) error {
+	fks, err := dbh.getForeignKeys()
+	if err != nil {
+		return err
+	}
+
+	for fkName, fk := range fks {
+		ti, ok := tableMap[fkName]
+		if !ok {
+			ti = &TableInfo{}
+			tableMap[fkName] = ti
+		}
+		ti.ForeignKeys = fk
+	}
+	return nil
+}
+
 func Get(cfg Config) (*Info, error) {
 	vtParams := &mysql.ConnParams{
 		Host:   cfg.VTParams.Host,
@@ -74,46 +169,38 @@ func Get(cfg Config) (*Info, error) {
 		DbName: cfg.VTParams.DbName,
 	}
 
+	var tableInfo []*TableInfo
+	tableMap := make(map[string]*TableInfo)
+
 	dbh := NewDBHelper(vtParams)
-	ts, err := dbh.getTableSizes()
+
+	globalVariables, err := dbh.getGlobalVariables()
 	if err != nil {
 		return nil, err
 	}
 
-	var tableInfo []*TableInfo
-	tableMap := make(map[string]*TableInfo)
+	if err := getTableSizes(dbh, tableMap); err != nil {
+		return nil, err
+	}
 
-	for tableName, tableRows := range ts {
-		tableMap[tableName] = &TableInfo{
-			Name: tableName,
-			Rows: tableRows,
-		}
+	if err := getPrimaryKeys(dbh, tableMap); err != nil {
+		return nil, err
 	}
 
-	tc, err := dbh.getColumnInfo()
-	if err != nil {
+	if err := getIndexes(dbh, tableMap); err != nil {
 		return nil, err
 	}
 
-	for tableName, columns := range tc {
-		ti, ok := tableMap[tableName]
-		if !ok {
-			ti = &TableInfo{
-				Name: tableName,
-			}
-			tableMap[tableName] = ti
-		}
-		ti.Columns = columns
+	if err := getForeignKeys(dbh, tableMap); err != nil {
+		return nil, err
 	}
 
 	for tableName := range tableMap {
 		tableInfo = append(tableInfo, tableMap[tableName])
 	}
-
-	globalVariables, err := dbh.getGlobalVariables()
-	if err != nil {
-		return nil, err
-	}
+	sort.Slice(tableInfo, func(i, j int) bool {
+		return tableInfo[i].Name < tableInfo[j].Name
+	})
 
 	dbInfo := &Info{
 		FileType:        "dbinfo",

go/dbinfo/dbinfo_test.go

@@ -54,10 +54,11 @@ func TestDBInfoLoad(t *testing.T) {
 
 	t.Run("validateGlobalVariables", func(t *testing.T) {
 		require.NotEmpty(t, si.GlobalVariables)
-		require.Len(t, si.GlobalVariables, 3)
+		require.Len(t, si.GlobalVariables, 4)
 		expected := map[string]string{
 			"binlog_format":    "ROW",
 			"binlog_row_image": "FULL",
+			"gtid_mode":        "OFF",
 			"log_bin":          "ON",
 		}
 		require.EqualValues(t, expected, si.GlobalVariables)
@@ -144,4 +145,65 @@ func TestDBInfoGet(t *testing.T) {
 		require.NoError(t, err)
 		require.NotEmpty(t, gv)
 	})
+
+	t.Run("primary keys", func(t *testing.T) {
+		pks, err := dbh.getPrimaryKeys()
+		require.NoError(t, err)
+		require.Len(t, pks, 16)
+		want := map[string][]string{
+			"actor":         {"actor_id"},
+			"film":          {"film_id"},
+			"language":      {"language_id"},
+			"film_category": {"film_id", "category_id"},
+			"film_actor":    {"actor_id", "film_id"},
+		}
+		for tableName, Columns := range want {
+			pk, ok := pks[tableName]
+			require.True(t, ok)
+			require.Equal(t, Columns, pk.columns)
+		}
+	})
+
+	t.Run("indexes", func(t *testing.T) {
+		idxs, err := dbh.getIndexes()
+		require.NoError(t, err)
+		require.Len(t, idxs, 16)
+		idx, ok := idxs["film_actor"]
+		require.True(t, ok)
+		require.Len(t, idx.indexes, 2)
+		require.Equal(t, "idx_fk_film_id", idx.indexes["idx_fk_film_id"].Name)
+		require.Equal(t, []string{"film_id"}, idx.indexes["idx_fk_film_id"].Columns)
+		idx, ok = idxs["rental"]
+		require.True(t, ok)
+		require.Len(t, idx.indexes, 5)
+		require.Equal(t, "rental_date", idx.indexes["rental_date"].Name)
+		require.Equal(t, []string{"rental_date", "inventory_id", "customer_id"}, idx.indexes["rental_date"].Columns)
+		require.Equal(t, "PRIMARY", idx.indexes["PRIMARY_KEY"].Name)
+		require.Equal(t, []string{"rental_id"}, idx.indexes["PRIMARY_KEY"].Columns)
+	})
+
+	t.Run("foreign keys", func(t *testing.T) {
+		fks, err := dbh.getForeignKeys()
+		require.NoError(t, err)
+		require.Len(t, fks, 11)
+		fk, ok := fks["city"]
+		require.True(t, ok)
+		require.Len(t, fk, 1)
+		require.Equal(t, "country_id", fk[0].ColumnName)
+		require.Equal(t, "fk_city_country", fk[0].ConstraintName)
+		require.Equal(t, "country", fk[0].ReferencedTableName)
+		require.Equal(t, "country_id", fk[0].ReferencedColumnName)
+
+		fk, ok = fks["store"]
+		require.True(t, ok)
+		require.Len(t, fk, 2)
+		require.Equal(t, "address_id", fk[0].ColumnName)
+		require.Equal(t, "fk_store_address", fk[0].ConstraintName)
+		require.Equal(t, "address", fk[0].ReferencedTableName)
+		require.Equal(t, "address_id", fk[0].ReferencedColumnName)
+		require.Equal(t, "manager_staff_id", fk[1].ColumnName)
+		require.Equal(t, "fk_store_staff", fk[1].ConstraintName)
+		require.Equal(t, "staff", fk[1].ReferencedTableName)
+		require.Equal(t, "staff_id", fk[1].ReferencedColumnName)
+	})
 }

go/dbinfo/utils.go

@@ -130,3 +130,106 @@ func (dbh *DBHelper) getGlobalVariables() (map[string]string, error) {
 	}
 	return gv, nil
 }
+
+type primaryKey struct {
+	tableName string
+	columns   []string
+}
+type primaryKeys map[string]*primaryKey
+
+func (dbh *DBHelper) getPrimaryKeys() (primaryKeys, error) {
+	vtConn, cancel, err := dbh.GetConnection()
+	if err != nil {
+		return nil, err
+	}
+	defer cancel()
+
+	pks := make(primaryKeys)
+	queryPrimaryKeys := "select table_name, column_name from information_schema.key_column_usage where constraint_name = 'PRIMARY' and table_schema = '%s' order by table_name"
+	query := fmt.Sprintf(queryPrimaryKeys, dbh.vtParams.DbName)
+	qr, err := vtConn.ExecuteFetch(query, -1, false)
+	if err != nil {
+		return nil, err
+	}
+	for _, row := range qr.Rows {
+		tableName := row[0].ToString()
+		columnName := row[1].ToString()
+		pk, ok := pks[tableName]
+		if !ok {
+			pk = &primaryKey{tableName: tableName}
+			pks[tableName] = pk
+		}
+		pk.columns = append(pk.columns, columnName)
+	}
+	return pks, nil
+}
+
+type tableIndex struct {
+	tableName string
+	indexes   map[string]*Index
+}
+
+func (dbh *DBHelper) getIndexes() (map[string]*tableIndex, error) {
+	vtConn, cancel, err := dbh.GetConnection()
+	if err != nil {
+		return nil, err
+	}
+	defer cancel()
+
+	idxs := make(map[string]*tableIndex)
+	queryIndexes := "select table_name, index_name, column_name, non_unique from information_schema.statistics where table_schema = '%s' order by table_name, index_name"
+	query := fmt.Sprintf(queryIndexes, dbh.vtParams.DbName)
+	qr, err := vtConn.ExecuteFetch(query, -1, false)
+	if err != nil {
+		return nil, err
+	}
+	for _, row := range qr.Rows {
+		tableName := row[0].ToString()
+		indexName := row[1].ToString()
+		columnName := row[2].ToString()
+		nonUnique, _ := row[3].ToBool()
+		tidx, ok := idxs[tableName]
+		if !ok {
+			tidx = &tableIndex{tableName: tableName, indexes: make(map[string]*Index)}
+			idxs[tableName] = tidx
+		}
+		idxName := indexName
+		if idxName == "PRIMARY" {
+			idxName = "PRIMARY_KEY"
+		}
+		idx, ok := tidx.indexes[idxName]
+		if !ok {
+			idx = &Index{Name: indexName, NonUnique: nonUnique}
+			tidx.indexes[idxName] = idx
+		}
+		idx.Columns = append(idx.Columns, columnName)
+	}
+	return idxs, nil
+}
+
+func (dbh *DBHelper) getForeignKeys() (map[string][]*ForeignKey, error) {
+	vtConn, cancel, err := dbh.GetConnection()
+	if err != nil {
+		return nil, err
+	}
+	defer cancel()
+
+	fks := make(map[string][]*ForeignKey)
+	queryForeignKeys := "select table_name, column_name, constraint_name, referenced_table_name, referenced_column_name from information_schema.key_column_usage where table_schema = '%s' and referenced_table_name is not null"
+	query := fmt.Sprintf(queryForeignKeys, dbh.vtParams.DbName)
+	qr, err := vtConn.ExecuteFetch(query, -1, false)
+	if err != nil {
+		return nil, err
+	}
+	for _, row := range qr.Rows {
+		tableName := row[0].ToString()
+		fk := &ForeignKey{
+			ColumnName:           row[1].ToString(),
+			ConstraintName:       row[2].ToString(),
+			ReferencedTableName:  row[3].ToString(),
+			ReferencedColumnName: row[4].ToString(),
+		}
+		fks[tableName] = append(fks[tableName], fk)
+	}
+	return fks, nil
+}