Skip to content

Commit

Permalink
Merge pull request #77 from planetscale/dbinfo2
Browse files Browse the repository at this point in the history 
dbinfo: add info for primary keys, indexes, foreign keys
  • Loading branch information
@rohit-nayak-ps
rohit-nayak-ps authored Dec 2, 2024
2 parents a2d740d + 87a7484 commit 6eda568
Show file tree
Hide file tree
Showing 4 changed files with 1,216 additions and 703 deletions.
141 changes: 114 additions & 27 deletions go/dbinfo/dbinfo.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@ import (
"encoding/json"
"io"
"os"
"sort"

"vitess.io/vitess/go/mysql"
)
Expand Down Expand Up @@ -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 {
Expand All @@ -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,
Expand All @@ -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",
Expand Down
64 changes: 63 additions & 1 deletion go/dbinfo/dbinfo_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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)
})
}
103 changes: 103 additions & 0 deletions go/dbinfo/utils.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
}
Loading

0 comments on commit 6eda568

Please sign in to comment.