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

Signed-off-by: Rohit Nayak <[email protected]>
vitessio · Nov 30, 2024 · e76c125 · e76c125
1 parent a2d740d
commit e76c125
Show file tree

Hide file tree

Showing 2 changed files with 202 additions and 0 deletions.
diff --git a/go/dbinfo/dbinfo_test.go b/go/dbinfo/dbinfo_test.go
@@ -144,4 +144,68 @@ 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.Equal(t, 16, idxs.len())
+		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"].indexName)
+		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"].indexName)
+		require.Equal(t, []string{"rental_date", "inventory_id", "customer_id"}, idx.indexes["rental_date"].columns)
+		require.Equal(t, "PRIMARY", idx.indexes["PRIMARY_KEY"].indexName)
+		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, "city", fk[0].tableName)
+		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, "store", fk[0].tableName)
+		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, "store", fk[1].tableName)
+		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)
+	})
 }
diff --git a/go/dbinfo/utils.go b/go/dbinfo/utils.go
@@ -130,3 +130,141 @@ 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
+}
+type index struct {
+	indexName string
+	columns   []string
+	nonUnique bool
+}
+
+type indexes map[string]*tableIndex
+
+func (i *indexes) len() int {
+	return len(*i)
+}
+
+func (dbh *DBHelper) getIndexes() (indexes, error) {
+	vtConn, cancel, err := dbh.GetConnection()
+	if err != nil {
+		return nil, err
+	}
+	defer cancel()
+
+	idxs := make(indexes)
+	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{indexName: indexName, nonUnique: nonUnique}
+			tidx.indexes[idxName] = idx
+		}
+		idx.columns = append(idx.columns, columnName)
+	}
+	return idxs, nil
+}
+
+// Foreign Keys with Dependency Information
+// SELECT TABLE_SCHEMA, TABLE_NAME, COLUMN_NAME, CONSTRAINT_NAME, REFERENCED_TABLE_SCHEMA, REFERENCED_TABLE_NAME, REFERENCED_COLUMN_NAME
+// FROM INFORMATION_SCHEMA.KEY_COLUMN_USAGE
+// WHERE TABLE_SCHEMA = 'your_database_name' AND REFERENCED_TABLE_NAME IS NOT NULL;
+
+type foreignKey struct {
+	tableName            string
+	columnName           string
+	constraintName       string
+	referencedTableName  string
+	referencedColumnName string
+}
+
+func (fk *foreignKey) String() string {
+	return fmt.Sprintf("Table: %s, Column: %s, Constraint: %s, RefTable: %s, RefColumn: %s", fk.tableName, fk.columnName, fk.constraintName, fk.referencedTableName, fk.referencedColumnName)
+}
+
+type foreignKeys map[string][]*foreignKey
+
+func (dbh *DBHelper) getForeignKeys() (foreignKeys, error) {
+	vtConn, cancel, err := dbh.GetConnection()
+	if err != nil {
+		return nil, err
+	}
+	defer cancel()
+
+	fks := make(foreignKeys)
+	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()
+		columnName := row[1].ToString()
+		constraintName := row[2].ToString()
+		referencedTableName := row[3].ToString()
+		referencedColumnName := row[4].ToString()
+		fk := &foreignKey{
+			tableName:            tableName,
+			columnName:           columnName,
+			constraintName:       constraintName,
+			referencedTableName:  referencedTableName,
+			referencedColumnName: referencedColumnName,
+		}
+		fks[tableName] = append(fks[tableName], fk)
+	}
+	return fks, nil
+}