Add database roles to tsh db ls -v (#35582)

* refactor and add --exclude-column

* add Database Roles column

* add database_roles to json,yaml

* remove --exclude-column and refactor databaseTableRow

* change slices import

tool/tsh/common/db.go

@@ -94,7 +94,7 @@ func onListDatabases(cf *CLIConf) error {
 	}
 
 	sort.Sort(types.Databases(databases))
-	return trace.Wrap(showDatabases(cf.Stdout(), cf.SiteName, databases, activeDatabases, accessChecker, cf.Format, cf.Verbose))
+	return trace.Wrap(showDatabases(cf, databases, activeDatabases, accessChecker))
 }
 
 func accessCheckerForRemoteCluster(ctx context.Context, profile *client.ProfileStatus, proxy *client.ProxyClient, clusterName string) (services.AccessChecker, error) {
@@ -230,7 +230,7 @@ func listDatabasesAllClusters(cf *CLIConf) error {
 	format := strings.ToLower(cf.Format)
 	switch format {
 	case teleport.Text, "":
-		printDatabasesWithClusters(cf.SiteName, dbListings, active, cf.Verbose)
+		printDatabasesWithClusters(cf, dbListings, active)
 	case teleport.JSON, teleport.YAML:
 		out, err := serializeDatabasesAllClusters(dbListings, format)
 		if err != nil {
@@ -1647,7 +1647,7 @@ func formatAmbiguousDB(cf *CLIConf, selectors resourceSelectors, matchedDBs type
 	var checker services.AccessChecker
 	var sb strings.Builder
 	verbose := true
-	showDatabasesAsText(&sb, cf.SiteName, matchedDBs, activeDBs, checker, verbose)
+	showDatabasesAsText(cf, &sb, matchedDBs, activeDBs, checker, verbose)
 
 	listCommand := formatDatabaseListCommand(cf.SiteName)
 	fullNameExample := matchedDBs[0].GetName()

tool/tsh/common/db_print.go

@@ -0,0 +1,135 @@
+/*
+ * Teleport
+ * Copyright (C) 2023  Gravitational, Inc.
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU Affero General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU Affero General Public License for more details.
+ *
+ * You should have received a copy of the GNU Affero General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ */
+
+package common
+
+import (
+	"fmt"
+	"io"
+	"reflect"
+	"regexp"
+	"slices"
+
+	"github.com/gravitational/teleport/api/types"
+	"github.com/gravitational/teleport/lib/asciitable"
+	"github.com/gravitational/teleport/lib/services"
+)
+
+type databaseTableRow struct {
+	Proxy         string
+	Cluster       string
+	DisplayName   string `title:"Name"`
+	Description   string
+	Protocol      string
+	Type          string
+	URI           string
+	AllowedUsers  string
+	DatabaseRoles string
+	Labels        string
+	Connect       string
+}
+
+func makeTableColumnTitles(row any) (out []string) {
+	// Regular expression to convert from "DatabaseRoles" to "Database Roles" etc.
+	re := regexp.MustCompile(`([a-z])([A-Z])`)
+
+	t := reflect.TypeOf(row)
+	for i := 0; i < t.NumField(); i++ {
+		field := t.Field(i)
+		title := field.Tag.Get("title")
+		if title == "" {
+			title = re.ReplaceAllString(field.Name, "${1} ${2}")
+		}
+		out = append(out, title)
+	}
+	return out
+}
+
+func makeTableRows[T any](rows []T) [][]string {
+	out := make([][]string, 0, len(rows))
+	for _, row := range rows {
+		var columnValues []string
+		v := reflect.ValueOf(row)
+		for i := 0; i < v.NumField(); i++ {
+			columnValues = append(columnValues, fmt.Sprintf("%v", v.Field(i)))
+		}
+		out = append(out, columnValues)
+	}
+	return out
+}
+
+type printDatabaseTableConfig struct {
+	writer              io.Writer
+	rows                []databaseTableRow
+	showProxyAndCluster bool
+	verbose             bool
+}
+
+func (cfg printDatabaseTableConfig) excludeColumns() (out []string) {
+	if !cfg.showProxyAndCluster {
+		out = append(out, "Proxy", "Cluster")
+	}
+	if !cfg.verbose {
+		out = append(out, "Protocol", "Type", "URI", "Database Roles")
+	}
+	return out
+}
+
+func printDatabaseTable(cfg printDatabaseTableConfig) {
+	allColumns := makeTableColumnTitles(databaseTableRow{})
+	rowsWithAllColumns := makeTableRows(cfg.rows)
+	excludeColumns := cfg.excludeColumns()
+
+	var printColumns []string
+	printRows := make([][]string, len(cfg.rows))
+	for columnIndex, column := range allColumns {
+		if slices.Contains(excludeColumns, column) {
+			continue
+		}
+
+		printColumns = append(printColumns, column)
+		for rowIndex, row := range rowsWithAllColumns {
+			printRows[rowIndex] = append(printRows[rowIndex], row[columnIndex])
+		}
+	}
+
+	var t asciitable.Table
+	if cfg.verbose {
+		t = asciitable.MakeTable(printColumns, printRows...)
+	} else {
+		t = asciitable.MakeTableWithTruncatedColumn(printColumns, printRows, "Labels")
+	}
+	fmt.Fprintln(cfg.writer, t.AsBuffer().String())
+}
+
+func formatDatabaseRolesForDB(database types.Database, accessChecker services.AccessChecker) string {
+	if database.SupportsAutoUsers() && database.GetAdminUser().Name != "" {
+		// may happen if fetching the role set failed for any reason.
+		if accessChecker == nil {
+			return "(unknown)"
+		}
+
+		autoUser, roles, err := accessChecker.CheckDatabaseRoles(database)
+		if err != nil {
+			log.Warnf("Failed to CheckDatabaseRoles for database %v: %v.", database.GetName(), err)
+		} else if autoUser.IsEnabled() {
+			return fmt.Sprintf("%v", roles)
+		}
+	}
+	return ""
+}

tool/tsh/common/db_print_test.go

@@ -0,0 +1,198 @@
+/*
+ * Teleport
+ * Copyright (C) 2023  Gravitational, Inc.
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU Affero General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU Affero General Public License for more details.
+ *
+ * You should have received a copy of the GNU Affero General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ */
+
+package common
+
+import (
+	"strings"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	apidefaults "github.com/gravitational/teleport/api/defaults"
+	"github.com/gravitational/teleport/api/types"
+	"github.com/gravitational/teleport/lib/services"
+)
+
+func Test_printDatabaseTable(t *testing.T) {
+	t.Parallel()
+
+	rows := []databaseTableRow{
+		databaseTableRow{
+			Proxy:        "proxy",
+			Cluster:      "cluster1",
+			DisplayName:  "db1",
+			Description:  "describe db1",
+			Protocol:     "postgres",
+			Type:         "self-hosted",
+			URI:          "localhost:5432",
+			AllowedUsers: "[*]",
+			Labels:       "Env=dev",
+			Connect:      "tsh db connect db1",
+		},
+		databaseTableRow{
+			Proxy:         "proxy",
+			Cluster:       "cluster1",
+			DisplayName:   "db2",
+			Description:   "describe db2",
+			Protocol:      "mysql",
+			Type:          "self-hosted",
+			URI:           "localhost:3306",
+			AllowedUsers:  "[alice]",
+			DatabaseRoles: "[readonly]",
+			Labels:        "Env=prod",
+		},
+	}
+
+	tests := []struct {
+		name   string
+		cfg    printDatabaseTableConfig
+		expect string
+	}{
+		{
+			name: "tsh db ls",
+			cfg: printDatabaseTableConfig{
+				rows:                rows,
+				showProxyAndCluster: false,
+				verbose:             false,
+			},
+			// os.Stdin.Fd() fails during go test, so width is defaulted to 80 for truncated table.
+			expect: `Name Description  Allowed Users Labels   Connect             
+---- ------------ ------------- -------- ------------------- 
+db1  describe db1 [*]           Env=dev  tsh db connect d... 
+db2  describe db2 [alice]       Env=prod                     
+
+`,
+		},
+		{
+			name: "tsh db ls --verbose",
+			cfg: printDatabaseTableConfig{
+				rows:                rows,
+				showProxyAndCluster: false,
+				verbose:             true,
+			},
+			expect: `Name Description  Protocol Type        URI            Allowed Users Database Roles Labels   Connect            
+---- ------------ -------- ----------- -------------- ------------- -------------- -------- ------------------ 
+db1  describe db1 postgres self-hosted localhost:5432 [*]                          Env=dev  tsh db connect db1 
+db2  describe db2 mysql    self-hosted localhost:3306 [alice]       [readonly]     Env=prod                    
+
+`,
+		},
+		{
+			name: "tsh db ls --verbose --all",
+			cfg: printDatabaseTableConfig{
+				rows:                rows,
+				showProxyAndCluster: true,
+				verbose:             true,
+			},
+			expect: `Proxy Cluster  Name Description  Protocol Type        URI            Allowed Users Database Roles Labels   Connect            
+----- -------- ---- ------------ -------- ----------- -------------- ------------- -------------- -------- ------------------ 
+proxy cluster1 db1  describe db1 postgres self-hosted localhost:5432 [*]                          Env=dev  tsh db connect db1 
+proxy cluster1 db2  describe db2 mysql    self-hosted localhost:3306 [alice]       [readonly]     Env=prod                    
+
+`,
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			var sb strings.Builder
+
+			cfg := test.cfg
+			cfg.writer = &sb
+
+			printDatabaseTable(cfg)
+			require.Equal(t, test.expect, sb.String())
+		})
+	}
+}
+
+func Test_formatDatabaseRolesForDB(t *testing.T) {
+	t.Parallel()
+
+	db, err := types.NewDatabaseV3(types.Metadata{
+		Name: "db",
+	}, types.DatabaseSpecV3{
+		Protocol: "postgres",
+		URI:      "localhost:5432",
+	})
+	require.NoError(t, err)
+
+	dbWithAutoUser, err := types.NewDatabaseV3(types.Metadata{
+		Name:   "dbWithAutoUser",
+		Labels: map[string]string{"env": "prod"},
+	}, types.DatabaseSpecV3{
+		Protocol: "postgres",
+		URI:      "localhost:5432",
+		AdminUser: &types.DatabaseAdminUser{
+			Name: "teleport-admin",
+		},
+	})
+	require.NoError(t, err)
+
+	roleAutoUser := &types.RoleV6{
+		Metadata: types.Metadata{Name: "auto-user", Namespace: apidefaults.Namespace},
+		Spec: types.RoleSpecV6{
+			Options: types.RoleOptions{
+				CreateDatabaseUserMode: types.CreateDatabaseUserMode_DB_USER_MODE_KEEP,
+			},
+			Allow: types.RoleConditions{
+				Namespaces:     []string{apidefaults.Namespace},
+				DatabaseLabels: types.Labels{"env": []string{"prod"}},
+				DatabaseRoles:  []string{"roleA", "roleB"},
+				DatabaseNames:  []string{"*"},
+				DatabaseUsers:  []string{types.Wildcard},
+			},
+		},
+	}
+
+	tests := []struct {
+		name          string
+		database      types.Database
+		accessChecker services.AccessChecker
+		expect        string
+	}{
+		{
+			name:     "nil accessChecker",
+			database: dbWithAutoUser,
+			expect:   "(unknown)",
+		},
+		{
+			name:     "roles",
+			database: dbWithAutoUser,
+			accessChecker: services.NewAccessCheckerWithRoleSet(&services.AccessInfo{
+				Username: "alice",
+			}, "clustername", services.RoleSet{roleAutoUser}),
+			expect: "[roleA roleB]",
+		},
+		{
+			name:     "db without admin user",
+			database: db,
+			accessChecker: services.NewAccessCheckerWithRoleSet(&services.AccessInfo{
+				Username: "alice",
+			}, "clustername", services.RoleSet{roleAutoUser}),
+			expect: "",
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			require.Equal(t, test.expect, formatDatabaseRolesForDB(test.database, test.accessChecker))
+		})
+	}
+}