nats-upload: use cobra and add clean subcommand

main.go

@@ -3,66 +3,162 @@ package main
 import (
 	"context"
 	"errors"
-	"flag"
 	"fmt"
 	"log"
 	"os"
 	"path/filepath"
 	"sort"
-	"strconv"
 	"strings"
+	"time"
 
 	"github.com/nats-io/nats.go"
 	"github.com/nats-io/nats.go/jetstream"
+	"github.com/spf13/cobra"
+	"github.com/spf13/pflag"
+	"github.com/spf13/viper"
 	"golang.org/x/mod/semver"
 )
 
-func main() {
-	var (
-		natsURL     = flag.String("nats", getEnv("INPUT_NATS_URL", "nats://localhost:4222"), "NATS server URL")
-		bucketName  = flag.String("bucket", getEnv("INPUT_BUCKET", "binaries"), "Object store bucket name")
-		directory   = flag.String("dir", getEnv("INPUT_SOURCE", "upload"), "Directory containing binaries to upload")
-		prefix      = flag.String("prefix", getEnv("INPUT_STRIP_PREFIX", ""), "Prefix to strip from paths (like 'upload/')")
-		binaryName  = flag.String("binary", getEnv("INPUT_BINARY", ""), "Binary name (defaults to first binary found)")
-		notifyTopic = flag.String("notify", getEnv("INPUT_NOTIFY_TOPIC", "binaries.update"), "NATS topic to publish update notification")
-		skipNotify  = flag.Bool("skip-notify", getEnvBool("INPUT_SKIP_NOTIFY", false), "Skip publishing update notification")
-		cleanup     = flag.Int("cleanup", getEnvInt("INPUT_CLEANUP", 0), "Keep only N most recent versions (0 disables cleanup)")
-		cleanupAll  = flag.Bool("cleanup-all", getEnvBool("INPUT_CLEANUP_ALL", false), "Cleanup all binaries, not just current one")
-		justClean   = flag.Bool("just-clean", getEnvBool("INPUT_JUST_CLEAN", false), "Dont upload, just cleanup old versions")
-	)
-	flag.Parse()
+type Config struct {
+	NatsURL     string `mapstructure:"nats"`
+	BucketName  string `mapstructure:"bucket"`
+	Directory   string `mapstructure:"dir"`
+	Prefix      string `mapstructure:"prefix"`
+	BinaryName  string `mapstructure:"binary"`
+	NotifyTopic string `mapstructure:"notify"`
+	SkipNotify  bool   `mapstructure:"skip-notify"`
+	Cleanup     int    `mapstructure:"cleanup"`
+	CleanupAll  bool   `mapstructure:"cleanup-all"`
+}
 
-	if *directory == "" && *cleanup == 0 && *justClean == false {
-		log.Fatal("Directory path is required or cleanup must be enabled")
+var rootCmd = &cobra.Command{
+	Use:   "nats-upload",
+	Short: "Upload binaries to NATS object store and cleanup old versions",
+	RunE: func(cmd *cobra.Command, args []string) error {
+		var cfg Config
+		if err := viper.Unmarshal(&cfg); err != nil {
+			return fmt.Errorf("failed to unmarshal config: %w", err)
+		}
+		if cfg.Directory == "" && cfg.Cleanup == 0 {
+			return errors.New("directory path is required or cleanup must be enabled")
+		}
+		return runUploadAndCleanup(cmd.Context(), &cfg)
+	},
+}
+
+var cleanCmd = &cobra.Command{
+	Use:   "clean",
+	Short: "Cleanup old versions in NATS object store",
+	RunE: func(cmd *cobra.Command, args []string) error {
+		var cfg Config
+		if err := viper.Unmarshal(&cfg); err != nil {
+			return fmt.Errorf("failed to unmarshal config: %w", err)
+		}
+		if cfg.Cleanup == 0 {
+			return errors.New("cleanup count must be greater than 0")
+		}
+		return runCleanupOnly(cmd.Context(), &cfg)
+	},
+}
+
+func init() {
+	cobra.OnInitialize(initConfig)
+
+	rootCmd.PersistentFlags().String("nats", "nats://localhost:4222", "NATS server URL")
+	rootCmd.PersistentFlags().String("bucket", "binaries", "Object store bucket name")
+	rootCmd.PersistentFlags().String("binary", "", "Binary name (defaults to first binary found)")
+	rootCmd.PersistentFlags().Int("cleanup", 2, "Keep only N most recent versions (0 disables cleanup)")
+	rootCmd.PersistentFlags().Bool("cleanup-all", false, "Cleanup all binaries, not just current one")
+	rootCmd.PersistentFlags().Bool("clean-all", false, "Alias for --cleanup-all")
+
+	rootCmd.Flags().String("dir", "upload", "Directory containing binaries to upload")
+	rootCmd.Flags().String("prefix", "", "Prefix to strip from paths (like 'upload/')")
+	rootCmd.Flags().String("notify", "binaries.update", "NATS topic to publish update notification")
+	rootCmd.Flags().Bool("skip-notify", false, "Skip publishing update notification")
+}
+
+func bindPFlag(fs *pflag.FlagSet, key string, flagNames ...string) {
+	name := key
+	if len(flagNames) > 0 {
+		name = flagNames[0]
+	}
+	if err := viper.BindPFlag(key, fs.Lookup(name)); err != nil {
+		log.Fatalf("error binding %s flag: %v", key, err)
+	}
+}
+
+func init() {
+	rootPersistentFlags := rootCmd.PersistentFlags()
+	for _, name := range []string{"nats", "bucket", "binary", "cleanup", "cleanup-all"} {
+		bindPFlag(rootPersistentFlags, name)
 	}
 
-	ctx := context.Background()
+	rootFlags := rootCmd.Flags()
+	for _, name := range []string{"dir", "prefix", "notify", "skip-notify"} {
+		bindPFlag(rootFlags, name)
+	}
 
-	nc, err := nats.Connect(*natsURL)
+	rootCmd.AddCommand(cleanCmd)
+}
+
+func initConfig() {
+	viper.SetEnvPrefix("INPUT")
+	viper.AutomaticEnv()
+
+	viper.RegisterAlias("nats_url", "nats")
+	viper.RegisterAlias("source", "dir")
+	viper.RegisterAlias("strip_prefix", "prefix")
+	viper.RegisterAlias("notify_topic", "notify")
+	viper.RegisterAlias("clean_all", "cleanup-all")
+}
+
+type NATSClient struct {
+	Conn  *nats.Conn
+	JS    jetstream.JetStream
+	Store jetstream.ObjectStore
+}
+
+func getNATSConnection(ctx context.Context, cfg *Config) (*NATSClient, error) {
+	nc, err := nats.Connect(cfg.NatsURL)
 	if err != nil {
-		log.Fatalf("Failed to connect to NATS: %v", err)
+		return nil, fmt.Errorf("failed to connect to NATS: %w", err)
 	}
-	defer nc.Close()
 
 	js, err := jetstream.New(nc)
 	if err != nil {
-		log.Fatalf("Failed to create JetStream context: %v", err)
+		nc.Close()
+		return nil, fmt.Errorf("failed to create JetStream context: %w", err)
 	}
 
-	store, err := js.ObjectStore(ctx, *bucketName)
+	store, err := js.ObjectStore(ctx, cfg.BucketName)
 	if err != nil {
 		store, err = js.CreateObjectStore(ctx, jetstream.ObjectStoreConfig{
-			Bucket:      *bucketName,
+			Bucket:      cfg.BucketName,
 			Description: "Binary storage for self-update",
 		})
 		if err != nil {
-			log.Fatalf("Failed to get/create object store: %v", err)
+			nc.Close()
+			return nil, fmt.Errorf("failed to get/create object store: %w", err)
 		}
-		log.Printf("Created object store: %s", *bucketName)
+		log.Printf("Created object store: %s", cfg.BucketName)
 	}
 
-	if *directory != "" && *justClean == false {
-		err = filepath.Walk(*directory, func(path string, info os.FileInfo, err error) error {
+	return &NATSClient{
+		Conn:  nc,
+		JS:    js,
+		Store: store,
+	}, nil
+}
+
+func runUploadAndCleanup(ctx context.Context, cfg *Config) error {
+	client, err := getNATSConnection(ctx, cfg)
+	if err != nil {
+		return err
+	}
+	defer client.Conn.Close()
+
+	if cfg.Directory != "" {
+		err := filepath.Walk(cfg.Directory, func(path string, info os.FileInfo, err error) error {
 			if err != nil {
 				return err
 			}
@@ -76,28 +172,28 @@ func main() {
 				return fmt.Errorf("failed to read %s: %w", path, err)
 			}
 
-			relPath, err := filepath.Rel(*directory, path)
+			relPath, err := filepath.Rel(cfg.Directory, path)
 			if err != nil {
 				return fmt.Errorf("failed to get relative path: %w", err)
 			}
 
 			objectKey := relPath
-			if *prefix != "" {
-				objectKey = strings.TrimPrefix(relPath, *prefix)
+			if cfg.Prefix != "" {
+				objectKey = strings.TrimPrefix(relPath, cfg.Prefix)
 			}
 
 			objectKey = filepath.ToSlash(objectKey)
 
-			if *binaryName == "" {
+			if cfg.BinaryName == "" {
 				parts := strings.Split(objectKey, "/")
 				if len(parts) >= 2 {
-					*binaryName = parts[0]
+					cfg.BinaryName = parts[0]
 				}
 			}
 
 			log.Printf("Uploading %s as %s (%d bytes)", path, objectKey, len(data))
 
-			_, err = store.PutBytes(ctx, objectKey, data)
+			_, err = client.Store.PutBytes(ctx, objectKey, data)
 			if err != nil {
 				return fmt.Errorf("failed to upload %s: %w", path, err)
 			}
@@ -107,64 +203,53 @@ func main() {
 		})
 
 		if err != nil {
-			log.Fatalf("Failed to upload files: %v", err)
+			return fmt.Errorf("failed to upload files: %w", err)
 		}
 
-		log.Printf("Successfully uploaded all files from %s to NATS object store '%s'", *directory, *bucketName)
+		log.Printf("Successfully uploaded all files from %s to NATS object store '%s'", cfg.Directory, cfg.BucketName)
 	}
 
-	if *cleanup > 0 {
-		log.Printf("Cleaning up old versions, keeping %d most recent", *cleanup)
-		err = cleanupOldVersions(ctx, store, *binaryName, *cleanup, *cleanupAll)
+	if cfg.Cleanup > 0 {
+		log.Printf("Cleaning up old versions, keeping %d most recent", cfg.Cleanup)
+		err := cleanupOldVersions(ctx, client.Store, cfg.BinaryName, cfg.Cleanup, cfg.CleanupAll)
 		if err != nil {
-			log.Fatalf("Failed to cleanup old versions: %v", err)
+			return fmt.Errorf("failed to cleanup old versions: %w", err)
 		}
 	}
 
-	if !*skipNotify && *notifyTopic != "" {
-		log.Printf("Publishing update notification to topic: %s", *notifyTopic)
+	if !cfg.SkipNotify && cfg.NotifyTopic != "" {
+		log.Printf("Publishing update notification to topic: %s", cfg.NotifyTopic)
 
-		message := fmt.Sprintf("binaries updated in %s", *bucketName)
-		err = nc.Publish(*notifyTopic, []byte(message))
+		message := fmt.Sprintf("binaries updated in %s", cfg.BucketName)
+		err := client.Conn.Publish(cfg.NotifyTopic, []byte(message))
 		if err != nil {
-			log.Fatalf("Failed to publish notification: %v", err)
+			return fmt.Errorf("failed to publish notification: %w", err)
 		}
 
 		// Flush to ensure message is sent
-		err = nc.Flush()
+		err = client.Conn.Flush()
 		if err != nil {
-			log.Fatalf("Failed to flush notification: %v", err)
+			return fmt.Errorf("failed to flush notification: %w", err)
 		}
 
 		log.Printf("✓ Published update notification")
 	}
+	return nil
 }
 
-func getEnv(key, defaultValue string) string {
-	if value := os.Getenv(key); value != "" {
-		return value
+func runCleanupOnly(ctx context.Context, cfg *Config) error {
+	client, err := getNATSConnection(ctx, cfg)
+	if err != nil {
+		return err
 	}
-	return defaultValue
-}
+	defer client.Conn.Close()
 
-func getEnvBool(key string, defaultValue bool) bool {
-	if value := os.Getenv(key); value != "" {
-		b, err := strconv.ParseBool(value)
-		if err == nil {
-			return b
-		}
+	log.Printf("Cleaning up old versions, keeping %d most recent", cfg.Cleanup)
+	err = cleanupOldVersions(ctx, client.Store, cfg.BinaryName, cfg.Cleanup, cfg.CleanupAll)
+	if err != nil {
+		return fmt.Errorf("failed to cleanup old versions: %w", err)
 	}
-	return defaultValue
-}
-
-func getEnvInt(key string, defaultValue int) int {
-	if value := os.Getenv(key); value != "" {
-		i, err := strconv.Atoi(value)
-		if err == nil {
-			return i
-		}
-	}
-	return defaultValue
+	return nil
 }
 
 func cleanupOldVersions(ctx context.Context, store jetstream.ObjectStore, currentBinary string, keepCount int, cleanAll bool) error {
@@ -184,12 +269,12 @@ func cleanupOldVersions(ctx context.Context, store jetstream.ObjectStore, curren
 			continue
 		}
 
-		binaryName := parts[0]
+		binName := parts[0]
 		arch := parts[1]
-		pathKey := binaryName + "/" + arch
+		pathKey := binName + "/" + arch
 
 		// If not cleaning all and this isn't the current binary, skip
-		if !cleanAll && currentBinary != "" && binaryName != currentBinary {
+		if !cleanAll && currentBinary != "" && binName != currentBinary {
 			continue
 		}
 
@@ -239,3 +324,13 @@ func cleanupOldVersions(ctx context.Context, store jetstream.ObjectStore, curren
 
 	return nil
 }
+
+func main() {
+	ctx, cancel := context.WithTimeout(context.Background(), 15*time.Minute)
+	defer cancel()
+
+	if err := rootCmd.ExecuteContext(ctx); err != nil {
+		_, _ = fmt.Fprintln(os.Stderr, err)
+		os.Exit(1)
+	}
+}