package main import ( "context" "errors" "flag" "fmt" "log/slog" "os" "path/filepath" "strings" "syscall" "databasus-agent/internal/config" "databasus-agent/internal/features/api" "databasus-agent/internal/features/restore" "databasus-agent/internal/features/start" "databasus-agent/internal/features/upgrade" "databasus-agent/internal/logger" ) var Version = "dev" func main() { if len(os.Args) < 2 { printUsage() os.Exit(1) } switch os.Args[1] { case "start": runStart(os.Args[2:]) case "_run": runDaemon(os.Args[2:]) case "stop": runStop() case "status": runStatus() case "restore": runRestore(os.Args[2:]) case "version": fmt.Println(Version) default: fmt.Fprintf(os.Stderr, "unknown command: %s\n", os.Args[1]) printUsage() os.Exit(1) } } func runStart(args []string) { fs := flag.NewFlagSet("start", flag.ExitOnError) isSkipUpdate := fs.Bool("skip-update", false, "Skip auto-update check") cfg := &config.Config{} cfg.LoadFromJSONAndArgs(fs, args) if err := cfg.SaveToJSON(); err != nil { fmt.Fprintf(os.Stderr, "Failed to save config: %v\n", err) } log := logger.GetLogger() isDev := checkIsDevelopment() runUpdateCheck(cfg.DatabasusHost, *isSkipUpdate, isDev, log) if err := start.Start(cfg, Version, isDev, log); err != nil { if errors.Is(err, upgrade.ErrUpgradeRestart) { reexecAfterUpgrade(log) } fmt.Fprintf(os.Stderr, "Error: %v\n", err) os.Exit(1) } } func runDaemon(args []string) { fs := flag.NewFlagSet("_run", flag.ExitOnError) if err := fs.Parse(args); err != nil { os.Exit(1) } log := logger.GetLogger() cfg := &config.Config{} cfg.LoadFromJSON() if err := start.RunDaemon(cfg, Version, checkIsDevelopment(), log); err != nil { if errors.Is(err, upgrade.ErrUpgradeRestart) { reexecAfterUpgrade(log) } log.Error("Agent exited with error", "error", err) os.Exit(1) } } func runStop() { log := logger.GetLogger() if err := start.Stop(log); err != nil { fmt.Fprintf(os.Stderr, "Error: %v\n", err) os.Exit(1) } } func runStatus() { log := logger.GetLogger() if err := start.Status(log); err != nil { fmt.Fprintf(os.Stderr, "Error: %v\n", err) os.Exit(1) } } func runRestore(args []string) { fs := flag.NewFlagSet("restore", flag.ExitOnError) pgDataDir := fs.String("target-dir", "", "Target pgdata directory (required)") backupID := fs.String("backup-id", "", "Full backup UUID (optional)") targetTime := fs.String("target-time", "", "PITR target time in RFC3339 (optional)") isSkipUpdate := fs.Bool("skip-update", false, "Skip auto-update check") cfg := &config.Config{} cfg.LoadFromJSONAndArgs(fs, args) if err := cfg.SaveToJSON(); err != nil { fmt.Fprintf(os.Stderr, "Failed to save config: %v\n", err) } log := logger.GetLogger() isDev := checkIsDevelopment() runUpdateCheck(cfg.DatabasusHost, *isSkipUpdate, isDev, log) if *pgDataDir == "" { fmt.Fprintln(os.Stderr, "Error: --target-dir is required") os.Exit(1) } if cfg.DatabasusHost == "" || cfg.Token == "" { fmt.Fprintln(os.Stderr, "Error: databasus-host and token must be configured") os.Exit(1) } if cfg.PgType != "host" && cfg.PgType != "docker" { fmt.Fprintf(os.Stderr, "Error: --pg-type must be 'host' or 'docker', got '%s'\n", cfg.PgType) os.Exit(1) } apiClient := api.NewClient(cfg.DatabasusHost, cfg.Token, log) restorer := restore.NewRestorer(apiClient, log, *pgDataDir, *backupID, *targetTime, cfg.PgType) ctx := context.Background() if err := restorer.Run(ctx); err != nil { fmt.Fprintf(os.Stderr, "Error: %v\n", err) os.Exit(1) } } func printUsage() { fmt.Fprintln(os.Stderr, "Usage: databasus-agent [flags]") fmt.Fprintln(os.Stderr, "") fmt.Fprintln(os.Stderr, "Commands:") fmt.Fprintln(os.Stderr, " start Start the agent (WAL archiving + basebackups)") fmt.Fprintln(os.Stderr, " stop Stop a running agent") fmt.Fprintln(os.Stderr, " status Show agent status") fmt.Fprintln(os.Stderr, " restore Restore a database from backup") fmt.Fprintln(os.Stderr, " version Print agent version") } func runUpdateCheck(host string, isSkipUpdate, isDev bool, log *slog.Logger) { if isSkipUpdate { return } if host == "" { return } apiClient := api.NewClient(host, "", log) isUpgraded, err := upgrade.CheckAndUpdate(apiClient, Version, isDev, log) if err != nil { log.Error("Auto-update failed", "error", err) os.Exit(1) } if isUpgraded { reexecAfterUpgrade(log) } } func checkIsDevelopment() bool { dir, err := os.Getwd() if err != nil { return false } for range 3 { if data, err := os.ReadFile(filepath.Join(dir, ".env")); err == nil { return parseEnvMode(data) } if _, err := os.Stat(filepath.Join(dir, "go.mod")); err == nil { return false } dir = filepath.Dir(dir) } return false } func parseEnvMode(data []byte) bool { for line := range strings.SplitSeq(string(data), "\n") { line = strings.TrimSpace(line) if line == "" || strings.HasPrefix(line, "#") { continue } parts := strings.SplitN(line, "=", 2) if len(parts) == 2 && strings.TrimSpace(parts[0]) == "ENV_MODE" { return strings.TrimSpace(parts[1]) == "development" } } return false } func reexecAfterUpgrade(log *slog.Logger) { selfPath, err := os.Executable() if err != nil { log.Error("Failed to resolve executable for re-exec", "error", err) os.Exit(1) } log.Info("Re-executing after upgrade...") if err := syscall.Exec(selfPath, os.Args, os.Environ()); err != nil { log.Error("Failed to re-exec after upgrade", "error", err) os.Exit(1) } }