pgsql-broker/cmd/broker/db.go

package main

import (
	"bufio"
	"fmt"
	"os"
	"strconv"
	"strings"
	"time"

	"github.com/spf13/cobra"
	"golang.org/x/term"

	"git.warky.dev/wdevs/pgsql-broker/pkg/broker/config"
)

// Flags shared by the `db` subcommands.
var (
	dbConfigPath string

	dbAddFrom            string
	dbAddName            string
	dbAddHost            string
	dbAddPort            int
	dbAddDatabase        string
	dbAddUser            string
	dbAddPassword        string
	dbAddSSLMode         string
	dbAddMaxOpenConns    int
	dbAddMaxIdleConns    int
	dbAddConnMaxLifetime string
	dbAddConnMaxIdleTime string
	dbAddQueueCount      int
	dbAddTenantID        string
	dbAddAutoMigrate     bool
	dbAddDisabled        bool
	dbAddNonInteractive  bool
	dbYes                bool
)

var dbCmd = &cobra.Command{
	Use:   "db",
	Short: "Manage databases in the broker config file",
	Long:  `Add, remove, enable, disable, or list the databases configured in the broker config file.`,
}

var dbAddCmd = &cobra.Command{
	Use:   "add [name]",
	Short: "Add a database to the config file",
	Long: `Add a database entry to the config file.

Every field can be supplied via flags for fully non-interactive use. Any
field left unset falls back to an interactive prompt (when stdin is a
terminal) that is pre-filled with the value from --from (a named existing
instance) or, if --from is not given, the last database currently in the
file. Press Enter at a prompt to accept the shown value.`,
	Args: cobra.MaximumNArgs(1),
	RunE: func(cmd *cobra.Command, args []string) error {
		if len(args) == 1 && dbAddName == "" {
			dbAddName = args[0]
		}
		return runDBAdd(cmd)
	},
}

var dbRemoveCmd = &cobra.Command{
	Use:   "remove <name>",
	Short: "Remove a database from the config file",
	Args:  cobra.ExactArgs(1),
	RunE: func(cmd *cobra.Command, args []string) error {
		return runDBRemove(args[0])
	},
}

var dbEnableCmd = &cobra.Command{
	Use:   "enable <name>",
	Short: "Enable a database (removes the disabled flag)",
	Args:  cobra.ExactArgs(1),
	RunE: func(cmd *cobra.Command, args []string) error {
		return runDBSetDisabled(args[0], false)
	},
}

var dbDisableCmd = &cobra.Command{
	Use:   "disable <name>",
	Short: "Disable a database (excluded from `start` until re-enabled)",
	Args:  cobra.ExactArgs(1),
	RunE: func(cmd *cobra.Command, args []string) error {
		return runDBSetDisabled(args[0], true)
	},
}

var dbListCmd = &cobra.Command{
	Use:   "list",
	Short: "List databases configured in the config file",
	Args:  cobra.NoArgs,
	RunE: func(cmd *cobra.Command, args []string) error {
		return runDBList()
	},
}

func init() {
	rootCmd.AddCommand(dbCmd)
	dbCmd.AddCommand(dbAddCmd, dbRemoveCmd, dbEnableCmd, dbDisableCmd, dbListCmd)

	dbCmd.PersistentFlags().StringVar(&dbConfigPath, "config", "broker.yaml", "config file to edit")

	dbAddCmd.Flags().StringVar(&dbAddFrom, "from", "", "prime defaults from this existing database name (default: last database in the file)")
	dbAddCmd.Flags().StringVar(&dbAddName, "name", "", "unique name for the new database (required)")
	dbAddCmd.Flags().StringVar(&dbAddHost, "host", "", "PostgreSQL host")
	dbAddCmd.Flags().IntVar(&dbAddPort, "port", 0, "PostgreSQL port (default 5432)")
	dbAddCmd.Flags().StringVar(&dbAddDatabase, "database", "", "database name")
	dbAddCmd.Flags().StringVar(&dbAddUser, "user", "", "database user")
	dbAddCmd.Flags().StringVar(&dbAddPassword, "password", "", "database password")
	dbAddCmd.Flags().StringVar(&dbAddSSLMode, "sslmode", "", "SSL mode (disable, require, verify-ca, verify-full)")
	dbAddCmd.Flags().IntVar(&dbAddMaxOpenConns, "max-open-conns", 0, "max open connections (default 25)")
	dbAddCmd.Flags().IntVar(&dbAddMaxIdleConns, "max-idle-conns", 0, "max idle connections (default 5)")
	dbAddCmd.Flags().StringVar(&dbAddConnMaxLifetime, "conn-max-lifetime", "", "connection max lifetime (e.g. 5m, default 5m)")
	dbAddCmd.Flags().StringVar(&dbAddConnMaxIdleTime, "conn-max-idle-time", "", "connection max idle time (e.g. 10m, default 10m)")
	dbAddCmd.Flags().IntVar(&dbAddQueueCount, "queue-count", 0, "number of concurrent queues (default 4)")
	dbAddCmd.Flags().StringVar(&dbAddTenantID, "tenant-id", "", "RLS tenant id (default \"default\")")
	dbAddCmd.Flags().BoolVar(&dbAddAutoMigrate, "auto-migrate", false, "apply pending migrations automatically on connect")
	dbAddCmd.Flags().BoolVar(&dbAddDisabled, "disabled", false, "add the database in a disabled state")
	dbAddCmd.Flags().BoolVar(&dbAddNonInteractive, "non-interactive", false, "fail instead of prompting for missing fields")
	dbAddCmd.Flags().BoolVarP(&dbYes, "yes", "y", false, "skip the confirmation prompt")

	dbRemoveCmd.Flags().BoolVarP(&dbYes, "yes", "y", false, "skip the confirmation prompt")
}

// runDBAdd builds a new DatabaseConfig by layering (in priority order)
// explicit flags, then interactive prompts primed from --from / the last
// database in the file, then falls back to leaving optional fields unset
// so config.applyDatabaseDefaults keeps handling defaults at load time.
func runDBAdd(cmd *cobra.Command) error {
	doc, err := config.LoadFileDoc(dbConfigPath)
	if err != nil {
		return err
	}

	var prime config.DatabaseConfig
	if dbAddFrom != "" {
		found, ok, err := doc.FindDatabase(dbAddFrom)
		if err != nil {
			return err
		}
		if !ok {
			return fmt.Errorf("--from instance %q not found in %s", dbAddFrom, dbConfigPath)
		}
		prime = found
	} else if last, ok, err := doc.LastDatabase(); err != nil {
		return err
	} else if ok {
		prime = last
	}
	// Never prime the new entry's name from an existing one.
	prime.Name = ""

	interactive := !dbAddNonInteractive && term.IsTerminal(int(os.Stdin.Fd()))
	reader := bufio.NewReader(os.Stdin)

	if interactive {
		fmt.Fprintln(os.Stderr, "Adding a new database. Press Enter to accept the shown value.")
	}

	db := config.DatabaseConfig{}

	if dbAddName != "" {
		db.Name = dbAddName
	} else {
		db.Name, err = resolveStringField(cmd, reader, interactive, "name", "", true, dbAddName)
		if err != nil {
			return err
		}
	}
	if db.Name == "" {
		return fmt.Errorf("name is required (--name, a positional argument, or run interactively)")
	}
	db.Host, err = resolveStringField(cmd, reader, interactive, "host", prime.Host, true, dbAddHost)
	if err != nil {
		return err
	}
	db.Database, err = resolveStringField(cmd, reader, interactive, "database", prime.Database, true, dbAddDatabase)
	if err != nil {
		return err
	}
	db.User, err = resolveStringField(cmd, reader, interactive, "user", prime.User, true, dbAddUser)
	if err != nil {
		return err
	}
	db.Password, err = resolveStringField(cmd, reader, interactive, "password", prime.Password, false, dbAddPassword)
	if err != nil {
		return err
	}
	db.SSLMode, err = resolveStringField(cmd, reader, interactive, "sslmode", prime.SSLMode, false, dbAddSSLMode)
	if err != nil {
		return err
	}
	db.TenantID, err = resolveStringField(cmd, reader, interactive, "tenant_id", prime.TenantID, false, dbAddTenantID)
	if err != nil {
		return err
	}

	port, err := resolveIntField(cmd, reader, interactive, "port", prime.Port, "port", dbAddPort)
	if err != nil {
		return err
	}
	db.Port = port

	maxOpen, err := resolveIntField(cmd, reader, interactive, "max_open_conns", prime.MaxOpenConns, "max-open-conns", dbAddMaxOpenConns)
	if err != nil {
		return err
	}
	db.MaxOpenConns = maxOpen

	maxIdle, err := resolveIntField(cmd, reader, interactive, "max_idle_conns", prime.MaxIdleConns, "max-idle-conns", dbAddMaxIdleConns)
	if err != nil {
		return err
	}
	db.MaxIdleConns = maxIdle

	queueCount, err := resolveIntField(cmd, reader, interactive, "queue_count", prime.QueueCount, "queue-count", dbAddQueueCount)
	if err != nil {
		return err
	}
	db.QueueCount = queueCount

	connMaxLifetime, err := resolveDurationField(cmd, reader, interactive, "conn_max_lifetime", prime.ConnMaxLifetime, "conn-max-lifetime", dbAddConnMaxLifetime)
	if err != nil {
		return err
	}
	db.ConnMaxLifetime = connMaxLifetime

	connMaxIdleTime, err := resolveDurationField(cmd, reader, interactive, "conn_max_idle_time", prime.ConnMaxIdleTime, "conn-max-idle-time", dbAddConnMaxIdleTime)
	if err != nil {
		return err
	}
	db.ConnMaxIdleTime = connMaxIdleTime

	if cmd.Flags().Changed("auto-migrate") {
		db.AutoMigrate = dbAddAutoMigrate
	} else {
		db.AutoMigrate = prime.AutoMigrate
	}
	if cmd.Flags().Changed("disabled") {
		db.Disabled = dbAddDisabled
	} else {
		db.Disabled = prime.Disabled
	}

	for field, value := range map[string]string{
		"host":     db.Host,
		"database": db.Database,
		"user":     db.User,
	} {
		if strings.TrimSpace(value) == "" {
			return fmt.Errorf("%s is required (--%s or an interactive value)", field, strings.ReplaceAll(field, "_", "-"))
		}
	}

	if err := doc.AddDatabase(db); err != nil {
		return err
	}

	if !dbYes && interactive {
		fmt.Fprintf(os.Stderr, "\nAbout to add database %q to %s:\n", db.Name, dbConfigPath)
		printDBSummary(os.Stderr, db)
		fmt.Fprint(os.Stderr, "Proceed? [Y/n]: ")
		line, _ := reader.ReadString('\n')
		line = strings.TrimSpace(strings.ToLower(line))
		if line == "n" || line == "no" {
			fmt.Println("Aborted; config file not modified.")
			return nil
		}
	}

	if err := doc.Save(); err != nil {
		return err
	}
	fmt.Printf("Added database %q to %s\n", db.Name, dbConfigPath)
	return nil
}

func runDBRemove(name string) error {
	doc, err := config.LoadFileDoc(dbConfigPath)
	if err != nil {
		return err
	}
	db, ok, err := doc.FindDatabase(name)
	if err != nil {
		return err
	}
	if !ok {
		return fmt.Errorf("database %q not found in %s", name, dbConfigPath)
	}

	if !dbYes {
		if !term.IsTerminal(int(os.Stdin.Fd())) {
			return fmt.Errorf("refusing to remove %q non-interactively without --yes", name)
		}
		fmt.Fprintf(os.Stderr, "Remove database %q (%s@%s:%d/%s) from %s? [y/N]: ",
			name, db.User, db.Host, db.Port, db.Database, dbConfigPath)
		reader := bufio.NewReader(os.Stdin)
		line, _ := reader.ReadString('\n')
		line = strings.TrimSpace(strings.ToLower(line))
		if line != "y" && line != "yes" {
			fmt.Println("Aborted; config file not modified.")
			return nil
		}
	}

	if _, err := doc.RemoveDatabase(name); err != nil {
		return err
	}
	if err := doc.Save(); err != nil {
		return err
	}
	fmt.Printf("Removed database %q from %s\n", name, dbConfigPath)
	return nil
}

func runDBSetDisabled(name string, disabled bool) error {
	doc, err := config.LoadFileDoc(dbConfigPath)
	if err != nil {
		return err
	}
	found, err := doc.SetDisabled(name, disabled)
	if err != nil {
		return err
	}
	if !found {
		return fmt.Errorf("database %q not found in %s", name, dbConfigPath)
	}
	if err := doc.Save(); err != nil {
		return err
	}
	verb := "Enabled"
	if disabled {
		verb = "Disabled"
	}
	fmt.Printf("%s database %q in %s\n", verb, name, dbConfigPath)
	return nil
}

func runDBList() error {
	doc, err := config.LoadFileDoc(dbConfigPath)
	if err != nil {
		return err
	}
	dbs, err := doc.ListDatabases()
	if err != nil {
		return err
	}
	if len(dbs) == 0 {
		fmt.Printf("No databases configured in %s\n", dbConfigPath)
		return nil
	}
	fmt.Printf("%-20s %-20s %-6s %-16s %-16s %-10s\n", "NAME", "HOST", "PORT", "DATABASE", "USER", "STATUS")
	for _, db := range dbs {
		status := "enabled"
		if db.Disabled {
			status = "disabled"
		}
		port := db.Port
		if port == 0 {
			port = 5432
		}
		fmt.Printf("%-20s %-20s %-6d %-16s %-16s %-10s\n", db.Name, db.Host, port, db.Database, db.User, status)
	}
	return nil
}

func printDBSummary(w *os.File, db config.DatabaseConfig) {
	password := ""
	if db.Password != "" {
		password = "(set)"
	}
	fmt.Fprintf(w, "  host:              %s\n", db.Host)
	fmt.Fprintf(w, "  port:              %s\n", intOrDefault(db.Port, 5432))
	fmt.Fprintf(w, "  database:          %s\n", db.Database)
	fmt.Fprintf(w, "  user:              %s\n", db.User)
	fmt.Fprintf(w, "  password:          %s\n", password)
	fmt.Fprintf(w, "  sslmode:           %s\n", stringOrDefault(db.SSLMode, "disable"))
	fmt.Fprintf(w, "  max_open_conns:    %s\n", intOrDefault(db.MaxOpenConns, 25))
	fmt.Fprintf(w, "  max_idle_conns:    %s\n", intOrDefault(db.MaxIdleConns, 5))
	fmt.Fprintf(w, "  conn_max_lifetime: %s\n", durationOrDefault(db.ConnMaxLifetime, 5*time.Minute))
	fmt.Fprintf(w, "  conn_max_idle_time:%s\n", durationOrDefault(db.ConnMaxIdleTime, 10*time.Minute))
	fmt.Fprintf(w, "  queue_count:       %s\n", intOrDefault(db.QueueCount, 4))
	fmt.Fprintf(w, "  tenant_id:         %s\n", stringOrDefault(db.TenantID, "default"))
	fmt.Fprintf(w, "  auto_migrate:      %t\n", db.AutoMigrate)
	fmt.Fprintf(w, "  disabled:          %t\n", db.Disabled)
}

func intOrDefault(v, def int) string {
	if v == 0 {
		return fmt.Sprintf("%d (default)", def)
	}
	return strconv.Itoa(v)
}

func stringOrDefault(v, def string) string {
	if v == "" {
		return fmt.Sprintf("%s (default)", def)
	}
	return v
}

func durationOrDefault(v, def time.Duration) string {
	if v == 0 {
		return fmt.Sprintf("%s (default)", def)
	}
	return v.String()
}

// resolveStringField returns, in priority order: the flag value if the
// flag was explicitly set, otherwise an interactive prompt (pre-filled
// with defaultVal) when interactive, otherwise defaultVal. When required
// and still empty after all of that, an error is returned by the caller.
func resolveStringField(cmd *cobra.Command, reader *bufio.Reader, interactive bool, label, defaultVal string, required bool, flagVal string) (string, error) {
	flagName := flagNameFor(label)
	if cmd.Flags().Changed(flagName) {
		return flagVal, nil
	}
	if !interactive {
		return defaultVal, nil
	}
	for {
		fmt.Fprintf(os.Stderr, "%s [%s]: ", label, defaultVal)
		line, _ := reader.ReadString('\n')
		line = strings.TrimSpace(line)
		if line == "" {
			line = defaultVal
		}
		if line == "" && required {
			fmt.Fprintln(os.Stderr, "value is required")
			continue
		}
		return line, nil
	}
}

func resolveIntField(cmd *cobra.Command, reader *bufio.Reader, interactive bool, label string, defaultVal int, flagName string, flagVal int) (int, error) {
	if cmd.Flags().Changed(flagName) {
		return flagVal, nil
	}
	if !interactive {
		return defaultVal, nil
	}
	for {
		fmt.Fprintf(os.Stderr, "%s [%d]: ", label, defaultVal)
		line, _ := reader.ReadString('\n')
		line = strings.TrimSpace(line)
		if line == "" {
			return defaultVal, nil
		}
		n, err := strconv.Atoi(line)
		if err != nil {
			fmt.Fprintln(os.Stderr, "must be a whole number")
			continue
		}
		return n, nil
	}
}

func resolveDurationField(cmd *cobra.Command, reader *bufio.Reader, interactive bool, label string, defaultVal time.Duration, flagName, flagVal string) (time.Duration, error) {
	if cmd.Flags().Changed(flagName) {
		d, err := time.ParseDuration(flagVal)
		if err != nil {
			return 0, fmt.Errorf("--%s: %w", flagName, err)
		}
		return d, nil
	}
	if !interactive {
		return defaultVal, nil
	}
	for {
		fmt.Fprintf(os.Stderr, "%s [%s]: ", label, defaultVal)
		line, _ := reader.ReadString('\n')
		line = strings.TrimSpace(line)
		if line == "" {
			return defaultVal, nil
		}
		d, err := time.ParseDuration(line)
		if err != nil {
			fmt.Fprintln(os.Stderr, "must be a duration like 5m, 30s")
			continue
		}
		return d, nil
	}
}

// flagNameFor maps a prompt label to its cobra flag name for the string
// fields prompted in runDBAdd.
func flagNameFor(label string) string {
	switch label {
	case "name":
		return "name"
	case "host":
		return "host"
	case "database":
		return "database"
	case "user":
		return "user"
	case "password":
		return "password"
	case "sslmode":
		return "sslmode"
	case "tenant_id":
		return "tenant-id"
	default:
		return label
	}
}