✨ Add shell completion for --dbname and --table

clevyr · Mar 24, 2022 · e42e3ab · e42e3ab
1 parent a617986
commit e42e3ab
Show file tree

Hide file tree

Showing 5 changed files with 84 additions and 3 deletions.
diff --git a/internal/config/database.go b/internal/config/database.go
@@ -12,6 +12,8 @@ type Databaser interface {
 	DefaultPort() uint16
 
 	DatabaseEnvNames() []string
+	ListDatabasesQuery() string
+	ListTablesQuery() string
 	DefaultDatabase() string
 
 	UserEnvNames() []string

diff --git a/internal/config/exec.go b/internal/config/exec.go
@@ -2,4 +2,5 @@ package config
 
 type Exec struct {
 	Global
+	DisableHeaders bool
 }
diff --git a/internal/config/flags/database.go b/internal/config/flags/database.go
@@ -1,6 +1,11 @@
 package flags
 
-import "github.com/spf13/cobra"
+import (
+	"github.com/clevyr/kubedb/internal/config"
+	"github.com/clevyr/kubedb/internal/util"
+	"github.com/spf13/cobra"
+	"strings"
+)
 
 func Grammar(cmd *cobra.Command) {
 	cmd.PersistentFlags().String("grammar", "", "database grammar. detected if not set. (postgres, mariadb)")
@@ -28,6 +33,10 @@ func Format(cmd *cobra.Command) {
 
 func Database(cmd *cobra.Command) {
 	cmd.PersistentFlags().StringP("dbname", "d", "", "database name to connect to")
+	err := cmd.RegisterFlagCompletionFunc("dbname", listDatabases)
+	if err != nil {
+		panic(err)
+	}
 }
 
 func Username(cmd *cobra.Command) {
@@ -56,12 +65,57 @@ func NoOwner(cmd *cobra.Command, p *bool) {
 
 func Tables(cmd *cobra.Command, p *[]string) {
 	cmd.Flags().StringSliceVarP(p, "table", "t", []string{}, "dump the specified table(s) only")
+	err := cmd.RegisterFlagCompletionFunc("table", listTables)
+	if err != nil {
+		panic(err)
+	}
 }
 
 func ExcludeTable(cmd *cobra.Command, p *[]string) {
 	cmd.Flags().StringSliceVarP(p, "exclude-table", "T", []string{}, "do NOT dump the specified table(s)")
+	err := cmd.RegisterFlagCompletionFunc("exclude-table", listTables)
+	if err != nil {
+		panic(err)
+	}
 }
 
 func ExcludeTableData(cmd *cobra.Command, p *[]string) {
 	cmd.Flags().StringSliceVar(p, "exclude-table-data", []string{}, "do NOT dump data for the specified table(s)")
+	err := cmd.RegisterFlagCompletionFunc("exclude-table-data", listTables)
+	if err != nil {
+		panic(err)
+	}
+}
+
+func listTables(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) {
+	conf := config.Exec{DisableHeaders: true}
+	err := util.DefaultSetup(cmd, &conf.Global)
+	if err != nil {
+		return nil, cobra.ShellCompDirectiveNoFileComp
+	}
+	return queryInDatabase(cmd, args, conf, conf.Grammar.ListTablesQuery())
+}
+
+func listDatabases(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) {
+	conf := config.Exec{DisableHeaders: true}
+	err := util.DefaultSetup(cmd, &conf.Global)
+	if err != nil {
+		return nil, cobra.ShellCompDirectiveNoFileComp
+	}
+	return queryInDatabase(cmd, args, conf, conf.Grammar.ListDatabasesQuery())
+}
+
+func queryInDatabase(cmd *cobra.Command, args []string, conf config.Exec, query string) ([]string, cobra.ShellCompDirective) {
+	r := strings.NewReader(query)
+	var buf strings.Builder
+	sqlCmd := conf.Grammar.ExecCommand(conf)
+	err := conf.Client.Exec(conf.Pod, []string{"sh", "-c", sqlCmd.Join()}, r, &buf, false)
+	if err != nil {
+		return nil, cobra.ShellCompDirectiveNoFileComp
+	}
+
+	cobra.CompDebugln(buf.String(), true)
+
+	names := strings.Split(buf.String(), "\n")
+	return names, cobra.ShellCompDirectiveNoFileComp
 }
diff --git a/internal/database/grammar/mariadb.go b/internal/database/grammar/mariadb.go
@@ -22,6 +22,14 @@ func (MariaDB) DatabaseEnvNames() []string {
 	return []string{"MARIADB_DATABASE"}
 }
 
+func (MariaDB) ListDatabasesQuery() string {
+	return "show databases"
+}
+
+func (MariaDB) ListTablesQuery() string {
+	return "show tables"
+}
+
 func (MariaDB) DefaultDatabase() string {
 	return "db"
 }
@@ -61,10 +69,14 @@ func (MariaDB) PasswordEnvNames() []string {
 }
 
 func (MariaDB) ExecCommand(conf config.Exec) *command.Builder {
-	return command.NewBuilder(
+	cmd := command.NewBuilder(
 		command.NewEnv("MYSQL_PWD", conf.Password),
 		"mysql", "--host=127.0.0.1", "--user="+conf.Username, "--database="+conf.Database,
 	)
+	if conf.DisableHeaders {
+		cmd.Push("--skip-column-names")
+	}
+	return cmd
 }
 
 func (MariaDB) DumpCommand(conf config.Dump) *command.Builder {

diff --git a/internal/database/grammar/postgres.go b/internal/database/grammar/postgres.go
@@ -26,6 +26,14 @@ func (Postgres) DatabaseEnvNames() []string {
 	return []string{"POSTGRES_DB"}
 }
 
+func (Postgres) ListDatabasesQuery() string {
+	return "SELECT datname FROM pg_database WHERE datistemplate = false"
+}
+
+func (Postgres) ListTablesQuery() string {
+	return "SELECT table_name FROM information_schema.tables WHERE table_schema='public' AND table_type='BASE TABLE'"
+}
+
 func (Postgres) DefaultDatabase() string {
 	return "db"
 }
@@ -103,10 +111,14 @@ func (Postgres) PasswordEnvNames() []string {
 }
 
 func (Postgres) ExecCommand(conf config.Exec) *command.Builder {
-	return command.NewBuilder(
+	cmd := command.NewBuilder(
 		command.NewEnv("PGPASSWORD", conf.Password),
 		"psql", "--host=127.0.0.1", "--username="+conf.Username, "--dbname="+conf.Database,
 	)
+	if conf.DisableHeaders {
+		cmd.Push("--tuples-only")
+	}
+	return cmd
 }
 
 func quoteParam(param string) string {