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 ", 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 ", 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 ", 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 } }